Fast-WAM 方法原理详解

Fast-WAM 方法原理详解:把"想象未来"留在训练中

0. 先用一分钟抓住全貌

Fast-WAM 要解决的不是"怎样更快地生成未来视频",而是一个更根本的问题:

World Action Model(WAM)的收益,到底主要来自训练时学习预测未来,还是来自推理时真的把未来视频生成出来?

传统 WAM 常采用"先想象、再行动":先用视频生成模型迭代去噪出未来画面,再根据这些画面生成机器人动作。形式上,它近似计算

\p(a_{1:H}\\mid o,l) = \\int p(v_{1:T}\\mid o,l)\\,p(a_{1:H}\\mid o,l,v_{1:T})\\,dv_{1:T}, \\

其中:

  • \(o\) 是当前视觉观测;
  • \(l\) 是语言指令;
  • \(v_{1:T}\) 是未来视频;
  • \(a_{1:H}\) 是未来 \(H\) 步动作块(action chunk)。

Fast-WAM 的观察是:"学会预测未来"和"测试时生成未来"是两件可以拆开的事。于是它采用以下策略:

  • 训练时:同时做未来视频预测和动作生成,让视频预测任务把 Video DiT 训练成懂运动、接触和物体变化的"世界表征器"。
  • 推理时:只输入当前一帧,不创建未来视频噪声,也不做未来视频去噪;当前帧经过 Video DiT 一次,所得逐层特征直接供 Action DiT 生成动作。

因此,更完整地写,Fast-WAM 推理时建模

\p_\\theta(a_{1:H}\\mid o,l,s) = p_\\theta(a_{1:H}\\mid z_\\theta(o;l,s),l,s), \\

其中 \(s\) 是当前本体状态;\(z_\theta(o;l,s)\) 是当前观测在语言和本体上下文调制下得到的隐式世界表征,而不是显式生成的未来画面。后文为了简洁,有时仍把全部条件统称为上下文 \(c\)。

一句话概括:

Fast-WAM 用"预测未来"教会模型理解世界,但部署时不让模型真的画出未来,只让它直接行动。

这带来论文中最重要的实验现象:保留视频协同训练但删除测试时未来生成,性能基本不掉;反过来,删除视频协同训练,性能明显下降。论文报告 Fast-WAM 在单张 RTX 5090D V2 上延迟约 \(190\text{ ms}\),而两阶段 Fast-WAM-IDM 约为 \(810\text{ ms}\)。


1. 为什么 WAM 原本要"先想象、再行动"

一个普通视觉-语言-动作策略(VLA)直接学习

\o,l \\longrightarrow a_{1:H}. \\

它的优点是直接,缺点是训练目标只要求"动作答对",未必迫使视觉骨干理解:手臂推一下物体会往哪里走、夹爪闭合后物体是否会被抓住、布料受力后会怎样变形。

WAM 增加未来视频预测任务,希望模型通过预测"接下来会看到什么"来学习物理动态。典型做法有两类:

  1. 联合生成:未来视频 token 和动作 token 一起从噪声迭代去噪,两者在每一步互相通信。
  2. 视频后接动作:先完整生成未来视频,再把生成结果送进逆动力学模型(IDM)或动作头。

两种方法都在测试时显式生成未来。问题在于视频 latent 很长,Video DiT 又很大;迭代生成未来视频会成为机器人闭环控制的主要延迟。

Fast-WAM 进一步追问:如果视频预测真正有价值的部分,是它在训练期塑造了更好的表示,那么推理期的未来视频也许只是昂贵的"中间产品",并非动作预测所必需。

可以把它类比成学生训练:做大量"预测下一步会发生什么"的练习,是为了形成物理直觉;考试时不必先画一整套未来分镜,才能回答下一步该怎么做。

1.1 Fast-WAM 真正解耦的两个因素

论文通过统一框架分别控制两个开关:

方法 训练时视频建模 推理时显式生成未来 动作如何使用视觉
Fast-WAM 只看当前帧的世界特征
Fast-WAM-Joint 动作与正在去噪的未来视频联合生成
Fast-WAM-IDM 先生成未来视频,动作再看完整未来
Fast-WAM w.o. video co-train 架构和 Fast-WAM 相同,但没有视频损失

这组对照回答了两个不同问题:

  • 比较 Fast-WAM 与 Joint/IDM:测试时显式未来想象是否必要?
  • 比较 Fast-WAM 与 no-video-co-train:训练时的视频预测监督是否重要?

2. 一条样本到底包含什么

先忽略网络,把一条训练轨迹切片看清楚。Fast-WAM 的一个训练样本包含:

\\\mathcal D_i= \\left( I_{0:32}\^{(1:C)}, s_{0:32}, a_{0:31}, l \\right). \\

其中 \(C\) 是相机数,\(I\) 是图像,\(s\) 是机器人本体状态,\(a\) 是动作,\(l\) 是任务指令。

当前开源配置固定取连续 \(33\) 个控制时刻,因此有 \(32\) 个状态转移和 \(32\) 个动作。模型不是把全部 \(33\) 帧都送进视频分支,而是按 \(4\) 倍时间间隔抽帧:

\\\{I_0,I_4,I_8,\\ldots,I_{32}\\}, \\

总共 \(9\) 帧。动作仍保留 \(32\) 步,所以相邻两张训练视频帧之间对应 \(4\) 个控制动作。

2.1 为什么是 33、32 和 9

  • 观察窗口有 \(33\) 个时刻:\(0\) 到 \(32\)。
  • 相邻时刻之间有 \(32\) 次动作:\(a_0\) 到 \(a_{31}\)。
  • 每隔 \(4\) 个时刻取一帧:\(32/4+1=9\) 帧。
  • Wan VAE 要求输入帧数满足 \(T\bmod 4=1\),而 \(9\bmod 4=1\),正好合法。

这里的"9 个视频帧"是 VAE 输入端的 RGB 帧数。Wan VAE 还会做 \(4\) 倍时间压缩,所以 \(9\) 张 RGB 帧编码后只剩

