MaskWAM — 统一 Mask 提示与预测的世界动作模型

先厘清梯队：VLA → WAM

VLA(vision-language-action) 直接把"观测+指令"映射到动作；WAM(world-action model) 在中间加一步"预测未来观测"作为代理任务(proxy task)，用视频预测来捕捉物理动态和时空结构，让表征更利于下游控制。MaskWAM 属于 WAM 阵营(基线里的 Motus、FastWAM 也是 WAM，π0/π0.5/GR00T-N1 是 VLA)。

动机：WAM 的两大空间瓶颈

论文指出，现有 WAM 在"如何编码目标物体及其空间位置"上有两个根本缺陷：

核心洞察：用视觉信号(mask)注入空间引导

作者的关键判断是：既然 WAM 本质上是视觉驱动架构，那么空间引导就应该通过视觉模态注入，而不是塞进文本。 mask 正是这样一种密集、以物体为中心的视觉空间信号。于是他们提出两个协同机制：

• 通过 mask 预测实现语义监督：在预测 RGB 的同时预测未来 mask，提供高度语义化的监督信号，强制策略优先关注任务相关区域(抑制背景噪声)。
• 通过视觉提示实现空间锚定：集成首帧目标 mask，建立精确空间锚点，大幅减少指代歧义——哪怕在严重杂乱场景也能给策略可靠引导。

问题设定

给定当前 RGB 观测 I₀∈ℝ^3×H×W、本体感觉 s₀∈ℝ^D、语言指令 ℓ、以及可选的首帧目标 mask M₀∈{0,1}^H×W，MaskWAM 联合预测动作块 a_1:K、未来 RGB 帧 I_1:T 和未来 mask M_1:T：

其中 K 是动作块大小，D 是动作维度，T 是预测 horizon(实现中 T=8)。模型建立在 Wan 2.2 视频生成模型之上，把 mask 作为任务相关区域的显式空间表示，将传统 RGB-space WAM 扩展为 RGB-mask-action 联合建模。

整体数据流

s₀ + 带噪动作经投影层送入动作专家；语言经 T5 以 cross attention 注入视觉分支。

关键张量维度

机制 ①：统一 RGB-Mask 编码 —— 把 mask 渲染成 RGB 图，复用同一个 VAE

最自然的想法是给 mask 单独训一个编码器，但作者发现这会引入表征不匹配。他们的做法：把任务相关 mask 渲染成一张三通道 RGB 兼容图像(固定调色板标目标、另一种颜色标背景)，让它与 RGB 帧有相同的时间 horizon 和空间分辨率，然后用同一个冻结的 Wan 2.2 Causal 3D VAE 编码器 𝓔 编码二者：

沿通道维拼接后送入扩散骨架，使每个 latent 位置同时编码外观(RGB)和任务相关空间引导(mask)。

机制 ②：Patch Embedding 扩展 —— 零初始化优雅接入

为吃下 2C 通道输入，把预训练视频骨架的输入 patch embedding 从 C 扩到 2C：原始视觉通道继承预训练权重，新增 mask 通道初始化为零。这保证微调开始时模型行为与预训练完全一致，mask 信息在训练中被逐渐吸收，不破坏视觉先验。输出头同样 C→2C，预测 RGB 和 mask 两路 latent 速度场。

机制 ③：首帧 mask 作为策略条件 + Mask Dropout 统一提示

部署时 M₀ 是一个可选的首帧锚点，可由 SAM-3 从文本短语、点击、边界框或粗糙 mask 生成。它只在 episode 开头提供一次，无需实时跟踪——首帧 mask 与"预测的未来 mask"之间的对应关系，让世界模型自动把目标规格传播到未来 RGB、未来 mask 和动作中。

Mask Dropout：训练时以概率 p=0.5 把 M₀ 替换为零张量。这让同一个模型既能在语言歧义时吃首帧 mask 做精确消歧，也能在语言清晰时纯靠 RGB+语言运行，无需训两套模型。

机制 ④：联合 Flow Matching + 解耦噪声调度

