年份 2024

作者Gaoyue Zhou1,∗Hengkai Pan1, Yann LeCun1,2, Lerrel Pinto1

1Courant Institute, New York University, 2Meta-FAIR

摘要

根据控制动作预测未来结果的能力是物理推理的基础。然而，这种预测模型（通常称为世界模型）已被证明难以学习，并且通常是为具有在线策略学习的特定任务解决方案而开发的。我们认为，世界模型的真正潜力在于它们仅使用被动数据来推理和规划各种问题的能力。具体来说，我们要求世界模型具有以下三个属性：1） 可在离线、预先收集的轨迹上进行训练，2） 支持测试时行为优化，以及 3） 促进与任务无关的推理。为了实现这一目标，我们提出了 DINO 世界模型 （DINO-WM），这是一种无需重建视觉世界即可对视觉动态进行建模的新方法。DINO-WM 利用 DINOv2 预先训练的空间补丁特征，使其能够通过预测未来的补丁特征来从离线行为轨迹中学习。这种设计允许 DINO-WM 通过动作序列优化来实现观察目标，通过将所需的目标补丁特征视为预测目标来促进与任务无关的行为规划。我们在各个领域评估了 DINO-WM，包括迷宫导航、桌面推送和粒子操纵。我们的实验表明，DINO-WM 可以在测试时生成零镜头行为解决方案，而无需依赖专家演示、奖励建模或预先学习的逆向模型。值得注意的是，与之前最先进的工作相比，DINOWM 表现出强大的泛化能力，适应不同的任务系列，例如任意配置的迷宫、具有不同对象形状的推送操作和多粒子场景。

学习的内容

论文背景

• 模仿学习和强化学习方法主要是基于策略，泛化性太差
• 世界模型效果好，但是在解决通用任务方面仍然存在重大挑战。
• 世界建模的挑战，online的世界模型需要对每个新任务重新训练，offline的世界模型可以在环境收集的轨迹离线数据集上训练，当需要解决任务时，需要用到很强的辅助信息（可以有采用专家演示的形式 Pathak et al. （2018），结构化关键点 Ko et al. （2023）;温 et al. （2024）， 访问预训练逆模型 Du et al. （2023）;Ko et al. （2023） 或密集奖励函数，所有这些都降低了使用离线世界模型的普遍性。）构建更好的离线世界模型的核心问题是，是否有不影响其通用性的替代辅助信息？

我们的目标是从预先收集的离线数据集中学习任务无关的世界模型，并使用这些世界模型在测试时进行视觉推理和控制。在测试时，我们的系统从任意环境状态开始，并以 RGB 图像的形式提供目标观察，并被要求执行一系列动作 a0， ...， aT，以便实现目标状态。

这种方法不同于在线强化学习 （RL） 中使用的世界模型，在线强化学习 （RL） 的目标是优化手头一组固定任务的奖励（Dreamer）。

方法

作者提出了DINO-WM，这是一种从离线轨迹数据集构建与任务无关的世界模型的新方法。

1、对于嵌入，我们使用了 DINOv2 模型中的预训练补丁特征，该模型提供了空间和以对象为中心的先验表示。

3、鉴于这些视觉嵌入和操作，DINO-WM 使用 ViT 架构来预测未来的嵌入。

2、一旦这个模型被训练出来，解决任务的计划就被构建为视觉目标的达成，即根据当前的观察达到未来的预期目标。用规划的方法进行决策。

DINO-BASED WORLD MODELS (DINO-WM)

Observation Model

为了学习通用世界模型要求：1） 独立于任务和环境，以及 2） 包含丰富的空间信息，这在导航和操作任务中至关重要。

我们选择开箱即用的预训练 DINOv2 模型作为我们世界模型的观察模型，因为它已被证明在需要大量空间理解的对象检测、语义分割和深度估计任务方面表现出色。

DINO-WM 的架构。给定观察值 ot−k：t，我们优化 at：T −1 的动作顺序，以最小化对预期目标 og 的预测损失。所有前向计算都是在潜在空间 z 中完成的。这里 pθ 表示 DINO-WM 的动力学模型，用于进行未来预测。