\T_z=\\frac{9-1}{4}+1=3 \\

个 latent 时间片。不要把"9 个 RGB 帧"和"3 个 latent 帧"混为一谈。

2.2 多相机如何处理

Fast-WAM 没有为每个相机建立独立视觉编码器,而是先把同一时刻的多相机图像拼成一张大图,再送入 VAE。

数据集 相机 单相机预处理 拼接方式 最终图像 动作维度 本体状态维度 归一化
LIBERO 主视角 + 腕部 各缩放到 \(224\times224\) 水平拼接 \(224\times448\) \(7\) \(8\) min-max 到约 \(-1,1\)
RoboTwin 头部 + 左腕 + 右腕 头部 \(256\times320\);双腕各 \(128\times160\) 头部在上,双腕并排在下 \(384\times320\) \(14\) \(14\) z-score

最终 RGB 会按

\x_{\\mathrm{img}}=2\\frac{x_{\\mathrm{uint8}}}{255}-1 \\

映射到 \(-1,1\)。

2.3 动作、状态和 padding

LIBERO 的 \(7\) 维动作可理解为 \(6\) 维末端执行器位姿变化加 \(1\) 维夹爪;其 \(8\) 维状态包含 \(6\) 维末端位姿和 \(2\) 维夹爪状态。RoboTwin 的 \(14\) 维动作/状态对应双臂关节控制表示。

轨迹靠近 episode 末尾时,未来长度可能不够。数据集会复制末端数据补齐,同时生成:

  • action_is_pad:哪些动作步是补齐的;
  • image_is_pad:哪些图像帧是补齐的;
  • proprio_is_pad:哪些状态步是补齐的。

训练 loss 会屏蔽补齐的动作和视频时间片。LIBERO 对 delta 位姿维度还有一个细节:若某个补齐动作是相对位移,它会被置为 \(0\);夹爪维度不是 delta,不按同一规则处理。

视频 padding 从 RGB 帧映射到 VAE latent 时间片时遵循 VAE 的 \(4\) 倍时间分组。对于 \(9\) 个 RGB 标志 \((m_0,\ldots,m_8)\),两个未来 latent 标志为

\m\^{z}_1=\\bigwedge_{j=1}\^{4}m_j,\\qquad m\^{z}_2=\\bigwedge_{j=5}\^{8}m_j. \\

也就是说,一组 \(4\) 帧全部是 padding 时,对应 latent 时间片才视作 padding。第一 latent 帧是条件,本来就不计视频 loss。

2.4 文本与本体状态如何进入模型

原始任务文本会套入固定模板:

text 复制代码
A video recorded from a robot's point of view executing the following instruction: {task}

训练前,T5 会把每个去重后的指令编码为最多 \(128\) 个、每个 \(4096\) 维的文本向量,并缓存到磁盘。训练时直接读取缓存,因此不反复运行大文本编码器。

本体状态没有单独形成一条 Transformer 分支。模型只取窗口第一个时刻的状态 \(s_0\),通过线性层投影到 \(4096\) 维,再作为一个额外的上下文 token拼在文本 token 后面:

\c=\[E_{\\mathrm{T5}}(l);W_s s_0+b_s. \]

Video DiT 和 Action DiT 都通过 cross-attention 读取这组上下文。于是语言告诉模型"要做什么",本体 token 告诉模型"机器人现在在哪里"。

2.5 数据构造伪代码

text 复制代码
函数 构造训练样本(轨迹, 起点):
    # 1. 取 33 个控制时刻;越过 episode 结尾时复制最后一个时刻并记录 padding
    图像序列, 状态序列, 动作块, padding掩码 = 滑窗读取(轨迹, 起点, 长度=33)

    # 2. 每个相机独立缩放,再按数据集规定拼成一张图
    对 t = 0...32:
        多相机大图[t] = 拼接相机(图像序列[t])

    # 3. 视频分支每 4 个控制时刻取一帧,得到 9 帧;动作仍是完整 32 步
    视频 = 多相机大图[[0, 4, 8, ..., 32]]
    动作 = 动作块[0:32]

    # 4. 归一化图像、动作和本体状态
    视频 = 图像归一化到负一至一(视频)
    动作 = 数据集统计量归一化(动作)
    状态 = 数据集统计量归一化(状态序列)

    # 5. 读取预计算的 T5 文本 embedding
    文本上下文 = 读取缓存(固定模板.format(任务指令))

    返回 视频, 动作, 状态, 文本上下文, padding掩码

3. 模型结构:两个专家,一次受控的信息交换

Fast-WAM 总参数量约 \(6\text{B}\):

  • Video DiT :来自预训练 Wan2.2-TI2V-5B,约 \(5\text{B}\),负责视频 latent 和当前帧世界表示。
  • Action DiT :约 \(1\text{B}\),负责生成 \(32\) 步动作块。
  • 视频 VAE:把 RGB 视频压缩为 latent,训练中冻结。
  • T5 文本编码器:编码任务指令,训练中冻结且通常提前缓存。

两条 DiT 分支组成 Mixture-of-Transformer(MoT)。这里的 "Mixture" 不是稀疏 MoE 路由:每一层始终有一个视频专家和一个动作专家,两者参数分开,但在 self-attention 处把 token 放在一起做一次带掩码的 mixed attention。

3.1 关键尺寸

项目 Video DiT Action DiT
层数 \(30\) \(30\)
hidden size \(3072\) \(1024\)
FFN size \(14336\) \(4096\)
attention heads \(24\) \(24\)
每头维度 \(128\) \(128\)
Q、K、V 各自的投影维度 \(24\times128=3072\) \(24\times128=3072\)
输入 token 视频 patch 每个动作一步一个 token

虽然两个专家的 hidden size 不同,但注意力头数和 head dimension 相同,所以两边投影后的 \(Q\)、\(K\)、\(V\) 各自 都是 \(3072\) 维,可以在序列维拼接后共同计算注意力。注意力输出切回各自序列后,再由各专家自己的输出投影、残差、cross-attention 和 FFN 继续处理。

