26年6月来自小米EV的论文“Discrete-WAM: Unified Discrete Vision-Action Token Editing for World-Policy Learning”。

自动驾驶需要推理自身行为如何塑造周围环境的演变。然而,大多数端到端方法依赖于从状态到动作的直接映射,仅捕捉相关性,而未显式建模受动作影响的动态特性。相比之下,基于连续潜空间的“世界模型”往往缺乏针对反事实未来进行因果推理所需的组合式结构。为此,提出 Discrete-WAM——一种统一的潜视觉-动作世界策略。该模型将未来的视觉状态与自身动作表示为对齐的离散 Token,从而支持跨不同未来场景的组合式因果推理。基于这种统一的离散对齐机制,Discrete-WAM 构建一个具有统一生成任务的共享离散扩散框架,将世界建模、世界-动作策略及支持分层决策的策略整合在一起,从而在多样化的驾驶场景中实现组合式泛化。

如图 1 Discrete-WAM 概述。 Discrete-WAM 在统一的离散空间中联合编辑视觉、决策和动作token,通过统一的预训练和奖励引导的训练后提供可编辑的未来观察和规划轨迹。
请添加图片描述


0 世界策略建模

世界模型(World Models, WMs)已成为物理智能领域的核心范式,旨在建模外部世界如何随智体(agent)的动作而演变 [23]。在自动驾驶领域,早期受世界模型启发的方法通常结合规划与中间表征(如占用预测 [53, 71, 80] 或联合运动预测 [15, 25, 31, 52]),其中学习未来动态主要是为了支持下游的决策制定。随着大规模视觉基础模型 [38] 的出现,近期的研究日益转向利用未来的视觉信号进行监督 [21, 57, 77, 89],即利用图像和视频生成作为丰富的监督信号,用于表征学习 [75]、面向规划的推理 [91] 乃至显式或隐式的奖励建模 [44, 79]。与此同时,另一类研究侧重于仿真与数据生成引擎。这些方法通过结构化世界建模 [90](如行为仿真 [50]、三维场景重建 [20, 70] 或生成式视频模型 [77, 88])构建交互式场景或长尾场景。通过综合生成多样化的未来结果与反事实交互,它们为策略评估、闭环训练及安全性验证提供了可扩展的环境。尽管取得了一定成效,但现有大多数方法仍将世界建模与规划视为松散耦合的组件;在这些方法中,世界模型主要充当辅助预测器或仿真器,而非在统一的因果框架下与决策制定过程进行联合优化。

这一局限性推动近期“世界-动作建模”(World-Action Modeling, WAM)[104] 的发展,该方法将未来观测与策略生成统一建模为一个联合学习问题 [9, 43]。部分研究将未来世界与智体动作建模为单一的生成过程 [59, 82, 95],而另一些研究则强调统一预训练 [96]——涵盖共享的世界-策略表征 [5]、潜动作建模 [22] 以及价值-觉察世界模型 [37]——旨在减少显式观测生成的冗余,并增强环境动态与控制之间的耦合度。然而,现有的建模方案主要构建于连续潜空间之上,这往往会导致表征模糊性问题 [66, 98]。这些局限性促使人们近期开始采用基于离散token空间的并行生成范式,其中包括掩码建模(mask modeling)和离散扩散(discrete diffusion)等方法 [73, 84, 99]。尽管这些方法提高了生成效率并支持迭代优化,但大多数方法仍将决策过程建模为从观测到动作的直接映射,而未显式学习涵盖未来世界状态与策略的统一生成先验。相比之下,Discrete-WAM 将观测与动作统一置于共享的离散表征空间内。统一生成式预训练为世界演变与策略生成建立了共同先验,而联合离散扩散的建模方式则为世界建模与决策制定提供了一个统一的框架。

如图 2 所示:Discrete-WAM模型架构
请添加图片描述

1 Discrete-WAM 基础架构

Discrete-WAM 架构包含四个主要组件:(1) 一个视觉 VQ Tokenizer(将视觉观测编码为离散语义 Token)和一个投影器(将视觉特征与 Transformer 隐藏层维度对齐);(2) 一个上下文编码器(将自车状态和导航指令注入潜序列中);(3) 一个仅包含解码器(decoder-only)的 Transformer 主干网络(在统一的 Token 空间内联合建模观测、动作及未来演变);以及 (4) 用于世界建模、策略生成和“世界-动作”序列生成的多任务预测头。

