DINO-WM: World Models on Pre-trained Visual Featuresenable Zero-shot Planning-CSDN博客

本文链接：https://blog.csdn.net/qq_33673253/article/details/147069382

DINO-WM：基于预训练视觉特征的世界模型实现零样本规划

Gaoyue Zhou 1 Hengkai Pan 1 Yann LeCun 1 2 Lerrel Pinto 1

摘要能够预测给定控制动作的未来结果对于物理推理至关重要。然而，这种预测模型（通常称为世界模型）一直很难学习，并且通常用于特定任务的在线策略学习解决方案。为了释放世界模型的真正潜力，我们认为它们应该 1）能够在离线预收集的轨迹上进行训练，2）支持测试时的行为优化，3）促进与任务无关的推理。为此，我们提出了DINO世界模型（DINO-WM），这是一种新的方法，用于在不重建视觉世界的情况下对视觉动态进行建模。DINO-WM利用使用DINOv2预训练的空间补丁特征，使其能够通过预测未来补丁特征从离线行为轨迹中学习。这使得DINO-WM能够通过动作序列优化实现观测目标，通过将目标特征作为预测目标来促进与任务无关的规划。我们证明DINO-WM在六个环境中实现了零样本行为解决方案，这些环境包括任意配置的迷宫、不同物体形状的推操纵和多粒子场景等多样化任务家族，无需专家演示、奖励建模或预训练的逆模型，超越了以前的最新技术。

1. 引言

近年来，机器人技术和具身人工智能取得了巨大进展。模仿学习和强化学习的进步使智能体能够在各种任务中学习复杂的行为（Agarwal等，2022；Zhao等，2023；Lee等，2024；Ma等，2024；Hafner等，2024；Hansen等，2024；Haldar等，2024；Jia等，2024）。尽管取得了这些进展，但泛化仍然是一个主要挑战（Zhou等，2023）。现有方法主要依赖于一旦训练完成，在部署期间以前馈方式运行的策略——将观测映射到动作，而无需进一步的优化或推理。在这种框架下，成功泛化本质上要求智能体在训练完成后具备所有可能任务和场景的解决方案，这只有在智能体在训练期间见过类似场景时才有可能（Reed等，2022；Brohan等，2023b；Brohan等，2023a；Etukuru等，2024）。然而，提前为所有潜在任务和环境学习解决方案既不现实，也不高效。

与其在训练期间学习所有可能任务的解决方案，另一种方法是在训练数据上拟合动态模型，并在运行时优化特定任务的行为。这些动态模型，也称为世界模型（Ha & Schmidhuber，2018），在机器人学和控制领域有着悠久的历史（Sutton，1991；Todorov & Li，2005；Williams等，2017）。最近，几项工作表明世界模型可以基于原始感官数据进行训练（Hafner等，2019；Micheli等，2023；Robine等，2023；Hansen等，2024；Hafner等，2024）。这使得能够灵活地使用基于模型的优化来获得策略，因为它避免了显式状态估计的需要。尽管如此，使用世界模型解决通用任务仍面临重大挑战。

为了理解世界建模中的挑战，让我们考虑学习世界模型的两种广泛范式：在线和离线。在在线设置中，通常需要访问环境以便持续收集数据以改进世界模型，这反过来又改进了策略和后续数据收集。然而，在线世界模型仅在所优化策略的覆盖范围内是准确的。因此，虽然它可以用于训练强大的特定任务策略，但每个新任务即使在相同环境中也需要重新训练。相反，在离线设置中，世界模型是在环境中收集的离线轨迹数据集上进行训练的，这消除了其对特定任务的依赖，前提是数据集提供了足够的覆盖。然而，当需要解决任务时，该领域的方需要强大的辅助信息，这些信息可以是专家演示（Pathak等，2018；Wang等，2023）、结构化关键点（Ko等，2023；Wen等，2024）、预训练逆模型的访问（Du等，2023；Ko等，2023）或密集奖励函数（Ding等，2024），所有这些都降低了使用离线世界模型的通用性。构建更好的离线世界模型的核心问题是是否存在其他辅助信息，而不会影响其通用性。

在本文中，我们提出了DINO-WM，这是一种新的简单方法，用于从离线轨迹数据集中构建与任务无关的世界模型。DINO-WM在紧凑的嵌入空间中对世界动态进行建模，而不是对原始观测本身进行建模。对于嵌入，我们使用DINOv2预训练的补丁特征，它提供了空间和以对象为中心的表示先验。我们假设这种预训练表示能够实现稳健和一致的世界建模，从而减少了对特定任务数据覆盖的需求。给定这些视觉嵌入和动作，DINO-WM使用ViT架构来预测未来的嵌入。一旦该模型在离线数据集上训练完成，规划以解决任务就被构建为视觉目标到达，即给定当前观测，到达未来的期望目标。由于DINO-WM的预测质量很高（见图4），我们只需使用模型预测控制和推理时间优化来达到期望的目标，而无需在测试期间提供任何额外信息。