3.2 视频怎样变成 token

Wan VAE 把 RGB 视频编码为

\z\\in\\mathbb R\^{B\\times48\\times T_z\\times H/8\\times W/8}. \\

Video DiT 再用大小为 \((1,2,2)\) 的 3D patch embedding,也就是时间上不继续压缩,空间上每 \(2\times2\) 个 VAE latent 像素形成一个 token。

以 LIBERO 为例:

  • 输入拼接图为 \(224\times448\);
  • VAE latent 空间为 \(28\times56\);
  • DiT patch 后每个 latent 时间片有 \(14\times28=392\) 个 token;
  • \(9\) 个 RGB 帧经时间压缩成为 \(3\) 个 latent 时间片;
  • 训练视频序列共有 \(3\times392=1176\) 个 token;
  • 推理只保留第一个 latent 时间片,因此只有 \(392\) 个视觉 token。

RoboTwin 最终图为 \(384\times320\),对应每个 latent 时间片 \(24\times20=480\) 个 token;训练为 \(1440\) 个视频 token,Fast-WAM 推理只用 \(480\) 个。

把常用的无 batch 形状集中列出如下。文本缓存固定补齐到 \(128\) 个 token;拼入一个本体 token 后,上下文最多是 \(129\times4096\)。

数据集 RGB 视频 \(T,C,H,W\) VAE latent \(C_z,T_z,H_z,W_z\) 视频 token 动作 \(H,D_a\) 状态窗口 实际取的状态
LIBERO \(9,3,224,448\) \(48,3,28,56\) \(1176\) \(32,7\) \(33,8\) \(s_0\in\mathbb R^8\)
RoboTwin \(9,3,384,320\) \(48,3,48,40\) \(1440\) \(32,14\) \(33,14\) \(s_0\in\mathbb R^{14}\)

动作 \(a_t\) 与从控制时刻 \(t\) 到 \(t+1\) 的转移对齐。具体控制周期以及 RoboTwin \(14\) 个分量的机器人学语义由原始数据/环境定义,论文和当前 Fast-WAM 配置没有给出足够信息,本文不作猜测。

3.3 动作怎样变成 token

归一化动作块为

\a\\in\\mathbb R\^{B\\times H\\times D_a}, \\

其中默认 \(H=32\),LIBERO 的 \(D_a=7\),RoboTwin 的 \(D_a=14\)。一个线性层把每一步动作投影到 \(1024\) 维,因此恰好有 \(32\) 个 action token。动作位置通过一维 RoPE 表示。

Action DiT 的骨干不是完全随机初始化。代码先把 Video DiT 中形状相同的参数直接复制;形状不同的参数沿各维做线性插值缩放。若最后一维从视频宽度 \(d_v\) 缩到动作宽度 \(d_a\),还乘

\\\alpha=\\sqrt{\\frac{d_v}{d_a}} \\

以调整幅值。动作输入层和输出头保持随机初始化。这使动作专家从一开始就继承视频 Transformer 的结构先验,又允许其宽度缩到 \(1024\)。

3.4 每一层 MoT 实际做什么

设第 \(m\) 层的视频和动作表示分别为 \(X_v^{(m)}\)、\(X_a^{(m)}\)。每个专家先独立产生

\(Q_v,K_v,V_v)=g_v\^{(m)}(X_v\^{(m)}), \\

\(Q_a,K_a,V_a)=g_a\^{(m)}(X_a\^{(m)}). \\

然后在序列维拼接:

\Q=\[Q_v;Q_a,\quad K=K_v;K_a,\quad V=V_v;V_a, \]

并执行

\Y=\\operatorname{softmax}\\left(\\frac{QK\^\\top}{\\sqrt{d_h}}+M\\right)V, \\

其中 \(M\) 是结构化 attention mask。得到的 \(Y\) 再按视频/动作长度切开,分别走各自专家的输出投影、语言 cross-attention、FFN 和残差连接。

因此"共享注意力"准确地说是:共享一次带掩码的注意力计算和可见键值集合,不共享整套 Transformer 参数。

位置编码也仍由各专家按自己的数据结构生成:视频 token 使用覆盖时间、高度和宽度的 3D RoPE,动作 token 使用动作序列上的 1D RoPE。两者都把旋转编码应用到各自的 \(Q/K\),随后才把同维度的 \(Q/K/V\) 沿序列拼接;并不是把动作位置硬映射到某个视频空间坐标。

3.5 结构化注意力掩码是整个方法的核心

令:

  • \(f_0\):干净的当前帧 latent token;
  • \(f_+\):加噪后的未来帧 latent token;
  • \(a\):加噪后的动作 token。

行表示 query,列表示它能读取的 key。Fast-WAM 训练时的可见关系是:

Query \(\downarrow\) / Key \(\rightarrow\) 当前帧 \(f_0\) 未来帧 \(f_+\) 动作 \(a\)
当前帧 \(f_0\)
未来帧 \(f_+\)
动作 \(a\)

这张表有四个必须理解的含义:

  1. 当前帧是共享视觉锚点,但它不读取未来或动作,所以自身不会被未来信息污染。
  2. 未来视频 token 可以双向读取整个视频分支,用当前帧和其他未来 token 完成视频去噪。
  3. 动作 token 彼此双向可见,因此一次生成完整动作块,而不是自回归地逐动作生成。
  4. 动作永远不能读取未来视频 token,所以训练时不存在"看答案":部署时删除未来分支后,动作的条件分布没有突然少掉一类输入。

所有视频和动作 token 还会通过各自的 cross-attention 读取语言与本体上下文。