MoT(Mixture of Transformers) 含两个交互分支：视觉分支联合去噪 RGB+mask，轻量动作专家去噪动作块。训练用统一 flow matching，但解耦两个域的噪声时间步：

• 视觉域：RGB 和 mask 共享时间步 τ_v，保证外观与结构 mask 时空对齐；
• 动作域：采样独立时间步 τ_a，迫使动作专家在各种视觉噪声水平下学习去噪——这为推理时的 partial denoising 奠定理论基础。

三项分别是 RGB 视频、mask 表征、动作轨迹的 flow-matching 损失。仓库 config 中三者权重默认各为 1.0。

训练设置(论文 Implementation Details)

• Video VAE 与 T5 文本编码器全程冻结，只训 DiT 骨架 + MoT 动作专家 + 扩展的 patch embed / 输出头。
• 输入分辨率 384×320(T-shape)，预测 horizon T=8。
• Mask dropout 概率 p=0.5。
• 损失：𝓛_video+𝓛_mask+𝓛_act(联合 flow matching，解耦 τ_v/τ_a)。
• 注意力：block-wise causal attention mask，让 RGB / mask / action token 在统一序列里保持正确因果关系(图 2b)。

数据标注 pipeline(附录 B，图 6)

为拿到训练所需的每帧 mask，作者搭了自动标注流程，且区分两类任务以避免 oracle 泄漏：

• 语言清晰任务：Qwen3-VL 解析指令识别任务相关物体 → SAM3 分割并跨帧跟踪。
• 语言歧义任务：文本无法消歧，由人在首帧点一个 point prompt → SAM3 传播成完整 mask 并跟踪。

所有首轮标注人工验证，低质量 episode 用关键帧点提示修正后 SAM3 重传播。效果：91% 的 episode 无需人工修正，标注 50 个 episode 约 3 分钟；语言歧义任务的初始点提示每 episode 仅多花 5-10 秒。

推理：Partial Denoising —— 不生成完整视频

完整视频生成太慢。MaskWAM 的推理只做一步去噪：

动作专家联合关注这些部分去噪的 latent、语言、本体感觉来生成动作块，再配合 KV-caching 进一步加速。首帧 mask 只在 episode 开头由 SAM3 生成一次。

评测协议矩阵：作者到底测了什么

LIBERO(表 1) —— 即使不用 mask 提示也涨点

联合 mask 预测把 RGB-only(97.3) 提到 98.4——部署时不用任何 mask 提示，仅辅助训练目标就带来提升，Spatial +2.0 / Goal +1.2。注意力图显示 RGB-only 常盯背景，mask 监督让模型精确关注任务区域。

RoboTwin 2.0(表 2) —— 双臂多任务随机化

92.2，超 π0 达 19.4、超 FastWAM 达 4.5。Mask-only(88.8) 原生就比 RGB-only(87.3) 强，联合两者最大化到 92.2。

真机(Dual-arm Xtrainer)：语言歧义是主战场

语言清晰任务(表 3，4 任务)

MaskWAM 平均 84.3，超 FastWAM(79.0)、π0.5(72.3)。整合未来 mask 预测比自家 RGB-only 一致涨点(Task1 86→91，Task3 76→81)——mask 不损害清晰任务，反而稳定增强。

语言歧义任务(图 4 / 表 5，整体 84.9)

整体 84.9，超最强 baseline π0-mask(51.7) 达 33.2。三个 OOD 轴均保持高性能——干扰物几乎不掉、新物体类别级迁移、光照强不变性。

核心消融(表 4)

• Q1：RGB-only(97.3) / Mask-only(97.6) / 联合(98.4)——成对预测是必要的表征正则化器。
• Q2：无未来 mask 预测 → 21.6%。未来 mask 预测是"锚定视觉提示"的关键。
• Q3：mask 提示(84.9) 远胜坐标文本(18.2)。密集视觉先验 vs 稀疏文本坐标。

诚实提示：官方仓库当前是"骨架"，核心实现未发布

① 配置里能看到的架构意图 🟢 已发布

maskwam/models/wam.py 的 MaskWAMConfig 把论文的三大要素落成了显式开关：

