清华王岩团队提出StageWAM：世界动作模型不应该只预测短期未来

论文团队来自清华大学产业研究院（AIR）、北京航空航天大学、清华大学 AIR 无锡创新中心、中国人民大学和 TARS Robotics。本文主要介绍该团队的最新工作 StageWAM:Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation。

近年来，World-Action Model（WAM）开始成为机器人操作中的一条重要技术路线。相比直接从观测生成动作的 VLA，WAM 会同时预测未来视频和机器人动作，因此能够显式学习物体运动、接触以及局部场景变化。但现有 WAM 通常将“未来”表示成一个固定长度的视频—动作片段，这种短期未来很适合描述机械臂接下来怎么运动，却没有明确告诉模型：当前任务接下来真正应该进入什么状态。

StageWAM 从这个问题出发，将机器人操作中的未来拆成两个互补的尺度：一个是描述局部运动和接触变化的 short-term physical future，另一个是描述任务下一阶段应该达到什么状态的 stage-level semantic future。具体来说，StageWAM 使用一个基于 V-JEPA2 的 Stage-JEPA 预测下一任务阶段的 latent representation，再将这个 latent representation 注入原有的WAM，共同指导未来视频和动作的生成。在 RoboTwin 2.0 的 50 个任务上，StageWAM 最终取得了 90.25% 的整体成功率，同时成功轨迹的平均执行步数相比 Motus 减少了 5.97%。

• 论文题目： StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation

• 合作机构： 清华大学智能产业研究院（AIR）、北京航空航天大学、清华大学 AIR 无锡创新中心、中国人民大学、它石智航

• 论文链接： https://arxiv.org/abs/2608.10780

考虑一个很常见的机器人操作过程：机械臂需要先抓起一个物体，再把它移动到指定位置。对于一个 WAM 来说，它可以不断预测未来几帧视频和接下来的一段动作，例如机械臂向物体靠近、夹爪闭合、物体被抬起。只要每一个局部预测足够准确，机器人就有机会通过不断重规划完成整个任务。

但这种方式存在一个问题：模型所看到的未来始终是一个固定的短时间窗口。它能够知道“接下来几步怎么走”，却并没有显式表示“这一阶段结束以后应该到哪里”。因此，即便生成的局部运动本身合理，策略仍然可能在某个状态附近进行大量重复修正，而没有快速进入下一个真正与任务进度相关的状态。

论文中的 Figure 1 对这一点进行了很直观的区分。传统 VLA 从当前观测和语言指令直接预测动作；WAM 在此基础上进一步预测短期视觉未来；而 StageWAM 额外预测一个 stage guidance，表示当前任务接下来应该达到的进度状态，再利用这个状态指导局部动作生成。

这也是 StageWAM 最核心的观点：机器人操作中的“未来”并不是只有一种。 短期未来适合表达物体运动、接触和局部物理变化，而任务进度关注的则是更抽象的对象状态和关系变化。例如对于一个放置任务，短期未来需要知道机械臂怎样移动，但阶段级未来关心的是物体是否已经被抓取、是否已经移动到目标区域，以及下一步应该形成怎样的物体—目标关系。

直接把 WAM 的视频预测长度不断拉长并不是一个理想的解决方案。随着预测距离增加，像素级生成需要同时处理大量外观、运动和时间细节，而“下一任务阶段”本身并不需要规定这些细节。因此，StageWAM 没有尝试生成一张更远的目标图像，而是把阶段未来放在latent空间中进行预测。

StageWAM 的整体结构如下图所示。方法首先通过 Stage-JEPA 从当前观测和任务指令中预测下一阶段的 latent representation，然后把这一latent representation 作为内部条件加入 WAM。这样一来，JEPA 负责提供任务进度方向，而 WAM 继续负责它本来擅长的局部视频和动作建模。

这里首先要解决的问题是：训练数据中并没有人工标注好的任务阶段。为此，StageWAM 不预先定义“抓取”“移动”“放置”等阶段，而是利用冻结的 V-JEPA2 表征，从 demonstration 中自动寻找状态变化明显的位置。

