ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

Huimin Wang; Yue Wang; Bihao Cui; Pengxiang Li; Ben Lu; Mingqian Wang; Tong Wang; Chuan Tang; Teng Zhang; Kun Zhan

Source-linked AI summary

ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan

arXiv:2605.04647v2May 6, 2026cs.RO

TL;DR

Autonomous-driving planners need trajectory representations that support correction along common longitudinal and lateral error axes, while self-editing stages remain difficult to couple effectively. ReflectDrive-2 uses masked discrete diffusion with token-space AutoEdit and full-rollout reinforcement learning, achieving 91.0 PDMS camera-only and 94.8 PDMS with best-of-6 oracle selection on NAVSIM.

Problem
Driving-planning errors concentrate along longitudinal progress and lateral heading, while supervised self-editing leaves drafting and correction stages decoupled.
Method
ReflectDrive-2 combines goal-conditioned masked discrete-diffusion trajectory drafting with same-model token rewriting and reinforcement learning over the complete draft-and-edit rollout.
Results
91.0 PDMS camera-only and 94.8 PDMS with best-of-6 oracle selection are achieved on NAVSIM.
Takeaways & Limitations
The results suggest masked discrete diffusion can provide an editable and deployable foundation for VLA driving policies.
Takeaways & Limitations
Fixed-resolution BEV coordinate tokens bound waypoint spatial precision, and the RL stage optimizes a lightweight closed-loop planning score that proxies real-world driving objectives.

Abstract

from arXiv · show

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.

1 Introduction

ReflectDrive-2 is a goal-conditioned masked discrete-diffusion planner that uses AutoEdit to revise trajectory tokens within a decision–draft–reflect process. Structure-aware perturbations and full-rollout reinforcement learning couple drafting with editing, while a reflection-aware decoding stack improves efficiency.

Motivation and training: Planning errors concentrate along longitudinal speed and lateral heading axes, motivating structure-aware perturbations for training AutoEdit.Examples include speed overshoot, under-progress, late braking, lane deviation, clipped turns, and drivable-area violations.
Planner and AutoEdit: Discrete diffusion enables any subset of trajectory tokens to be re-masked and rewritten by the same model without an auxiliary network or separate inference mode.This native revision operates in the shared discrete token space while conditioning on the remaining tokens.
Planner and AutoEdit: ReflectDrive-2 conditions on panoramic cameras, navigation, and ego state, then generates discrete trajectory tokens through goal proposal, parallel drafting, and in-place AutoEdit revision.Goal points represent candidate behavioral hypotheses, while final waypoint tokens anchor the behavior hypothesis and remaining tokens realize the 4-second plan.
Results and efficiency: 91.0 PDMS camera-only and 94.8 PDMS under best-of-6 oracle selection are achieved on NAVSIM, while the reflective decoding stack reaches 30.2 ms average latency on NVIDIA Thor.The stack combines shared-prefix KV cache reuse, Alternating Step Decode, and fused on-device unmasking.
Motivation and training: Full-rollout RL co-adapts the drafter and editor, addressing the weak coupling produced by supervised training alone and amplifying inference-time AutoEdit.With supervised training alone, the drafter and editor optimize separate losses and neither stage accounts for the other’s closed-loop behavior.

2 Related Work

Prior work spans end-to-end and VLA trajectory planners, discrete diffusion generation, editable token decoding, embodied refinement, and reinforcement learning for diffusion. ReflectDrive-2 is positioned as a closed-loop driving approach that jointly trains drafting and editing with post-edit reward.