@dataclass
class MaskWAMConfig:
    rgb_encoder: str = "vit_base_patch16"
    mask_encoder: str = "vit_small_patch16"
    hidden_dim: int = 768
    num_layers: int = 12
    action_dim: int = 7
    action_chunk_size: int = 8
    pred_horizon: int = 4              # 预测的未来 RGB/mask 帧数
    use_mask_prompt: bool = True       # ← 机制③ mask 作为输入提示
    predict_future_mask: bool = True   # ← 机制① mask 作为预测目标
    loss_weights: dict = {"rgb":1.0, "mask":1.0, "action":1.0}  # ← 式3 三项损失

use_mask_prompt 与 predict_future_mask 两个布尔量，正好对应论文"mask 既是输入、又是输出"的双重角色；loss_weights 对应式(3)的 𝓛_video+𝓛_mask+𝓛_act。

② 类结构印证"统一编码 + 三个预测目标" 🟢 API 结构

MaskWAM.forward 的签名与 docstring 明确了输入输出契约(实现待发布)：

def forward(self, rgb, language, mask_prompt=None) -> dict:
    """
    rgb: (B, T, 3, H, W) 观测帧
    mask_prompt: 可选首帧目标 mask (B, 1, H, W)
    Returns: dict with 'actions', 'future_rgb', 'future_mask'
    """
    raise NotImplementedError("MaskWAM.forward is not released yet.")

返回 actions / future_rgb / future_mask 三路输出，与论文"联合预测动作块、未来 RGB、未来 mask"(式1)一致。MaskEncoder 的 docstring 也点明"同一编码器复用于首帧 mask 提示 与 未来 mask 预测目标"。

③ 骨架与论文最终设计的差异(重要) 🩷 核对提醒

评价与启发

它把什么问题问清楚了

"世界模型该预测什么"是 WAM 路线的隐含假设。主流答案是"预测未来 RGB(what happens)"。MaskWAM 给出补充答案：还要预测"什么重要(what matters)"——用未来 mask 把物体中心的语义监督注入 WAM。同时它用一个漂亮的实证反转说明：在视觉驱动的 WAM 里，空间信息该走视觉通道(mask)而非文本通道(坐标)。

方法论上的三个亮点

• 零成本的注意力正则化：即使部署不用 mask 提示，仅把未来 mask 预测当辅助目标就能提升 RGB-only 策略 1-5%——这是可直接搬到任意端到端策略的训练技巧。
• 共享编码器处理异构信号：把 mask 渲染成 RGB 兼容格式后复用预训练 VAE，避开独立编码器的表征不匹配。这思路可迁移到深度图、触觉图等其它模态。
• 解耦噪声调度：视觉/动作用不同 flow-matching 时间步，让动作分支在带噪视觉上下文学习，从而支持推理时 partial denoising 省算力。

值得追问 / 存疑处

• 效率完全没量化：反复强调省延迟，却不给 ms/Hz/参数量/显存。基于 Wan 2.2 的方法，真机实时性是关键悬念。
• 归因不够干净：声称"WAM 里 mask 预测特别有效"，却没把同样的 mask 预测目标加到 π0(VLA) 做对照——现有 π0-mask 只是加法融合 mask 输入，未做 mask 预测。无法完全区分收益来自 mask 预测本身还是 WAM 框架。
• 清晰任务边际收益小：LIBERO 仅比 FastWAM 高 0.8，主要价值集中在歧义/OOD 场景——若你的场景足够简单，mask 标注成本未必划算。
• 依赖 SAM3 与单一平台：歧义场景需 SAM3 提供首帧 mask(透明/反光物体可能失败)；真机只在 Xtrainer 双臂平台验证。

与本站其它页联动

配合 ImageWAM 页对比两种 WAM 效率哲学(取 KV cache 不解码图 vs partial denoising 不生成完整视频)；配合 DiT4DiT 页看视频动态+动作联合建模的另一条路；配合"世界模型"谱系页理解 WAM 与 VLA 的分野；配合 DP·DiT·Flow Matching 页理解这里的 flow-matching 动作专家。

自测：7 题检验理解