4D-WAM: Infusing Spatiotemporal Awareness into World Action Models through Trajectory Fields
Lishan Yang, Wenxuan Song, Xi Wang, Pingyue Sheng, Zheng Fang, Ziyang Zhou, Junjie He, Haodong Yan, Jiayi Chen, Nan Sun, Qiao Sun, Pengwei Wang, Lingqiao Liu, Yan Wang, Yuxiang Gao, Feras Dayoub, Haoang Li
cs.RO
2026-08-08
训练期把 Trace Anything 的 3D 轨迹场对齐进 WAM 的中间特征,LIBERO-Plus 七种扰动均值从 62.21% 提到 71.01%,推理架构不变。
世界动作模型(World Action Model,WAM)把「预测下一秒画面会怎么变」和「机械臂该怎么动」放在一个视频生成骨架里联合训练。视觉骨干看的是 2D 像素,机械臂却活在 3D 空间,中间隔着一层表示鸿沟:抓杯子要知道杯子在哪个方位、多高、离夹爪多远,这些在像素里都是间接的。
近一拨工作往模型里塞 3D 信号:输入端加深度图、点云,或把深度图当辅助重建目标让模型自己学。这篇指出两条都不彻底:深度图是逐帧独立的 3D 快照,3D 结构怎么随时间演化没有显式表示,而操作的「动态」恰恰在这上面;监督重建像素级深度,注意力又会被低层次细节吸走,学不到高层 4D 理解(3D 结构+时间演化)。
训练期借一个 4D 教师模型做表示对齐,推理零开销。教师是 Trace Anything,把视频里每个像素表示成一条随时间变化的参数化 3D 轨迹,同时编码空间结构和时间演化。
此前在 VLA 模型上有效的做法,是把学生模型的中间特征直接对齐到 3D 基础模型特征。搬到 WAM 上稳定变差,论文给了诊断:π0.5 的中间特征与 VGGT 的 3D 特征在各层稳定正相似,而 WAM(如 Lingbot-VA)的中间特征与 4D 特征在多数层不相关甚至负相关,硬对齐带来优化冲突。但把两边各自的相邻帧特征做差再比相似度,相似度明显升高:绝对特征对不上,时间变化对得上。
两个损失都加在 DiT 骨干第 20 层中间表示上:
总损失是原 WAM 损失(视频隐空间去噪+动作预测)加两项对齐项,权重建议都低于 0.1,不动原目标。训练完 Trace Anything、投影器、对齐损失全拆掉,推理架构和参数量与原模型相同。
分布内基准以 FastWAM-Joint 为基座:
| 设置 | FastWAM-Joint | 4D-WAM |
| RoboTwin 2.0 Clean(%) | 90.84 | 91.92 |
| RoboTwin 2.0 Random(%) | 90.32 | 90.76 |
| LIBERO 四套件均值(%) | 98.2 | 98.6 |
LIBERO 接近饱和,这 0.4 个点是渐进改进,真正的差距在分布外。
LIBERO-Plus 七种扰动(相机、机器人、语言、光照、背景、噪声、布局)下,均值从 62.21% 提到 71.01%;相机扰动单项从 27.89% 拉到 45.15%,17.26 个点。相机扰动改的是观测几何,打在空间理解上,增益分布跟设计对得上。
RoboTwin Clean2Rand(干净数据训练、域随机化评测)以 Lingbot-VA 为基座:随机化设置从 34.6% 提到 41.8%,是对比模型里唯一两种设置均值都过 60% 的;同场 FastWAM 掉到 1.2%。
视频预测(LIBERO-Plus,100 回合):PSNR 从 19.22 到 21.44,SSIM 从 0.75 到 0.83。消融上去掉任一对齐项,LIBERO 均值都落到 98.0%(完整版 98.6%);对齐层 16 到 24 层里第 20 层最好。
真机实验用 ARX LIFT2 双臂平台,四个长程任务、550 条演示、每任务 50 次评测:基座 Lingbot-VA 四个任务三个成功率为 0,4D-WAM 四个都能偶尔做完,总成功率 5.5%。不好看,但三个全零和四个都能走通是两种局面。
附录探针实验给机制证据:冻结骨干、只训线性探针,用第 20 层特征预测 token 在下一帧/末帧的位置。LIBERO 上终点预测 Top-1 从 0.4150 到 0.4487,平均排名从 4.66 降到 4.07;LIBERO-Plus 上趋势保持。
机器人学习里「VLA 模型对 3D 不敏感」是公认短板,主流补法是改输入(加深度/点云通道)或改输出(加几何预测头),都要动架构,推理期多一份开销。这篇走纯训练期知识注入:架构不动,推理零开销,任何 DiT-based WAM 都能照方抓药,对做机器人基座模型的人是可叠在其他改进上的增益,扰动和域偏移场景尤甚。
「绝对特征对不上就对齐差分」不限于机器人,表示空间不兼容的蒸馏场景都用得上;query-key 分布的终点约束,示范了怎么把长程对应关系做成可监督的量。
分布内增益渐进,LIBERO 那 0.4 个点在饱和基准上说明不了太多,方法的真实卖点全在 OOD 侧。真机成功率 5.5%,绝对水平低,只能说比基座的全零强,离能用还有距离,550 条演示的数据量也可能是瓶颈。对齐层选第 20 层是在这个骨干上搜的,换骨干要不要重搜,论文没给规律。Trace Anything 自身的误差会直接传进监督信号,它在训练分布外的轨迹质量如何,没有讨论。Clean2Rand 实验里 4D-WAM 换了基座(FastWAM 到 Lingbot-VA),理由是 FastWAM 在该设置下崩得厉害;选择合理,但两个主打结果用不同基座,横向读数字要留意。