2 视觉 Token 化

视觉 Tokenizer 将连续的相机观测转换为紧凑的离散视觉 Token,从而在原始图像与 Transformer 主干网络之间提供了一种 Token 级别的接口。这种离散表示使得视觉观测能够以与动作 Token 相同的序列格式进行处理,同时保留了下游世界建模和策略建模所需的场景语义。为了获得紧凑的视觉表示,沿用先前工作 [93] 的方法,预训练一个基于 VQ-VAE 的 Tokenizer [72],将前视相机图像编码为离散视觉 Token。

3 动作 Token 化

在动作表示方面,将连续的未来运动转换为一种兼容离散 Token 的表示形式,该形式基于加速度词表构建。给定跨越 H 个时间步的未来轨迹,首先利用三次样条(cubic spline)对离散轨迹进行拟合,以获得平滑的连续曲线。随后,利用二阶有限差分法计算每个未来时间步的自车坐标系下二维加速度(记为 a_x, a_y),其中 a_x 和 a_y 分别对应纵向和横向加速度。样条拟合步骤改善所得加速度序列的连续性,并确保对恢复出的加速度进行积分后,能与原始离散轨迹高度吻合。

通过将 a_x 和 a_y 的有效范围分别独立划分为 N_x 和 N_y 个区间(bin),构建一个均匀分布的二维加速度词表。这些区间的笛卡尔积构成了一个大小为 N_x * N_y 的网格结构动作词表。不将每个连续加速度向量简单地分配给单个最近的区间(这会引入确定性的硬量化误差),而是采用“软目标”(soft target)表示法,即利用相邻词表项的组合来表示该向量。具体而言,对于每个加速度分量,找出夹住该连续值的两个相邻区间中心,并根据其在两者间的相对位置分配插值权重。在二维加速度词表中,这种方法针对坐标 (ax, ay) 周围的四个相邻原型生成一个软标签(soft label)。在这种网格插值机制下,连续加速度值在每个加速度网格单元内均能以精确插值的形式表示。

在训练阶段,动作预测头(action head)针对该软目标分布(而非独热标签)通过交叉熵损失进行优化。在推理阶段,通过对加速度词表中的原型进行加权求和,可将预测的动作分布映射回连续加速度值。若能精确恢复软目标分布,则连续加速度便可精确重构为相邻词表原型的加权和。因此,所提出的软标签表示法消除确定性硬分配量化带来的误差,而剩余的重构误差则归因于分布预测的不匹配。

由此,未来的动作序列 A_t+1:t+H 得以在离散加速度词表上进行表示,同时通过软标签插值保留连续控制值。这种动作token化(tokenization)方案使得视觉token与动作 token能够在同一个 Transformer token空间内进行联合建模,同时避免因连续动作硬量化而引入的确定性误差。

4 统一的世界-策略(Unified World-Policy)

Discrete-WAM 将自动驾驶建模为一个基于离散视觉、决策和动作 Token 的统一世界-策略建模问题。在时间步 t,将场景上下文记为 C_t,其中包含历史视觉观测、自车状态信息和导航指令。未来的视觉观测表示为离散视觉 Token 序列 V_t+1:t+H,未来的动作策略表示为离散动作 Token 序列 A_t+1:t+H,其中每个动作 Token 对应于定义的量化加速度原型。进一步将高层决策条件记为 D_t,它捕捉稀疏的低频驾驶结构,例如机动意图、目标车道、粗略航点、速度趋势或交互优先级。

在此表示法下,Discrete-WAM 在一个共享的 Token 编辑接口下整合三种训练模式。这些模式的区别在于使用哪些 Token 流作为条件输入,以及将哪些 Token 流视为预测目标。

第一种模式是世界建模,其目标是视觉预测。给定场景上下文 C_t 和未来动作 Token A_t+1:t+H,模型预测未来的视觉 Token V_t+1:t+H。该任务训练模型理解场景在特定动作序列下的演变方式,并为下游策略学习提供受动作条件约束的世界动态。

第二种模式是世界-策略建模,它联合训练视觉预测和动作预测。在此模式下,模型在同一个 Token 序列内推理未来的动作及其引发的视觉后果。动作 Token 代表未来策略,而视觉 Token 代表相应的未来世界状态。这种表述方式促使模型将策略生成与世界演变耦合起来,而不是将它们作为独立的目标进行学习。

