ReflectDrive-2: Autoedição Alinhada por Aprendizagem por Reforço para Condução de Difusão Discreta

Resumo

Apresentamos o ReflectDrive-2, um planejador de difusão discreta mascarada com especialista de ação separado para condução autónoma, que representa planos como *tokens* de trajetória discretos e os gera através de descodificação mascarada paralela. Este espaço discreto de *tokens* permite a revisão in-situ da trajetória: o AutoEdit reescreve *tokens* selecionados usando o mesmo modelo, sem exigir uma rede de refinamento auxiliar. Para treinar esta capacidade, utilizamos um procedimento em duas fases. Primeiro, construímos perturbações com consciência da estrutura das trajetórias especialistas ao longo das direções de progressão longitudinal e de rumo lateral e supervisionamos o modelo para recuperar a trajetória especialista original. De seguida, afinamos o *rollout* completo de decisão-rascunho-reflexão com aprendizagem por reforço (RL), atribuindo a recompensa final de condução à trajetória pós-edição final e propagando o crédito do gradiente de política através de transições de *rollout* completo. O RL de *rollout* completo revela-se crucial para acoplar a geração de rascunhos e a edição: apenas com treino supervisionado, o AutoEdit em tempo de inferência melhora o PDMS no máximo 0.3, enquanto o RL aumenta o seu ganho para 1.9. Também co-projetamos uma *stack* de descodificação reflexiva eficiente para o *pipeline* decisão-rascunho-reflexão, combinando a reutilização de KV de prefixo partilhado, a Descodificação por Passos Alternados e o desmascaramento fundido no dispositivo. No NAVSIM, o ReflectDrive-2 atinge 91.0 de PDMS com entrada apenas de câmara e 94.8 de PDMS numa configuração *oráculo* de melhor-de-6, enquanto executa com uma latência média de 31.8 ms no 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.