一个常见疑问是:动作不能看未来 token,视频协同训练如何帮助动作?答案在共享的 Video DiT 参数和 \(f_0\) 特征。未来 token 为了预测物体运动,会读取 \(f_0\);视频 loss 的梯度因而会训练产生 \(f_0\) 逐层键值的 Video DiT。动作分支在每一层又读取这些 \(f_0\) 键值,所以它得到的是被未来预测任务塑造过的当前世界表示。

这里"学到运动、接触等动态信息"是由训练目标和信息路径给出的机制解释。论文用任务成功率和消融结果证明这种协同训练对控制有帮助,但没有用表征探针直接测量某个 token 是否编码了某一种物理概念。

3.6 一个需要特别注意的开源默认设置

当前发布配置的 action_conditioned 为 false,也就是视频分支的未来预测并不显式以真实动作作为条件。代码虽然把动作参数传给 Video DiT,但默认配置下不会使用它。视频预测的条件是当前帧、语言和本体上下文。

因此,对当前开源 release 最准确的说法是:它通过"任务条件下的未来视频预测"学习动态先验,而不是严格意义上的 \(p(v_{1:T}\mid o,l,a)\) 动作条件视频模型。论文的核心对照结论不依赖把视频分支解释成动作条件生成器,但读者不应擅自假定未来帧看到了真实动作。


4. Flow Matching:视频和动作共用的生成原理

Fast-WAM 的两个输出都不是直接回归,而是用 flow matching 从高斯噪声逐渐运输到真实数据。

对任意目标 \(y\),它可以是视频 latent \(z\),也可以是动作块 \(a\)。采样噪声

\\\epsilon\\sim\\mathcal N(0,I), \\

在数据与噪声之间作线性插值:

\y_\\sigma=(1-\\sigma)y+\\sigma\\epsilon,\\qquad \\sigma\\in\[0,1. \]

当 \(\sigma=0\) 时,\(y_\sigma=y\) 是干净数据;当 \(\sigma=1\) 时,\(y_\sigma=\epsilon\) 是纯噪声。这条直线轨迹对 \(\sigma\) 的速度是常数:

\\\frac{d y_\\sigma}{d\\sigma}=\\epsilon-y. \\

模型 \(f_\theta\) 的任务,就是看到 \(y_\sigma\)、噪声时间和条件后,预测这个速度:

\f_\\theta(y_\\sigma,\\sigma,o,l)\\approx\\epsilon-y. \\

论文给出的基础目标是

\\\mathcal L_{\\mathrm{FM}}(y) = \\mathbb E_{y,\\epsilon,\\sigma} \\left\[ \\left\\\|f_\\theta(y_\\sigma,\\sigma,o,l)-(\\epsilon-y)\\right\\\|_2\^2 \\right. \]

直观上,模型不是直接猜终点 \(y\),而是在任意噪声位置学习一个"往哪个方向走"的向量场。

4.1 代码中的时间采样不是简单均匀分布

当前实现先采样

\u\\sim\\mathcal U(0,1), \\

再用默认 \(s=5\) 的 shift 函数

\\\sigma=\\phi_s(u)=\\frac{su}{1+(s-1)u}. \\

\(s>1\) 会把更多样本推向高噪声区。训练代码内部还把时间写成 \(t=1000\sigma\),但真正混合数据时仍会除以 \(1000\),所以数学本质仍是上面的 \(\sigma\in0,1\)。

此外,代码还给每条样本的 MSE 乘时间权重。令 \(N=1000\),并定义离散网格

\u_j=1-\\frac{j}{N},\\qquad t_j=N\\phi_s(u_j),\\qquad j=0,\\ldots,N-1. \\

先计算

\q(t)= \\exp\\left\[-2\\left(\\frac{t-N/2}{N}\\right)\^2\\right-y_{\min}, \]

其中

\y_{\\min}=\\min_j \\exp\\left\[-2\\left(\\frac{t_j-N/2}{N}\\right)\^2\\right. \]

最后将其除以该离散网格上的均值,使平均权重大约为 \(1\):

\w(t)=\\frac{q(t)}{\\frac{1}{N}\\sum_{j=0}\^{N-1}q(t_j)+\\varepsilon_w}, \\qquad \\varepsilon_w=10\^{-10}. \\

它降低靠近两端时间的权重,更强调中间噪声区域。

论文与代码差异:论文正文将时间分布描述为 logit-normal,并给出未加权的标准 flow-matching MSE;当前仓库实际采用上述 shift-uniform 采样和时间加权。理解论文思想时可用标准公式,复现当前代码时应以这里的真实调度器为准。

4.2 为什么推理是从 \(\sigma=1\) 走到 \(0\)

训练学到的是正向速度 \(\epsilon-y\),而生成时从纯噪声出发,沿相反的时间方向积分。为避免与动作的物理时间下标混淆,下面用括号上标表示数值积分步。代码取

\u_k=1-\\frac{k}{K},\\qquad \\sigma_k=\\phi_s(u_k),qquad k=0,\\ldots,K, \\

其中默认 \(K=10\),初值 \(x^{(0)}\sim\mathcal N(0,I)\) 位于 \(\sigma_0=1\)。由于 \(\Delta\sigma_k=\sigma_{k+1}-\sigma_k<0\),Euler 更新

\x\^{(k+1)}=x\^{(k)}+f_\\theta(x\^{(k)},\\sigma_k,c)\\Delta\\sigma_k \\

自然会从噪声端走向数据端。

举一个一维理想例子。真实动作 \(y=0.2\),初始噪声 \(\epsilon=1.0\),完美模型始终输出

\v=\\epsilon-y=0.8. \\

若只用一步从 \(\sigma_0=1\) 走到 \(\sigma_1=0\),则

\x\^{(1)}=1.0+0.8(0-1)=0.2, \\

恰好回到真实动作。实际模型不完美且向量场会随位置变化,所以用多步积分更稳。


5. 完整训练 pipeline