第三种模式是策略建模,其目标是先进行决策预测,随后进行动作预测。模型首先根据场景上下文 C_t 预测高层决策条件 D_t。在以该决策为条件的情况下,动作 Token 规划器随后预测未来的动作序列 A_t+1:t+H。这种层级分解将策略学习划分为两个层面:高层决策任务负责捕捉多模态驾驶选择,而底层动作预测任务则侧重于根据所选决策生成平滑且在时间上一致的轨迹。

这三种模式共同构成一个统一的架构,支持动作条件下的世界预测、世界与策略的联合学习,以及决策条件下的动作生成。


1 统一预训练

遵循统一的世界-策略(world-policy)表述,预训练过程针对多个任务族实例化条件化 Token 编辑目标。各任务使用相同的离散视觉和动作 Token 接口,但在哪些 Token 流被视为条件化输入以及哪些受损 Token 被作为编辑目标进行监督方面存在差异。这种设计使得世界建模、策略预测以及联合世界-策略建模能够共享同一个训练框架,同时保留各自任务特有的条件化结构。

训练任务。通过三个任务族来实现统一预训练:世界建模、策略建模以及联合世界-策略建模。

Token 化。将视觉观测和未来动作均 Token 化为离散 Token 序列,以便在统一的 Transformer 架构中进行联合建模。

对于图像 Token 化,采用预训练的 Token 化器。参考先前工作 [93](其 Token 化器与 [67] 对齐),该量化器包含一个大小为 KV 的码本(codebook)。每张输入图像被划分为互不重叠的 H_V × W_V 个图像块(patch),并编码为离散视觉 Token 序列。该过程独立应用于 H 个输入帧中的每一帧,从而生成视觉 Token 序列 V_t+1:t+H。

对于动作 Token 化,利用前述基于加速度的量化方法,将连续的未来轨迹离散化。具体而言,平滑后的轨迹被转换为以自我为中心(ego-centric)的纵向和横向加速度,随后量化为网格结构的加速度词表。由此产生离散动作 Token 序列 A_t+1:t+H,其中每个 Token 对应一个二维加速度原型(prototype)。

为了将动作 Token 纳入下游的离散扩散建模中,通过动作嵌入表(action embedding table)将每个动作 Token 与一个可学习的嵌入向量关联起来。在扩散训练期间,受损或部分掩码的动作 Token 经嵌入处理后,与视觉 Token 一同输入 Transformer;同时,模型在离散扩散目标下进行训练,以预测动作 Token。通过联合优化,动作嵌入学习离散运动原型的紧凑表征,并在共享隐空间中与视觉 Token 表征实现隐式对齐。因此,视觉和动作 Token 可由统一的 Transformer 联合处理,从而实现基于视觉条件的、通过离散扩散进行的未来动作生成。

世界建模。世界建模任务旨在学习受动作条件控制的未来视觉预测。给定场景上下文 C_t 和未来动作序列 A_t+1:t+H,该模型预测未来的视觉 Token 序列 V_t+1:t+H。

在训练过程中,未来的动作tokens通过“教师强制”(teacher forcing)机制作为条件输入提供,而监督信号则作用于未来的视觉tokens。该任务旨在训练模型捕捉不同的动作序列如何导致不同的未来世界演变。

策略建模。策略建模任务学习基于决策条件的层级化动作生成。模型首先根据场景上下文 C_t 预测高层决策骨架 D_t,随后在上下文和决策tokens的共同条件下,预测未来的动作序列 A_t+1:t+H:

世界-策略建模。该任务在不使用显式决策tokens的情况下,联合学习动作预测与基于动作的世界状态预测。给定场景上下文 C_t,未来的动作tokens与视觉tokens将沿预测时域排列成交错序列。在每一个未来时间步,动作预测基于已有的历史动作与视觉tokens进行;而视觉预测则基于已有的历史动作与视觉tokens以及当前的动作 tokens进行。这种机制与任务特定的注意掩码(attention mask)相匹配:每个预测模块均可关注其被允许访问的历史动作-视觉上下文,但无法获取其对应的干净目标tokens。

在训练过程中,相应的注意掩码(attention mask)基于“Token 编辑”范式实现这种依赖模式:带有噪声的动作 Token 和视觉 Token 利用允许的历史动作-视觉上下文进行编辑,而干净的目标 Token 则与对应的噪声 Token 保持隔离。

