StageWAM：面向机器人操作世界-动作模型的阶段联合嵌入预测

StageWAM：面向机器人操作世界-动作模型的阶段联合嵌入预测

论文链接： arXiv:2608.10780 | 作者： Xiao Liu, Yuguang Yang, Xi Wang, Kai Jiang, Cheng Chi, Yong Xu, Wenchao Ding, Yilun Chen, Yan Wang | 代码状态： 未公开

一句话总结

StageWAM 在世界-动作模型（WAM）的短期视频-动作预测之上，叠加一个基于 JEPA 的阶段级语义未来预测器，用预测的下一任务阶段潜在表征门控注入 WAM 视频流，在不增加推理架构的前提下将 50 项双臂操作任务的成功率从 87.85% 提升至 90.25%，并减少成功回合的执行步数。

研究背景与动机

通用机器人策略的目标是将多模态视觉观测、机器人状态和语言任务指令映射为跨任务的 executable 动作。当前主流方案分为两大阵营：视觉-语言-动作模型（VLA）和世界-动作模型（WAM）。VLA 直接从当前观测和指令生成动作，语义理解和指令遵循能力强，但不显式建模环境的未来变化。WAM 则联合预测未来视频帧和动作块，提供了对物体运动、接触和局部场景演化的密集监督，但预测一段合理的视频-动作轨迹并不意味着策略理解了任务应该到达的下一个状态。

这一区分至关重要。现有 WAM 通常将未来表示为一个固定长度的密集时间块，能够捕捉局部运动和接触动力学，但将下一个任务相关的物体状态变化或关系变化隐含在密集 rollout 中。例如，在"拿起杯子放到架子上"这个任务中，WAM 可以预测未来若干帧的杯子运动轨迹，但它并不显式知道"杯子应该被放置到架子上"这个阶段目标。策略可能在局部做出正确的运动，却在整体任务进度上偏离方向。

这引出一个核心问题：如何在生成局部动作之前，让机器人策略预测任务的下一个有意义阶段？这种阶段未来应编码任务相关的物体状态和关系变化（如空间对齐或机器人-物体交互），而不承诺像素细节或过渡时长。直接预测远距离目标图像与这一目标不匹配，因为像素级预测会强调外观而非任务完成。

StageWAM 的出发点正是区分两种互补的未来：短期物理未来描述即时的视频-动作演化，阶段级语义未来指定下一个应到达的物体或关系状态。前者由 WAM 负责，后者由 JEPA 负责。两者的结合使得策略既知道"下一步该到哪"，又知道"怎么到那"。

联合嵌入预测架构（JEPA）为建模这种阶段级未来提供了天然机制。JEPA 不重建每个像素，而是从当前上下文预测目标状态的潜在表征，保留任务相关的语义和物理结构，同时对低级外观不敏感。StageWAM 利用目标条件 JEPA 预测器，从机器人演示中推断的下一个任务阶段的潜在表征，作为内部进度目标来条件化短期的世界和动作生成。

预备知识

V-JEPA2 是 Meta 提出的自监督视频表示学习模型，通过对视频片段的掩码预测任务学习通用的视频表征。StageWAM 使用冻结的 V-JEPA2 编码器提取观测的潜在表征，这一表征对视频中的语义和物理结构敏感而对像素级外观变化不敏感，正是阶段级预测所需要的。模型 facebook/vjepa2-vitl-fpc64-256 接受 64 帧输入。

Motus 是一种基于混合 Transformer（Mixture-of-Transformers）的世界-动作模型，包含视频生成专家、动作专家和视觉-语言理解专家三个流。视频专家从 Wan2.2 初始化，建模局部视觉演化；动作专家处理机器人状态和密集动作块；理解专家用冻结的 Qwen3-VL 编码器提取指令和观测相关特征。三者通过三模态联合注意力交换信息，同时保持各自的模态特异性归一化、前馈层和输出头。

方法详解