5.1 训练前初始化

  1. 加载 Wan2.2-5B 的 Video DiT、视频 VAE 和 T5。
  2. 将 Video DiT 的骨干权重插值到较窄的 Action DiT;动作输入/输出层随机初始化。
  3. 预计算所有任务指令的 T5 embedding。
  4. 计算或加载动作与状态的 dataset statistics。
  5. 冻结 VAE 和 T5;训练 Video DiT、Action DiT 以及本体状态投影层。

5.2 单个 batch 的前向过程

给定视频 \(I_{0:8}\)、动作 \(a_{0:31}\)、当前本体状态 \(s_0\) 和文本 \(l\):

  1. 冻结的 VAE 编码全部 \(9\) 帧,得到干净视频 latent \(z\)。
  2. 保存第一 latent 帧 \(z_0\) 作为干净视觉锚点。
  3. 视频和动作独立 采样噪声与时间:\((\epsilon_v,\sigma_v)\)、\((\epsilon_a,\sigma_a)\)。
  4. 构造 \(z_{\sigma_v}\) 和 \(a_{\sigma_a}\)。
  5. 把 \(z_{\sigma_v}\) 的第一帧重新替换成干净 \(z_0\);该帧的 DiT 时间 embedding 也强制设为 \(0\)。
  6. Video DiT 和 Action DiT 分别做输入投影,并在每层通过结构化 mixed attention 交换允许的信息。
  7. 视频头预测 \(\epsilon_v-z\),动作头预测 \(\epsilon_a-a\)。
  8. 计算未来视频 loss 和有效动作 loss,加权求和后反向传播。

5.3 代码真实损失

对第 \(i\) 个样本,动作每一步先对维度求平均:

\e\^{a}_{i,h} = \\frac{1}{D_a} \\left\\\| \\hat v\^a_{i,h}-(\\epsilon\^a_{i,h}-a_{i,h}) \\right\\\|_2\^2. \\

设 \(m^a_{i,h}=1\) 表示有效动作、\(0\) 表示 padding,则

\\\mathcal L_{\\mathrm{act}} = \\frac{1}{B}\\sum_{i=1}\^{B} w(t_i\^a) \\frac{\\sum_h m\^a_{i,h}e\^a_{i,h}} {\\max(1,\\sum_hm\^a_{i,h})}. \\

视频 loss 对通道和空间维求平均,再对有效 latent 时间片平均。第一帧是条件,不计算视频重建 loss。设未来 latent 时间片的有效掩码为 \(m^v_{i,\tau}\),则

\e\^v_{i,\\tau} = \\operatorname{mean}_{c,x,y} \\left( \\hat v\^v_{i,c,\\tau,x,y} -(\\epsilon\^v_{i,c,\\tau,x,y}-z_{i,c,\\tau,x,y}) \\right)\^2, \\

\\\mathcal L_{\\mathrm{vid}} = \\frac{1}{B}\\sum_{i=1}\^{B} w(t_i\^v) \\frac{\\sum_{\\tau\>0}m\^v_{i,\\tau}e\^v_{i,\\tau}} {\\max(1,\\sum_{\\tau\>0}m\^v_{i,\\tau})}. \\

总损失为

\\\mathcal L = \\lambda_v\\mathcal L_{\\mathrm{vid}} + \\lambda_a\\mathcal L_{\\mathrm{act}}. \\

当前配置只显式写了 \(\lambda_a=1\),代码对缺省的 \(\lambda_v\) 使用 \(1\),因此默认是

\\\lambda_v=\\lambda_a=1. \\

5.4 训练伪代码

text 复制代码
函数 FastWAM单步训练(batch):
    视频, 动作, 状态, 文本, padding = batch

    # 冻结模块:VAE 和 T5 不参与梯度更新
    干净视频latent = VAE编码(视频)
    当前帧latent = 干净视频latent[:, :, 第0个latent时间片]
    上下文 = 拼接(文本embedding, 线性投影(状态[:, 0]))

    # 视频与动作独立采样 flow-matching 时间和噪声
    σ视频 = shift采样(均匀随机数, shift=5)
    σ动作 = shift采样(均匀随机数, shift=5)
    噪声视频 = 标准高斯(与干净视频latent同形状)
    噪声动作 = 标准高斯(与动作同形状)

    加噪视频 = (1-σ视频) * 干净视频latent + σ视频 * 噪声视频
    加噪动作 = (1-σ动作) * 动作 + σ动作 * 噪声动作

    # 当前帧始终是干净条件,不是待生成目标
    加噪视频[:, :, 0] = 当前帧latent

    视频token = VideoDiT输入层(加噪视频, 视频时间, 上下文)
    动作token = ActionDiT输入层(加噪动作, 动作时间, 上下文)

    对 layer = 1...30:
        # 各专家产生自己的 Q/K/V,再在统一 mask 下做 mixed attention
        视频token, 动作token = MoT层(
            视频token, 动作token,
            mask="动作只看当前帧和动作;未来视频不向动作泄漏",
            上下文=上下文
        )

    预测视频速度 = VideoDiT输出头(视频token)
    预测动作速度 = ActionDiT输出头(动作token)

    视频目标 = 噪声视频 - 干净视频latent
    动作目标 = 噪声动作 - 动作

    # 当前帧、padding 视频时间片、padding 动作步均不计入对应 loss
    L视频 = 时间加权且带有效掩码的MSE(预测视频速度, 视频目标)
    L动作 = 时间加权且带有效掩码的MSE(预测动作速度, 动作目标)
    L总 = λ视频 * L视频 + λ动作 * L动作

    反向传播并更新(VideoDiT, ActionDiT, 本体状态投影层)
    返回 L总, L视频, L动作

5.5 优化设置

论文与配置共同给出的主要设置为:

  • AdamW,\(\mathrm{lr}=10^{-4}\),\(\beta=(0.9,0.95)\);
  • weight decay 为 \(0.01\);
  • cosine 学习率,代码带前 \(5\%\) step 的线性 warmup;
  • bf16 混合精度;
  • gradient clipping 为 \(1.0\);
  • 每卡 batch size 为 \(16\);
  • LIBERO 论文按 \(20\text{k}\) step 训练;RoboTwin 按 \(30\text{k}\) step 训练;
  • 开源 task 配置以 epoch 数控制上限:LIBERO 为 \(10\) epoch,RoboTwin 为 \(5\) epoch,实际 step 数取决于数据集长度和 GPU 数。

