ReflectDrive-2 : Auto-édition alignée par apprentissage par renforcement pour la conduite de diffusion discrète
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
Résumé
Nous présentons ReflectDrive-2, un planificateur par diffusion discrète masquée doté d'un expert d'action distinct pour la conduite autonome, qui représente les plans sous forme de jetons de trajectoire discrets et les génère par un décodage masqué parallèle. Cet espace de jetons discrets permet une révision de trajectoire *in-place* : AutoEdit réécrit les jetons sélectionnés en utilisant le même modèle, sans nécessiter de réseau de raffinement auxiliaire. Pour entraîner cette capacité, nous utilisons une procédure en deux étapes. Premièrement, nous construisons des perturbations structurellement conscientes des trajectoires expertes selon les directions de progression longitudinale et de cap latéral, et supervisons le modèle pour qu'il retrouve la trajectoire experte originale. Nous affinons ensuite le déroulement complet décision--ébauche--réflexion par apprentissage par renforcement (RL), en attribuant la récompense terminale de conduite à la trajectoire finale post-édition et en propageant le crédit du gradient de politique à travers les transitions du déroulement complet. Le RL sur déroulement complet s'avère crucial pour coupler l'ébauche et l'édition : avec un entraînement supervisé seul, AutoEdit en inférence n'améliore le PDMS que de 0.3 au maximum, tandis que le RL porte ce gain à 1.9. Nous co-concevons également une pile de décodage réflexif efficace pour le pipeline décision--ébauche--réflexion, combinant la réutilisation KV de préfixe partagé, le décodage par étapes alternées et le démasquage fusionné sur puce. Sur NAVSIM, ReflectDrive-2 atteint un PDMS de 91.0 avec une entrée caméra uniquement et un PDMS de 94.8 dans un scénario oracle best-of-6, tout en fonctionnant avec une latence moyenne de 31.8 ms sur 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.