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
分类: cs.RO, cs.AI, cs.LG
发布日期: 2025-04-03 (更新: 2025-05-23)
🔗 代码/项目: PROJECT_PAGE
💡 一句话要点
提出统一世界模型(UWM)，通过视频和动作扩散预训练提升机器人泛化能力。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 机器人学习 世界模型 扩散模型 模仿学习 视频预测 动作预测 Transformer 预训练
📋 核心要点
- 模仿学习依赖高质量专家数据，难以扩展到大型机器人模型，而大量无动作标注的视频数据蕴含丰富信息但难以利用。
- UWM通过统一的Transformer架构整合动作和视频扩散过程，独立控制模态扩散时间步，灵活表示策略、动力学和视频生成。
- 实验表明，UWM预训练能有效提升策略的泛化性和鲁棒性，并能从无动作视频中学习，优于传统模仿学习。
📝 摘要（中文）
模仿学习是构建通用机器人的一种有前景的方法。然而，由于其依赖于高质量的专家演示，扩展模仿学习以适应大型机器人基础模型仍然具有挑战性。同时，大量描绘各种环境和不同行为的视频数据很容易获得。这些数据提供了关于真实世界动态和智能体-环境交互的丰富信息。然而，由于缺乏动作标注，直接利用这些数据进行模仿学习已被证明是困难的。在这项工作中，我们提出了统一世界模型（UWM），这是一个允许利用视频和动作数据进行策略学习的框架。具体来说，UWM在一个统一的Transformer架构中集成了动作扩散过程和视频扩散过程，其中独立的扩散时间步长控制每个模态。通过控制每个扩散时间步长，UWM可以灵活地表示策略、前向动力学、逆向动力学和视频生成器。通过模拟和真实世界的实验，我们表明：（1）UWM能够通过动力学和动作预测对大规模多任务机器人数据集进行有效的预训练，从而产生比模仿学习更具泛化性和鲁棒性的策略，（2）UWM通过独立控制特定模态的扩散时间步长，自然地促进了从无动作视频数据中学习，进一步提高了微调策略的性能。我们的结果表明，UWM为利用大型异构数据集进行可扩展的机器人学习提供了一个有希望的步骤，并提供了模仿学习和世界建模这两种通常不同的范例之间的简单统一。
🔬 方法详解
问题定义：论文旨在解决机器人学习中，如何有效利用大规模无动作标注的视频数据，克服模仿学习对高质量专家数据的依赖，从而提升机器人策略的泛化性和鲁棒性的问题。现有方法要么依赖于昂贵的专家数据，要么难以直接从无动作视频中学习到有效的策略。
核心思路：论文的核心思路是将动作扩散过程和视频扩散过程集成到一个统一的Transformer架构中，通过独立控制每个模态的扩散时间步长，使得模型能够灵活地表示策略、前向动力学、逆向动力学和视频生成器。这种统一的框架允许模型同时从有动作标注的数据和无动作标注的视频数据中学习。
技术框架：UWM的技术框架包含以下几个主要模块：1) 一个统一的Transformer架构，用于处理视频和动作数据；2) 一个动作扩散过程，用于学习策略；3) 一个视频扩散过程，用于学习环境动态；4) 独立的扩散时间步长控制机制，用于灵活地控制每个模态的学习过程。整体流程是，首先使用大规模的机器人数据集（包括有动作标注和无动作标注的数据）对UWM进行预训练，然后使用少量特定任务的数据对UWM进行微调，最终得到一个能够执行特定任务的机器人策略。
关键创新：UWM最重要的技术创新点在于它将动作扩散和视频扩散集成到一个统一的框架中，并允许独立控制每个模态的扩散时间步长。这种设计使得模型能够同时从有动作标注的数据和无动作标注的视频数据中学习，从而克服了传统模仿学习对高质量专家数据的依赖。此外，UWM还能够灵活地表示策略、前向动力学、逆向动力学和视频生成器，从而为机器人学习提供了更多的可能性。
关键设计：UWM的关键设计包括：1) 使用Transformer架构作为统一的表示学习器；2) 使用扩散模型来建模动作和视频的生成过程；3) 设计独立的扩散时间步长控制机制，允许模型灵活地控制每个模态的学习过程；4) 使用合适的损失函数来训练模型，例如，使用动作预测损失来训练动作扩散过程，使用视频预测损失来训练视频扩散过程。
🖼️ 关键图片
📊 实验亮点
实验结果表明，UWM在模拟和真实机器人任务中均优于传统模仿学习方法。通过在包含动力学和动作预测的大规模多任务机器人数据集上进行预训练，UWM获得了更具泛化性和鲁棒性的策略。此外，UWM能够从无动作视频数据中学习，进一步提高了微调策略的性能。
🎯 应用场景
UWM可应用于各种机器人任务，如家庭服务机器人、工业机器人和自动驾驶。通过利用大量无标注视频数据，降低了机器人学习对高质量专家数据的依赖，加速了机器人技术的普及。未来，UWM有望成为机器人通用基础模型，赋能更智能、更灵活的机器人。
📄 摘要（原文）
Imitation learning has emerged as a promising approach towards building generalist robots. However, scaling imitation learning for large robot foundation models remains challenging due to its reliance on high-quality expert demonstrations. Meanwhile, large amounts of video data depicting a wide range of environments and diverse behaviors are readily available. This data provides a rich source of information about real-world dynamics and agent-environment interactions. Leveraging this data directly for imitation learning, however, has proven difficult due to the lack of action annotation. In this work, we present Unified World Models (UWM), a framework that allows for leveraging both video and action data for policy learning. Specifically, a UWM integrates an action diffusion process and a video diffusion process within a unified transformer architecture, where independent diffusion timesteps govern each modality. By controlling each diffusion timestep, UWM can flexibly represent a policy, a forward dynamics, an inverse dynamics, and a video generator. Through simulated and real-world experiments, we show that: (1) UWM enables effective pretraining on large-scale multitask robot datasets with both dynamics and action predictions, resulting in more generalizable and robust policies than imitation learning, (2) UWM naturally facilitates learning from action-free video data through independent control of modality-specific diffusion timesteps, further improving the performance of finetuned policies. Our results suggest that UWM offers a promising step toward harnessing large, heterogeneous datasets for scalable robot learning, and provides a simple unification between the often disparate paradigms of imitation learning and world modeling. Videos and code are available at https://weirdlabuw.github.io/uwm/.