跳转至

Fast-WAM: Do World Action Models Need Test-Time Future Imagination?

会议: NeurIPS 2026
arXiv: 2603.16666
代码: https://yuantianyuan01.github.io/FastWAM/
领域: 机器人
关键词: 世界动作模型, 具身智能, 视频扩散 Transformer, 测试时计算, 流匹配

一句话总结

针对世界动作模型(WAM)推理时迭代去噪未来视频带来的高延迟瓶颈,本文提出解耦训练期视频建模与测试期显式未来生成的 Fast-WAM,仅保留训练阶段视频联合训练作为物理表征监督,在推理时单次前向直接生成动作,以 190ms 实时延迟和无具身预训练的极高数据效率达到主流 SOTA 水平。

研究背景与动机

构建通用具身智能体需要策略不仅能够建立从视觉观察到动作序列的端到端映射,还能理解并推理物理环境在机器人交互下的动态演化。相比于主要依赖静态图文预训练、缺乏物理演化先验的标准视觉-语言-动作(VLA)模型,世界动作模型(World Action Models, WAMs)将未来视觉预测与动作生成联合建模,被视为迈向更高阶物理交互理解的关键方向。

然而,现有 WAM 大多遵循“先想象后执行”(imagine-then-execute)范式,即在测试时必须首先迭代去噪生成未来观测画面,再基于生成的未来特征条件化预测动作。这种设计由于视频扩散去噪的多步迭代,带来了巨大的推理延迟(通常在 800ms 以上),严重制约了机器人在真实物理场景中的闭环控制与实时响应。更为本质的是,学术界一直未能厘清:WAM 带来的性能收益,究竟源自训练阶段视频预测目标赋予模型的强物理动力学表征先验,还是源于测试时显式“想象”未来所提供的额外前瞻线索?现存工作将这两者高度缠绕,缺乏受控的解耦验证。

本文的核心思考在于解耦训练目标与推理机制:如果世界建模的核心价值在于训练期塑造高质量的物理潜在表征空间,那么测试时完全可以剥离昂贵的未来视频生成开销。核心 idea:将训练期的视频协同预测与测试期的显式未来生成彻底解耦,提出 Fast-WAM 架构——训练时保留未来视频潜变量联合流匹配以注入物理动力学先验,测试时直接去除未来视频分支,仅通过视频 DiT 的单次前向编码提取当前世界潜在表征并生成动作,实现 190ms 毫秒级低延迟的实时控制。

方法详解

整体框架

Fast-WAM 旨在保留世界建模对表征学习的增益,同时完全消除推理阶段显式生成未来画面的昂贵计算。系统由预训练视频扩散 Transformer(Wan2.2-5B DiT 骨干、T5 文本编码器与 3D 因果视频 VAE)和轻量动作专家 DiT 组成,采用共享注意力的混合 Transformer(Mixture-of-Transformer, MoT)架构组织。

在训练阶段,模型接收当前首帧观测的干净潜在 token、未来时序视频帧的加噪潜在 token 以及动作 chunk token,利用统一的流匹配(Flow Matching)目标进行联合多任务优化;结构化注意力掩码确保动作 token 能够吸收当前视觉表征,但无法前向窥探未来视频 token,防止未来信息泄露。在推理阶段,Fast-WAM 直接将未来视频分支剔除,首帧观测潜变量仅需在视频 DiT 中执行单次前向传播编码为紧凑的世界表征,动作专家据此在 10 步流匹配去噪下直接生成连续动作 chunk,大幅收缩延迟。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:首帧观测图像 + 任务语言指令"] --> B["预训练视频 VAE 与 T5 文本编码"]
    B --> C["共享注意力 MoT 骨干网络<br/>Wan2.2-5B 视频 DiT + 1B 动作专家"]
    C --> D["单向信息隔离掩码:动作无法窥探未来视频"]
    D --> E["双分支协同流匹配训练<br/>动作速度场 + 未来视频潜在速度场"]
    D --> F["单次前向测试时推断<br/>移除未来分支,单遍编码提取世界表征"]
    F --> G["输出:实时连续动作块 Chunk"]