StageWAM 的训练分为两个阶段。阶段一训练目标条件 JEPA 预测器，从当前观测和任务目标推断下一阶段的未来表征。阶段二冻结该预测器，训练 WAM 在预测的阶段条件下生成短期视觉未来和动作块。推理时，两个组件以滚动时域循环顺序运行。

阶段一：JEPA 预测器训练

阶段级预测提供任务演化的指导而非详细实现，捕捉物体状态和关系的语义变化，将局部外观和运动留给 WAM。因此，这一阶段的学习使用表征空间监督，在引入视觉和动作目标之前完成。

阶段对构造。 阶段一从机器人轨迹构造监督，无需手动阶段标签。对于一条轨迹，每 8 帧采样候选中心 $\mathcal{C}=\{c_j\}_{j=1}^{M}$，用冻结的 V-JEPA2 编码每个中心处 64 帧片段，得到 token 表征 $U_j=\{u_{j,n}\}_n$。通过对 token 池化得到片段级表征 $v_j$。然后计算局部窗口 $w=2$ 内的表征变化分数：

$$q_j = \max\left\{\mathcal{N}\left(\lVert v_{j^+} - v_{j^-} \rVert_2\right),\; \mathcal{N}\left(\max_n \lVert u_{j^+,n} - u_{j^-,n} \rVert_2\right)\right\}$$

其中 $\mathcal{N}$ 是每条 episode 的 min-max 归一化，$j^- = \max(1, j-w)$，$j^+ = \min(M, j+w)$。分数 $q_j$ 同时考虑片段级和 token 级的表征变化，捕捉从粗到细的阶段转移信号。

随后用半径 $r=3$ 的时序非极大值抑制（NMS）保留至多 $K=5$ 个转移关键帧：

$$\mathcal{K} = \{c_j \mid j \in \operatorname{NMS}(\{q_j\}_{j=1}^{M};\; r,\; K)\}$$

关键帧与首尾帧共同定义边界 $\mathcal{B} = \operatorname{sort}(\{0\} \cup \mathcal{K} \cup \{T-1\})$。对每个阶段区间 $[b_i, b_{i+1})$，当前帧 $t$ 与其下一边界 $b_{i+1}$ 配对，构造密集训练对：

$$\mathcal{P} = \bigcup_{i=0}^{L-1}\{(t, b_{i+1}) \mid b_i \leq t < b_{i+1}\}$$

关键帧仅定义边界，阶段一使用密集的当前到下一阶段对 $(t, b_{i+1})$ 进行训练。这种自动阶段检测方法避免了对人工标注的依赖，同时利用 V-JEPA2 的语义表征能力识别有意义的任务转移点。

目标条件阶段预测。 StageWAM 将未来预测模型 $f_\theta$ 实现为 Stage-JEPA，使用冻结的 V-JEPA2 编码器和从 V-JEPA2 预测器初始化的可训练目标条件预测器。由于 V-JEPA2 需要 64 帧输入，密集对 $(t, b_{i+1})$ 被转换为以 $t$ 为中心的 64 帧当前观测片段 $\mathcal{S}_t$ 和以 $b_{i+1}$ 为中心的 64 帧目标观测片段 $\mathcal{S}_t^{\text{target}}$。冻结编码器 $E$ 分别产生当前潜在 $z_t = E(\mathcal{S}_t)$ 和目标潜在 $z_t^{\text{target}} = E(\mathcal{S}_t^{\text{target}})$。预测器产生 token 级下一阶段预测 $\widehat{Z}_t^{\text{target}}$，池化和归一化后得到 $\widehat{z}_t^{\text{target}}$ 用于监督。

对于一批 $B$ 个归一化的预测-目标对，优化余弦损失：

$$\mathcal{L}_{\text{stage}} = \frac{1}{B}\sum_{i=1}^{B}\left(1 - (\widehat{z}_i^{\text{target}})^\top z_i^{\text{target}}\right)$$

该损失直接在表征空间对齐预测潜在与目标潜在，避免像素重建带来的外观偏差，使预测器专注于任务进度的语义目标。

阶段二：阶段条件 WAM 训练