DINO-WM在六个环境套件上进行了实验评估，这些环境包括迷宫导航、滑动操纵、机械臂控制和可变形物体操纵任务。我们的实验揭示了以下发现：

DINO-WM产生了高质量的未来世界建模，这可以通过从训练有素的解码器中改进的视觉重建来衡量。在我们最困难的任务的LPIPS指标上，这比以前的最新技术工作提高了56%（见第4.7节）。
给定使用DINO-WM训练的潜在世界模型，我们展示了在我们最困难的任务中实现任意目标的高成功率，平均比以前的工作提高了45%（见第4.3节）。
DINO-WM可以在任务家族内的环境变化中进行训练（例如，导航的不同迷宫布局或操纵的不同物体形状），并与以前的工作相比取得了更高的成功率（见第4.5节）。

DINO-WM的代码和模型已开源，以确保可重复性，并在我们的匿名项目网站上提供规划视频：https://dino-wm.github.io。

图1。我们提出了Dino-WM，这是一种使用图像框架retrained DINOv2 embeddings 训练视觉模型的方法。经过训练后，给定目标观察，我们可以通过使用模型预测控制（B）通过Dino-WM进行计划直接优化代理行为。pretrained 嵌入可显着提高先前最新世界模型（C）的性能。

2. 相关工作

我们的工作建立在开发世界模型、从模型中优化行为以及利用紧凑视觉表示的几项工作之上。为了简洁起见，我们只讨论与DINO-WM最相关的部分。

基于模型的学习：

从动态模型中学习有着丰富的文献，涉及控制、规划和机器人学领域（Sutton，1991；Todorov & Li，2005；Astolfi等，2008；Holkar & Waghmare，2010；Williams等，2017）。最近的工作表明，建模动态和预测未来状态可以显著增强具身智能体的基于视觉的学习，涵盖各种应用，包括在线强化学习（Micheli等，2023；Robine等，2023；Hansen等，2024；Hafner等，2024）、探索（Sekar等，2020；Mendonca等，2021；2023a）、规划（Watter等，2015；Finn & Levine，2017；Ebert等，2018；Hafner等，2019）和模仿学习（Pathak等，2018）。这些方法最初集中在状态空间动态（Deisenroth & Rasmussen，2011；Lenz等，2015；Chua等，2018；Nagabandi等，2019），并且已经扩展到处理基于图像的输入，这在本文中得到解决。这些世界模型可以预测像素空间中的未来状态（Finn & Levine，2017；Ebert等，2018；Ko等，2023；Du等，2023）或潜在表示空间（Yan等，2021）。在像素空间中进行预测由于需要图像重建和使用扩散模型的开销而计算成本高昂（Ko等，2023）。另一方面，潜在空间预测通常与重建图像的目标相关联（Hafner等，2019；Micheli等，2023；Hafner等，2024），这引发了对学习特征是否包含足够任务信息的担忧。此外，这些模型中的许多都纳入了奖励预测（Micheli等，2023；Robine等，2023；Hafner等，2024），或使用奖励预测作为辅助目标来学习潜在表示（Hansen等，2022；2024），这使得世界模型具有任务特定性。在本文中，我们旨在将任务依赖信息与潜在空间预测解耦，努力开发一种多功能且与任务无关的世界模型，能够泛化到不同场景。

作为世界模型的生成模型：

随着大规模基础模型的兴起，已经有 initiatives 在自动驾驶（Hu等，2023）、控制（Yang等，2023；Bruce等，2024）和通用视频生成（Liu等，2024）领域构建基于代理动作条件的大型视频生成世界模型。这些模型旨在在文本或高级动作序列的条件下生成视频预测。尽管这些模型在数据增强等下游任务中显示出实用性，但它们对语言条件的依赖限制了其在需要精确视觉目标的应用中的应用。此外，使用扩散模型进行视频生成使其计算成本高昂，进一步限制了其在测试时优化技术（如MPC）中的应用。在本文中，我们旨在构建一个潜在空间中的世界模型，而不是原始像素空间，以实现更精确的规划和控制。

预训练视觉表示：