与此同时,世界-策略模型(world-policy model)始终对未来的视觉 Token 施加 Token 编辑监督,促使模型在动作条件上下文下,从受损的视觉 Token 恢复出未来的世界状态。此外,动作流也可采用同样的 Token 编辑策略进行训练,即受损的未来动作 Token 被编辑并向真实动作序列靠拢。这种联合监督使模型不仅能学习如何根据当前上下文生成合理的策略,还能学习生成的动作如何塑造后续的视觉演变。因此,世界-策略建模为动作生成、世界预测以及基于动作的反事实推理提供一个统一的优化目标。

任务特定的注意掩码。用任务特定的注意掩码来适配各训练目标所需的信息流。其基本原则是:历史上下文 Token 作为“教师强制”(teacher-forced)的条件信息可见,而当前或未来的受损 Token 则通过 Token 编辑进行预测,且在此过程中不泄露其对应的干净目标信息。

在世界建模(world modeling)中,未来的动作 Token 被视为纯净的条件输入,而非编辑目标。因此,该任务中的动作流不采用“纯净-噪声”双重填充机制。未来的视觉流采用 Token 编辑格式构建;模型基于历史观测、自我状态(ego-state)、导航 Token 以及给定的未来动作序列,来预测受损的未来视觉 Token。

在策略建模(policy modeling)中,上下文 Token 遵循因果注意结构,使模型能够基于历史观测和可用状态信息进行推理,而无需获取未来的上下文信息。未来的动作 Token 构成编辑目标块,并允许在动作块内部进行双向注意交互。这种双向动作注意机制使模型能够联合优化完整的未来动作序列,同时因果上下文掩码(causal context mask)防止来自观测端的未来信息泄露。

在世界-策略联合建模中,未来的视觉 Token 和动作 Token 均可采用 Token 编辑监督进行训练。对于每个可编辑流,采用双路径填充顺序,将纯净 Token 和噪声 Token 分置于不同的块中。注意掩码强制在纯净目标块与相应的噪声预测块之间实现双向隔离:噪声 Token 可以利用允许的上下文和特定任务的可见条件,但不能直接关注其对应的纯净目标。在训练过程中,教师强制(teacher forcing)机制提供历史真值(ground-truth)信息作为上下文,模型则学习将噪声视觉和动作 Token 编辑为纯净目标。

训练目标。统一的预训练目标结合 Token 级分类损失与连续运动重构损失。对于视觉 Token 编辑,模型在离散视觉词表上受交叉熵损失监督。与仅监督受损位置的目标函数不同,将 Token 分类损失应用于所有可编辑的视觉 Token 位置(包括纯净 Token 和噪声 Token)。对于受损位置,该损失训练模型从噪声输入中恢复原始的纯净目标;对于纯净位置,同样的损失鼓励恒等映射,要求模型保留已接近真实值的 Token,而非进行不必要的编辑。这种全位置监督为 Token 编辑提供隐式的停止信号:模型不仅学习如何修正噪声 Token,还学习何时无需进行编辑。当潜视觉嵌入(latent visual embeddings)可用时,进一步引入潜重构损失,以保留超越离散 Token 索引的细粒度视觉语义。针对动作Token(action token)的编辑任务,模型在离散动作词表上采用交叉熵损失进行监督。除了Token分类外,还引入运动层面的监督,以确保离散动作分布在物理上与连续的未来运动保持一致。首先,对与预测动作分布相关的解码加速度进行监督,促使离散动作Token保留准确的底层运动语义。其次,将预测的动作分布转换为连续加速度,并对其进行两次时间积分以重构未来的自车轨迹,进而对该轨迹施加轨迹层面的回归损失。此外,还针对积分得到的未来位置引入一种辅助的分解式位置分类损失。该辅助任务独立监督纵向和横向坐标,并对由预测动作Token生成的轨迹提供细粒度的空间约束。

获取用于运动重构的连续加速度的一种简单方法,是将预测的动作概率与加速度词表相乘,并利用所得的全分布期望值。然而,动作预测往往具有内在的多模态特性:不同的加速度模态可能对应于不同的合理驾驶决策。直接对整个预测分布求期望可能会导致不兼容模态的平均化,从而产生物理上不合理的中间加速度值。这一问题称为“解码引起的模态平均化”(decoding-induced mode averaging)。

