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

Paper: arXiv:2603.16666 Code: yuantianyuan01/FastWAM Code reference: main @ 45d8e145 (2026-04-03)

1. Motivation (研究动机)

这篇论文讨论的是 World Action Model (WAM) 在机器人控制中的一个核心设计问题：WAM 是否真的需要在测试时显式生成未来视频，还是只需要在训练时利用视频预测目标来塑造更好的世界表示。现有 WAM 通常把两件事绑在一起：训练时学习未来视觉状态，推理时也先 imagine future observations，再根据想象出的未来生成动作。这个范式直观，但代价很高，因为视频生成通常依赖多步 diffusion / flow denoising，测试时延迟会直接进入闭环控制周期。

传统 VLA policy 直接建模 $p(a_{1:H}\mid o,l)$，也就是从当前视觉观察 $o$ 和语言指令 $l$ 预测未来动作 chunk $a_{1:H}$。它的优点是推理接口简单、速度较快；缺点是 web-scale 视觉语言预训练主要来自静态 image-text 数据，对物理动态、接触变化、物体形变和长时序交互没有显式建模。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})d{v}_{1:T}.$ 这带来两个瓶颈。第一，未来视频生成慢，尤其当模型需要反复 denoise latent video tokens 时，闭环机器人系统会出现明显延迟。第二，现有方法难以回答因果问题：性能提升到底来自训练时的 video prediction objective，还是来自测试时显式 imagine 的未来视频。如果二者被同一个系统同时使用，实验结果只能说明 WAM 有效，不能说明哪一部分真正必要。

Fast-WAM 的研究价值就在于拆开这个耦合。作者提出一个更直接的问题：如果保留视频共训练，让 video DiT 学到物理动态表示，但测试时不再生成未来帧，动作性能是否仍然接近 imagine-then-execute WAM？如果答案是肯定的，那么 WAM 的主要收益可能不是“测试时想象未来”，而是“训练时通过未来视频预测学习更好的 latent world representation”。这对实际机器人很重要，因为闭环控制更关心稳定低延迟；如果可以用 190 ms 级别的单步 world encoding 替代 580-810 ms 的未来视频生成，就能把 WAM 从离线规划式系统推向更实时的部署。

2. Idea (核心思想)

核心洞察：WAM 的收益可以被分解为两个因素，训练时的 video co-training 和推理时的 explicit future imagination。Fast-WAM 保留前者，去掉后者，用 pretrained video DiT 作为单次前向的 world encoder，再让 action expert 在这个 latent world representation 上做动作 denoising。

这与标准 imagine-then-execute WAM 的根本区别在于，Fast-WAM 推理时不采样 $v_{1:T}$，也不对未来视频 latent 做多步 denoising。它只编码当前观察帧的 latent tokens，缓存 video branch 的 K/V 表示，然后 action tokens 在 denoising 过程中 attend 到这些已经被视频共训练塑造过的表示。也就是说，模型在测试时看起来像直接 policy，但它的视觉 backbone 在训练时仍然承担 WAM 式未来建模任务。

作者还构造了三类 controlled variants：Fast-WAM-Joint 对应 joint video-action denoising，Fast-WAM-IDM 对应先生成未来视频再做 inverse dynamics/action prediction，Fast-WAM w.o. video co-train 则只移除训练时视频建模目标。这样实验能比较“有无测试时未来想象”和“有无视频共训练”两个变量。结果显示，Fast-WAM 与两个 imagine-then-execute 变体差距很小，而移除 video co-training 会大幅掉点，支持论文的中心结论。

3. Method (方法)

3.1 Overall framework

Figure 1 解读：图中对比了三种 WAM 范式。oc。(B) future video prediction followed by action prediction 先 denoise future video，再让动作模型 conditioned on generated future representation，代表 causal WAM / inverse dynamics 路线。(C) Fast-WAM 的关键变化是训练时仍保留 video denoising objective，但推理时只对当前帧做一次 video DiT forward，再用 action DiT denoise action tokens；未来视频分支不再被实例化。

Fast-WAM 的形式化目标是把 WAM 中的中间未来视频变量去掉。标准 WAM 写成上面的积分形式，而 Fast-WAM 在推理时直接建模： $p_{\theta }(a_{1:H}\mid o,l),$ 同时引入 video backbone 产生的 latent world representation $z(o,l)$： $p_{\theta }(a_{1:H}\mid o,l)=p_{\theta }(a_{1:H}\mid z(o,l)).$ 这里 $z(o,l)$ 不是通过采样未来视频得到的，而是由当前 observation latent 经 video DiT 单次编码得到。训练阶段，video DiT 还要预测未来 video latents，因此 $z(o,l)$ 被迫携带运动、接触和任务相关的时序结构；推理阶段，这些结构以隐式表示的形式服务 action denoising。