在视觉表示学习领域取得了重大进展，其中紧凑的特征可以捕获空间和语义信息，以便直接用于下游任务。像ImageNet预训练的ResNet（He等，2016）、I-JEPA（Assran等，2023）、DINO（Caron等，2021；Oquab等，2024）等预训练模型，以及用于视频的V-JEPA（Bardes等，2024），以及用于机器人的R3M（Nair等，2022）、MVP（Xiao等，2022）等，允许快速适应下游任务，因为它们包含丰富的空间和语义信息。虽然这些模型中的许多使用单一全局特征来表示图像，但引入Vision Transformers（ViTs）（Dosovitskiy等，2021）使得预训练补丁特征的使用成为可能，如DINO（Caron等，2021；Oquab等，2024）所示。DINO采用自蒸馏损失，使模型能够有效地学习表示，捕获语义布局并提高图像内的空间理解。在本文中，我们利用DINOv2的补丁嵌入来训练我们的世界模型，并证明它作为一个多功能编码器，能够处理各种精确任务。

3. DINO世界模型

概述和问题表述：我们的工作遵循基于视觉的控制任务框架，将环境建模为部分可观测马尔可夫决策过程（POMDP）。POMDP由元组定义，其中O表示观测空间，A表示动作空间。环境的动态由转移分布建模，该分布根据过去的动作和观测预测未来的观测。

在本文中，我们的目标是从预收集的离线数据集中学习与任务无关的世界模型，并在测试时使用这些世界模型进行视觉推理和控制。在测试时，我们的系统从任意环境状态开始，并以RGB图像的形式提供目标观测，并被要求执行一系列动作a0, ..., aT以达到目标状态。这种方法与在线强化学习中使用的世界模型不同，在线强化学习的目标是针对一组固定任务优化奖励（Hafner等，2024；Hansen等，2024），也与通过文本提示指定目标的文本条件世界模型不同（Du等，2023；Ko等，2023）。

3.1 DINO基础的世界模型（DINO-WM）

我们在潜在空间中建模环境的动态。更具体地说，在每个时间步t，我们的世界模型包括以下组件：

其中观测模型将图像观测编码为潜在状态zt，转移模型接收长度为H的过去潜在状态历史。解码器模型接收潜在zt，并重建图像观测ot。我们用θ表示这些模型的参数。注意，我们解码器的训练目标与世界模型其余部分的训练目标是独立的。这消除了在训练和测试期间重建图像的需求，与Micheli等（2023）和Hafner等dreamerv3（2024）中将观测模型和解码器的训练耦合在一起的方法相比，减少了计算成本。我们在附录A.4.3中对此选择进行了消融实验。

DINO-WM仅建模离线轨迹数据集中从环境中获得的信息，与最近的在线RL世界模型不同，后者还需要任务相关信息，如奖励（Hafner等，2020；Hansen等，2022；2024）、折扣因子（Hafner等，2022；Robine等，2023）和终止条件（Micheli等，2023；Hafner等，2024）。

3.1.1 观测模型

为了在许多环境和现实世界中学习通用世界模型，我们认为观测模型应该

1）与任务和环境无关，

2）捕获丰富的空间信息以用于导航和操纵。

与以往工作不同，以往工作中观测模型总是为手头的任务学习（Hafner等，2024），我们认为从头开始学习良好的观测模型可能是低效的，通常在面对新环境时也不可能，因为感知是一项从大规模互联网数据中受益的一般任务。因此，我们使用预训练的DINOv2模型作为我们世界模型的观测模型，利用其强大的空间理解能力，用于对象检测、语义分割和深度估计等任务（Oquab等，2024）。在训练和测试期间，观测模型保持冻结。在每个时间步t，它将图像ot编码为补丁嵌入，其中N表示补丁 patches 数量，E表示嵌入维度。这一过程在图2中进行了可视化。

图2。Dino-WM的结构。给定观察结果，我们优化了动作序列最小化所需目标 $o_g$ 的预测损失。所有正向计算都是在潜在空间z中完成的。这里的 $p_{\theta}$ 表示Dino-WM的动力学模型，该模型用于做出未来的预测

3.1.2 转移模型

我们采用ViT架构（Dosovitskiy等，2021）作为转移模型，因为它适合处理补丁特征。我们移除了标记化层，因为它在补丁嵌入上操作，有效地将其转换为仅解码器的变换器。我们进一步对架构进行了一些修改，以允许额外的本体感知和控制器动作条件。

