← 返回 2026-05-08

ReflectDrive-2: 强化学习对齐的自编辑离散扩散驾驶规划 ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

Huimin Wang, Yue Wang, Bihao Cui, Pengxiang Li, Ben Lu, Mingqian Wang, Tong Wang, Chuan Tang, Teng Zhang, Kun Zhan 📅 2026-05-06 👍 9 2026-07-13 08:36
VLA规划 强化学习 离散扩散 自动驾驶 自纠错

用RL联合优化起草与编辑,让离散扩散驾驶规划真正学会自纠错。

前置知识

掩码离散扩散

一种把序列生成建模为按比例逐步把 [MASK] 替换为真实 token 的过程。训练时按比例 $t\in[0,1]$ 随机遮蔽 token,模型从部分被遮蔽序列 $x_t$ 预测原始 token $x_0$。推理时从全 [MASK] 出发,少数几轮并行去掩码即可得到完整序列。

ReflectDrive-2 的轨迹生成完全建立在这一范式之上,并利用其'任意子集可重掩码并重写'的性质实现 AutoEdit 同模型内编辑。

VLA 规划器

Vision-Language-Action 规划器把视觉、导航语言指令与自车状态一起作为 token 输入到统一 Transformer,用同一骨干直接生成动作轨迹。本文的具体输入是环视相机 + 路线指令 + 自车状态,输出是 8 个 BEV 离散坐标 waypoint。

它是本文对比的基线范式(AutoVLA、DriveVLA、ReCogDrive),理解它才能体会 ReflectDrive-2 的'自车状态 + 视觉 + 语言 → 离散 token 轨迹'统一表示的关键。

模仿学习的协变量偏移

模仿学习在专家分布上训练,测试时若策略访问到专家很少去的状态,误差会沿因果链放大。论文指出驾驶策略误差集中在纵向速度误判与横向航向漂移两个轴上,正好是结构化失败最易发生的方向。

它是论文设计 AutoEdit 结构化扰动(纵向进度缩放、横向旋转)的动机来源,也是 RL 后训练被引入的根本原因。

策略梯度与组相对优势

把同组采样轨迹的回报 $R(\tau_g)$ 与组均值之差作为优势 $A_g$,再用 PPO clip 目标更新策略。本文把这一目标应用到离散 token 转移,并加入 $1\{x^{s+1}_{g,p}\neq x^s_{g,p}\}$ 指示,只给真正更新的 token 授信。

它是 ReflectDrive-2 在 SFT 之后做 RL 微调、并将信用同时分给起草与编辑两阶段的数学基础。

NAVSIM 闭环规划基准

基于 nuPlan 构建的非交互式闭环规划基准,给定环视传感器预测未来 4 秒 2 Hz 的 8 点轨迹。主要指标 PDMS 由无责碰撞 NC、可行驶区域合规 DAC、TTC、舒适性、进度 EP 加权得到。

论文所有数字都汇报于 NAVSIM,理解 PDMS 的合成方式才能看懂 Table 3 中 EP 从 82.2 飙到 89.4 这一单一维度的爆炸性提升。

研究动机

现有自动驾驶规划器在模仿学习设定下产生的错误并非随机分布,而是高度集中在两个轴上:纵向进度失误(超调、欠进度、刹车过晚)与横向航向漂移(车道偏离、转弯被裁剪、驶出可行驶区域)。经典模块化栈、端到端规划器只能给出一条轨迹,无法在原地修正;自回归 VLA 规划器需要串行 token 解码,纠错必须重新跑完整序列;连续扩散规划器虽然并行生成,但反向的是高斯噪声过程,无法对应驾驶员的结构化失败模式。更关键的是,即使把训练好的编辑器叠加在 drafter 之上,SFT-only 设定下推理期自编辑对 PDMS 贡献仅有 $\leq +0.3$——编辑能力在权重里沉睡。

本文的目标是本文要回答两个问题:(1) 能否设计一个表示,让'原地修正'天然成为同模型内的一种 token-to-token 重写;(2) 能否让 drafter 与 editor 在同一闭环奖励下联合优化,使推理期自编辑从 +0.3 PDMS 跃升到接近 +2 PDMS。最终目标是 NAVSIM 上 camera-only 设定达到 91.0 PDMS 单轨迹、94.8 PDMS best-of-6 oracle,并在 NVIDIA Thor 上跑出 30.2 ms 平均延迟。

与已有工作不同的是,此前最接近的 DriveFine 把 refiner 作为独立训练的第二阶段,drafter 与 editor 没有共享奖励;LLaDA 2.1 的 T2T 只用置信度启发式做解码期编辑,未针对下游控制任务训练;HDP、DriveFine 的 RL 都只优化单阶段 rollout。本文切入点是把'决策-起草-反思'三段统一进同一个共享离散 token 空间,再用 RL 把 terminal driving reward 同时授信给起草与编辑两个阶段的 token 转移,让二者共同适应同一个闭环目标。

核心方法