Transition Model

• 我们采用 ViT架构作为过渡模型，因为它是处理补丁特征的自然选择。但是，需要对架构进行一些修改，以允许对本体感觉和控制器操作进行额外的调节。
• 基于前H步预测下一步,H是超参数
• 我们在 ViT 模型中实现了一个因果注意力机制，使模型能够在帧级别自回归地预测潜在因素。具体而言patch vector
• 我们认为，在帧级别进行预测并将一个观察的补丁向量视为一个整体，可以更好地捕获全局结构和时间动态，对整个观察的依赖关系进行建模，而不是孤立的标记，从而提高时间泛化。
• 为了模拟代理体的动作对环境的影响，作者使用MLP，从原始动作表示映射的K维动作向量连接到中，其中i=1,,N

损失函数：

Decoder for interpretability

为了帮助可视化和可解释性，我们使用一堆转置卷积层将补丁表示解码回图像像素

解码器的训练完全独立于transition model 训练，具有几个优点：1） 解码器的质量不会影响世界模型解决下游任务的推理和规划能力，以及 2） 在规划过程中，无需重建原始像素图像，从而降低计算成本。尽管如此，解码器仍然很有价值，因为它增强了世界模型预测的可解释性。

VISUAL PLANNING WITH DINO-WM

• 采用模型预测控制（MPC），世界模型接收当前观测值和目标观测值，planning的过程表示为agent采取一系列的动作从到达的过程。
• 我们利用随机优化算法交叉熵法 （CEM） 来优化每次迭代的动作顺序。损失函数是当前潜在状态与目标潜在状态之间的均方误差 （MSE）

作者尝试过梯度下降（GD）来优化，但是发现CEM效果更好

CME优化过程

1. 给定当前观测值 o0 和目标观测值 og，都表示为 RGB 图像，观测值首先被编码为潜伏状态：
2. 规划目标定义为最终时间步 T 的预测潜在状态与目标潜在状态之间的均方误差 （MSE）：
3. 在每次计划迭代中，CEM 从分布中抽取 N 个动作序列的总体，每个动作序列的长度为 T 。初始分布设置为 Gaussian。
4. 对于每个采样的动作序列，世界模型用于预测潜在空间中的结果轨迹：,并且成本 C 是针对每个轨迹计算的。
5. 选择成本最低的前 K 个动作序列，并相应地更新分布的均值和协方差。
6. 从更新的分布中抽取一组新的 N 个操作序列，并重复该过程，直到获得成功或经过我们设置为超参数的固定迭代次数后。
7. 优化过程完成后，前 k 个操作 在环境中执行。然后，该过程将在下一个时间步中对新的观测值重复。

GD优化过程

实验

任务环境

观察空间都是大小为 （224， 224） 的 RGB 图像。

• Point Maze
• Push-T
• Wall
• Rope Manipulation
• Granular Manipulation

baselines

• IRIS：IRIS 使用离散自动编码器将视觉输入转换为标记，并使用 GPT Transformer 预测未来观察的标记。它结合了这些组成部分，通过富有想象力的程序来学习政策和价值功能。
• DreamerV3：DreamerV3 学习一个世界模型，将视觉输入解释为分类表示。它根据给定的行动预测未来的表现和奖励，并根据其想象的轨迹训练演员-批评家政策。
• TD-MPC2：TD-MPC2 在潜在空间中学习无解码器的世界模型，并使用奖励信号来优化潜在空间。它是 reconstructionfree 世界建模的强大基线。
• AVDC：AVDC 利用扩散模型，根据初始观察和文本目标描述生成任务执行的想象视频。然后，它会估计帧之间的光流以捕获物体的运动并生成机械臂命令。

使用 DINO-WM 优化行为

预训练的视觉表示是否重要

2、

疑问

1、反向传播需要传播到Encoder吗？好像是不需要的

2、GT也就是o_g是怎么获得的？数据集自带的吗，如果在真实场景中没有o_g怎么办？