我们的转移模型接收过去潜在状态的历史和动作，其中H是表示模型上下文长度的超参数，并预测下一步时间的潜在状态zt。为了正确捕获时间依赖关系，即时间t的世界状态应仅依赖于过去的观测和动作，我们在ViT模型中实现了因果注意力机制，使模型能够以帧级别自回归地进行预测。具体来说，每个补丁向量对于潜在状态zt接收。这与过去的工作IRIS（Micheli等，2023）不同，IRIS同样将每个观测表示为一系列向量，但在标记 token 级别自回归地预测，接收以及。我们认为，在帧级别进行预测并将一个观测的补丁向量作为一个整体来处理，更好地捕获了全局结构和时间动态，建模了整个观测的依赖关系，而不是孤立的标记，从而实现了更好的时间泛化。我们在附录A.4.2的消融实验中展示了这种注意力掩码的有效性。

为了模拟代理动作对环境的影响，我们通过这些动作来条件化世界模型的预测。具体来说，我们将K维动作向量（使用多层感知器（MLP）从原始动作表示映射而来）连接到每个补丁向量，i = 1, ..., N。当有本体感知信息时，我们通过将其连接到观测潜在变量来类似地整合它，从而将其整合到潜在状态中。

我们在教师强迫下训练世界模型。在训练期间，我们将轨迹切分为长度为H + 1的段，并在每个H个预测帧上计算潜在一致性损失。对于每个帧，我们计算：

其中ϕ是动作编码器模型，可以将动作映射到更高维度。注意，我们的世界模型训练完全在潜在空间中进行，无需重建原始像素图像。

3.1.3 用于解释性的解码器

为了帮助可视化和解释性，我们使用一系列转置卷积层将补丁表示解码回图像像素，类似于Razavi等（2019）的方法。给定一个预收集的数据集，我们通过简单的重建损失来优化解码器qθ的参数θ：

解码器的训练完全独立于转移模型的训练，提供了几个优势：

1）解码器不影响世界模型在解决下游任务中的推理和规划能力，

2）在规划期间无需重建原始像素图像，从而减少了计算成本。

然而，解码器仍然有价值，因为它增强了世界模型预测的可解释性。尽管将这种解码器损失反向传播到预测器是可能的，但我们对此选择进行了消融实验，并发现与省略解码器损失相比，它对性能有负面影响。详细信息请参阅附录A.4.3。

3.2 使用DINO-WM进行视觉规划

为了评估世界模型的质量，我们在测试时执行轨迹优化并测量性能。虽然规划方法本身相当标准，但它们作为强调世界模型质量的手段。为此，我们的世界模型接收当前观测o0和目标观测og，两者都表示为RGB图像。我们将规划制定为搜索代理将采取的一系列动作以达到og。我们采用模型预测控制（MPC），通过考虑未来动作的结果来促进规划。

我们使用交叉熵方法（CEM）在每次迭代中优化动作序列。规划成本定义为当前潜在状态与目标潜在状态之间的均方误差（MSE），如下所示：

MPC框架和CEM优化过程的详细信息在附录A.5.1中提供。由于我们的世界模型是可微分的，一种可能更有效的优化方法是通过梯度下降（GD）优化这个目标，允许世界模型直接指导代理朝着特定目标前进。GD的详细信息在附录A.5.2中提供。然而，我们在实验中观察到CEM在我们的实验中优于GD，完整结果在附录A.5.3中。我们假设在训练期间和规划目标中引入正则化可以进一步提高性能，并将其留作未来工作。

4. 实验

我们的实验旨在解决以下关键问题：

1）我们是否可以有效地使用预收集的离线数据集训练DINO-WM？

2）一旦训练完成，DINO-WM是否可以用于视觉规划？

3）世界模型的质量在多大程度上依赖于预训练的视觉表示？

4）DINO-WM是否能够泛化到新配置，例如空间布局和物体排列的变化？

我们在六个环境套件上训练和评估DINO-WM（完整描述请参阅附录A.1），并将其与多种最新世界模型进行比较，这些模型在潜在空间或原始像素空间中进行预测。

4.1 环境和任务

我们在六个环境套件上进行评估，这些环境具有不同的动态复杂性，其中一些取自标准机器人基准测试，如D4RL（Fu等，2021）和DeepMind Control Suite（Tassa等，2018），如图3所示。这些环境包括迷宫导航（Maze, Wall）、桌面推送的精细控制（PushT）和机械臂控制（Reach），以及使用XArm进行的可变形物体操纵（Rope, Granular）。

在所有环境中，任务是从任意初始状态开始，达到由目标观测指定的随机采样目标状态。对于PushT，目标配置被采样以确保在25步内可行。对于Granular，目标需要将所有粒子聚集到具有随机位置和大小的方形中。所有环境中的观测都是224×224的RGB图像。环境的完整描述请参阅附录A.1。

4.2 基线