6. Fast-WAM 推理:不生成视频,但动作仍会去噪

这是最容易被一句"single-pass inference"误导的地方。

Fast-WAM 删除的是未来视频的迭代去噪 ,不是所有迭代。动作仍然从高斯噪声通过默认 \(10\) 个 flow-matching 步生成。"单次前向"准确地指当前帧只通过 Video DiT 一遍,并缓存其逐层 K/V;后续每个动作去噪步只重跑 Action DiT。

6.1 推理的精确流程

给定当前多相机观测 \(o\)、指令 \(l\) 和当前本体状态 \(s\):

  1. 按训练时相同布局拼接多相机图像,并归一化到 \(-1,1\)。

  2. VAE 只编码当前一帧,得到 \(z_0\)。

  3. T5 编码语言,将本体状态投影成一个 token 并拼到文本上下文。

  4. 从高斯分布初始化动作噪声

    \x_a\^{(0)}\\sim\\mathcal N(0,I),\\qquad x_a\^{(0)}\\in\\mathbb R\^{1\\times32\\times D_a}. \\

  5. 令视频时间为 \(0\),让干净 \(z_0\) 通过 Video DiT 的 \(30\) 层一次;在第 \(m\) 层缓存该层视频 token 产生的 \((K_v^{(m)},V_v^{(m)})\)。

  6. 对 \(K=10\) 个动作去噪时间:Action DiT 在每层重新计算当前动作的 \(Q_a,K_a,V_a\),其中 \(Q_a\) 同时注意缓存的视频 \(K_v,V_v\) 和本步动作的 \(K_a,V_a\)。

  7. 用 Euler 公式更新动作样本,直至 \(\sigma=0\)。

  8. 把动作从归一化空间映射回机器人控制空间。

缓存有效的原因是:当前图像、语言和本体状态在一个 action chunk 的去噪过程中不变;Video DiT 的输入时间固定为 \(0\),其逐层表示也不变。反复计算它只是在浪费算力。

6.2 推理伪代码

text 复制代码
函数 FastWAM生成动作块(当前多相机图像, 指令, 当前本体状态):
    拼接图 = 按训练布局拼接并缩放(当前多相机图像)
    当前帧latent = VAE编码(拼接图)              # 只编码一帧
    文本token = T5编码(固定模板.format(指令))
    本体token = 线性投影(归一化后的当前本体状态)
    上下文 = 拼接(文本token, 本体token)

    动作样本 = 标准高斯(shape=[1, 32, 动作维度])

    # 当前帧只过一次 30 层 Video DiT;保存每层供动作查询的 K/V
    视频KV缓存 = VideoDiT逐层编码并缓存KV(
        当前帧latent, 视频时间=0, 上下文=上下文
    )

    # 默认 10 步,只重复较小的 Action DiT
    时间点, 时间差 = 构造从σ=1到σ=0的shift调度(步数=10, shift=5)
    对 (σ, Δσ) 属于 (时间点, 时间差):
        预测速度 = ActionDiT(
            动作样本, 动作时间=σ,
            可读取KV=视频KV缓存 + 当前动作KV,
            上下文=上下文
        )
        动作样本 = 动作样本 + 预测速度 * Δσ

    动作块 = 反归一化(动作样本[0])
    返回 动作块                              # [32, 动作维度]

6.3 CFG 在默认推理中实际上没有额外计算

论文写推理 CFG scale 为 \(1.0\)。当前 infer_action 接口虽然保留 negative_prompttext_cfg_scale 参数,但动作路径没有运行正/负条件双分支;scale 为 \(1.0\) 本来也等价于不做额外 classifier-free guidance。因此默认 \(10\) 步中,每步只运行一次 Action DiT。


7. 从动作块到闭环机器人控制

模型一次预测 \(H=32\) 步,并不意味着机器人盲目执行完 \(32\) 步才看环境。实际部署采用 receding-horizon control(滚动时域控制):

  1. 获取最新观测并预测一个 \(32\) 步动作块。
  2. 只执行前 \(R\) 步。
  3. 丢弃剩余旧动作,重新观察并规划。

默认配置中:

  • LIBERO:\(R=10\);
  • RoboTwin:\(R=24\)。

这在反应速度和推理开销之间折中。\(R\) 小,闭环纠错更及时,但模型调用更频繁;\(R\) 大,吞吐更高,但更依赖较远期动作预测。

LIBERO 推理还会把夹爪值转换回环境约定,并可将夹爪命令二值化。RoboTwin 将预测的 \(14\) 维关节动作直接放入队列逐步执行。可选的 action ensembler 会把不同重规划时刻对同一未来时刻的动作预测取平均,但当前 LIBERO 默认关闭。

text 复制代码
函数 闭环执行任务(环境, 策略, 每次执行步数R):
    动作队列 = 空
    当任务未结束且未超时:
        如果动作队列为空:
            观测 = 环境.获取观测()
            候选动作块 = 策略.生成动作块(观测)       # 共 32 步
            动作队列 = 候选动作块的前 R 步

        当前动作 = 动作队列.弹出队首()
        环境.执行(当前动作)

注意:RoboTwin 的加速评估可在动作队列未空时跳过 RGB 渲染,这会让保存的视频看起来帧率很低,但不改变实际动作队列逻辑。


8. 两个"显式想象"对照变体

8.1 Fast-WAM-Joint:视频和动作一起去噪

Joint 版本与 Fast-WAM 的主要结构差别只有 attention mask:动作 query 不再只看 \(f_0\),而是可以看所有未来视频 token。

Query \(\downarrow\) / Key \(\rightarrow\) 当前帧 \(f_0\) 未来帧 \(f_+\) 动作 \(a\)
当前帧 \(f_0\)
未来帧 \(f_+\)
动作 \(a\)