3.2 Model architecture

Figure 2a 解读：Fast-WAM 基于 Wan2.2-5B 的 video DiT、T5 text encoder 和 video VAE。语言指令先由 text encoder 得到 context，并通过 cross-attention 提供给 video tokens 和 action tokens。当前帧和未来帧经 VAE 编码为 latent video tokens；动作 chunk 经 Action Encoder 投影为 action tokens。模型主体是 Mixture-of-Transformer (MoT)：video branch 复用 pretrained video DiT，action branch 是 reduced hidden dimension 的 action expert DiT，二者通过 shared / mixed attention 交互。

模型输入分三组。第一组是 clean latent tokens of the first observation frame，作为视觉锚点。第二组是 future video latent tokens，只在训练时加入，并被加噪后用于 video flow matching。第三组是 action tokens，用于动作 chunk generation。训练时 video branch 和 action branch 同时经过 MoT，但注意力 mask 保证动作 token 不能偷看 future video token；它只能 attend 到 clean first-frame tokens 和 action branch 内部上下文。这样 video co-training 会塑造 shared video backbone，但不会把 ground-truth future video 信息泄漏给动作预测。

Figure 2b 解读：mask 图展示了训练和推理的信息流。训练时，$f_0$ 是可见的当前帧 token，未来帧 $f_1,\dots ,f_h$ 与动作 $a_1,\dots ,a_h$ 被加噪；video tokens 可以做 video denoising，action tokens 可以做 action denoising，但 action tokens 不直接读取 future video tokens。推理时只保留 $f_0$ 与动作 token，future frames 分支消失。released code 中 _build_mot_attention_mask() 明确实现了这个约束：video-to-video 使用 video expert 自己的 mask，action-to-action 全可见，action-to-video 只开放 first-frame video tokens。

Action expert 的结构不是一个独立小 MLP，而是与 video DiT 层数、head 数、head dim 对齐的 ActionDiT。configs/model/fastwam.yaml 中 action branch 使用 hidden_dim=1024、ffn_dim=4096、num_heads=24、attn_head_dim=128、num_layers=30，而 video branch 使用 Wan2.2-5B 的 hidden_dim=3072、ffn_dim=14336、num_layers=30。代码在 FastWAM.from_wan22_pretrained() 中检查 action expert 的 num_heads、attn_head_dim、num_layers 必须与 video expert 一致，因为 MoT mixed attention 需要对齐 attention geometry。

3.3 Flow matching objective

Fast-WAM 使用同一套 flow matching 形式训练动作和未来视频。给定目标变量 $y$，其中 $y$ 可以是动作 chunk $a_{1:H}$，也可以是 VAE future video latent $z_{1:T}$，采样 Gaussian noise $ϵ\sim \mathcal{N}(0,I)$ 和时间 $t\in (0,1)$，构造插值： $y_t=(1-t)y+tϵ.$ 模型预测 velocity field，目标为： ${\mathcal{L}}_{\mathrm{FM}}(y)={\mathbb{E}}_{y,ϵ,t}\left[{\Vert f_{\theta }(y_t,t,o,l)-(ϵ-y)\Vert }_2^2\right].$ 动作损失和视频损失分别为： ${\mathcal{L}}_{\mathrm{act}}={\mathcal{L}}_{\mathrm{FM}}(a_{1:H}),{\mathcal{L}}_{\mathrm{vid}}={\mathcal{L}}_{\mathrm{FM}}(z_{1:T}).$ 总损失： $\mathcal{L}={\mathcal{L}}_{\mathrm{act}}+\lambda {\mathcal{L}}_{\mathrm{vid}}.$ released code 对这个公式的实现更细一些。training_loss() 分别为 video latents 和 action tokens 采样噪声与 timestep，通过 WanContinuousFlowMatchScheduler.add_noise() 得到 noisy input，再用 training_target() 得到 velocity target。video loss 使用 F.mse_loss(...).mean(dim=(1,3,4)) 聚合通道和空间维度，并根据 image_is_pad 排除 padding frame；action loss 使用 token-level MSE，并根据 action_is_pad 排除 padding action。最后代码做 scheduler weight： ${\mathcal{L}}_{code}={\lambda }_v{w}_v(t_v){\mathrm{MSE}}_{video}+{\lambda }_a{w}_a(t_a){\mathrm{MSE}}_{action}.$ 论文公式与 released code 实现差异：论文只写 ${\mathcal{L}}_{act}+\lambda {\mathcal{L}}_{vid}$，代码中显式有 loss_lambda_video 和 loss_lambda_action 两个权重，默认都为 1.0；configs/model/fastwam.yaml 只列出 lambda_action: 1.0，未显式列出 lambda_video，因此 video 权重来自代码默认值 1.0。论文正文说 AdamW weight decay 为 0.01，但 released code 的 configs/train.yaml 写的是 weight_decay: 0.0；这里应以论文实验报告为主，同时在复现实验时检查作者是否另有未公开 override。论文说推理使用 10 denoising steps；代码 infer_action() 默认 num_inference_steps=20，但 configs/train.yaml 的 eval_num_inference_steps: 10 和 configs/sim_robotwin.yaml 的 EVALUATION.num_inference_steps=${eval_num_inference_steps} 与论文一致。

