Amortizing Trajectory Diffusion with Keyed Drift Fields

1. Problem Statement

Diffusion-based trajectory planners have emerged as powerful tools for offline reinforcement learning (RL) due to their ability to model complex, multimodal distributions of future behaviors. However, they suffer from a critical bottleneck: inference latency.

• Iterative Denoising: Standard diffusion models require $T$ sequential denoising steps to generate a single trajectory. In a receding-horizon control loop (where the planner runs at every time step), this results in high computational costs and slow planning rates, making them impractical for real-time, high-frequency control.
• The "One-Step" Challenge: The authors ask: Can we achieve diffusion-like planning behavior (multimodality, conditional controllability) with a single-step inference?
• The Failure Mode of Naïve Approaches: Simply training a one-step generator to match the data distribution in full trajectory space fails. In high-dimensional state-action sequences, distance metrics are dominated by unconstrained future dimensions. This causes the model to collapse toward "average" trajectories, losing action diversity and resulting in static, inert behavior (mode collapse).

2. Methodology: Keyed Drifting Policies (KDP)

The authors propose Keyed Drifting Policies (KDP), a one-step trajectory generator that amortizes the iterative refinement of diffusion models into the training process. The core innovation is a conditioning-aware neighborhood definition.

Key Concepts

1. Key Space vs. Full Space:

   • Key Space ($\kappa$): A compact space aligned with the condition (e.g., the current state $s_0$). Similarity between trajectories is measured here to ensure the model focuses on condition-matched data.
   • Full Trajectory Space: The actual updates (drifts) are applied to the full trajectory window (states and actions over time $H$).
   • Insight: By measuring "neighborhood" in the key space, the model avoids being misled by unconstrained future dimensions that dominate Euclidean distance in the full space.

2. Keyed Drift Field:
   KDP trains a generator $g_\psi$ using a drift field $V(x)$ that performs two simultaneous operations:

   • Attraction: Pulls generated trajectories toward dataset windows that match the current condition (key).
   • Repulsion: Pushes generated trajectories away from other samples generated by the current model to maintain diversity and prevent mode collapse.
   • Mathematical Formulation: The drift is a weighted difference between the mean of condition-matched dataset windows ($\mu^+$) and the mean of current model samples ($\mu^-$), where weights are computed via a softmax kernel based on key distance ($d_\kappa$).

3. Amortized Refinement (Stop-Gradient Target):
   Instead of applying the drift field iteratively at inference time, KDP "amortizes" this refinement into the model parameters during training:

   • A drifted target $\tilde{x}_i$ is computed as the current sample plus the drift field: $\tilde{x}_i = \text{sg}(x_i + V_i)$.
   • The generator is trained to regress toward this target using a weighted squared error loss.
   • Hard Constraints: The initial state (and optionally goal states) are clamped to the condition. The drift field is masked so it cannot modify these clamped coordinates, ensuring the plan starts exactly from the current observation.

4. Inference:
   At runtime, KDP performs a single forward pass. It generates $K$ candidate trajectories in parallel from noise and the current condition. A learned scorer selects the best candidate, and the first action is executed. This eliminates the $T$-step sequential loop.

3. Key Contributions

1. Keyed Drifting Policies (KDP): A novel one-step conditional trajectory generator that avoids the mode-collapse failure of naïve distribution matching by using a key-space metric for neighborhood selection.
2. Amortized Refinement: A training recipe that uses a stop-gradient drifted target to internalize the iterative refinement process of diffusion models, enabling single-step inference without sacrificing multimodal diversity.
3. Practical Deployment: An end-to-end evaluation demonstrating that KDP can run on real-time hardware (drones and robotic arms) with significantly lower latency than diffusion baselines while maintaining or improving task success.

4. Experimental Results

The authors evaluated KDP on standard D4RL benchmarks (locomotion, maze navigation, dexterous manipulation) and real-world hardware (Crazyflie drone, SO-100 arm).

• Performance vs. Latency Trade-off:

  • Locomotion (D4RL): KDP achieves performance comparable to strong offline RL baselines and significantly outperforms Diffuser (a standard diffusion planner) while reducing planning latency by orders of magnitude (e.g., ~4ms vs ~150ms per step).
  • Long-Horizon Planning (Maze2D/AntMaze): KDP outperforms Diffuser in success rates. Crucially, it allows for larger candidate sets ($K$) without increasing sequential depth, enabling better selection without latency penalties.
  • Dexterous Manipulation (Adroit): KDP shows substantial gains over Diffuser (e.g., 53.4 vs 10.7 score on Pen-clone). The one-step drift objective aligns better with the need for precise, contact-rich actions than iterative denoising.

• Real-World Hardware:

  • On a Crazyflie drone and an SO-100 arm, KDP achieved ~38Hz and ~17Hz replanning rates, respectively, compared to ~4Hz and ~3Hz for Diffuser.
  • Despite the massive speedup, KDP maintained or slightly improved task success rates (e.g., 94% vs 92% for navigation). This validates that faster replanning compensates for the lack of iterative refinement in dynamic environments.

• Ablation Studies:

  • No Keying: Removing the key-space distance (using full-window distance) causes performance to collapse (e.g., score drops from 90.7 to 2.91 in locomotion), confirming that conditioning-aware neighborhoods are essential.
  • No Repulsion: Removing the repulsion term leads to uncontrolled drift and action diversity collapse.
  • Self-Negative Masking: Excluding self-samples from repulsion is necessary to prevent the drift from becoming a no-op.

5. Significance

This paper addresses a fundamental limitation in applying generative models to real-time robotics: the latency of iterative sampling.

• Paradigm Shift: It moves generative planning from an "iterative refinement at inference" paradigm to an "amortized refinement at training" paradigm.
• Real-Time Viability: By reducing planning latency from hundreds of milliseconds to single-digit milliseconds, KDP makes diffusion-like multimodal planning viable for high-frequency control loops (e.g., drone flight, dexterous manipulation) where previous diffusion methods were too slow.
• Robustness: The method proves that a single-step generator can retain the benefits of trajectory-level uncertainty and diversity, provided the training signal is correctly aligned with the conditioning constraints.

In summary, KDP successfully bridges the gap between the expressive power of diffusion models and the strict latency requirements of real-world robotic control.