直观上,作者把一次驾驶决策拆成'先想好要去哪、再快速起草路径、最后在原地微调'三步。技术路线是一条统一离散 token 流:视觉、路线指令、自车状态三种条件信号与未来轨迹一起送进同一个双向 Transformer 骨干;先由目标点后验选出 $N_g$ 个行为假设(goal token),每个 goal 触发一次并行掩码扩散起草(4 轮左右把全 [MASK] 填成 8 个 waypoint = 16 个 BEV 坐标 token),最后 AutoEdit 在同一 token 空间对低置信度 token 做直接替换,无需重新掩码。训练分两阶段:先以结构化扰动(纵向 $\beta$ 缩放 + 横向 $\alpha$ 旋转)的 SAP 损失教会编辑能力,再以 PDMS 作为 terminal reward、组相对优势 + token 级 PPO clip 目标对完整 draft→edit 链路做 RL 微调,指示函数 $1\{x^{s+1}_{g,p}\neq x^s_{g,p}\}$ 让信用只流向真正发生变化的 token。

本质区别是把'自编辑'从解码期启发式或独立第二阶段升级为同一策略 rollout 的一个被 RL 显式优化的相位:terminal reward 落在 post-edit 轨迹上,policy gradient 同时穿过 drafting 与 AutoEdit 两个相位,二者因此在同一闭环目标下共同适应。配合同 token 空间的目标点后验提供行为级多模态、共享前缀 KV 复用与 Alternating Step Decode 让反射结构同时定义了一个高效运行时栈。AutoEdit 不是用 [MASK] 中转,而是把当前 concrete token 序列直接映射到替换 token,从根本上区别于 LLaDA 2.1 的 T2T 与 DriveFine 的独立 refiner。

方法步骤详情

输入 $o_t=(v_t,\ell_t,s_t)$ 经共享 Transformer 编码后分四步:(1) goal head 预测 BEV 离散后验,用 top-k + NMS 选 $N_g$ 个行为假设;(2) 对每个 goal 固定 goal token、其余轨迹 token 置 [MASK],经 $S_{draft}=3\sim5$ 轮并行去掩码生成 draft;(3) AutoEdit 取 concrete 序列预测替换 token,按 commit mask 直接原地写入,循环 $K$ 轮得 final;(4) 训练分两阶段:先以 $\mathcal{L}_{SAP}$(同 DLM 形式但输入为扰动 $\tilde x^0$)教编辑,再用 PDMS 作 reward、组相对优势 $A_g=R(\tau_g)-\frac{1}{G}\sum_j R(\tau_j)$,对全程 $S_{draft}+S_{edit}$ 个 token 转移用 clip 目标更新;可行驶区域 field loss 与 DLM 损失并入 SFT 总目标 $\mathcal{L}_{sup}$。

技术新颖性

三项贡献彼此耦合:(a) 决策-起草-反思的 goal-conditioned 掩码扩散 VLA,目标点作为行为级先验被 NMS 解耦为不同机动假设;(b) Reward-coupled AutoEdit,把结构化扰动预训练与全 rollout RL 微调结合,让同一策略优化起草与编辑;(c) Reflective decoding stack,把 AutoEdit 复用作跨帧时序细化器(ASD),配合共享前缀 KV 复用、Action-Expert FFN(hidden 4096→1024)、fused CUDA unmasking kernel,把单帧决策-起草-反射从 71.2 ms 压到 41.1 ms(full-step)或 19.3 ms(lite-step)。技术新颖性集中在'如何让 self-correction 真正成为被优化对象'这一以前被绕开的开放问题上。

ReflectDrive-2 architecture.
Figure 2: ReflectDrive-2 architecture.
Training, inference, and deployment of ReflectDrive-2.
Figure 3: Training, inference, and deployment of ReflectDrive-2.

实验结果

Table 3 揭示最关键发现:SFT-only 时推理期 AutoEdit 仅给 $\leq+0.3$ PDMS (DLM 84.8→85.0、+DACF 87.2→87.3、+SAP 训练 87.7→88.0),加 RL 后跃升到 +1.9(89.1→91.0)。Table 4 NAVSIM 单轨迹:end-to-end 方法 83.4–90.3、camera-only VLA 89.1–90.8,ReflectDrive-2 camera-only 达 91.0 PDMS(NC 97.3、DAC 98.1、TTC 92.5、Comf. 100.0、EP 89.4),EP 89.4 vs ReCogDrive 87.3 是单维度最大提升。Table 5 best-of-6 oracle 94.8 PDMS 与人类参考持平;Table 1 推理优化 71.2 → +KV 56.0 → +FFN 43.4 → +CUDA 41.1 ms,lite-step 19.3 ms,平均 30.2 ms。