我们将DINO-WM与以下最新模型进行比较，这些模型通常用于控制。对于IRIS、DreamerV3和TD-MPC2，我们在没有奖励或任务信息的情况下使用离线数据集训练模型，并在学习的世界模型上执行MPC以解决下游任务。

a) IRIS（Micheli等，2023）：IRIS通过离散自编码器将视觉输入编码为标记，并使用GPT变换器预测未来标记，从而通过想象进行策略和价值学习。

b) DreamerV3（Hafner等，2024）：DreamerV3将视觉输入编码为分类表示，预测未来状态和奖励，并从想象的轨迹中训练演员-评论家策略。

c) TD-MPC2（Hansen等，2024）：TD-MPC2在潜在空间中学习无解码器的世界模型，并使用奖励信号优化潜在变量。

d) AVDC（Ko等，2023）：AVDC使用扩散模型从初始观测和文本目标生成任务执行视频。我们在附录A.6中为动作条件变体提供了定性评估和MPC规划结果。

4.3 使用DINO-WM优化行为

使用训练有素的世界模型，我们研究DINO-WM是否可以直接在潜在空间中用于零样本规划。

对于Maze、Reach、PushT和Wall环境，我们采样50个初始和目标状态，并测量所有实例的成功率。由于Rope和Granular环境的环境步进时间，我们在这些环境上评估了10个实例的Chamfer距离（CD）。在Granular中，我们从验证集中采样一个随机配置，目标是将材料推入一个在随机选择的位置和比例的方形形状。

如表1所示，在诸如Wall和PointMaze等较简单的环境中，DINO-WM与DreamerV3等最新世界模型相当。然而，DINO-WM在需要准确推断丰富接触信息和物体动态以完成任务的操纵环境中显著优于以前的工作。我们注意到，对于TD-MPC2，缺乏奖励信号使其难以学习良好的潜在表示，从而导致性能不佳。所有环境的规划可视化的详细信息可以在附录A.10中找到。

随着更多数据的可用性，DINO-WM是否能学习到更好的环境动态？我们在附录A.4.1中进行了一系列消融实验，表明规划性能随着训练数据量的增加而正向扩展。我们还在附录A.8中展示了DINO-WM的完整推理和规划时间，与传统模拟相比显示出显著加速，特别是在计算密集型的可变形环境中。

4.4 预训练视觉表示是否重要？

我们使用不同的预训练通用目的编码器作为世界模型的观测模型，并评估它们的下游规划性能。具体来说，我们使用了在机器人控制和通用感知中常用的以下编码器：R3M（Nair等，2022）、ImageNet预训练的ResNet-18（Russakovsky等，2015；He等，2016）和DINO CLS（Caron等，2021）。这些编码器的详细描述在附录A.3中。

在PointMaze任务中，涉及简单的动态和控制，我们观察到具有各种观测编码器的世界模型都接近完美成功率。然而，随着环境复杂性的增加——需要更精确的控制和空间理解——将观测编码为单一潜在向量的世界模型表现出显著的性能下降。我们认为，基于补丁的表示更好地捕获了空间信息，与R3M、ResNet和DINO CLS等模型相比，这些模型将观测简化为单一全局特征向量，丢失了操纵任务所需的关键空间细节。

4.5 泛化到新环境配置

我们评估世界模型的泛化能力，不仅跨不同目标，还跨各种环境配置。我们构建了三个环境家族——WallRandom、PushObj和GranularRandom——在这些环境中，模型在未见过的配置上进行测试，具有随机目标。环境的详细描述可以在附录A.2中找到。

从表5中，我们观察到DINO-WM在WallRandom中表现出显著更好的性能，表明模型有效地学习了墙壁和门的一般概念，即使它们被放置在训练中未见过的位置。相比之下，其他方法难以准确定位门的位置并导航通过它。PushObj任务对所有方法来说仍然是一个挑战，因为该模型仅在四种物体形状上进行了训练，这使得精确推断相关物理参数变得困难。在GranularRandom中，代理遇到的粒子数量比训练时少，与训练实例相比，这导致了分布外的图像。然而，DINO-WM准确地编码了场景，并成功地将粒子聚集到指定的方形位置，与基线相比具有最低的Chamfer距离（CD），显示出更好的场景理解。我们假设这是由于DINO-WM的观测模型将场景编码为补丁特征，使得粒子数量的变化仍然在每个图像补丁的分布内。

4.6 与生成视频模型的定性比较

鉴于生成视频模型的突出地位，人们自然会认为它们可以作为世界模型。我们比较了DINO-WM与AVDC（Ko等，2023），这是一种基于扩散的生成模型。如图6所示，尽管AVDC可以生成视觉上逼真的未来图像，但这些图像缺乏物理合理性。

