ReflectDrive-2: Verstärkungslernen-angepasste Selbstbearbeitung für diskrete Diffusionssteuerung
ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving
May 6, 2026
Autoren: Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan
cs.AI
Zusammenfassung
Wir stellen ReflectDrive-2 vor, einen maskierten diskreten Diffusionsplaner mit separatem Aktions-Experten für das autonome Fahren, der Pläne als diskrete Trajektorien-Tokens repräsentiert und sie durch paralleles maskiertes Decodieren erzeugt. Dieser diskrete Token-Raum ermöglicht eine direkte Trajektorienrevision: AutoEdit überschreibt ausgewählte Tokens mit demselben Modell, ohne dass ein zusätzliches Verfeinerungsnetzwerk erforderlich ist. Um diese Fähigkeit zu trainieren, verwenden wir ein zweistufiges Verfahren. Zuerst konstruieren wir strukturierte Störungen von Expertentrajektorien entlang der Längsrichtung (Fortschritt) und der Querrichtung (Gierwinkel) und beaufsichtigen das Modell dabei, die ursprüngliche Expertentrajektorie wiederherzustellen. Anschließend feintunen wir den vollständigen Decision-Draft-Reflect-Rollout mit Reinforcement Learning (RL), wobei wir die finale Belohnung für das Fahrverhalten der endgültigen, bearbeiteten Trajektorie zuweisen und die Policy-Gradient-Credits über die vollständigen Rollout-Transitionen zurückverteilen. Vollständiges Rollout-RL erweist sich als entscheidend für die Kopplung von Entwurf und Bearbeitung: Unter rein beaufsichtigtem Training verbessert AutoEdit zur Inferenzzeit PDMS höchstens um 0,3, während RL den Gewinn auf 1,9 steigert. Wir entwickeln zudem einen effizienten Reflective-Decoding-Stack für die Decision-Draft-Reflect-Pipeline, der Shared-Prefix-KV-Wiederverwendung, Alternating Step Decode und fusioniertes Unmasking auf dem Endgerät kombiniert. Auf NAVSIM erreicht ReflectDrive-2 91,0 PDMS mit rein kamerabasierter Eingabe und 94,8 PDMS in einer Best-of-6-Orakel-Einstellung, bei einer durchschnittlichen Latenz von 31,8 ms auf 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.