Trajectory planning: End-to-end planners map sensor inputs directly to trajectories, while SMART tokenizes multi-agent trajectories for autoregressive next-token prediction.VLA planners inherit language priors but decode token-by-token, making latency scale with trajectory length.
Discrete diffusion: D3PM extends diffusion modeling to discrete variables, and MaskGIT enables parallel masked-token generation through confidence-based unmasking.Recent masked-diffusion language models include LLaDA and Seed Diffusion.
Editable planning: Discrete diffusion supports editable planning because arbitrary token subsets can be re-masked and regenerated.Existing token-editing methods are often decoding-time heuristics or independently trained refinement stages; LLaDA 2.1 T2T revises tokens according to model confidence.
Driving refinement: DriveFine introduces refinement-augmented masked-diffusion driving, but trains and optimizes its refiner separately from the drafter.ReflectDrive-2 instead treats drafting and editing as a single composed rollout.
Driving refinement: ReflectDrive-2 assigns terminal driving reward to the post-edit trajectory and applies policy-gradient credit to token transitions from both drafting and editing stages.This couples the drafter and editor under a shared closed-loop objective, unlike refinement work focused on multimodal understanding rather than closed-loop control.
Diffusion reinforcement learning: DDPO and DPPO apply policy gradients to continuous diffusion, while d1, d2, and SPG develop reinforcement-learning methods for discrete diffusion with differing treatments of multi-step structure and bounds.The passage distinguishes d1’s GRPO-style approach, d2’s step-aware gradients with group-relative advantage, and SPG’s tighter ELBO/EUBO bounds.

3 Preliminaries

The planner maps multimodal driving observations and route instructions to safe, comfortable, rule-compliant future trajectories using masked discrete diffusion. Its in-place editing interface supports trained trajectory revision, while reinforcement learning aligns drafting and editing with closed-loop driving rewards.

Problem formulation: At each time step, the ego vehicle receives panoramic visual tokens, route-level linguistic instructions, and ego-state information, then generates a future waypoint trajectory.The trajectory is represented as τ = {(x_k, y_k)}_K and should satisfy safety, comfort, rule compliance, and route consistency.
Masked diffusion: Masked discrete diffusion represents trajectories as BEV coordinate tokens, independently masks them, and uses a bidirectional Transformer to reconstruct originals from multimodal context.Inference starts from a fully masked sequence and performs parallel denoising in a small number of steps.
Selective re-generation: Selective re-generation enables arbitrary in-place rewriting by masking chosen tokens and denoising the resulting partial sequence from an edit-dependent effective time.AutoEdit retains this interface but replaces decoding-time heuristics with a trained operator coupled to the drafter through a shared RL reward.
Reinforcement learning: Reinforcement learning treats trajectory generation as a Markov decision process and optimizes expected closed-loop driving reward rather than only imitating the data distribution.The objective is J(θ) = E_τ∼π_θ(·|o)[R(τ)] with group-relative advantages and a discrete-diffusion policy gradient.
Reinforcement learning: The policy-gradient credit assignment restricts credit to tokens updated at each denoising step and jointly credits drafting and AutoEdit transitions.The total rollout length is S = S_draft + S_edit, coupling both stages through the same reward.

4 Method

ReflectDrive-2 unifies goal proposal, masked discrete trajectory drafting, and in-place token-space correction within a shared action representation. Its structure-aware supervision and full draft-and-edit reinforcement-learning rollout jointly train AutoEdit as part of the planning policy rather than as postprocessing.

Unified planning pipeline: The planner predicts discrete BEV goal-point hypotheses, drafts a goal-conditioned trajectory with masked diffusion, and selectively rewrites drafted tokens with AutoEdit.Goals represent behavior-level alternatives, while each selected goal conditions a separate masked-diffusion drafting branch.
Reinforcement learning: The RL stage assigns terminal driving reward to the final post-edit trajectory and propagates policy-gradient credit through both drafting and AutoEdit transitions.This rollout objective optimizes drafting for trajectories that can be improved and editing for corrections that improve closed-loop score.
Masked trajectory drafting: 8 waypoints yield a length-L = 16 trajectory token sequence, whose masked positions are filled in parallel over denoising rounds.Generation cost depends on denoising rounds rather than trajectory-token count, and the same masked-token interface supports later rewriting.
AutoEdit: AutoEdit directly predicts and commits replacement tokens from the concrete trajectory sequence, without re-masking, re-denoising, or an auxiliary refinement network.Low-confidence non-goal tokens are edited while goal tokens remain fixed as the behavior anchor.
AutoEdit training: Structure-aware perturbations target longitudinal progress errors and lateral heading deviations, training the shared model to map perturbed token sequences back to clean trajectories.Longitudinal scaling produces under-progress or overshoot, while lateral rotation creates coherent deviation while preserving smoothness.

5 Efficient Inference for Reflective Masked Planning

ReflectDrive-2 reduces reflective masked-planning latency through a stacked inference optimization chain combining shared-prefix caching, an action-expert FFN, alternating temporal decoding, and on-device token updates. The final deployment stack achieves 30.2 ms per frame on NVIDIA Thor.