推理时同时初始化未来视频噪声和动作噪声,在每个时间步共同更新:

\z_{k+1}=z_k+\\hat v_v\\Delta\\sigma_k, \\

\a_{k+1}=a_k+\\hat v_a\\Delta\\sigma_k. \\

当前帧 \(z_{\mathrm{cur}}\) 每一步都重新钉回干净值。视频和动作拥有各自的 scheduler;当前默认二者都是 \(K=10\)、shift \(5\),所以时间点重合并在同一循环中同步更新,但实现上可以分别配置。动作能读取正在成形的未来视频,因此属于显式未来想象;代价是每个去噪步都必须运行庞大的 Video DiT。

8.2 Fast-WAM-IDM:先视频,后动作

IDM 版本把生成拆成两个阶段:

  1. 运行 Video DiT,经过 \(K=10\) 步得到完整未来视频 latent。
  2. 将最终视频以时间 \(0\) 额外 prefill 一次,缓存它在 \(30\) 层中的 K/V;再经过 \(K=10\) 步 Action DiT 去噪动作。

它比 Joint 更像字面意义上的"先想象,再执行",串行依赖也最强。

训练时若总是让动作看到干净真实视频,会产生 teacher-forcing gap:训练条件完美,推理条件却是带误差的生成视频。为减轻这个问题,代码以 \(p=0.5\) 的概率对动作条件视频添加随机强度噪声。与此同时,另有一份加噪视频分支负责标准视频去噪 loss;只有这份分支计入视频 loss,动作则读取"干净或增强后"的 teacher-forcing 视频。

8.3 三者计算路径对比

方法 Video DiT 推理次数 Action DiT 推理次数 是否产生未来 latent 串行阶段
Fast-WAM \(1\) 次 prefill \(K\) 视觉编码 -> 动作去噪
Joint \(K\) \(K\),与视频同轮 联合去噪
IDM \(K+1\) \(K\) 先视频 \(K\) 步,再视频 prefill 一次,最后动作 \(K\) 步

这里 \(K=10\) 是论文和评估配置的默认值。表中的次数指完整分支通过其全部 Transformer 层的次数。


9. 用一个具体例子串起全流程

假设 LIBERO 指令是"把红色杯子放到左边盘子上"。当前输入包括主相机、腕部相机和 \(8\) 维本体状态。

9.1 训练时

数据窗口提供控制时刻 \(0\) 到 \(32\):

  • 视频目标选 \(I_0,I_4,\ldots,I_{32}\),主视角与腕部视角水平拼为 \(224\times448\);
  • 动作目标保留 \(a_0,\ldots,a_{31}\),每步 \(7\) 维;
  • \(I_0\) 是当前帧,其余帧展示靠近杯子、夹住、抬起和移动的过程。

VAE 将 \(9\) 帧编码成 \(3\) 个 latent 时间片。第一片保持干净,后两片和整个动作块分别加噪。模型同时学习:

  • 从当前场景、任务语义和带噪未来恢复视频速度场;
  • 从当前场景、任务语义、本体状态和带噪动作恢复动作速度场。

动作 token 看不到后两个 latent 时间片,所以即使训练视频里已经显示杯子最终到了盘子上,它也不能偷看这个结果。它只能借助经过视频任务训练的当前帧表示,自己推断怎样移动机械臂。

9.2 推理时

机器人现在只拍到主相机和腕部相机的当前画面:

  1. 两图拼接并经 VAE 得到单帧 latent。
  2. Video DiT 结合"把红色杯子放到左边盘子上"和当前本体状态编码该帧,缓存 \(30\) 层 K/V。
  3. Action DiT 从一个 \(32\times7\) 的随机高斯矩阵开始,经过 \(10\) 次更新得到动作块。
  4. LIBERO 只执行前 \(10\) 步,然后重新拍照再规划。

如果第一个动作块执行后杯子发生滑动,新观测会反映这个偏差,下一次重规划便能纠正。Fast-WAM 不靠一次生成很远的显式未来视频来闭环,而靠"有世界建模训练的动作策略 + 高频重新观测"完成控制。


10. 实验究竟支持了什么结论

论文不使用 embodied pretraining,对比结果为:

方法 LIBERO 平均成功率 RoboTwin 平均成功率
Fast-WAM \(97.6\%\) \(91.8\%\)
Fast-WAM-Joint \(98.5\%\) \(90.6\%\)
Fast-WAM-IDM \(98.0\%\) \(91.3\%\)
无视频协同训练 \(93.5\%\) \(83.8\%\)

可以看到:

  • 显式想象版本相对 Fast-WAM 没有稳定优势。LIBERO 上 Joint 略高,RoboTwin 上反而 Fast-WAM 略高。
  • 删除视频协同训练造成的下降明显更大:LIBERO 下降 \(4.1\) 个百分点,RoboTwin 下降 \(8.0\) 个百分点。
  • 真实毛巾折叠任务使用约 \(60\) 小时遥操作数据;无视频协同训练版本只剩 \(10\%\) 成功率。论文将其作为视频协同训练有益于该任务的证据,但正文没有给出各 Fast-WAM 变体的精确成功率和试验次数,不能仅凭这个数字量化效果大小。
  • Fast-WAM 延迟约 \(190\text{ ms}\),IDM 约 \(810\text{ ms}\),说明删除视频生成带来实际部署收益。

LIBERO 结果来自 \(40\) 个任务共 \(2000\) 次试验;RoboTwin 对每个任务评估 \(100\) 次,覆盖 \(50\) 余个任务。论文没有报告方差或置信区间,因此 Fast-WAM、Joint 和 IDM 之间 \(0.5\) 到 \(1.2\) 个百分点的差异不能据此宣称具有统计显著性。"显式想象没有稳定优势"是比"它们完全相同"更准确的表述。