3.4 Inference: no future video generation

Fast-WAM 推理路径对应代码中的 infer_action()。它先检查 video_attention_mask_mode='first_frame_causal'，说明推理只允许 first-frame causal video representation。然后把输入图像经 VAE 编码为 first_frame_latents，调用 video expert 的 pre_dit() 得到 video tokens；这些 token 不是每个 denoising step 都重新计算，而是通过 MoT.prefill_video_cache() 逐层缓存 video K/V。之后 action latents 从 Gaussian noise 初始化，按 action scheduler 的 denoising schedule 迭代，每一步只重算 action branch，并通过 forward_action_with_video_cache() 让 action queries attend 到 cached video K/V 和当前 action K/V。

这个实现解释了为什么 Fast-WAM 能低延迟：video branch 是一次 prefill，未来视频 latent 不参与推理；每个 action denoising step 主要更新 action tokens。与 Fast-WAM-IDM 的差别是，IDM 需要先完整生成 future video，再进行动作预测；与 Fast-WAM-Joint 的差别是，Joint 在每一步同时 denoise video 和 action，因此 latency 更高。

3.5 Controlled variants

Fast-WAM-Joint：打开 video-action token 的交互，让未来视频和动作一起 denoise。这模拟 joint-generation WAM：动作预测在推理时仍与 future video denoising 绑定，因此更接近“边想象边执行”的范式。

Fast-WAM-IDM：先从当前观察和语言上下文生成 future video token，再让 action prediction conditioned on generated future representation。论文还说明跟随 LingBot-VA 的设置，对 ground-truth video tokens 做概率 $p=0.5$ 的 noise augmentation。

Fast-WAM w.o. video co-train：保持架构和推理过程不变，只移除 video modeling objective。它是最关键的对照组，因为它能回答：如果没有训练时的未来视频预测，仅保留 single-pass architecture 是否足够。实验显示它掉点最多，因此支持 video co-training 是主要收益来源。

3.6 Pseudocode from released code

Code reference: main @ 45d8e145 (2026-04-03) — pseudocode and mapping based on this commit

# source: src/fastwam/models/wan22/fastwam.py:_build_mot_attention_mask
# paper: §3.2, Figure 2b
def build_fastwam_attention_mask(video_seq_len, action_seq_len, tokens_per_frame, device):  # 定义 Fast-WAM 的 MoT 可见性规则，把论文 Figure 2b 的信息流约束落到代码 mask
    total_seq_len = video_seq_len + action_seq_len  # 将 video tokens 和 action tokens 拼成同一个 MoT 序列，对应论文 Figure 2b 的联合 attention mask
    mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device)  # 初始化所有 token 之间都不可见，后续只打开论文允许的信息流
    mask[:video_seq_len, :video_seq_len] = build_video_to_video_mask(video_seq_len, tokens_per_frame, device)  # 复用 video DiT 的视频内部 mask，让未来视频 token 只按 WAM 视频建模规则读取当前帧和视频上下文
    mask[video_seq_len:, video_seq_len:] = True  # 允许 action tokens 彼此双向注意力，支撑 action chunk 内部的联合 denoising
    first_frame_tokens = min(tokens_per_frame, video_seq_len)  # 计算当前观测帧 f0 覆盖的 token 数，作为 action branch 唯一可读的视觉锚点
    mask[video_seq_len:, :first_frame_tokens] = True  # 只让 action tokens attend 到 first-frame video tokens，阻断 future video token 泄漏到动作预测
    return mask  # 返回结构化 mask，用代码实现论文“训练有 video co-training、动作不偷看未来帧”的解耦设计

