提升56%！纽约大学DINO-WM：预训练世界模型实现零样本规划

最新推荐文章于 2024-12-04 17:14:25 发布

自动驾驶之心

最新推荐文章于 2024-12-04 17:14:25 发布

阅读量411

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=Mzg2NzUxNTU1OA==&mid=2247641860&idx=2&sn=35601034f0cb761fd01e82ab3699c136&chksm=cf13bf03a6312338ce10cd79ea7d5b7f2a9edc9959808e0eca9c768e4fc294e3f167f998f285&scene=126&sessionid=0

版权

点击下方卡片，关注“具身智能之心”公众号

作者 | 具身智能之心编辑 | 具身智能之心

本文只做学术分享，如有侵权，联系删文

>>点击进入→具身智能之心技术交流群

写在前面&出发点

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

内容出自国内首个具身智能全栈学习社区：具身智能之心知识星球，这里包含所有你想要的。

一些介绍

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

与在训练期间学习所有可能任务的解决方案不同，另一种方法是使用训练数据拟合一个动力学模型，并在运行时优化特定任务的行为。这些动力学模型也被称为世界模型，在机器人技术和控制领域有着悠久的历史。最近，一些研究表明，可以使用原始观测数据训练世界模型。这使得能够灵活地使用基于模型的优化来获得策略，因为它避免了显式状态估计的需求。尽管如此，在使用世界模型解决通用任务方面仍然存在重大挑战。

为了理解世界建模中的挑战，让我们考虑学习世界模型的两种主要范式：在线和离线。在在线设置中，通常需要访问环境，以便可以持续收集数据来改进世界模型，进而改进策略以及后续的数据收集。然而，在线世界模型仅在所优化策略覆盖的范围内准确。因此，虽然它可用于训练强大的特定任务策略，但即使在同一环境中，对于每个新任务都需要重新训练。相比之下，在离线设置中，世界模型是在环境中收集的轨迹的离线数据集上进行训练的，这消除了其对任务特定性的依赖，前提是数据集具有足够的覆盖范围。然而，当需要解决任务时，该领域的方法需要强大的辅助信息来克服任务特定域缺乏密集覆盖的问题。这种辅助信息可以是专家演示，结构化关键点，访问预训练的逆模型或密集奖励函数，所有这些都会降低离线世界模型的通用性。构建更好离线世界模型的核心问题是，是否存在不损害其通用性的替代辅助信息？

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

DINO-WM在涵盖迷宫导航、滑动操作和粒子操作任务的四个环境套件上进行了实验评估。我们的实验得出了以下发现：

DINO-WM能够生成高质量的未来世界模型，这可以通过训练后的解码器改进的视觉重建来衡量。在最困难的任务的LPIPS指标上，这比先前最先进的工作提高了56%。
利用DINO-WM训练的潜在世界模型，在最困难的任务上实现了任意目标的高达成率，平均比先前的工作提高了45%。
DINO-WM可以在任务家族内的不同环境变体（例如，导航中的不同迷宫布局或操作中的不同物体形状）上进行训练，并且与先前的工作相比，实现了更高的成功率。

DINO-WM的代码和模型将开源，以确保可重复性，代码：https://dino-wm.github.io。

DINO世界模型

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

本工作的目标是从预先收集的离线数据集中学习任务无关的世界模型，并在测试时间使用这些世界模型进行视觉推理和控制。在测试时间，系统从任意环境状态开始，并被提供一张RGB图像形式的目标观测，这与先前的工作保持一致。系统被要求执行一系列动作，以便达到目标状态。这种方法与在线强化学习（RL）中使用的世界模型不同，后者的目标是优化手头固定任务集的奖励，也与通过文本提示指定目标的文本条件世界模型不同。

1）基于DINO的世界模型（DINO-WM）

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

其中，观测模型将图像观测编码为潜在状态，而转移模型则接收长度为H的过去潜在状态的历史记录。解码器模型接收一个潜在状态，并重构图像观测。我们用θ来表示这些模型的参数。请注意，解码器是完全可选的，因为解码器的训练目标与训练世界模型的其他部分无关。这消除了在训练和测试期间都需要重构图像的需求，与Hafner等人（2024）和Micheli等人（2023）中将观测模型的训练和解码器的训练耦合在一起相比，这降低了计算成本。

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