在单个时间步内可能发生大的不现实变化，并且该模型难以达到确切的目标状态。未来生成模型的进步可能有助于解决这些问题。

我们进一步比较了DINO-WM与AVDC的一个变体，其中扩散模型被训练为在给定当前观测和动作的情况下生成下一个观测ot+1，而不是一次性生成整个观测序列，以文本目标为条件。如附录A.6所述，动作条件扩散模型在长期预测中偏离了真实观测，使其不足以进行准确的任务规划。

4.7 解码和解释潜在变量

尽管DINO-WM在潜在空间中运行，且观测模型没有使用像素重建目标进行训练，但训练解码器有助于解释预测。我们评估了所有模型预测未来的图像质量，发现我们的方法优于其他方法，即使那些编码器是与环境特定的重建目标一起训练的。图4中的开环展开展示了DINO-WM的稳健性，尽管缺乏明确的像素监督。

我们在所有模型预测的未来帧上报告了学习感知图像补丁相似性 Learned Perceptual Image Patch Similarity（LPIPS）（Zhang等，2018），通过比较图像的深度表示来评估感知相似性，较低的分数反映了更接近的视觉相似性。附录A.7提供了包括结构相似性指数（SSIM）（Wang等，2004）在内的额外结果。

5. 结论

我们介绍了DINO-WM，这是一种简单而有效的方法，用于在潜在空间中建模视觉动态，无需像素空间重建。我们证明了DINO-WM能够捕捉环境动态并泛化到未见过的配置，独立于任务规范，实现测试时的视觉推理，并通过规划生成下游任务的零样本解决方案。DINO-WM在弥合任务无关世界建模与推理和控制之间的差距方面迈出了一步，为现实世界应用中的通用世界模型提供了有希望的前景。

局限性和未来工作：首先，DINO-WM假设可以访问具有足够状态-动作覆盖的离线数据集，这在高度复杂的环境中可能难以获得。这可能通过将DINO-WM与探索策略相结合，并在获得新经验时更新模型来解决。其次，DINO-WM仍然依赖于从代理那里获得的真实动作，这在使用互联网上的大量视频数据进行训练时可能并不总是可行的。最后，虽然我们目前在下游任务解决中在动作空间中进行规划，但这项工作的扩展可能涉及开发一个分层结构，将高层规划与低层控制策略相结合，以解决更精细的控制任务。

A. 附录

A.1 环境和数据集生成

a) PointMaze：在这个由（Fu等，2021）引入的环境中，任务是一个在笛卡尔方向x和y上受力驱动的2自由度球达到目标点。智能体的动态结合了物理属性，如速度、加速度和惯性，使运动更加真实。我们生成了2000个完全随机的轨迹来训练我们的世界模型。我们简称这个任务为Maze，以简化表格中的表示。

b) Wall：这个定制的2D导航环境由一堵墙分隔的两个房间组成。智能体的任务是从一个房间中的随机起始位置导航到另一个房间中的目标，穿过门。我们提供了一个变体，其中墙和门的位置是随机的，测试模型对新配置的泛化能力。对于固定墙设置，我们在一个完全随机的数据集上训练了1920个轨迹，每个轨迹有50个时间步。对于具有多种训练环境配置的变体，我们生成了10240个随机轨迹。

c) Reacher：这是一个来自DeepMind Control Suite（Tassa等，2018）的连续控制任务，其中一个2关节机械臂在2D空间中达到目标。我们通过要求整个机械臂而不仅仅是末端执行器匹配任意目标姿态来增加难度。为了训练世界模型，我们生成了3000个有100步的轨迹。我们简称这个任务为Reach，以简化表格中的表示。

d) Push-T：这个由（Chi等，2024）引入的环境包含一个推动智能体与T形块交互。目标是在25步内将智能体和T形块从随机初始化状态引导到已知的可行目标配置。与以前的设置不同，固定的绿色T形块不再代表T形块的目标位置，而仅作为视觉锚点。成功需要精确理解智能体和物体之间的接触丰富动态，使其成为视觉运动控制和物体操纵的具有挑战性的测试。我们使用各种噪声级别的原始发布的专家轨迹生成了18500个样本的数据集。此外，我们通过改变物体的形状和颜色来引入变化，以评估模型适应新任务的能力。为此变体，我们生成了20000个有100步的随机采样轨迹。

e) Rope Manipulation：这个任务由（Zhang等，2024）引入，使用Nvidia Flex（Zhang等，2024）进行模拟，包括一个X臂与桌面上的软绳交互。目标是从任意起始配置将绳子移动到测试时指定的目标配置。为了训练，我们生成了一个从随机起始位置开始的20步随机动作的1000个轨迹的随机数据集，而测试则在从不同起始位置设置的目标配置上进行，包括随机变化的方向和空间位移。