Optimization chain: 30.2 ms per frame is achieved by the final stack, combining full decision–draft–reflect passes on full-step frames with lightweight temporal AutoEdit on lite-step frames.The optimization chain is applied cumulatively rather than as independent serving tricks.
Optimization chain: 56.8 ms to 42.2 ms: shared-prefix KV reuse avoids recomputing visual, route-instruction, and ego-state prefixes across proposal, drafting, and AutoEdit phases.The cache switches between Single and Batch states according to the serving phase.
Optimization chain: 4096 to 1024 hidden dimensions: the compact action-expert FFN lowers decode latency from 42.2 ms to 29.6 ms for constrained trajectory-token decoding.The action branch uses a short fixed-length token block and constrained action vocabulary.
On-device token updates: Fused on-device CUDA updates combine confidence ranking, token selection, and token-state updates for both masked drafting and token-to-token AutoEdit.The fusion avoids CPU synchronization at every update step.
Alternating Step Decode: 9.5 ms versus 27.3 ms: Alternating Step Decode uses full decision–draft–reflect passes on full-step frames and temporal token-to-token AutoEdit on lite-step frames.Lite-step frames transform the previous plan into the current ego frame and apply a short update instead of rebuilding the trajectory.

6 Experiments

Experiments on NAVSIM show that full-rollout RL substantially strengthens inference-time AutoEdit, yielding strong standard and oracle planning performance. Ablations and visualizations further support structured behavioral diversity, targeted trajectory revision, and efficient decoding.

Evaluation setup: ReflectDrive-2 is evaluated on NAVSIM using PDMS for 4-second ego-trajectory prediction from 1,192 training and 136 test scenes.PDMS aggregates collision, drivable-area, time-to-collision, comfort, and ego-progress criteria.
RL and AutoEdit: +1.9 PDMS is the inference-time AutoEdit gain after full-rollout RL, versus at most +0.3 PDMS before RL.This isolates the coupling between drafting and editing that supervised training alone does not achieve.
Standard results: 91.0 PDMS is achieved in standard single-trajectory evaluation with camera-only input, exceeding ReCogDrive (90.8) and GoalFlow (90.3).Ego progress reaches EP = 89.4, while DAC is 98.1 and comfort is 100.0; the camera-only advantage over listed VLA peers ranges from +0.2 to +1.9 PDMS.
Best-of-N results: 94.8 PDMS is reached in best-of-6 evaluation with oracle selection, recovering 3.8 PDMS over the single-trajectory result and matching the NAVSIM human reference.The recovered headroom indicates that goal proposals expose multiple behavior hypotheses rather than duplicate endpoint samples.
Qualitative analysis: Goal points generate distinct turning and interaction behaviors, while AutoEdit rewrites drafts around drivable boundaries and nearby agents in the same trajectory-token space.The revisions are structured rewrites rather than cosmetic smoothing.
Ablations and efficiency: +2.4 PDMS comes from field loss, RL raises EP from 82.2 to 89.3, and performance plateaus around 3–5 decoding or AutoEdit steps.More goal proposals expose additional hypotheses, with an NMS threshold near 1.2 m optimal because smaller thresholds retain duplicates and larger thresholds remove genuine alternatives.

7 Conclusion

ReflectDrive-2 frames autonomous driving as joint decision making, trajectory drafting, and self-correction in a shared discrete token space. Its editable masked-diffusion foundation supports lossless planning quality while fixed-resolution coordinates limit waypoint precision.

Core approach: ReflectDrive-2 combines goal-point hypotheses, parallel masked diffusion drafting, and same-policy AutoEdit within a shared discrete token space.AutoEdit rewrites drafts without an auxiliary repair network.
Conclusion: The central finding is lossless planning quality, suggesting masked discrete diffusion can provide an editable and deployable foundation for VLA driving policies.The conclusion presents editability and deployability as properties of the masked discrete diffusion foundation.
Limitations and future work: Fixed-resolution BEV coordinate tokens make trajectories interpretable and editable but bound waypoint spatial precision by the coordinate-bin size.Proposed remedies include finer coordinate vocabularies, residual offsets, or hybrid discrete-continuous action heads while retaining token-space editability.