观测模型

我们的目标是学习一个适用于多种环境和现实世界的通用世界模型，我们认为观测模型应该，1）与任务和环境无关，2）包含丰富的空间信息，这对于导航和操作任务至关重要。与以往总是为当前任务学习观测模型的工作不同，我们认为，面对新环境时，世界模型不可能总是从零开始学习观测模型，因为感知是一个可以从大量互联网数据中学习的通用任务。因此，我们选择现成的预训练DINOv2模型作为我们世界模型的观测模型，因为它在需要深刻理解空间的目标检测、语义分割和深度估计任务中表现出色。在训练和测试期间，观测模型都保持不变。在每个时间步t，它将图像编码为block嵌入，其中N表示块的数量，E表示嵌入维度。该过程如图2所示。

转移模型

我们为转移模型采用了ViT架构，因为它是处理block特征的自然选择。然而，需要对架构进行一些修改，以便对本体感受和控制器动作进行额外的条件控制。

我们的转移模型接收过去潜在状态和动作的历史记录，其中H是表示模型上下文长度的超参数，并预测下一个时间步的潜在状态。为了正确捕获时间依赖性，即时间t处的世界状态应仅依赖于之前的观测和动作，我们在ViT模型中实现了因果注意力机制，使模型能够在帧级别上自回归地预测潜在状态。对于潜在状态的每个block向量，它关注于。这与过去的工作IRIS不同，后者同样将每个观测表示为向量序列，但在标记级别上自回归地预测，同时关注于以及$[z^i_t]^<k_{i=1}$。我们认为，在帧级别上进行预测，并将一个观测的block向量视为一个整体，能更好地捕获全局结构和时间动态，对整个观测而非孤立的标记进行依赖建模，从而实现更好的时间泛化。< p=""></k_{i=1}$。我们认为，在帧级别上进行预测，并将一个观测的block向量视为一个整体，能更好地捕获全局结构和时间动态，对整个观测而非孤立的标记进行依赖建模，从而实现更好的时间泛化。<>

为了模拟agent动作对环境的影响，将世界模型的预测建立在这些动作的基础上。将从原始动作表示使用多层感知器（MLP）映射得到的K维动作向量与每个块向量（对于i = 1, ..., N）进行拼接。当本体感受信息可用时，同样通过将其与观测潜在状态进行拼接来整合它，从而将其纳入潜在状态。

我们使用教师强制（teacher forcing）方法训练世界模型。在训练过程中，将轨迹切分为长度为H + 1的片段，并在每个预测的H帧上计算潜在一致性损失。对于每一帧，我们计算：

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

用于可解释性的解码器

为了辅助可视化和可解释性，我们使用转置卷积层的堆叠来将block表示解码回图像像素，这与Razavi等人（2019）的方法类似。给定一个预先收集的数据集，通过一个简单的重建损失来优化解码器θ的参数θ，该损失定义为：

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

2）使用DINO-WM进行视觉规划

可以说，为了评估世界模型的质量，它需要能够支持下游的推理和规划。一个标准的评估指标是在测试时间使用这些世界模型进行轨迹优化并测量性能。虽然规划方法本身相当标准，但它作为强调世界模型质量的一种方式。为此，我们的世界模型接收当前观测值和目标观测值，两者均以RGB图像表示。我们将规划定义为寻找一系列动作的过程，这些动作是agent为到达将采取的。为了实现这一点，采用了模型预测控制（MPC），它通过考虑未来动作的结果来促进规划。

我们利用交叉熵方法（CEM），一种随机优化算法，来优化每次迭代中的动作序列。规划成本定义为当前潜在状态与目标潜在状态之间的均方误差（MSE）。

实验对比

我们的实验旨在解决以下关键问题：1）是否可以使用预先收集的离线数据集有效地训练DINO-WM？2）一旦训练完成，DINO-WM是否可用于视觉规划？3）世界模型的质量在多大程度上依赖于预训练的视觉表示？4）DINO-WM是否适用于新的配置，例如空间布局和物体排列的变化？为了回答这些问题，在五个环境套件中训练和评估了DINO-WM，并将其与多种最先进的在潜在空间和原始像素空间中建模世界的世界模型进行了比较。

