ReflectDrive-2: Zelfbewerking Afgestemd op Reinforcement Learning voor Discrete Diffusion Driving
ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving
May 6, 2026
Auteurs: Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan
cs.AI
Samenvatting
Wij introduceren ReflectDrive-2, een gemaskeerde discrete diffusie-planner met een gespecialiseerde actie-expert voor autonoom rijden, die plannen representeert als discrete trajecttokens en deze genereert via parallelle gemaskeerde decodering. Deze discrete tokenruimte maakt *in-place* trajectrevisie mogelijk: AutoEdit herschrijft geselecteerde tokens met hetzelfde model, zonder dat een hulpverfijningsnetwerk nodig is. Om deze capaciteit te trainen, gebruiken we een procedure in twee fasen. Eerst construeren we structuurbewuste perturbaties van expert-trajecten langs de longitudinale voortgangs- en laterale richtingsassen en trainen we het model onder supervisie om het oorspronkelijke expert-traject te herstellen. Vervolgens verfijnen we de volledige *decision–draft–reflect* roll-out met reinforcement learning (RL), waarbij we de uiteindelijke rijbeloning toekennen aan het definitieve, na-bewerking traject en de *policy-gradient credit* doorvoeren via volledige-roll-out transities. Volledige-roll-out RL blijkt cruciaal voor het koppelen van ontwerpen en bewerken: bij enkel onder supervisie getrainde modellen verbetert AutoEdit tijdens inferentie PDMS met maximaal 0.3, terwijl RL de winst verhoogt naar 1.9. We ontwerpen ook een efficiënte reflectieve decoderingstack voor de *decision–draft–reflect* pijplijn, die *shared-prefix KV reuse*, *Alternating Step Decode* en gefuseerd *on-device unmasking* combineert. Op NAVSIM behaalt ReflectDrive-2 een PDMS van 91.0 met alleen camera-input en 94.8 PDMS in een *best-of-6 oracle* setting, terwijl het draait met een gemiddelde latentie van 31.8 ms op NVIDIA Thor.
English
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.