# source: src/fastwam/models/wan22/fastwam.py:training_loss
# paper: §3.2 Eq.5-Eq.9
def fastwam_training_loss(model, sample):  # 定义训练损失流程，同时优化动作 flow matching 和未来视频 latent co-training
    inputs = model.build_inputs(sample)  # 将原始 batch 编码成 VAE latent、动作 chunk、语言 context 和 padding mask，对齐论文的 observation/action/instruction 输入
    x_video = inputs["input_latents"]  # 取出未来视频 latent 目标 z_{1:T}，这是 video co-training 的 flow matching 学习对象
    action = inputs["action"]  # 取出动作 chunk a_{1:H}，这是 action expert 需要生成的控制序列
    context, context_mask = inputs["context"], inputs["context_mask"]  # 取出 T5 语言条件及其 mask，让 video/action tokens 都通过 cross-attention 对齐任务指令
 
    eps_v = torch.randn_like(x_video)  # 为视频 latent 采样 Gaussian noise epsilon，对应论文 Eq.5 的 flow matching 插值端点
    t_v = model.train_video_scheduler.sample_training_t(x_video.shape[0], device=x_video.device, dtype=x_video.dtype)  # 为每个样本采样视频分支的训练 timestep，用于构造不同噪声强度的 video target
    noisy_video = model.train_video_scheduler.add_noise(x_video, eps_v, t_v)  # 按 y_t=(1-t)y+tε 构造 noisy video latent，让 video DiT 学习从噪声场回到未来视频动态
    target_video = model.train_video_scheduler.training_target(x_video, eps_v, t_v)  # 生成视频 velocity target ε-y，对应论文 Eq.6 的 flow matching 监督信号
    if model.video_expert.fuse_vae_embedding_in_latents:  # 检查代码配置是否把 VAE first-frame embedding 直接融合进 latent，以匹配 Fast-WAM 的 clean observation anchor
        noisy_video[:, :, 0:1] = inputs["first_frame_latents"]  # 保持当前观测帧 f0 为 clean anchor，使 video/action 预测都基于真实当前视觉状态而非噪声帧
 
    eps_a = torch.randn_like(action)  # 为动作 chunk 采样 Gaussian noise，使 action branch 也使用和视频分支一致的 flow matching 形式
    t_a = model.train_action_scheduler.sample_training_t(action.shape[0], device=action.device, dtype=action.dtype)  # 为动作分支独立采样 timestep，允许动作 denoising 与视频 denoising 使用各自噪声进度
    noisy_action = model.train_action_scheduler.add_noise(action, eps_a, t_a)  # 构造 noisy action tokens，作为 action expert 的输入状态
    target_action = model.train_action_scheduler.training_target(action, eps_a, t_a)  # 生成动作 velocity target，让模型学习从 noisy action 恢复真实动作 chunk
 
    video_pre = model.video_expert.pre_dit(noisy_video, t_v, context, context_mask, action=action, fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"])  # 把 noisy video latent 投影成 video DiT tokens，并注入语言/时间条件以训练 world representation
    action_pre = model.action_expert.pre_dit(noisy_action, t_a, context, context_mask)  # 把 noisy action chunk 投影成 action tokens，使动作生成进入与 video branch 对齐的 MoT 表示空间
    mask = build_fastwam_attention_mask(video_pre["tokens"].shape[1], action_pre["tokens"].shape[1], video_pre["meta"]["tokens_per_frame"], video_pre["tokens"].device)  # 构造 Figure 2b 的结构化 mask，保证动作预测只看当前帧而不看未来视频 token
    out = model.mot({"video": video_pre["tokens"], "action": action_pre["tokens"]}, attention_mask=mask)  # 在共享 attention 的 MoT 中联合处理 video/action tokens，让 video objective 塑造可供动作使用的 latent world features
    pred_video = model.video_expert.post_dit(out["video"], video_pre)  # 将 MoT 输出的 video tokens 还原成视频 velocity 预测，用于未来视频 flow matching loss
    pred_action = model.action_expert.post_dit(out["action"], action_pre)  # 将 MoT 输出的 action tokens 还原成动作 velocity 预测，用于 action chunk flow matching loss
    loss_video = masked_weighted_mse(pred_video[:, :, 1:], target_video[:, :, 1:], inputs["image_is_pad"], t_v)  # 只在未来帧上计算视频损失并排除 padding，实现论文的 video co-training objective
    loss_action = masked_weighted_mse(pred_action, target_action, inputs["action_is_pad"], t_a)  # 在有效 action tokens 上计算动作损失，使模型学习真实机器人控制 chunk
    return model.loss_lambda_video * loss_video + model.loss_lambda_action * loss_action  # 合并 video/action 两个 flow matching 目标，对应论文 Eq.9 及代码中的两个 loss lambda

# source: src/fastwam/models/wan22/fastwam.py:infer_action; src/fastwam/models/wan22/mot.py:prefill_video_cache,forward_action_with_video_cache
# paper: §3.2 Eq.4 and Figure 1(C)
@torch.no_grad()
def fastwam_infer_action(model, image, prompt, action_horizon, steps=10):  # 定义 Fast-WAM 推理路径，只生成动作而不生成未来视频
    context, context_mask = model.encode_prompt(prompt)  # 将语言指令编码成 cross-attention 条件，让动作生成受任务语义约束
    first_frame = model.vae.encode(image)  # 只编码当前观测帧 f0，不实例化 future video latents，这是 Fast-WAM 低延迟推理的关键
    video_pre = model.video_expert.pre_dit(  # 用 video backbone 对当前帧做 single forward pass，形成 action denoising 可读取的 latent world representation
        x=first_frame,  # 输入只有 clean first-frame latent，对应论文中不显式采样 future observations 的 z(o,l)
        timestep=torch.zeros(first_frame.shape[0], device=first_frame.device),  # video branch 使用零 timestep 表示当前帧编码阶段不是未来视频去噪步骤
        context=context,  # 注入语言条件，使 latent world representation 与任务指令绑定
        context_mask=context_mask,  # 传入文本 mask，避免 padding token 参与 cross-attention
        action=None,  # 推理时 video branch 不接收 ground-truth action，保持 direct policy 的测试接口
    )  # 对当前帧做一次 video DiT pre-processing，得到论文 Eq.4 中 z(o,l) 的 token 表示而不是生成未来帧
    action = torch.randn(1, action_horizon, model.action_expert.action_dim)  # 从 Gaussian noise 初始化动作 chunk，后续只对动作 latent 做 iterative denoising
    mask = build_fastwam_attention_mask(video_pre["tokens"].shape[1], action.shape[1], video_pre["meta"]["tokens_per_frame"], video_pre["tokens"].device)  # 构造推理 mask，使 action tokens 只能读取当前帧 world representation 和动作自身上下文
    kv_cache = model.mot.prefill_video_cache(  # 预填充 video branch 的逐层 K/V cache，避免每个 action denoising step 重跑 video DiT
        video_tokens=video_pre["tokens"],  # 使用当前帧 video tokens 作为缓存内容，代表训练时 video co-training 学到的 world features
        video_freqs=video_pre["freqs"],  # 保留 video RoPE 频率，让缓存中的 K/V 与原 video DiT 位置编码一致
        video_t_mod=video_pre["t_mod"],  # 保留 video timestep modulation，使缓存与 video expert 的 block 输入格式一致
        video_context_payload={"context": video_pre["context"], "mask": video_pre["context_mask"]},  # 保留视频分支 cross-attention 条件，确保缓存已经融合语言任务信息
        video_attention_mask=mask[:video_pre["tokens"].shape[1], :video_pre["tokens"].shape[1]],  # 只取 video-to-video 子 mask，避免缓存阶段混入 action tokens
    )  # 逐层缓存 video branch 的 K/V，让每个动作 denoising step 复用同一个 single-pass world encoding
    for t, delta in model.infer_action_scheduler.build_inference_schedule(steps):  # 遍历动作 flow matching 的推理 schedule，逐步把 action noise 更新为动作 chunk
        action_pre = model.action_expert.pre_dit(action, t, context, context_mask)  # 将当前 noisy action 和 timestep 编码成 action tokens，用于本轮 velocity 预测
        action_tokens = model.mot.forward_action_with_video_cache(  # 用当前 action query 读取 cached video K/V 和 action K/V，复现论文 Figure 1(C) 的 KV-cache 推理
            action_tokens=action_pre["tokens"],  # 输入当前 timestep 的 action tokens，因为动作 latent 每轮都会更新
            action_freqs=action_pre["freqs"],  # 使用 action RoPE 频率保持 action chunk 内部位置关系
            action_t_mod=action_pre["t_mod"],  # 使用 action timestep modulation 告诉 action expert 当前去噪阶段
            action_context_payload={"context": action_pre["context"], "mask": action_pre["context_mask"]},  # 让 action branch 继续读取语言任务条件
            video_kv_cache=kv_cache,  # 复用 single-pass video cache，避免显式生成或反复重算未来视频 token
            attention_mask=mask,  # 使用联合 mask 限制 action queries 只读取允许的 video/action key-value
            video_seq_len=video_pre["tokens"].shape[1],  # 告诉 MoT joint 序列中 video prefix 的长度，便于切出 action query rows
        )  # 只重算 action branch，并让 action queries attend 到 cached video K/V，实现“无测试时未来想象”的动作去噪
        velocity = model.action_expert.post_dit(action_tokens, action_pre)  # 将更新后的 action tokens 解码成动作 velocity，作为 flow matching 反向更新方向
        action = model.infer_action_scheduler.step(velocity, delta, action)  # 沿 scheduler 步长更新动作 latent，逐步从噪声收敛到动作 chunk
    return action[0]  # 返回最终 action horizon，对应论文的 direct-policy inference interface

# source: src/fastwam/models/wan22/fastwam_joint.py:infer_action; src/fastwam/models/wan22/fastwam.py:_predict_joint_noise
# paper: §3.3 and Figure 1(A)
def fastwam_joint_variant_step(model, image, prompt, num_video_frames, action_horizon, steps=10):  # 定义 Fast-WAM-Joint 对照组，用显式未来视频 denoising 检验 test-time imagination 的价值
    context, context_mask = model.encode_prompt(prompt)  # 编码语言指令，使 joint WAM 变体的视频生成和动作生成共享同一任务条件
    first_frame = model.vae.encode(image)  # 编码当前观测帧 f0，作为 joint denoising 中固定不变的视觉起点
    video = torch.randn(1, model.vae.model.z_dim, latent_t(num_video_frames), latent_h, latent_w)  # 初始化整段未来视频 latent 噪声，显式保留 test-time future imagination 分支
    video[:, :, 0:1] = first_frame  # 将第一个 latent frame 固定为真实当前帧，防止 joint denoising 改写 observation anchor
    action = torch.randn(1, action_horizon, model.action_expert.action_dim)  # 初始化动作 chunk 噪声，使动作与未来视频在同一推理循环中共同去噪
    for t_v, d_v, t_a, d_a in joint_schedule(steps):  # 同步遍历 video/action 两个 scheduler，让未来视频和动作在每一步共同去噪
        pred_video, pred_action = model._predict_joint_noise(  # 调用共享 MoT 同时预测视频和动作 velocity，体现 joint-generation WAM 范式
            latents_video=video,  # 传入当前 noisy future video latents，使 Joint 变体测试时仍显式维护未来想象
            latents_action=action,  # 传入当前 noisy action latents，让动作与视频共享同一次 MoT forward
            timestep_video=t_v,  # 指定视频分支当前 denoising timestep，用于预测 future video velocity
            timestep_action=t_a,  # 指定动作分支当前 denoising timestep，用于预测 action velocity
            context=context,  # 注入同一语言条件，保证未来视频想象和动作生成服务于同一个任务
            context_mask=context_mask,  # 传入文本 mask，避免 padding 干扰两个分支的 cross-attention
            fuse_vae_embedding_in_latents=True,  # 保持 first-frame VAE embedding 融合，确保 joint rollout 仍以当前观测为锚点
        )  # 同时预测 video/action velocity，并且 Joint 变体允许 action attend 到全部 video tokens 来模拟 imagine-then-execute WAM
        video = model.infer_video_scheduler.step(pred_video, d_v, video)  # 更新未来视频 latent，因此该变体在测试时持续付出视频 denoising 成本
        action = model.infer_action_scheduler.step(pred_action, d_a, action)  # 更新动作 latent，使动作生成与未来视频想象在每一步保持耦合
        video[:, :, 0:1] = first_frame  # 每轮后重新钉住 first frame，保持生成视频以真实当前观察为条件
    return model.vae.decode(video), action[0]  # 返回显式生成的未来视频和动作，用作论文 §3.3 中 joint-generation 对照

3.7 Code-to-paper mapping

Code reference: main @ 45d8e145 (2026-04-03) — pseudocode and mapping based on this commit

Paper concept	Released code anchor	Notes
Fast-WAM MoT architecture	src/fastwam/models/wan22/fastwam.py, FastWAM; src/fastwam/models/wan22/mot.py, MoT	Video expert + action expert with mixed attention.
Action expert DiT	src/fastwam/models/wan22/action_dit.py, ActionDiT	hidden_dim=1024, num_layers=30, heads aligned with video DiT.
Structured attention mask	FastWAM._build_mot_attention_mask()	Action tokens attend to action tokens and first-frame video tokens only.
Training objective	FastWAM.training_loss()	Separate video/action noise, scheduler targets, weighted MSE, padding masks.
Single-pass inference	FastWAM.infer_action()	Encodes first frame, pre-fills video K/V cache, denoises action only.
Video K/V cache	MoT.prefill_video_cache() and MoT.forward_action_with_video_cache()	Avoids recomputing video branch at every action denoising step.
Model config	configs/model/fastwam.yaml	Wan2.2-5B backbone, action expert dimensions, scheduler shifts.
Evaluation config	configs/sim_robotwin.yaml, configs/train.yaml	eval_num_inference_steps=10, text_cfg_scale=1.0 in evaluation.

4. Experimental Setup (实验设置)

Datasets and tasks：LIBERO 包含 Spatial、Object、Goal、Long 四个 suites，每个 suite 500 demonstrations、10 tasks；论文按标准协议训练，并在 40 tasks、2000 trials 上报告 success rate。RoboTwin 2.0 是双臂操作 benchmark，超过 50 tasks；训练数据由 2,500 clean-scene demonstrations 和 25,000 heavy-randomization demonstrations 组成，每个任务在 clean 和 randomized setting 下各 100 trials。真实机器人实验是 Galaxea R1 Lite 平台上的 towel folding，作者收集 60 hours teleoperated demonstrations，任务涉及可形变物体和长时程闭环操作。

Figure 3 解读：该图展示真实 towel-folding benchmark。毛巾折叠难点不只是抓取位置，还包括布料形变后的视觉状态变化、双臂协调、长 horizon 的执行顺序，以及失败后的闭环修正。作者用这个任务测试 Fast-WAM 是否能在没有 embodied pretraining 的情况下获得足够的数据效率，并且是否能以更低 latency 达到真实系统可用的控制频率。

Baselines：RoboTwin 对比 ${\pi }_0$、${\pi }_{0.5}$、Motus、Motus from Wan2.2、LingBot-VA、LingBot-VA from Wan2.2，以及 Fast-WAM 三个 controlled variants。LIBERO 对比 OpenVLA、${\pi }_0$、${\pi }_{0.5}$、LingBot-VA、Motus 和 Fast-WAM variants。真实任务图中还比较了 ${\pi }_{0.5}$、${\pi }_{0.5}$ w.o. pretrain、Fast-WAM、Fast-WAM-Joint、Fast-WAM-IDM、Fast-WAM w.o. video co-train。

Metrics：simulation 主要报告 success rate。RoboTwin 分 clean、randomized 和 average success rate；LIBERO 分 Spatial/Object/Goal/Long 和 average success rate。真实任务报告 success rate、average completion time 和 inference latency；completion time 越短代表策略不是靠反复试错慢慢纠正，而是能更有效完成折叠。

Training config：论文正文报告 backbone 为 pretrained Wan2.2-5B，包括 video DiT、text encoder、video VAE；action expert hidden dimension $d_a=1024$，约 1B parameters，总模型约 6B parameters。action horizon $h=32$；视频帧 temporal downsample $4\times$，每个 chunk 得到 9 video frames；多相机图像先拼成 single image 再进入 VAE。推理使用 10 denoising steps、CFG scale 1.0；optimizer 为 AdamW，learning rate $1\times 10^{-4}$，cosine schedule，mixed precision，gradient clipping 1.0。LIBERO 训练 20k steps，RoboTwin 和真实任务训练 30k steps。latency 在单张 NVIDIA RTX 5090D V2 32GB GPU 上测量。

released code 中可核对的配置：configs/model/fastwam.yaml 使用 model_id=Wan-AI/Wan2.2-TI2V-5B、tokenizer_model_id=Wan-AI/Wan2.1-T2V-1.3B、tokenizer_max_len=128、video/action scheduler train_shift=5.0、infer_shift=5.0、num_train_timesteps=1000。configs/train.yaml 设置 batch_size=2、learning_rate=1.0e-4、lr_scheduler_type=cosine、mixed_precision=bf16、max_grad_norm=1.0、eval_num_inference_steps=10。需要注意论文和 released code 在 weight_decay 上不一致：论文写 0.01，configs/train.yaml 写 0.0。

5. Experimental Results (实验结果)

5.1 Simulation results

RoboTwin 主表结果：Fast-WAM 在无 embodied pretraining 条件下达到 Clean 91.88、Rand. 91.78、Average 91.8。相比之下，${\pi }_0$ 为 62.2 average，${\pi }_{0.5}$ 为 79.8，Motus with pretraining 为 87.8，Motus from Wan2.2 without pretraining 为 77.3，LingBot-VA with pretraining 为 92.2，LingBot-VA from Wan2.2 without pretraining 为 80.6。Fast-WAM 在不使用 embodied pretraining 的情况下接近 pretrained LingBot-VA 的 92.2，并超过 Motus 的 87.8。

Controlled variants 在 RoboTwin 上非常关键：Fast-WAM-Joint 为 Clean 90.84、Rand. 90.32、Average 90.6；Fast-WAM-IDM 为 Clean 91.16、Rand. 91.34、Average 91.3；Fast-WAM w.o. video co-train 为 Clean 82.76、Rand. 84.80、Average 83.8。也就是说，显式 future imagination 的两种变体并没有明显超过 Fast-WAM，而移除 video co-training 造成约 8.0 个 average points 的下降。这直接支持论文结论：training-time video objective 比 test-time future generation 更重要。

LIBERO 主表结果：Fast-WAM 在 Spatial 98.2、Object 100.0、Goal 97.0、Long 95.2，Average 97.6。OpenVLA average 76.5，${\pi }_0$ 94.1，${\pi }_{0.5}$ 96.9，LingBot-VA 98.5，Motus 97.7。Fast-WAM-Joint 在 LIBERO 上 average 98.5，Fast-WAM-IDM 98.0，Fast-WAM w.o. video co-train 93.5。这里 Joint 略高于 Fast-WAM，但差距只有 0.9 points；无 video co-training 则比 Fast-WAM 低 4.1 points，并且 Spatial 从 98.2 降到 89.2，Long 从 95.2 降到 90.0。

5.2 Real-world performance and latency

Figure 4 解读：左图横轴是 average completion time，纵轴是 success rate，越靠左上越好；右图是 inference latency。真实任务中 pretrained ${\pi }_{0.5}$ 仍是最强点，success rate 高且 completion time 短。Fast-WAM family 内部，Fast-WAM-IDM success rate 最高，Fast-WAM completion time 更好；但最重要的是，所有带 video co-training 的 Fast-WAM variants 都显著优于 ${\pi }_{0.5}$ w.o. pretrain。去掉 video co-training 后，Fast-WAM w.o. video co-train 只有 10% success rate，并且 completion time 最长，说明真实形变物体任务对 world-grounded representation 更敏感。

latency 结果体现部署价值：${\pi }_{0.5}$ 约 180 ms，Fast-WAM 和 Fast-WAM w.o. video co-train 均为 190 ms，Fast-WAM-Joint 为 580 ms，Fast-WAM-IDM 为 810 ms。Fast-WAM 相比最慢的 IDM 变体超过 $4\times$ 更快，同时保持相近或更好的总体性能。这个结果说明跳过 explicit future video generation 不是单纯工程加速，而是在性能接近的前提下改变 WAM 推理接口，使其更像可实时闭环的 direct policy。

5.3 Ablation interpretation

论文最有价值的不是单个 benchmark 的 SOTA 数字，而是 controlled comparison 的方向一致性。RoboTwin 上 Fast-WAM、Joint、IDM 三者都在 90.6-91.8 average 区间，差距小于 1.2 points；无 video co-training 降到 83.8。LIBERO 上三者在 97.6-98.5 区间，无 video co-training 为 93.5。真实任务中三者的 success/completion trade-off 有差异，但无 video co-training 直接崩到 10% success rate。这个模式说明，如果模型已经通过未来视频预测学到了动态表示，测试时显式生成未来帧并不是主要性能来源。

这也给 WAM 设计一个实际建议：未来工作可以把更多资源投入训练阶段的 world modeling objective、pretraining data scale、representation quality，而不是默认把视频生成带到测试时。对于机器人系统，低延迟和可控计算预算通常比可视化 future rollout 更重要。Fast-WAM 表明 WAM 的“world model”可以以隐式 latent representation 存在，不一定要以 explicit generated future frames 形式出现。

5.4 Limitations

作者在结论中明确提到未来方向是研究更大规模 pretraining data 和 model scaling 对该设计的影响。当前论文已经显示无 embodied pretraining 也能强，但它没有证明在大规模 embodied pretraining 后 explicit future imagination 一定不重要。真实任务也只报告 towel folding，一类长时程可形变物体任务，尚不能覆盖导航、工具使用、复杂 contact-rich assembly 等更多 embodied settings。

另一个复现风险来自 released code 与论文正文的轻微不一致：weight_decay、默认推理步数、loss weight 配置都需要在运行脚本时确认最终 Hydra override。论文的主要结论不依赖这些细节，但如果要复现表格数字，必须固定 commit、使用论文对应 config，并确认作者是否有私有训练脚本或 checkpoint 配置未在仓库中完全展开。

总体结论：Fast-WAM 通过保留 training-time video co-training、移除 test-time future video generation，证明 WAM 的关键价值可能在于训练阶段形成 world-grounded latent representation。它在 LIBERO 和 RoboTwin 上达到接近或超过强 baselines 的性能，在真实 towel-folding 上保持强表现，同时把 latency 控制在 190 ms。这篇论文对 WAM 的贡献更像一次设计归因实验：它不只是提出一个更快模型，而是指出“未来想象”可能应当主要作为训练信号，而不是必须作为推理过程。