ReflectDrive-2: Саморедактирование, согласованное с обучением с подкреплением, для дискретного диффузионного вождения

ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

May 6, 2026

Авторы: Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan

cs.AI

Аннотация

Мы представляем ReflectDrive-2 — планировщик на основе маскированного дискретного диффузионного преобразователя с отдельным экспертом по действиям для автономного вождения, который представляет планы в виде дискретных токенов траектории и генерирует их посредством параллельного маскированного декодирования. Это дискретное токенное пространство позволяет осуществлять локальный пересмотр траектории: метод AutoEdit перезаписывает выбранные токены с использованием той же модели, без необходимости во вспомогательной сети уточнения. Для обучения этой возможности мы используем двухэтапную процедуру. Сначала мы создаем структурно-осознанные возмущения экспертных траекторий вдоль продольного направления движения и латерального направления по курсу и обучаем модель восстанавливать исходную эксперную траекторию. Затем мы дообучаем полный цикл «принятие решения — черновое планирование — рефлексия» с помощью обучения с подкреплением (RL), назначая итоговое вознаграждение за вождение окончательной, отредактированной траектории и распространяя градиенты политики через переходы полного цикла. Обучение с подкреплением на полном цикле оказывается ключевым для связи этапов чернового планирования и редактирования: при использовании только обучения с учителем AutoEdit во время вывода улучшает метрику PDMS максимум на 0.3, тогда как RL увеличивает этот выигрыш до 1.9. Мы также совместно разработали эффективный стек рефлексивного декодирования для конвейера «решение — черновик — рефлексия», сочетающий повторное использование ключей и значений (KV) для общего префикса, попеременное пошаговое декодирование (Alternating Step Decode) и объединенное аппаратное снятие маскирования. В симуляторе NAVSIM ReflectDrive-2 достигает показателя PDMS 91.0 при использовании только камер и 94.8 PDMS в оракульной настройке best-of-6, работая со средней задержкой 31.8 мс на 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.