StageWAM：预判任务阶段 提升机器人操作成功率

近年来，世界动作模型（WAM）在机器人操作领域成为重要技术路线，相比直接从观测生成动作的VLA模型，WAM可以同时预测未来视频和机器人动作，能够显式学习物体运动、接触和局部场景变化。但现有WAM通常将未来表示为固定长度的视频-动作片段，虽然能描述机械臂接下来的局部运动，却无法明确告知模型当前任务接下来应该进入什么状态，即便局部运动预测准确，策略也可能在某个状态附近反复修正，无法快速推进到与任务进度相关的下一阶段。

针对这一问题，国内高校联合团队推出了StageWAM模型，该工作将机器人操作中的未来拆分为两个互补的尺度：一种是描述局部运动和接触变化的短期物理未来，另一种是描述任务下一阶段应达到状态的阶段语义未来。

StageWAM的核心设计是基于V-JEPA2构建Stage-JEPA模块，用于预测下一任务阶段的隐层表征，再将该表征注入原有WAM模型，共同指导未来视频和动作的生成。具体来说，模型首先通过Stage-JEPA从当前观测和任务指令中提取下一阶段的隐层信息，再将这些信息作为内部条件加入到WAM中，让JEPA负责提供任务进度方向，而WAM继续负责擅长的局部视频和动作建模。

在训练阶段，团队没有预先定义“抓取”“移动”“放置”等任务阶段，而是利用冻结的V-JEPA2表征，从演示轨迹中自动寻找状态变化明显的位置。对于连续轨迹，模型以固定间隔采样候选位置，截取64帧视频片段通过冻结编码器得到隐层表征，比较候选位置前后局部窗口的表征变化，综合得到过渡分数，经过平滑归一化和时间非极大值抑制后，保留变化最明显的关键帧，与轨迹首尾共同组成阶段边界。

基于这些阶段配对数据，Stage-JEPA分别编码当前观测片段和对应的下一阶段观测片段，预测模块接收当前观测表征和语言指令，预测下一阶段的隐层表征，并通过余弦损失与真实的下一阶段V-JEPA2表征对齐，模型不会直接生成未来图像或预测动作，而是学习从当前状态和任务目标出发，下一阶段应该达成的状态。

在具体的WAM模块设计中，StageWAM基于Motus构建短程未来建模的世界动作模型，包含视觉演化、动作生成和视觉语言理解三个分支，通过三模态联合注意力持续交换信息。Stage-JEPA输出的下一阶段隐层标记首先经过池化和MLP投影到视觉token特征空间，再通过门控残差的方式注入到视觉token中，门控参数用于控制阶段信息对原始WAM表征的影响，注入后的视觉token会进入后续的三模态联合注意力，因此阶段隐层不仅影响未来视频预测，还会通过分支间的信息交互进一步影响动作生成。训练时，已经训练好的Stage-JEPA会被冻结，仅优化WAM和对应的条件接口，形成明确的分工：Stage-JEPA负责预测“下一阶段应该发生什么”，WAM负责预测“为了让该阶段发生，短期世界应该如何变化、机器人现在应该如何动作”。

该研究在RoboTwin 2.0的50个双臂操作任务上进行了评测，在标准和随机两种设置下，StageWAM分别取得了91.42%和89.08%的成功率，整体成功率达到90.25%，相比基准模型Motus均有提升。同时，研究团队还观察到，在成功完成任务的场景中，StageWAM所需的动作步数更少，Motus的平均执行长度为79.57步，而StageWAM仅为74.82步，减少了5.97%。

消融实验验证了模型各模块的有效性，完全移除Stage-JEPA后，整体成功率下降到87.23%，证明真正发挥作用的是预测出的下一阶段信息参与了后续的世界动作生成，而非仅仅增加了额外的视觉编码器。

总的来说，现有WAM擅长短时间范围内的物理演化，能够支持运动、接触和动作执行的建模，但对于由多个阶段组成的复杂操作任务，仅依靠局部未来无法明确任务下一步的推进方向。StageWAM重新区分了机器人操作中的两种未来尺度，用短期物理未来描述局部环境变化，用阶段语义未来明确任务整体进度，通过将阶段隐层表征作为条件接入WAM，让模型无需用过长的视频同时承担局部物理预测和长程任务进度建模，实现了更高效的机器人操作任务完成。

论文相关信息：

• 论文题目：StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation
• 合作机构：清华大学智能产业研究院（AIR）、北京航空航天大学、清华大学AIR无锡创新中心、中国人民大学、它石智航
• 论文链接：https://arxiv.org/abs/2608.10780