具体来说，对于一条连续轨迹，StageWAM 首先以固定间隔采样候选位置，并以每个位置为中心截取 64 帧视频片段，通过冻结的 V-JEPA2 encoder 得到 latent representation。随后比较候选位置前后局部窗口的表征变化，同时考虑 pooled representation 的整体变化和 token-level 的局部变化，并将其综合为 transition score。对于候选位置j，其 transition score 定义为：

其中v表示 pooled representation，u表示 token-level representation。这个 score 经过平滑和归一化后，再进行 temporal NMS：按照 transition score 从高到低依次选择候选位置，当一个位置被选为 keyframe 后，就抑制其时间邻域内距离过近的其他候选点，避免同一次状态变化产生多个重复的阶段边界。最终最多保留 5 个变化最明显的 keyframes，并与轨迹开始帧和最后一帧共同组成阶段边界。

这些 stage boundary 只表示轨迹在这些位置附近发生了明显的状态转换。对于阶段区间[ b_i , b_i+1 )，其中每个当前时刻t都与下一阶段边界b_i+1组成一个训练 pair。因此，keyframe 只负责划分阶段，而 Stage-JEPA 实际使用的是大量 current-to-next-stage 的稠密监督。

有了这些 Stage-Pair 后，Stage-JEPA 分别编码当前 observation片段和对应的 next-stage observation片段。Stage Predictor 接收当前 observation representation 和语言 instruction，预测下一阶段的 latent representation，并通过 cosine loss 与真实下一阶段的 V-JEPA2 representation 对齐。也就是说，它不生成未来图像，也不直接预测动作，而是学习：从当前状态和任务目标出发，下一阶段应该变成什么样。

得到 predicted stage latent 后，StageWAM 基于 Motus 构建负责短程未来建模的 World-Action Model。其中 video、action 和 understanding 三个分支分别负责局部视觉演化、动作生成以及视觉语言理解，并通过 Tri-modal Joint Attention 持续交换信息。

Stage-JEPA 输出的 next-stage latent tokens 首先经过 pooling 和 MLP 投影到 video-token feature space，再通过 gated residual 的方式注入 video tokens：

其中 gate (α) 用于控制 stage information 对原始 WAM representation 的影响。注入后的 video tokens 再进入后续的 Tri-modal Joint Attention，因此 stage latent 不仅影响未来视频预测，也会通过不同分支之间的信息交互进一步影响动作生成。

训练时，已经训练好的 Stage-JEPA 被冻结，只优化 WAM 和对应的 conditioning interface。最终两个模块形成了明确的分工：Stage-JEPA 负责预测“下一阶段应该发生什么”，WAM 负责预测“为了让它发生，短期世界应该如何变化、机器人现在应该怎么动”。

论文在 RoboTwin 2.0 的 50 个双臂操作任务上进行了评测。在 clean 和 randomized 两种设置下，StageWAM 分别取得 91.42% 和 89.08% 的成功率，整体成功率达到 90.25%，相比 Motus 均有提升。

消融实验也说明，论文提出的各个模块都起到了作用。完全移除 JEPA 后，整体成功率下降到 87.23%。也就是说，StageWAM 中真正起作用的是预测出的下一阶段信息参与了后续 world-action generation，而不是简单增加一个额外的视觉 encoder。

除了成功率，论文还观察到了一个比较直接的现象：StageWAM 在成功完成任务时所需要的动作更少。在成功 rollout 中，Motus 的平均执行长度为 79.57 steps，而 StageWAM 为 74.82 steps，减少了 5.97%。

现有 WAM 擅长的是短时间范围内的物理演化，这对于建模运动、接触和动作执行非常重要；但对于一个由多个阶段组成的操作任务来说，仅有这种局部未来还不足以明确表示任务下一步应该向什么状态推进。StageWAM 重新区分了机器人操作中两种不同尺度的未来，短期未来告诉模型局部环境接下来会怎么变化，而阶段未来则进一步告诉它，这些局部动作最终应该把任务带向哪里。因此StageWAM增加了一个基于 JEPA 的阶段预测过程，用latent空间的future描述下一任务阶段，再将这一latent直接作为条件接入 WAM。这样，模型不需要用一段越来越长的视频同时承担局部物理预测和长程任务进度建模，而是分别用 short-term physical future 和 stage-level semantic future 表达两类信息。