Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets
作者: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025
许可协议
- 开放获取
- 论文: CC BY 4.0
以下内容均依据论文全文，按其 CC BY 4.0 许可协议使用；其中的图已注明作者。
一句话总结
单个扩散 transformer 为动作和未来图像分别设置独立的噪声时间步，可充当策略、动力学模型或视频预测器，并能在无动作视频上训练。
为何重要
一个模型既能在带动作标注的机器人数据上训练，也能在获取容易得多的无动作视频上训练。在采用 DROID 风格 Franka 装置的五项真实机器人任务测试中，预训练的 UWM 在每项任务上的成功率都最高，而在额外视频上联合训练使其进一步提升，在视觉干扰下也是如此。
问题
大型机器人模型的模仿学习依赖高质量的专家示范，而通过遥操作收集这些示范既昂贵又缓慢。展示真实世界动态的视频非常丰富，但当前大多数模仿学习方法无法使用它们，因为这些视频没有动作标签；而且行为克隆并不会显式地捕捉轨迹中存在的时间动态。
提出的方法
UWM 在一个以当前观测为条件的 transformer 中耦合了动作扩散过程和未来图像扩散过程，每种模态各有独立的扩散时间步。把某个模态的时间步设为 T 即将其边缘化，设为 0 即以其为条件，因此同一个模型可以采样策略 p(a|o)、正向动力学 p(o′|o,a)、逆向动力学 p(a|o,o′) 或视频预测 p(o′|o)。训练中使用无动作视频的方式是：将动作时间步固定为 T，并用高斯噪声填充缺失的动作。
关键见解
- 每种模态独立的扩散时间步起到部分掩码的作用，因此单个网络就能覆盖动作和未来帧上的边缘分布与条件分布。
- 带来提升的是预测未来帧，而不只是重建当前帧：在 Stack-Bowls/Block-Cabinet 上成功率为 0.86/0.76，而重建当前帧为 0.70/0.66，Diffusion Policy 为 0.48/0.60。
- 寄存器 token（默认 8 个）有助于在动作 token 与图像 token 之间交换信息，并且在 LIBERO 消融实验中，AdaLN 条件调制优于交叉注意力。
- 从零开始训练时，UWM 与 Diffusion Policy 表现大致相同；UWM 从预训练中获益更多。
- 在域内机器人视频上联合训练比在互联网人类活动视频（Kinetics-400 和 Something-Something v2）上联合训练帮助更大，不过两者都优于只用机器人数据。
主要结果
- 真实机器人，在 2,000 条 DROID 轨迹上预训练，分布内成功率：UWM 在 Stack-Bowls、Block-Cabinet、Paper-Towel、Hang-Towel 和 Rice-Cooker 上分别为 0.86、0.76、0.78、0.82 和 0.60，Diffusion Policy 为 0.48、0.60、0.52、0.64、0.35，GR1 为 0.66、0.66、0.60、0.66、0.40。
- 在另外 2,000 条作为无动作视频使用的 DROID 轨迹上联合训练后，UWM 的分布内成功率提高到 0.92、0.84、0.86、0.86 和 0.65，分布外成功率（前四项任务）从 0.76/0.60/0.78/0.64 提高到 0.84/0.72/0.84/0.76。
- 分类的 OOD 场景（光照、背景、杂物）：联合训练的 UWM 在 Stack-Bowls 和 Block-Cabinet 上为 21/30 和 15/30，预训练的 UWM 为 15/30 和 8/30，Diffusion Policy 为 12/30 和 6/30。
- LIBERO（在 LIBERO-90 上预训练，在五项存在分布偏移的 LIBERO-10 任务上微调）：UWM 平均成功率为 0.79 ± 0.11，DP 为 0.71 ± 0.12，PAD 为 0.57 ± 0.19，GR1 为 0.58 ± 0.14。
- 逆向动力学模式在轨迹长度内跟踪 50 条 LIBERO 轨迹：在 Book-Caddy 和 Soup-Cheese 上成功率为 0.73 和 0.45，策略模式为 0.37 和 0.33。
- 在 DROID 上训练 UWM 100K 步，需要在 4 张 NVIDIA A100 GPU 上耗时 24 小时。
局限
- 该模型尚未从大规模人类视频中学习，因此人与机器人之间的本体差异问题没有得到解决。
- 正向动力学预测可能包含伪影，这可能降低其在规划中的用处。
- 作者预期更密集的视频预测会带来进一步提升，而当前版本没有使用。
- 在 LIBERO 分布外设置中相对基线的提升远小于真实机器人上的提升；作者认为仿真动态更为简单。
- Rice-Cooker 只在 20 个接近分布内的初始化上评估，没有在 OOD 设置下评估，而且所有真实机器人测试都使用同一套 Franka 装置。
对领域的影响
该论文被引用 244 次。论文称视频和代码可在其项目页面获取，该页面链接了 GitHub 上的代码仓库 WEIRDLabUW/unified-world-model。
摘要
模仿学习已成为构建通用机器人的一种有前景的途径。然而，由于依赖高质量的专家示范，为大型机器人基础模型扩展模仿学习仍然充满挑战。与此同时，描绘广泛环境和多样行为的大量视频数据唾手可得。这些数据提供了关于真实世界动态以及智能体与环境交互的丰富信息来源。然而，由于缺少当今大多数方法所需的动作标注，直接将这些数据用于模仿学习已被证明是困难的。在本工作中，我们提出统一世界模型（Unified World Models，UWM），这是一个能够同时利用视频数据和动作数据进行策略学习的框架。具体而言，UWM 在一个统一的 transformer 架构中整合了动作扩散过程和视频扩散过程，每种模态由各自独立的扩散时间步控制。我们表明，只需控制各个扩散时间步，UWM 就能灵活地表示策略、正向动力学、逆向动力学和视频生成器。通过仿真和真实世界实验，我们表明：(1) UWM 能够在大规模多任务机器人数据集上同时以动力学预测和动作预测进行有效的预训练，得到比模仿学习更具泛化性和鲁棒性的策略；(2) 通过独立控制各模态的扩散时间步，UWM 自然地支持从无动作视频数据中学习，进一步提升微调后策略的性能。我们的结果表明，UWM 朝着利用大型异构数据集实现可扩展机器人学习迈出了有前景的一步，并在通常彼此割裂的模仿学习与世界建模两种范式之间提供了一种简单的统一。视频和代码见 https://weirdlabuw.github.io/uwm/
链接
- arXiv 2504.02792: https://arxiv.org/abs/2504.02792
- DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
- PDF: https://arxiv.org/pdf/2504.02792v1
- 代码: https://github.com/WEIRDLabUW/unified-world-model
- 项目主页: https://weirdlabuw.github.io/uwm/