Skip to content

Discrete Diffusion for Reflective Vision-Language-Action Models in Autonomous Driving

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Hang Zhao (Tsinghua); ICLR 2026; discrete diffusion VLA for end-to-end autonomous driving

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

ReflectDrive is an end-to-end autonomous driving framework that applies discrete diffusion to trajectory planning within a Vision-Language-Action (VLA) model, enabling gradient-free safety constraint enforcement at inference time. A reflective inference mechanism iteratively identifies unsafe trajectory tokens, corrects them via local discrete search, and regenerates surrounding segments via diffusion inpainting. On the NAVSIM benchmark, ReflectDrive achieves PDMS 91.1, approaching human-level performance (94.8).

Problem

Imitation learning-based VLA planners cannot inherently enforce hard safety constraints (collision avoidance, drivable-area compliance) because behavior cloning optimizes trajectory likelihood rather than constraint satisfaction. Existing fixes are unsatisfactory: rule-based post-refinement is brittle, reinforcement learning is largely confined to simulation, and continuous diffusion guidance requires expensive gradient computations that are slow and numerically unstable.

Method

ReflectDrive operates in two stages built on a discrete masked diffusion backbone.

Trajectory discretization. Continuous 2D waypoints are independently quantized along \(x\) and \(y\) axes into a uniform 1D codebook \(\mathcal{A}\) covering \([-100, 100]\) m with resolution \(\Delta_g = 0.3\) m (\(|\mathcal{A}| \approx 667\) tokens/dimension). A trajectory of \(N\) waypoints becomes a flat token sequence \(y = (y_{1,x}, y_{1,y}, \ldots, y_{N,x}, y_{N,y}) \in \mathcal{A}^{2N}\).

Discrete diffusion model. A pre-trained Diffusion Language Model (LLaDA-V) is fine-tuned as the VLA backbone via the masked denoising objective: $\(\mathcal{L}(\theta) = \mathbb{E}_{y,c,s,m^{(s)}}\!\left[-\sum_{i:\,m_i^{(s)}=1} \log p_\theta\!\left(y_i \mid \tilde{y}^{(s)}, c, s\right)\right]\)$ The masking-based forward process and parallel decoding at inference enable bidirectional inpainting natively.

Reflective inference. At test time, two stages run without any gradient computation: 1. Goal-conditioned generation: sample \(K'\) terminal waypoint candidates from \(p_\theta(y_N \mid c, s)\), apply NMS to obtain \(K\) spatially diverse goals \(G = \{G_1, \ldots, G_K\}\), generate \(K\) full trajectories via inpainting, and select \(\tau^* = \arg\max_k S_{\text{global}}(\tau_k)\). 2. Safety-guided regeneration: iteratively (up to 10 iterations) (a) detect unsafe waypoints with a Safety Scorer \(S_{\text{safe}}\); (b) for the earliest violation, run a local Manhattan-neighborhood search \(\mathcal{N}_\delta\) (\(\delta \le 10\) tokens) to find the token pair maximizing a Local Scorer \(S_{\text{local}}\); (c) fix corrected tokens as safety anchors and regenerate the surrounding sequence via diffusion inpainting. Falls back to best-seen candidate if budget exhausted.

Key Contributions

  • First application of discrete masked diffusion to end-to-end autonomous driving trajectory generation, integrated into a VLA model.
  • Gradient-free reflective inference that enforces hard safety constraints via discrete token search + diffusion inpainting, without modifying model weights.
  • Empirical validation on NAVSIM showing that the reflection mechanism resolves most safety violations in 1–3 iterations and achieves near-human closed-loop performance.

Results

All numbers from NAVSIM closed-loop evaluation (PDMS = weighted aggregate; higher is better):

  • ReflectDrive (w/o reflective inference): PDMS 84.8 — comparable to base E2E planners (UniAD 83.4, PARA-Drive 84.0).
  • ReflectDrive (full, constant-velocity agents): PDMS 91.1 — outperforms all camera-only baselines including GoalFlow (90.3) and AutoVLA (89.1).
  • Safety gains vs. ablation (w/o R.I. → full): DAC +3.9 pts (92.2 → 99.3 — highest among all methods), TTC +1.3 pts, NC +0.8 pts, EP +7.9 pts.
  • ReflectDrive† (ground-truth agent states oracle): PDMS 94.7 — matches human driving (94.8); NC 99.7, DAC 99.5, TTC 99.1, EP 88.9 (exceeds human EP of 87.5).
  • DAC 99.3 substantially outperforms all baselines including augmented E2E planners (GoalFlow: 98.3).

Limitations

  • Constant-velocity assumption for surrounding agents in the safety scorer degrades TTC and NC in interactive scenarios; accurate perception/prediction is a prerequisite for full gains.
  • Base model (w/o reflective inference) PDMS 84.8 is below augmented E2E planners, attributed to limited training data scale and headroom in the pre-trained VLM backbone.
  • Maximum 10-iteration budget caps worst-case safety correction; trajectories that cannot be repaired fall back to best-seen candidate rather than a guaranteed safe plan.
  • Local search neighborhood \(\mathcal{N}_\delta\) is intentionally small for efficiency — cannot accommodate large-scale intent changes (e.g., choosing a different turn at an intersection) during safety-guided regeneration; these require the goal-conditioned stage.
  • Evaluated only on NAVSIM; generalization to other benchmarks (nuPlan, Waymo Open) is not demonstrated.

Relevance to Vision-Language Models

ReflectDrive is a direct application of pre-trained VLMs (specifically masked diffusion LMs such as LLaDA-V) to embodied planning, demonstrating that VLM pre-training transfers to safety-critical trajectory generation via lightweight fine-tuning. The key insight — that discrete token representations enable constraint injection through search and inpainting without gradient access — is broadly applicable to any VLA system built on discrete generative models. This work extends the line of VLM-grounded autonomous driving (AutoVLA, DiffusionDrive) by showing that the inpainting property of masked diffusion specifically enables a new class of gradient-free inference-time alignment. For VLM researchers, it offers a concrete case study in how the architecture choice (discrete vs. continuous diffusion) shapes what kinds of post-hoc safety alignment are computationally feasible.