ReflectDrive-2: Autoedición Alineada con Aprendizaje por Refuerzo para la Conducción de Difusión Discreta
ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving
May 6, 2026
Autores: Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan
cs.AI
Resumen
Presentamos ReflectDrive-2, un planificador de difusión discreta enmascarada con un experto en acciones separado para la conducción autónoma, que representa los planes como tokens de trayectoria discretos y los genera mediante decodificación enmascarada paralela. Este espacio discreto de tokens permite la revisión in situ de trayectorias: AutoEdit reescribe tokens seleccionados utilizando el mismo modelo, sin necesidad de una red de refinamiento auxiliar. Para entrenar esta capacidad, utilizamos un procedimiento de dos etapas. Primero, construimos perturbaciones conscientes de la estructura de las trayectorias expertas a lo largo de las direcciones de progreso longitudinal y rumbo lateral, y supervisamos el modelo para que recupere la trayectoria experta original. Luego, afinamos el despliegue completo decisión-borrador-reflexión con aprendizaje por refuerzo (RL), asignando la recompensa final de conducción a la trayectoria posterior a la edición y propagando el crédito de gradiente de política a través de transiciones de despliegue completo. El RL de despliegue completo resulta crucial para acoplar el borrador y la edición: solo con entrenamiento supervisado, el AutoEdit en tiempo de inferencia mejora el PDMS como máximo en 0.3, mientras que el RL aumenta su ganancia a 1.9. También co-diseñamos una pila de decodificación reflexiva eficiente para la canalización decisión-borrador-reflexión, combinando la reutilización de KV de prefijo compartido, Decodificación por Pasos Alternados y el desenmascaramiento fusionado en el dispositivo. En NAVSIM, ReflectDrive-2 alcanza un PDMS de 91.0 con entrada solo de cámara y 94.8 de PDMS en una configuración oráculo best-of-6, mientras funciona con una latencia promedio de 31.8 ms en 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.