WAM 补充阶段指导所需的细粒度视觉动力学和动作。由于联合优化可能让细节导向的损失改变语义目标，阶段二冻结 JEPA，仅训练 WAM 和条件接口。当前观测片段和指令产生 $\widehat{z}_t^{\text{target}} \equiv \widehat{Y}_t^{\text{stage}}$。

StageWAM 在 Motus 的 Transformer 层之前对 WAM 进行增强。视频分支形成当前条件的含噪视频 token $X_t^v \in \mathbb{R}^{N_v \times d_v}$，Stage-JEPA 输出 $\widehat{Z}_t^{\text{target}} \in \mathbb{R}^{N_j \times d_j}$ 包含 $N_j$ 个预测阶段位置。对该输出做均值池化并用 MLP $P_\psi$ 匹配 Wan 视频 token 维度，然后做加性注入：

$$h_t = P_\psi\left(\operatorname{Pool}(\widehat{Z}_t^{\text{target}})\right) \in \mathbb{R}^{d_v}$$

$$\alpha = 0.2\,\sigma(\beta)$$

$$\widetilde{X}_t^v = X_t^v + \alpha\, h_t$$

其中 $\beta$ 是零维可训练参数，$\alpha$ 是所有样本、层、视频 token 和通道共享的全局有界标量门。初始值设为 0.02，确保训练初期阶段条件的注入是温和的残余更新。条件向量 $h_t$ 保持样本特异性并广播到该样本的所有视频 token。下游 WAM 随后预测局部视觉未来和动作块。

训练保留 WAM 原生的视觉和动作目标，并可选地包含条件接口的正则化项：

$$\mathcal{L} = \mathcal{L}_{\text{video}} + \mathcal{L}_{\text{action}} + \mathcal{L}_{\text{reg}}$$

$$\mathcal{L}_{\text{reg}} = \lambda_g\,\alpha^2 + \lambda_\rho(\rho - \rho_0)^2,\quad \rho = \frac{\lVert \alpha\, h_t \rVert}{\lVert X_t^v \rVert}$$

其中 $\rho$ 度量注入的阶段特征相对于原始视频 token 的幅度，$\rho_0$ 是手动选择的目标比率（设为 0.08）。第一项 $\lambda_g \alpha^2$ 偏置有界门趋向小的残余更新，第二项 $\lambda_\rho(\rho - \rho_0)^2$ 阻止相对更新偏离 $\rho_0$。两个正则化权重分别为 $\lambda_g = 10^{-4}$ 和 $\lambda_\rho = 10^{-3}$。

两个阶段因此学习互补信息：阶段一指定什么任务级变化应该发生，阶段二学习如何通过详细的视觉演化和动作实现它。

闭环推理

部署时，策略维护一个因果 64 帧观测缓冲 $\mathcal{H}_t$。在 64 个观测可用之前，重复最早观测帧填充时间片段。每次策略查询执行：

$$\widehat{z}_t^{\text{target}} = f_\theta(\mathcal{H}_t, l)$$

$$(\widehat{V}_t^{\text{short}}, \widehat{A}_t) = W_\phi(o_t, s_t, l, \widehat{z}_t^{\text{target}})$$

预测的阶段潜在通过与阶段二相同的条件接口注入，WAM 随后预测有限时域的视觉未来和动作块。环境执行动作块，将结果观测追加到缓冲，再次查询策略。阶段边界和目标观测片段仅用于训练时监督，推理时反复从观测历史中预测自身的阶段条件。

下面的流程图展示了 StageWAM 从阶段检测到推理的完整数据流：