为了缓解这一问题,仅在构建连续加速度、进行轨迹重构以及计算辅助位置监督损失时,采用模态-觉察(mode-aware)的解码策略。具体而言,利用多模态高斯混合模型(GMM)来拟合加速度词表上的预测分类分布;分别考察包含1个、2个和3个分量的GMM,并选择拟合误差最小的一个。随后,采用 top-p 采样选择一个高斯模态,对该模态内的动作概率进行重归一化,并基于重归一化后的分布与加速度词表计算期望加速度。最后,对所得的模态感知加速度进行时间积分,以重构轨迹。

这种模态-觉察(mode-aware)的解码方式,与用于动作token化的软标签插值相互独立。后者在构建动作目标时消除确定性硬分配带来的量化误差,而前者引入的目的则是为防止因对互不兼容的动作模态进行平均而导致连续重构损失。模态选择步骤引入额外的近似处理。
最终目标函数是各项依赖于任务的损失项的加权和。

训练流程。采用多阶段训练方案,以逐步实现视觉世界建模与动作生成的对齐。在第一阶段,同时针对“世界-策略建模”和“世界建模”任务进行视觉预训练。该阶段仅使用视觉预测损失,并通过“教师强制”(teacher forcing)机制将未来的动作 Token 作为条件输入。此阶段旨在训练出能够预测未来场景演变并与“动作条件下的动态特性”相一致的离散视觉表征。

在第二阶段,对视觉预测和动作预测进行联合训练。除了世界建模外,“世界-策略建模”任务还会激活动作预测损失,使模型能够在统一的 Token 编辑框架下,同时学习恢复未来的视觉 Token 和动作 Token。该阶段增强了策略生成与动作条件下的世界演变之间的耦合关系。
在第三阶段,用 LoRA 适配器进行动作微调。该阶段侧重于离散扩散策略模型,在保留早期阶段所学视觉世界表征的同时,专门针对未来动作的生成与精细化进行微调。

2 后训练阶段

为了进一步捕捉数据集中覆盖不足的罕见或安全关键行为,引入一个后训练微调阶段。该阶段利用基于模型的轨迹采样和强化学习,在保留既有行为能力的同时,针对具有挑战性的场景对策略进行优化。

策略采样。针对每个驾驶场景上下文 C_t,Discrete-WAM 利用 Token-Edit 规划器生成一组候选轨迹 {τ_1, . . . , τ_G},以实现高效探索。
训练目标。遵循分组相对策略优化(GRPO)范式,参考 [101] 中的方法,利用单步重构估计器计算当前策略 πA_ θ,i 下每个 token 的对数概率。


设置

数据集与基准。在 NAVSIM-v1 和 v2 基准 [8, 17] 上展示 Discrete-WAM 的端到端(E2E)生成与规划能力;该基准为端到端驾驶提供大规模驾驶场景。遵循标准协议,在 navtest 划分集上评估 Discrete-WAM,该集合包含以 2 Hz 采样的 12,000 个驾驶场景。NAVSIM 使用 PDMS 和扩展版 EPDMS 评估 4 秒时域轨迹的规划质量;其中,安全与规则相关的关键指标作为乘法约束纳入考量,而与进度及舒适度相关的项则通过加权聚合的方式进行组合。报告的指标包括:无责碰撞(NC)、可行驶区域合规性(DAC)、行驶方向合规性(DDC)、交通信号灯合规性(TLC)、自车进度(EP)、碰撞时间(TTC)、车道保持(LK)、历史舒适度(HC)以及扩展舒适度(EC)。

实现细节。Discrete-WAM 遵循多阶段训练流程。基于任务族的统一世界-策略预训练在完整的 nuPlan 训练集 [6] 上进行,训练步数为 200,000 步,学习率为 1 × 10⁻⁴。预训练完成后,Discrete-WAM 在 navtrain 数据集上进行 10 个 epoch 的监督微调(SFT)。最后,Discrete-WAM 进行另外 2 个 epoch 的强化学习(RL)后训练,以进一步优化规划行为。SFT 和 RL 后训练阶段均采用 LoRA 微调技术,学习率为 1 × 10⁻⁵。所有阶段均使用 32 块 NVIDIA H20 GPU 进行训练,采用 AdamW 优化器及余弦退火学习率调度策略。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