1）环境和任务

在我们的评估中，考虑了五个环境套件，涵盖了从简单的导航环境到具有不同动力学复杂性的操作环境。对于所有环境，观测空间均为大小为（224, 224）的RGB图像。

a) Point Maze：D4RL套件中的简单二维点迷宫导航环境。具有二维动作空间的点agent在U形迷宫中移动。agent的动力学结合了速度、加速度和惯性等物理属性，使移动更加真实。任务的目标是导航迷宫，从任意起始位置到达任意目标位置。

b) Push-T：环境中有一个推动者agent与T形块进行交互。目标是在25步内将agent和T形块从随机初始状态引导到已知的可行目标配置。任务要求agent和T形块都匹配目标位置。与之前的设置不同，固定的绿色T形块不再表示T形块的目标位置，而仅作为视觉参考锚点。成功完成任务需要对agent和物体之间丰富的接触动力学有精确的理解，这对视觉运动控制和物体操作构成了挑战。我们还引入了一个变体，其中存在多种物体形状。

c) Wall：此自定义二维导航环境包含两个由带门的墙隔开的房间。任务要求agent从一个房间的随机起始位置导航到另一个房间的目标位置，这需要agent穿过门。我们引入了一个环境变体，其中墙和门的位置是随机的，以评估模型对熟悉环境动力学的新配置的泛化能力。

d) Rope Manipulation：此任务使用Nvidia Flex模拟，包含一个XArm与放置在桌面上的绳子进行交互。目标是将绳子从任意起始配置移动到指定的目标配置。

e) Granular Manipulation：颗粒操作使用与绳子操作相同的设置，并操作约一百个颗粒以形成所需的形状。

2）基线（模型/方法）

将DINO-WM与以下常用于控制的最新模型进行了比较：

a) IRIS：IRIS使用离散自动编码器将视觉输入转换为标记（token），并使用GPT Transformer预测未来观测的标记。它结合这些组件，通过想象过程来学习策略和值函数。

b) DreamerV3：DreamerV3学习一个世界模型，将视觉输入解释为分类表示。它基于给定的动作预测未来的表示和奖励，并从其想象的轨迹中训练一个行动者-评论家策略。

c) TD-MPC2：TD-MPC2在潜在空间中学习一个无需解码器的世界模型，并使用奖励信号来优化潜在变量。它是无重建世界建模的强大基线。

d) AVDC：AVDC利用扩散模型根据初始观测和文本目标描述生成任务执行的想象视频。然后，它估计帧之间的光流以捕捉物体运动，并生成机械臂指令。

3）使用DINO-WM优化行为

使用一个训练好的世界模型，我们研究DINO-WM是否可以直接在潜在空间中进行零样本规划。对于PointMaze、Push-T和Wall环境，采样了50个初始状态和目标状态，以测量所有实例的成功率。由于Rope和Granular环境的步长时间因素，在这两个环境中对10个实例评估了Chamfer距离（CD）。在Granular环境中，从验证集中随机采样一个配置，目标是将材料推入一个随机选定位置和大小的方形形状。

如表1所示，在Wall和PointMaze等较简单的环境中，DINO-WM的表现与最先进的世界模型（如DreamerV3）相当。然而，在需要准确推断丰富的接触信息和物体动力学以完成任务的操作环境中，DINO-WM显著优于先前的工作。我们注意到，对于TD-MPC2而言，缺乏奖励信号使其难以学习良好的潜在表示，进而导致性能不佳。一些规划结果的可视化图像可见于图5。

4）预训练视觉表征是否重要？

使用不同的预训练通用编码器作为世界模型的观测模型，并评估它们在下游规划任务中的性能。具体来说，我们使用了以下在机器人控制和一般感知中常用的编码器：R3M、在ImageNet上预训练的ResNet-18以及DINO CLS。

在PointMaze任务中，该任务涉及简单的动力学和控制，观察到使用各种观测编码器的世界模型都实现了接近完美的成功率。然而，随着环境复杂性的增加，需要更精确的控制和空间理解，那些将观测编码为单个潜在向量的世界模型在性能上出现了显著下降。我们认为，与将观测简化为单个全局特征向量的R3M、ResNet和DINO CLS等模型相比，基于patch的表示方法能更好地捕捉空间信息，避免了在操纵任务中丢失至关重要的空间细节。