关键设计

1. 共享注意力的混合 Transformer(MoT)架构:实现高效的多模态统一表示与轻量动作解耦

为了在承接大规模视频生成先验的同时保持动作生成的灵活性与计算经济性,Fast-WAM 将 5B 参数的 Wan2.2 视频 DiT 作为世界建模底座,并引入隐藏层维度为 \(d_a = 1024\) 的 1B 参数动作专家 DiT。不同模态 token 在共享自注意力层中交互:视觉首帧图像被 VAE 编码为干净的潜在视觉锚点(visual anchor),语言任务指令经 T5 编码后通过跨注意力注入各分支。视频 DiT 和动作专家 DiT 各自承载对应模态的参数投影,避免了统一大网络在动作模态上的表达冗余,同时使得强大的预训练视频先验能无缝作为动作预测的物理特征供给底座。

2. 结构化单向信息隔离掩码:严格解耦协同训练与避免未来信息泄露

在联合训练过程中,token 被明确划分为三组:首帧干净潜在 token、未来帧加噪视频潜在 token,以及加噪动作 token。为了防止因果倒置或未来信息泄露到动作分支,模型施加了严格的结构化注意力掩码矩阵: - 未来加噪视频 token 仅在视频分支内部执行双向注意力,并允许单向关注首帧干净视觉锚点; - 动作 token 仅在动作专家内部双向交互,并单向关注首帧干净视觉锚点; - 核心设计在于动作 token 严禁关注任何未来加噪视频 token,首帧视觉锚点亦不反向关注未来视频与动作。 这种单向阻隔设计确保了动作策略学到的是由当前状态向物理隐空间投射的世界表征,使得模型在测试阶段直接移除未来视频分支时,动作分支的输入分布与注意力模式保持严格自洽,完全无需改变推理协议。

3. 联合流匹配目标函数与测试期未来剪枝:用训练动力学替代测试期迭代计算

Fast-WAM 对未来视频潜在表示与动作 chunk 均采用最优传输连续流匹配(Optimal Transport Flow Matching)建模。在时步 \(t \in (0, 1)\) 下,对真实目标 \(y\) 与高斯噪声 \(\epsilon \sim \mathcal{N}(0, I)\) 进行线性插值构建样本 \(y_t = (1-t)y + t\epsilon\)。网络参数化速度场以拟合边际向量场 \((\epsilon - y)\):

\[\mathcal{L}_{\mathrm{FM}}(y) = \mathbb{E}_{y, \epsilon, t} \left[ \| f_\theta(y_t, t, o, l) - (\epsilon - y) \|_2^2 \right]\]

总训练目标将动作速度场预测与未来视频速度场预测加权联合:

\[\mathcal{L} = \mathcal{L}_{\mathrm{act}} + \lambda \mathcal{L}_{\mathrm{vid}} = \mathcal{L}_{\mathrm{FM}}(a_{1:H}) + \lambda \mathcal{L}_{\mathrm{FM}}(z_{1:T})\]

其中 \(z_{1:T}\) 是经过 VAE 编码的未来时序视频潜变量。通过最小化 \(\mathcal{L}_{\mathrm{vid}}\),视频 DiT 的骨干参数被迫学习捕获时序交互和物体动力学演化特征;而在测试推断时,由于 \(\mathcal{L}_{\mathrm{vid}}\) 已经在权重中锁定了深层动力学感知能力,Fast-WAM 直接丢弃未来视频扩散链,仅由动作专家利用提取的潜在世界特征执行 10 步动作去噪,从机制上将测试时推理成本压缩至纯 VLA 级别。

损失函数 / 训练策略

模型采用 AdamW 优化器,学习率设为 \(1 \times 10^{-4}\),权重衰减为 0.01,并配合余弦退火学习率调度与 1.0 的梯度裁剪。训练过程采用混合精度以提升吞吐。动作预测界域设为 \(H = 32\),视频在时间维度进行 4 倍下采样(每个动作 chunk 对应 9 帧视频潜变量)。噪声调度遵循 logit-normal 分布,推理阶段去噪步数收缩至 10 步,无分类器引导(CFG)缩放系数设为 1.0。