f) Granular Manipulation：这个环境使用与Rope Manipulation相同的模拟设置，涉及操纵大约一百个粒子形成期望的形状。训练数据由从相同初始配置开始的20步随机动作的1000个轨迹组成，而测试则在从不同起始位置设置的特定目标形状上进行，同时粒子分布、间距和方向也有随机变化。

A.2 测试泛化的环境家族

WallRandom：基于Wall环境，但墙和门的位置是随机的。在测试时，任务要求从墙的一侧的随机起始位置导航到另一侧的随机位置，墙和门的位置在训练中未见过且不重叠。
PushObj：从Push-T环境派生，其中我们引入了新的块形状，包括类似Tetris的块和“+”形状。我们在四种形状上训练模型，并在两种未见过的形状上进行评估。任务涉及智能体和物体都达到目标位置。
GranularRandom：从Granular环境派生，其中我们用不同数量的粒子初始化场景。任务要求机器人将所有粒子聚集到一个在随机采样位置的方形形状。为此任务，我们直接使用在第4.3节中使用的固定数量材料训练的模型。

可视化可以在图5中找到。

A.3 预训练特征

a) R3M：一个在广泛的真实世界人类操作视频上预训练的ResNet-18模型（Nair等，2022）。 b) ImageNet：一个在ImageNet-1K数据集上预训练的ResNet-18模型（Russakovsky等，2015）。 c) DINO CLS：预训练的DINOv2模型提供了两种类型的嵌入：补丁和CLS。CLS嵌入是一个一维向量，封装了图像的全局信息。

A.4 消融实验

A.4.1 DINO-WM的扩展规律

为了分析DINO-WM的扩展行为，我们在PushT环境上使用不同大小的数据集（从200到18500个轨迹）训练世界模型并进行规划。我们的结果表明了一个清晰的趋势：随着数据集大小的增加，世界模型预测的质量和计划行为的性能都显著提高。更大的数据集使世界模型能够捕捉环境的更多样动态和细微差别，从而实现更准确的预测和更明智的规划。

A.4.2 DINO-WM的因果注意力掩码

我们在PushT上通过训练有和没有因果注意力掩码的DINO-WM来消融这个选择，历史长度h各不相同，使得模型接收输入ot−h+1, ot−h+2, ...ot，并输出ot−h+2, ...ot+1。对于带有掩码的模型，模型只能关注过去的观测来预测每个ot，而在没有掩码的情况下，预测输出序列中的任何观测都可以关注整个输入观测序列。我们在表6中展示了PushT设置上的规划成功率。当h=1时，带和不带此因果掩码的模型是等效的，两种模型都获得了不错的且相当的成功率。然而，随着历史长度的增加，没有掩码的情况下的成功率迅速下降，因为模型在训练期间可以通过关注未来帧来作弊，而这些帧在测试时是不可用的。添加因果掩码解决了这个问题，我们观察到随着历史长度的增加，性能得到提高，因为更长的历史可以更好地捕捉动态信息，如速度、加速度和物体动量。

A.4.3 带有重建损失的DINO-WM

虽然DINO-WM消除了使用像素重建损失训练世界模型的需要——避免了学习与下游任务无关的特征的风险——但我们进行了一项消融研究，其中预测器使用从解码器传播的重建损失进行训练。如表7所示，这种方法在PushT任务上表现相当不错，但略低于我们的方法，其中预测器完全独立于解码器进行训练。这强调了解耦特征学习与重建目标的优势。

A.5 规划优化

在本节中，我们详细描述了实验中的规划优化程序。

A.5.1 带有交叉熵方法的模型预测控制

a) 给定当前观测o0和目标观测og，两者都表示为RGB图像，首先将观测编码为潜在状态：

b) 规划目标定义为最终时间步T处的预测潜在状态与目标潜在状态之间的均方误差（MSE）：

c) 在每次规划迭代中，CEM从分布中采样N个动作序列，每个序列长度为T。初始分布设置为高斯分布。

d) 对于每个采样的动作序列，使用世界模型预测潜在空间中的结果轨迹：

并计算每个轨迹的成本C。

e) 选择成本最低的前K个动作序列，并相应地更新分布的均值和协方差。

f) 从更新后的分布中采样新的N个动作序列，并重复该过程，直到成功或经过我们设置为超参数的固定迭代次数。

g) 优化过程完成后，执行环境中的前k个动作a0, ..., ak。然后在下一个时间步重复该过程，使用新的观测。