5）推广到新的环境配置

我们希望能够衡量我们的世界模型不仅在环境中的不同目标之间具有泛化能力，而且在不同的环境本身之间也具有泛化能力。为此，我们构建了三个环境系列，其中世界模型将在未见过的环境中针对未见过的目标进行部署。我们的环境系列包括WallRandom、PushObj和GranularRandom。训练和测试示例的可视化图像如图6所示。

从表3中，我们观察到DINO-WM在WallRandom环境中表现出显著更好的性能，这表明世界模型已经有效地学习了墙壁和门的一般概念，即使它们位于训练期间未见过的位置。相比之下，其他方法难以准确识别门的位置并通过它进行导航。

PushObj任务对所有方法来说仍然具有挑战性，因为模型仅针对四种物体形状进行了训练，这使得难以精确推断物理参数，如重心和惯性。在GranularRandom中，智能体遇到的粒子数量少于训练期间的一半，与训练实例相比，这导致了分布外的图像。尽管如此，DINO-WM仍然准确地编码了场景，并成功地将粒子聚集到指定的方形位置，与基线相比具有最低的Chamfer距离（CD），从而表现出更好的场景理解能力。我们假设这是因为DINO-WM的观测模型将场景编码为补丁特征，使得粒子数量的变化仍然在每个图像补丁的分布范围内。

6）与生成式视频模型的定性比较

鉴于生成式视频模型的显著地位，有理由推测它们可以轻易地作为世界模型使用。为了探究DINO-WM相较于此类视频生成模型的实用性，我们将其与基于扩散的生成模型AVDC的想象轨迹进行了比较。如图7所示，我们发现，在基准数据集上训练的扩散模型生成的未来图像在视觉上大多很逼真，但在物理上却不太合理，因为我们可以看到在单个预测时间步内会发生很大的变化，并且可能难以达到确切的目标状态。未来潜在上更强大的生成模型可能会缓解这一问题。

还将DINO-WM与AVDC的一个变体进行了比较，在该变体中，扩散模型被训练为根据当前观测值和动作生成下一个观测值，而不是根据文本目标一次性生成整个观测序列。

7）解码和解释潜在变量

尽管DINO-WM在潜在空间中运行，且观测模型并未以像素重建为目标进行训练，但训练一个解码器对于解释模型的预测仍然很有价值。我们评估了所有模型预测的未来图像的质量，并发现我们的方法优于其他方法，甚至优于那些其编码器以特定于环境的重建目标进行训练的方法。在图4中展示了开环轨迹的可视化结果。这证明了DINO-WM的鲁棒性，尽管它缺乏明确的像素级监督。这里报告了两个关键指标：世界模型预测未来状态重建的结构相似性指数（SSIM）和学习的感知图像patch相似性（LPIPS）。SSIM通过评估预测图像和真实图像之间的结构信息和亮度一致性来衡量图像的感知质量，值越高表示相似性越大。而LPIPS则通过比较图像的深度表示来评估感知相似性，分数越低表示视觉相似性越近。

参考

[1] DINO-WM: WORLD MODELS ON PRE-TRAINED VISUAL FEATURES ENABLE ZERO-SHOT PLANNING

“具身智能之心”公众号持续推送具身智能领域热点：

【具身智能之心】技术交流群

具身智能之心是首个面向具身智能领域的开发者社区，聚焦大模型、视觉语言导航、机械臂、双足机器人、四足机器人、感知融合、强化学习、模仿学习、规控与端到端、机器人仿真、产品开发、自动标注等多个方向，目前近60+技术交流群，欢迎加入！扫码添加小助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）

【具身智能之心】知识星球

具身智能之心知识星球是国内首个具身智能开发者社区，近600人的社区。主要关注具身智能相关的数据集、开源项目、具身仿真平台、大模型、视觉语言模型、强化学习、具身智能感知定位、机器臂抓取、姿态估计、策略学习、轮式+机械臂、双足机器人、四足机器人、大模型部署、端到端、规划控制等方向。扫码加入星球，享受以下专有服务：