实验关键数据

主实验

论文在双臂协同操作基准 RoboTwin 2.0(50+ 任务,涵盖干净与重度随机化环境)及终身学习基准 LIBERO(4 个子集共 40 个任务)上进行了系统评测。值得注意的是,Fast-WAM 完全不依赖大规模具身先验预训练(Embodied Pretraining),直接从视频预训练模型初始化微调。

基准与设置 评估指标 Fast-WAM (本文, 无具身PT) 之前SOTA (含具身PT) 性能对比与差距
RoboTwin 2.0 (Clean) 任务平均成功率 (%) 91.88 92.90 (LingBot-VA) 差距仅 -1.02%,超越 Motus (88.66%)
RoboTwin 2.0 (Rand.) 强随机化成功率 (%) 91.78 91.50 (LingBot-VA) 反超 +0.28%,超越 Motus (87.02%)
RoboTwin 2.0 (Average) 全任务平均成功率 (%) 91.8 92.2 (LingBot-VA) 胜过未预训练 LingBot-VA (80.6%) 与 Motus (77.3%)
LIBERO-Spatial 子集成功率 (%) 98.2 98.5 (LingBot-VA) 接近 SOTA,超越 OpenVLA (84.7%)
LIBERO-Object 子集成功率 (%) 100.0 99.8 (Motus) / 99.6 (LingBot) 达成 100% 满分表现
LIBERO-Goal 子集成功率 (%) 97.0 97.2 (LingBot-VA) / 98.0 (\(\pi_{0.5}\)) 保持一流水平
LIBERO-Long 长视程成功率 (%) 95.2 98.5 (LingBot-VA) / 97.6 (Motus) 远超 OpenVLA (53.7%) 与 \(\pi_0\) (85.2%)
LIBERO (Overall) 4套件平均成功率 (%) 97.6 98.5 (LingBot-VA) 超越 \(\pi_{0.5}\) (96.9%) 与 OpenVLA (76.5%)

消融实验

为了严格解耦“测试时显式未来想象”与“训练时视频协同建模”的贡献,作者构建了严格受控的统一架构消融变体: - Fast-WAM:训练时双分支流匹配协同优化,测试时不生成未来视频,单次前向直接出动作; - Fast-WAM-Joint:代表联合去噪范式,测试时未来视频 token 与动作 token 共同迭代去噪; - Fast-WAM-IDM:代表先视频后动作逆动力学范式,测试时先完全去噪生成未来视频潜变量,再条件化预测动作; - Fast-WAM w/o video co-train:完全移除训练期的视频预测损失函数 \(\mathcal{L}_{\mathrm{vid}}\),退化为普通策略微调。

配置变体 RoboTwin Average (%) LIBERO Average (%) 真实世界折毛巾成功率 (%) 单步推理延迟 (ms) 说明
Fast-WAM (本文完整模型) 91.8 97.6 ~42 190 训练带视频协同,测试不生成未来;实现 4x+ 推理加速
Fast-WAM-Joint (联合去噪范式) 90.6 98.5 ~40 ~650 测试期迭代生成未来视频与动作,延迟显著升高
Fast-WAM-IDM (先生成后动作范式) 91.3 98.0 ~48 810 测试期完整跑完未来视频去噪,延迟最长(>800ms)
Fast-WAM w/o video co-train 83.8 93.5 10.0 190 彻底移除视频协同训练损失,性能全面剧烈暴跌

