ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving
:::info Stub — Full Engineering Breakdown Coming This paper was featured on Hugging Face Daily Papers on 2026-05-06 with 7 upvotes. A full breakdown with production viability rating, implementation notes, and honest limitations is being written. Subscribe to AI Letters → :::
| Authors | Huimin Wang et al. |
| Year | 2026 |
| HF Upvotes | 7 |
| arXiv | 2605.04647 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
We introduce ReflectDrive-2, a masked discrete diffusion planner with separate action expert for autonomous driving that represents plans as discrete trajectory tokens and generates them through parallel masked decoding. This discrete token space enables in-place trajectory revision: AutoEdit rewrites selected tokens using the same model, without requiring an auxiliary refinement network. To train this capability, we use a two-stage procedure. First, we construct structure-aware perturbations of expert trajectories along longitudinal progress and lateral heading directions and supervise the model to recover the original expert trajectory. We then fine-tune the full decision--draft--reflect rollout with reinforcement learning (RL), assigning terminal driving reward to the final post-edit trajectory and propagating policy-gradient credit through full-rollout transitions. Full-rollout RL proves crucial for coupling drafting and editing: under supervised training alone, inference-time AutoEdit improves PDMS by at most 0.3, whereas RL increases its gain to 1.9. We also co-design an efficient reflective decoding stack for the decision--draft--reflect pipeline, combining shared-prefix KV reuse, Alternating Step Decode, and fused on-device unmasking. On NAVSIM, ReflectDrive-2 achieves 91.0 PDMS with camera-only input and 94.8 PDMS in a best-of-6 oracle setting, while running at 31.8 ms average latency on NVIDIA Thor.
Engineering Breakdown
Plain English
ReflectDrive-2 is a discrete diffusion model for autonomous driving that generates driving plans as sequences of tokens and can edit selected tokens in-place without retraining. The key innovation is AutoEdit—a self-editing mechanism that rewrites trajectory tokens using the same model, trained first on supervised recovery of perturbed expert trajectories, then fine-tuned with RL to optimize actual driving rewards on the edited trajectories.
Key Engineering Insight
Using a discrete token representation of trajectories enables in-place editing without auxiliary networks, reducing model complexity and inference overhead. Training the edit capability through supervised perturbation recovery before RL fine-tuning creates a stable, compositional system where the model learns both to generate and refine plans with a single forward pass.
Why It Matters for Engineers
Production autonomous driving systems need to adapt plans quickly to changing conditions without expensive recomputation or additional model inference. This approach reduces latency by eliminating refinement networks and improves reliability by keeping edits within a learned, supervised token space rather than allowing unrestricted model regeneration.
Research Context
Prior diffusion-based driving planners required separate refinement networks or full regeneration to fix trajectory errors. ReflectDrive-2 advances this by showing that discrete tokenization + supervised edit training + RL alignment creates a single, efficient model that can both plan and self-correct. This opens the door to tighter planning loops in real autonomous systems where computational budget is fixed.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::