A.5.2 梯度下降

由于我们的世界模型是可微分的，我们还考虑了一种使用梯度下降（GD）的优化方法，该方法通过反向传播直接最小化成本来优化动作。

a) 首先，将当前观测o0和目标观测og编码为潜在空间： ˆz0 = enc(o0), zg = enc(og)。

b) 目标与CEM相同： C = ∥ˆzT − zg∥²，其中 ˆzt = p(ˆzt−1, at−1), ˆz0 = enc(o0)。

c) 使用成本相对于动作序列{a0, a1, ..., aT−1}的梯度，迭代更新动作：，其中η是学习率。

d) 该过程重复直到达到固定的迭代次数，然后我们在环境中执行前k个动作a0, ..., ak，其中k是一个预确定的超参数。

A.5.3 规划结果

这里我们展示了使用各种规划优化方法的完整规划性能。CEM表示我们使用CEM优化一系列动作，并在环境中执行这些动作而没有任何更正或重规划。同样，GD表示使用梯度下降进行优化，并以开环方式一次性执行所有计划的动作。MPC表示允许使用CEM进行重规划和滚动视野优化。

A.6 与动作条件生成模型的比较

我们比较了DINO-WM与AVDC的一个变体，其中扩散模型被训练为在给定当前观测ot和动作at的情况下生成下一个观测ot+1，而不是一次性生成整个观测序列，以文本目标为条件。然后我们在验证轨迹上使用此动作条件扩散模型进行开环展，可视化结果如图7所示。可以看出，动作条件扩散模型在长期预测中偏离了真实观测，使其不足以进行准确的任务规划。

A.7 解码潜在变量：LPIPS和SSIM指标

我们在四个更具挑战性的环境上报告了两个关键指标：结构相似性指数（SSIM）（Wang等，2004）和学习感知图像补丁相似性（LPIPS）（Zhang等，2018），用于世界模型预测的未来状态的重建。SSIM通过评估预测和真实图像之间的结构信息和亮度一致性来衡量图像的感知质量，较高值表示更大相似性。LPIPS通过比较图像的深度表示来评估感知相似性，较低分数反映了更接近的视觉相似性。

A.8 推理时间

推理时间是部署模型进行实时决策的关键因素。表10报告了在NVIDIA A6000 GPU上单次推理步骤所需的时间，模拟器推进一步所需的时间，以及使用交叉熵方法（CEM）生成最优动作序列的整体规划时间。DINO-WM的推理时间在不同环境中保持不变，因为模型大小和输入图像分辨率固定，与传统模拟推进相比显著加速。特别是在计算需求高的环境中，如可变形物体操纵，模拟推进每步需要几秒钟，而DINO-WM实现了快速推理和高效规划。规划时间是在CEM中每迭代100个样本和10次优化步骤的情况下测量的，表明DINO-WM可以在保持跨任务的准确性和适应性的同时实现可行的规划时间。

A.9 超参数和实现

我们在下面展示了DINO-WM的超参数和相关实现存储库。我们在所有环境中使用相同的超参数训练世界模型。

世界模型架构在所有环境中保持一致。我们使用基于DINOv2的编码器，它从调整为196×196像素的输入图像中提取形状为（14×14, 384）的特征。ViT骨干网络具有6层深度、16个注意力头和2048维的MLP，总共约1900万个参数。

为了确保预测任务是有意义的，因为附近的观测可能非常相似，我们在数据处理中引入了一个帧跳跃参数。该参数指定模型预测的未来有多远。每个环境的帧跳跃值在表11中提供。

DINOv2: https://github.com/facebookresearch/dinov2
DreamerV3: https://github.com/NM512/dreamerv3-torch
AVDC: https://github.com/flow-diffusion/AVDC
R3M: https://github.com/facebookresearch/r3m/

我们的预测器实现基于 https://github.com/lucidrains/vit-pytorch/

A.10 额外的规划可视化

我们在图8中展示了DINO-WM和我们的基线的规划实例的可视化。为了比较，我们展示了表现最佳的世界模型DINO CLS和DreamerV3。我们还在图9中展示了DINO-WM在所有任务上的规划可视化。对于每个环境，顶部（阴影）行显示执行计划动作后的环境观测，底部行显示世界模型的想象观测。

为了展示DINO-WM在测试时对不同目标的泛化能力，我们在图10和图11中展示了DINO-WM在提供相同初始观测但不同目标观测时的额外可视化。同样，我们展示了轨迹对，以比较执行一系列计划动作后环境的观测（顶部阴影行）与DINO-WM的想象轨迹（底部行）。最左侧列表示初始观测，最右侧阴影列表示目标观测。