关键发现

  • 训练期视频协同损失是决定性力量,测试期显式未来生成贡献极其有限:在 RoboTwin 上,完全移除视频协同损失使得成功率从 91.8% 暴跌至 83.8%(掉点 8.0%);在真实世界极其复杂的长视程可形变物体操作(折毛巾)中,去掉视频协同训练更导致成功率从 >40% 直接断崖式下跌至 10.0%。相反,对比测试期显式想象未来的 Fast-WAM-Joint 和 Fast-WAM-IDM,Fast-WAM 与它们处于完全同等的性能区间(RoboTwin 91.8% vs 90.6%/91.3%;LIBERO 97.6% vs 98.5%/98.0%)。这确凿证明了 WAM 的本质红利来自训练阶段视频预测目标塑造的物理世界潜在表征空间,而非测试期的显式画面发散。
  • 卓越的推理效率与工程部署优势:在单张 RTX 5090D GPU 测试下,Fast-WAM 仅需 190ms,相比 Fast-WAM-IDM 的 810ms 提速超过 4.2 倍,完美跨过真实机器人 5Hz 闭环控制门槛,且不需要维护庞大的去噪记忆显存。
  • 免具身预训练的高数据效率:在没有经历数千小时机器人预训练数据集洗礼的前提下,Fast-WAM 仅通过通用视频模型与少量演示微调,即逼近或超越了经历大量具身预训练的代表作(如 \(\pi_{0.5}\)、Motus 和 LingBot-VA)。

亮点与洞察

  • 颠覆 WAM 领域思维定势的反直觉洞察:先前学术界普遍深信具身智能体必须在脑海中逐步“脑补未来帧画面”才能做出前瞻决策。本文用极其干净的受控实验揭示:未来视频预测真正有价值的地方在于作为一种极佳的自监督表示学习手段,而非推理时的计算开销源。
  • 结构化掩码实现平滑剪枝:通过注意力掩码在训练期彻底阻断动作分支对未来帧的注意,实现了推理时物理分支的零副作用修剪,使得模型兼具世界模型的物理特征深度与 VLA 模型的简洁直推速度。
  • 泛化迁移启示:这种“重训练表征监督、轻测试时生成”的思想可以进一步推广到三维物理仿真、自动驾驶轨迹预测等任务中,用扩散或自回归生成作为辅助表征损失,而在部署端保持轻巧的单向推断头。

局限与展望

  • 作者承认的局限:目前的评测与架构主要围绕单次动作块预测(Single action chunk generation)展开,省略了更长视程下显式自回归滚动预测可能带来的长期误差积累分析;同时实验规模主要基于 6B 级别模型与特定微调场景。
  • 实验与设定观察:在极度依赖精确长程接触动力学的真实折毛巾任务中,Fast-WAM-IDM 仍保持微弱的成功率优势(约 48% 对 42%),说明在高度复杂的非刚体极限动态下,显式生成未来视觉特征可能仍保留少量尚未被完全蒸馏的信息,值得进一步探索特征蒸馏机制。
  • 未来改进方向:探索在更大规模具身数据预训练与模型扩展(Scaling Law)场景下,该结论的鲁棒性;探索能否引入轻量级隐空间动力学演化(Latent Dynamics Rollout)替代像素级去噪,以更低成本捕获超长视程推理增益。

相关工作与启发

  • vs Motus / Unified World Models: 后两者在推理时必须同时在共享潜空间中对未来视频与动作执行联合去噪,计算复杂度与时序长度深度绑定;本文证明了动作去噪无需绑定未来去噪,测试时单步视频前向即可释放同等能量。
  • vs LingBot-VA / Causal WAMs: 因果式 WAM 坚持先由扩散模型合成未来预测帧或未来特征,再经由逆动力学模块反解动作(耗时通常 >800ms);本文展示了只要注意力掩码设计得当,完全可以免除测试时显式生成未来,直接由动作专家端到端解码。
  • vs OpenVLA / \(\pi_{0.5}\): 标准 VLA 依赖静态互联网图文骨干,缺乏连续物理时空动力学建模;Fast-WAM 借鉴视频生成底座引入物理联合自监督损失,在无具身预训练条件下全面超越传统纯策略微调。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 犀利地挑战了主流 WAM 的惯性范式,解耦视角极具穿透力与启发性。
  • 实验充分度: ⭐⭐⭐⭐⭐ 从双臂模拟仿真、多任务套件到真实世界高难度非刚体折毛巾,对比消融变量控制极其严密。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑层层推进,问题陈述明确直接,图表呈现清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为具身控制在实时高频落地与大世界模型融合之间找到了极为实用且优雅的工程与理论折衷。