本文介绍本组NeurIPS2023世界模型与强化学习方向的最新工作：Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning。
针对世界模型的视频预训练方法难以从现实场景视频正迁移的问题，本文提出情境化世界模型（ContextWM），利用场景和运动解耦的思想，以促进多样化场景之间的知识迁移，并展示了其在多样的下游视觉控制任务中的性能收益。
论文标题：
Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning
论文链接：
https://arxiv.org/abs/2305.18499
代码链接：
https://github.com/thuml/ContextWM
一、引言
世界模型 (World Models) 是智能体通过与环境交互和观察所习得的演变模型，它赋予了智能体思考与推理的能力，有助于其快速学习新技能和有效解决复杂任务，在具身智能、自动驾驶、大语言模型等领域具备巨大潜力。
然而，现有的基于模型的强化学习 (Model-based Reinforcement Learning, MBRL) 算法通常从零开始学习领域特定的世界模型，这阻碍了通用世界知识的获取和迁移，未能充分释放基于世界模型的方法高效利用样本的潜力。
最近的研究探索了世界模型的视频预训练方法，以提升 MBRL 的样本效率，但却局限于领域特定或模拟的视频数据。鉴于无监督预训练在视觉和自然语言处理领域的成功，来自互联网的大规模现实场景视频 (In-the-wild videos) 有望成为获取和迁移通用世界知识的有效来源，但其天然的复杂性也为世界模型学习带来了挑战：
- 如何克服现实场景视频蕴含的视觉复杂性与多样性？
- 现实场景视频与视觉控制任务之间可迁移的共享知识是什么？
针对上述关键问题，我们认为即使在视觉上多样化的场景（contexts）之间，物体的运动（dynamics）仍然具有相似的语义结构。对场景和运动粗略的耦合建模可能会浪费大量的模型容量来处理低级的视觉细节，阻碍世界模型捕捉关于运动的共享知识。
因此，我们提出了情境化世界模型 (Contextualized World Models, ContextWM)，利用场景和运动解耦的思想，以促进多样化场景之间的共享知识获取。通过将情境化世界模型应用到预训练-微调范式中，我们解锁了现实场景视频到视觉控制任务的正迁移，取得了在多种下游任务中的显著性能收益。
二、背景
2.1 世界模型的实现
Dreamer 是 MBRL 中的经典方法，它的世界模型是一个典型的隐变量动力学模型（Latent Dynamics Model），包含以下四个部分：
2.2 世界模型的视频预训练
APV 在利用 RLBench 的视频数据进行预训练时，能够有效地正迁移到 Meta-world 任务，提升 MBRL 的样本效率（下图中的 APV (w/ Manipulation) 曲线）。然而，当 APV 尝试使用 Something-Something 这种现实场景视频的时候，却遇到了严重的欠拟合问题，并且在下游任务上也没有取得性能收益（APV (w/ Real-world) 曲线）。
我们认为，APV 对于现实场景视频的利用不够高效的根本原因在于它浪费了大量的模型容量用于建模难以迁移的场景信息。
三、情境化世界模型（ContextWM）
3.1 情境化隐变量动力学模型
观测序列中通常存在两组可分离的信息，即静态的场景信息（contexts）和动态的运动信息（dynamics）。场景信息表示了场景中物体的静态信息，如纹理、形状和颜色，而动态信息则编码了物体在时间上的变迁，如位置和运动方向。
3.2 情境化世界模型的实现
我们接下来提供一种面向视觉控制任务的情境化隐变量动力学模型的具体实现，称为情境化世界模型（Contextualized World Models, ContextWM）。ContextWM 基于 Dreamer 和 APV，通过引入一个平行的场景编码器，来增强图像解码器的场景重建能力，并鼓励原有的编码器专注于建模动态信息。
(1) 场景表达
我们假设场景信息均匀地分布在每一帧中，通过随机选择其中一帧，场景编码器应该学会与时间无关的特征。
(2) 多尺度交叉注意力
(3) 双重奖励预测
目标函数
四、实验
我们在多种视频数据集以及视觉控制任务上进行了实验，包括：
- 预训练数据集：Something-Something V2 物体交互视频、Human3.6M 人体运动视频、YouTubeDriving 现实驾驶视频
- 视觉控制任务：Meta-world 机械臂操作任务、DMC Remastered 运动控制的视觉泛化、CARLA 自动驾驶任务
4.1 Meta-world 实验结果
在 Meta-world 的六个任务上，Something-Something V2 预训练的 ContextWM 取得了样本效率上显著的正迁移，而基线方法 APV（图中的 Plain WM）却没有取得明显收益。
4.2 DMC Remastered 实验结果
在 DMC Remastered 视觉泛化任务中，从现实场景视频（Something-Something V2）进行预训练可以显著提高性能，而采用了场景与运动解耦的 ContextWM 则进一步释放了这一潜力。
4.3 不同预训练数据的影响
对于不同预训练数据集对下游任务的影响，我们的实验初步发现了以下结论：
- 使用物体交互视频（Something-Something V2）预训练普遍有所收益。
- 更接近下游任务的领域数据（例如 RLBench）收益更高，但更加多样化的数据集具有作为可扩展的替代方案的潜力。
- 预训练数据缺乏多样性（例如 Human3.6M）甚至可能具有负面影响。
4.4 可视化分析
(1) 视频表示
即使预训练的时候完全没有利用视频的标签信息，我们发现 ContextWM 学到的视频特征能够较清晰地区分不同标签代表的运动方向。
(2) 视频预测
对比 ContextWM 和基线方法的视频预测结果，我们发现 ContextWM 的预测很准确地刻画了水杯的形状和运动。并且，来自不同帧的交叉注意力成功地关注到了场景帧中的不同空间位置。
(3) 组合解码分析
通过从另一条轨迹中采样一帧来替代原始的场景信息，同时保持原轨迹中的运动信息不变，我们发现 ContextWM 正确地将新的场景与原始的运动信息重新组合。这表明我们的模型成功地学习了场景和运动的解耦表示。
五、总结
本文提出了情境化世界模型（ContextWM），并将其应用于现实场景视频预训练（IPV）范式，然后在下游任务上进行微调，以提高基于模型的强化学习（MBRL）的样本效率。实验证明了我们的方法在多种视觉控制任务上的有效性。我们的工作不仅突出了利用丰富的现实场景视频数据的潜力，还强调了利于知识迁移和可扩展学习的世界模型结构设计的重要性。
Illustration From IconScout By Delesign Graphics
快抢沙发！