延迟均在单张 NVIDIA RTX 5090D V2 32GB 上测得,但论文正文没有进一步拆分计时边界,也没有在文字中报告 Joint 的精确延迟。因此这些数字适合比较论文实现下的端到端设计点,不宜直接外推到其他硬件或复现。

严谨地说,这些实验支持的是:

在论文选择的 Wan2.2 骨干、数据规模、任务、动作块长度和 \(10\) 步生成设置下,训练时视频协同建模的贡献明显大于测试时显式生成未来视频的贡献。

它不等于证明"任何机器人任务都不需要规划未来"。以下情形仍可能从显式未来中获益:需要比较多个候选动作后果、必须满足长程约束、环境部分可观测且当前帧信息不足,或未来生成器能进行真正的反事实 action-conditioned rollout。Fast-WAM 的结论是一个有力的工程设计证据,不是对显式规划的普遍否定。


11. 最容易混淆的十个问题

11.1 Fast-WAM 还是 world model 吗?

训练范式上是:它有未来视频生成目标,并用该目标学习世界表征。推理接口上更像直接 VLA:输入当前观察、语言和本体状态,直接输出动作。它可以看作"用 world modeling 训练出来的直接策略"。

11.2 推理真的只有一次网络前向吗?

不是整个网络一次。Video DiT 对当前帧只前向一次并缓存 K/V;Action DiT 默认仍前向 \(10\) 次做动作 flow matching。

11.3 推理时会生成未来视频 latent,只是不解码 RGB 吗?

Fast-WAM 主路径不会。它根本不实例化未来视频噪声,只保留当前一帧 latent。Joint 和 IDM 才会生成未来 latent;是否解码成 RGB 只是可视化问题。

11.4 动作分支训练时是否偷看了真实未来?

没有。结构化 mask 禁止 action query 读取未来视频 key/value。它能读取当前帧和所有动作 token。

11.5 Action DiT 是自回归模型吗?

不是。\(32\) 个 action token 彼此双向可见,并行生成整个 chunk;迭代发生在噪声时间维,而不是动作序列从左到右。

11.6 当前帧为什么不加噪?

它是条件,不是生成目标。训练和推理都必须拥有稳定、干净的视觉锚点;其 token 时间设为 \(0\),视频 loss 也排除第一 latent 帧。

11.7 视频 loss 为什么能改变当前帧表示?

未来视频 token 会注意当前帧 token。未来预测误差的梯度会经过这些注意力连接和共享 Video DiT 参数,训练出对运动预测有用的当前帧逐层 K/V;动作恰好读取这些 K/V。

11.8 视频生成是否由动作条件控制?

当前 release 的默认配置不是。action_conditioned=false,未来视频依赖当前帧、语言和本体上下文。不能把它描述为用真实动作 rollout 出未来。

11.9 为什么输出 32 步却只执行 10 或 24 步?

长 chunk 让模型学习连贯动作并摊薄推理成本;提前重规划则恢复闭环反馈。执行长度是控制频率和计算成本之间的可调折中。

11.10 "没有 embodied pretraining"是否等于从零训练?

不是。Fast-WAM 使用大规模预训练的 Wan2.2 视频模型及其 VAE/T5,只是没有先在额外机器人轨迹上做 embodied pretraining。Action DiT 骨干也由 Video DiT 权重插值得到。


12. 论文描述与当前代码的对齐清单

主题 论文抽象 当前开源实现
核心输入 当前观察 \(o\)、语言 \(l\) 当前拼接 RGB、语言、首时刻本体状态 token;完整条件记为 \((o,l,s)\)
动作 horizon \(H=32\) \(32\)
视频帧 每 chunk \(9\) 帧 从 \(33\) 个控制时刻每 \(4\) 步取一帧
视频 latent 时间 未重点展开 VAE 后为 \(3\) 个时间片
主干 Wan2.2-5B + 1B Action DiT Video hidden \(3072\),Action hidden \(1024\),均 \(30\) 层
动作/未来隔离 结构化 mask 动作只看当前帧;未来视频不看动作
基础目标 标准 flow matching 同一线性插值与速度目标
时间采样 正文称 logit-normal shift-uniform,默认 shift \(5\)
loss 权重 正文未展开 显式 \(w(t)\),中间时间权重较大
loss 系数 \(\mathcal L_a+\lambda\mathcal L_v\) 默认两项系数均为 \(1\)
CFG scale \(1.0\) Fast-WAM 动作路径不运行额外负条件分支
Fast-WAM 视频推理 单次当前帧编码 单次 Video DiT prefill + 逐层 KV cache
动作推理 flow matching 默认 \(10\) 步 Euler 更新
闭环 论文省略外层自回归以便比较 评估中滚动重规划,LIBERO 执行 \(10\) 步、RoboTwin 执行 \(24\) 步

13. 最终心智模型

理解 Fast-WAM,只需牢牢记住下面这条因果链:

text 复制代码
训练视频预测
    ↓ 迫使 Video DiT 学习物体、运动、接触和任务进展
当前帧逐层 K/V 成为"隐式世界表征"
    ↓ Action DiT 在每层读取它们
动作块获得与视频协同训练相关的性能收益和动态先验
    ↓ 推理时未来 token 从未成为动作的必要输入
可以删除昂贵的未来视频去噪,只保留一次当前帧 prefill

所以,Fast-WAM 的创新重点不是新的生成损失,也不是新的机器人动作表示,而是通过严格的信息隔离和受控对照,证明并利用了一个设计原则:

世界模型可以主要作为一种训练监督和表示学习机制存在;是否在测试时显式展开未来,是另一个应由收益与延迟单独权衡的选择。

对实际系统而言,它给出的方案是:用强视频生成模型的预训练能力和未来预测目标训练一个"懂世界"的视觉骨干,再以 KV cache 把这个骨干变成动作生成器可复用的当前状态编码器。这样保留 WAM 的训练收益,同时把部署路径压缩成一次视觉编码、若干次较小动作专家去噪和周期性闭环重规划。