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)으로 미세 조정하며, 최종 편집 후 궤적에 종단 주행 보상을 할당하고 전체 롤아웃 전이를 통해 정책 그래디언트 신용을 전파한다. 전체 롤아웃 RL은 초안 작성과 편집의 결합에 중요하다는 것이 입증되었다: 지도 학습만으로는 추론 시점 AutoEdit가 PDMS를 최대 0.3까지 개선하지만, RL은 이를 1.9까지 증가시킨다. 또한 의사결정-초안-반성 파이프라인을 위한 효율적인 반성 디코딩 스택을 공동 설계하여 공유 프리픽스 KV 재사용, 교대 단계 디코딩(Alternating Step Decode), 그리고 융합된 온디바이스 언마스킹을 결합하였다. NAVSIM에서 ReflectDrive-2는 카메라 전용 입력으로 91.0 PDMS, 6개 오라클 베스트 설정(best-of-6 oracle setting)에서 94.8 PDMS를 달성했으며, NVIDIA Thor에서 평균 31.8ms의 지연 시간으로 실행된다.

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.