flowchart LR
    A["机器人演示轨迹"] --> B["V-JEPA2
特征提取"]
    B --> C["阶段边界检测
(NMS)"]
    C --> D["阶段对构造
(t, b_i+1)"]
    D --> E["Stage-JEPA
训练(余弦损失)"]
    E --> F["冻结JEPA
预测器"]
    F --> G["门控注入
alpha=0.2*sigma(beta)"]
    H["当前观测"] --> G
    G --> I["Motus WAM
视频+动作生成"]
    I --> J["执行动作块"]
    J --> H

图1：StageWAM 框架总览。冻结的 Stage-JEPA 指导注入 WAM 视频流，预测的下一阶段潜在 token 条件化短期视频流和动作块预测。

实验结果

实验在 RoboTwin 2.0 上进行，这是一个双臂操作基准，包含 50 项任务，在干净和随机化两种环境配置下评估。每个任务-配置对运行 100 次闭环 episode，禁用提前停止，共 10,000 次 episode。主要指标是任务成功率。

基线包括 GO-1（AgiBot 具身基础模型）、$\pi_{0.5}$（通用 VLA）、X-VLA（跨任务跨体预训练 VLA）和 Motus（联合建模理解、未来视频和动作的 WAM）。基线数据来自 Motus 论文的已发表结果，StageWAM 和 Motus 使用相同的闭环协议评估。

表1 展示了按语义操作技能分组的成功率。StageWAM 在干净设置下达到 91.42%，随机化设置下 89.08%，在干净设置的 7 个语义类别中的 5 个以及随机化设置的 6 个类别中优于 Motus。

语义类别(任务数)	GO-1干净	GO-1随机	pi0.5干净	pi0.5随机	X-VLA干净	X-VLA随机	Motus干净	Motus随机	StageWAM干净	StageWAM随机
获取与提升(5)	76.40	76.60	36.80	36.20	80.80	74.20	94.20	94.60	92.60	92.00
交接(2)	10.50	10.00	23.00	18.50	36.50	18.50	82.00	68.00	96.50	93.50
定向放置(13)	25.00	23.23	42.46	46.00	78.85	79.77	82.62	85.08	88.23	86.54
容器装载(8)	36.88	40.25	34.75	37.00	74.38	73.25	88.62	87.25	91.12	90.00
排列与堆叠(8)	17.62	17.25	43.62	42.38	63.00	64.75	90.12	88.00	91.88	89.62
铰链/设备交互(7)	61.43	54.71	47.86	47.43	78.43	82.43	93.57	89.86	90.71	90.00
工具使用(7)	42.71	37.71	57.86	58.43	70.57	73.71	91.29	86.43	95.57	87.86
全部(50)	37.86	36.24	42.98	43.84	72.88	72.84	88.66	87.02	91.42	89.08

表1：RoboTwin 2.0 按语义操作技能分组的成功率(%)。加粗为每行最高值。

最强的结果出现在执行依赖于显式任务进度变化的任务上。交接、定向放置和容器装载要求策略在转移和对齐过程中维持期望的物体归属或物体-目标关系。排列与堆叠需要有序的中间关系，工具使用和动态操作要求识别机器人应接近、对齐还是交互。这些任务直接匹配 StageWAM 的双尺度设计：Stage-JEPA 指定下一步应到达什么任务相关状态，WAM 建模如何通过局部视觉动力学和动作实现它。

相比之下，获取与提升主要依赖即时抓取几何和接触稳定性，留给 JEPA 条件化的互补阶段结构较少。铰链/设备交互包含有意义的状态变化，但成功还需要机构特异的接触、约束运动和精确时机，当前间接条件化的阶段信号可能不足。这说明阶段级指导对于"任务进度需要显式推理"的任务增益最大，对于"即时感知-运动足够"的任务增益有限。

表2 的消融实验隔离了 StageWAM 的四个主要组件：

配置	干净(%)	随机化(%)	总体(%)
移除JEPA	87.80	86.66	87.23
仅当前帧token未来指导	91.10	88.68	89.89
移除阶段一训练	91.34	88.62	89.98
移除门正则化	91.00	88.76	89.88
StageWAM(完整)	91.42	89.08	90.25

表2：StageWAM 在 RoboTwin 2.0 上的消融实验。成功率为 50 项任务的平均值。

移除 JEPA 产生最大的退化（从 90.25% 降至 87.23%），说明阶段条件通路提供了短期 WAM 主干之外的信息。保留 JEPA 分支但阻止未来潜在注入也降低了成功率，支持将预测的下一阶段潜在注入 WAM 表征而非仅用于修改当前观测。未做阶段一训练的模型仍有竞争力（89.98%），说明预训练的 V-JEPA2 预测器已提供有用的通用视觉结构，但在当前到下一阶段对上的任务特定训练给出了最佳的干净、随机化和总体结果。移除门正则化产生较小下降，说明 sigmoid 有界门已稳定残余更新，而两个正则化项提供了对其规模的额外约束。

定性分析（图2）提供了 StageWAM 与 Motus 在代表性成功回合上的比较。虽然两者最终都满足成功条件，但执行模式差异显著。在执行轨迹中，StageWAM 在短序列动作后即达到成功，而 Motus 持续更长时间并反复上下移动被操纵物体。这说明预测的阶段潜在作为内部进度目标，偏置局部动作趋向任务进度，减少了冗余的纠正运动。

图2：阶段指导执行的定性案例。StageWAM 使用预测的下一阶段潜在指导局部动作生成，以更少的执行步数达到成功条件。

图3：视觉 rollout 对比。StageWAM 快速从初始交互推进到下一任务相关状态，而 Motus 在物体周围进行多次局部纠正。

执行效率方面，StageWAM 将所有 rollout 的平均回合长度从 156.28 步降至 132.80 步，降幅 15.02%。这一汇总比较混合了成功和失败回合。仅看成功回合，平均长度从 79.57 步降至 74.82 步，降幅 5.97%，同时保持任务完成。这与阶段指导帮助策略更直接地趋向中间任务目标、减少冗余局部纠正的判断一致。

局限性分析

作者自述的局限。 论文在结论中明确指出两个未来方向，隐含了当前局限：一是 WAM 的推理延迟限制了闭环控制频率，StageWAM 虽然在成功率上提升显著，但两阶段顺序推理（先 JEPA 预测，再 WAM 生成）增加了延迟开销，在需要高频控制的精密操作场景中可能成为瓶颈；二是当前的多尺度潜在表征尚未联合编码粗粒度任务进度和细粒度接触、几何、运动信息，这意味着阶段条件目前是间接的，对于需要机构特异接触和精确时机的任务（如铰链/设备交互）增益有限。

独立判断的局限。 第一，阶段边界的自动检测依赖 V-JEPA2 的表征变化启发式，阈值参数（NMS 半径 $r=3$、最大关键帧 $K=5$、候选间隔 8 帧）是固定的，不同任务结构可能需要不同的边界粒度，论文未探讨这些参数的鲁棒性。第二，门控注入仅作用于视频 token 流，动作专家不直接接收阶段条件，阶段信息需要通过联合注意力间接影响动作预测，这一设计选择的理论正当性未被充分讨论。第三，所有实验限于仿真环境 RoboTwin 2.0，缺乏真实机器人验证，sim-to-real 的差距未被评估。

总结与展望

StageWAM 的核心贡献在于将机器人未来建模区分为两个互补的时间尺度，并用门控接口将阶段级语义预测直接注入世界-动作模型的视频流。这一设计不改变 WAM 的推理架构，仅在 Transformer 层之前增加一个有界残余注入，使得预训练的 Motus 模型可以无缝继承。50 项 RoboTwin 2.0 任务的实验证明，阶段级未来指导对需要显式任务进度推理的操作有实质增益，同时消融实验确认了 JEPA 预测器、潜在注入、阶段一训练和门正则化各自的作用。

从更广的视角看，StageWAM 代表了一种趋势：将表征学习中的预测思想（JEPA）与生成式世界模型（WAM）结合，用语义级别的未来目标引导局部生成。这不同于直接预测目标图像或视觉子目标的方法——后者受像素级外观影响，而 StageWAM 的潜在空间预测保留了任务相关结构的同时避免了对低级外观的敏感。未来工作沿两个机器人核心方向展开：降低 WAM 推理延迟以实现高频闭环控制，以及学习多尺度潜在表征以联合编码粗粒度任务进度和细粒度接触、几何、运动信息。

金句： 知道下一步该到哪，比知道怎么到那更重要——但在机器人操作中，你需要两者，而且你需要它们协同工作。