Inference optimization chain on NVIDIA Thor.
Table 1: Inference optimization chain on NVIDIA Thor.
Quality gates for deployment optimizations.
Table 2: Quality gates for deployment optimizations.
Effect of inference-time AutoEdit across training regimes.
Table 3: Effect of inference-time AutoEdit across training regimes.
Closed-loop planning results on NAVSIM.
Table 4: Closed-loop planning results on NAVSIM.
Best-of-N evaluation on NAVSIM.
Table 5: Best-of-N evaluation on NAVSIM.
Effect of training components.
Table 6: Effect of training components.
Decision diversity from goal points.
Figure 4: Decision diversity from goal points.
Inference-time reflection with AutoEdit.
Figure 5: Inference-time reflection with AutoEdit.
Sensitivity to diffusion steps.
Figure 6: Sensitivity to diffusion steps.
Sensitivity to goal-proposal parameters.
Figure 7: Sensitivity to goal-proposal parameters.
查看结构化数据
任务指标本文基线提升
NAVSIM 闭环规划(单轨迹) PDMS 91.0 ReCogDrive 90.8、GoalFlow 90.3、DriveVLA-W0 90.2 +0.2 vs ReCogDrive;camera-only 设定下 SOTA
NAVSIM Best-of-6 Oracle PDMS 94.8 人类参考 94.8、单轨迹 91.0 +3.8 vs 单轨迹,与人类参考持平
推理期 AutoEdit 增益 ΔPDMS +1.9 (with RL) +0.3 (SFT-only) +1.6,证明 RL × AutoEdit 耦合的必要性
EP(ego progress) EP 89.4 ReCogDrive 87.3、GoalFlow 85.0 +2.1 vs ReCogDrive,单维度最大提升
部署延迟(NVIDIA Thor) 平均延迟 ms 30.2 (full 41.1 / lite 19.3) Baseline 71.2 ms -57.6%,即 2.4× 加速

局限与改进

作者明确承认三点:(1) 轨迹用固定分辨率 BEV 坐标 token 表示,空间精度受限于坐标 bin 大小,要更高精度需更细的词表或残差/混合连续 head;(2) RL 奖励是轻量级闭环 PDMS,作为真实驾驶目标的代理,高保真交互式仿真器与更丰富的安全奖励对齐仍是开放问题;(3) 结构化扰动目前只覆盖纵向 $\beta$ 与横向 $\alpha$,尚未覆盖让行时机、cut-in 响应、间隙选择等交互级失败模式。从评测角度看,camera-only NC 97.3 与 TTC 92.5 低于 ReCogDrive 的 97.9/94.9、DriveVLA-W0 的 98.7/95.3,说明'为了进度牺牲碰撞安全裕度'的潜在权衡尚未完全消除;best-of-6 使用 oracle 选择,不构成标准 benchmark,且端到端闭环仿真仍是缺席的强假设。

独立分析的弱点

三个弱点值得指出。第一,Table 3 中 RL 把 EP 从 82.2 提到 89.3,但 NC 从 97.8 跌到 96.3、TTC 从 93.6 跌到 88.9,说明 reward 对'进度'偏好过强、对碰撞/TTC 边际惩罚不足,需把 DAC/TTC 子项拆进 reward shaping 而非整体 PDMS。第二,论文没有 ground-truth 闭环仿真结果,所有数字均依赖 NAVSIM 非交互式 log-replay,碰撞检查本质上是其他 agent 沿 log 原轨迹继续运动的假设,与实际多智能体交互差距明显。第三,AutoEdit 的 commit mask 默认基于置信度,缺少与下游驾驶语义的显式关联(例如对位于非可行驶区域边缘的 token 应优先重写),目前主要靠 RL 隐式学习,对 OOD 场景可能失效。改进方向:(1) 解耦 reward 并加入不确定性正则;(2) 在 nuPlan closed-loop sim 做最终验证;(3) 让 commit mask 由可学习的 risk head 显式生成。

未来方向

作者提出三条主线:更细的离散/连续混合动作头提升空间精度;交互级仿真器 + 安全导向奖励;把 SAP 扰动扩展到让行时机、cut-in 响应、间隙选择。基于本文结果可延伸的方向还有:(1) 把 decision-draft-reflect 三阶段范式泛化到具身 VLA,让机器人操控也获得可编辑的离散动作空间;(2) 在多智能体交互中引入对手建模,把其他 agent 的离散意图 token 也纳入上下文,让目标点后验与 AutoEdit 同时考虑他人行为;(3) 把共享前缀 KV 复用与 ASD 思想推广到 VLM/VLA 在线服务,让多帧流式推理普遍受益于反射结构;(4) 探索 sparse-reward / curiosity 形式的 RL,让 AutoEdit 在缺乏稠密 PDMS 标注的长尾场景下也能获得编辑信号。

复现评估

代码与权重未开源(论文未给出 GitHub 链接),但论文披露了较充分的训练配置:0.7B 掩码扩散语言骨干 + 0.1B ViT,均从 LiAuto 自有预训练权重初始化;输入为前后两帧 left-front/front/right-front 环视 + 路线指令 + 自车状态;输出为 8 waypoint × 2 = 16 个 BEV 离散坐标 token;规划 horizon 4 秒 / 2 Hz;训练集 navtrain 1192 场景,评测集 navtest 136 场景。SFT 先做,再用 PDMS 作 reward 做 RFT。部署侧披露了 NVIDIA Thor 上的全链路延迟分解。复现难度较高:私有预训练权重是最主要门槛;之后需要复现自定义的 SAP 扰动、Action-Expert FFN、fused CUDA unmasking kernel、ASD 跨帧调度等组件,其中 CUDA kernel 与 KV 缓存 rewind 是工程量最大两部分。