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

z_ai
RL笔记
2025-07-02
42热度
0评论

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

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

背景

与Dreamer等需要在在线交互中不断学习和优化的世界模型不同，DINO-WM旨在解决一个更具挑战性的问题：如何仅使用被动收集的离线数据，构建一个通用的、与具体任务无关的世界模型，并用它在测试时直接解决新任务？

现有挑战:
1. 在线世界模型: 虽然强大，但通常是为特定任务训练的。当任务改变时（即使在同一环境中），也需要重新训练或微调。它们与正在优化的策略高度耦合。
2. 离线世界模型: 现有方法通常需要额外的强力辅助信息，如专家演示、奖励函数或预训练的逆动力学模型，这限制了它们的通用性。
核心目标: DINO-WM 追求世界模型的三个理想属性：
1. 可离线训练: 能从预先收集的轨迹数据中学习。
2. 支持测试时优化: 能够在测试阶段为新任务规划行为。
3. 任务无关的推理: 模型本身不依赖于特定任务的奖励或目标。

方法

DINO-WM的核心思想是：不在像素空间进行重建或预测，而是在一个强大的、预训练的、紧凑的视觉特征空间中对世界动态进行建模。

模型组成

DINO-WM由三个主要部分构成，但核心是前两个：

观测模型 (Observation Model)
- 是什么: 一个冻结的、预训练的DINOv2模型。
- 作用: 将高维的图像观测 (o_t) 编码为一系列低维的、包含丰富空间信息的patch特征 (z_t \in \R^{N \times E})，其中N是patch数量，E是每个patch的维度。
- 优势: 直接利用了大型模型在海量数据上学到的通用视觉理解能力（如物体检测、语义分割、深度估计），避免了从零开始学习感知，使得模型一开始就拥有强大的视觉基础。
转移模型 (Transition Model)
- 是什么: 一个带有因果自注意力机制的Vision Transformer (ViT) 架构。
- 作用: 接收历史的patch特征序列 (z_{t-H:t-1}) 和动作序列 (z_{t-H:t-1})，预测下一时刻的patch特征 (\hat{z}_t)。
- 关键设计:
  - 因果注意力: 确保预测未来时只依赖过去的信息。
  - 帧级别(Frame-level)预测: 将一帧图像的所有patch视为一个整体进行处理和预测，更好地捕捉全局结构和时间动态。
  - 动作融合: 将动作向量通过一个小型MLP映射后，拼接到每个patch特征上，从而让模型理解动作对世界状态的影响。
解码器 (Decoder - 可选)
- 是什么: 一个由转置卷积层堆叠而成的网络。
- 作用: 仅用于可视化和解释，可以将预测的patch特征解码回像素图像，让我们能直观地看到世界模型的“想象”是否准确。
- 关键点: 解码器的训练与转移模型完全独立，在规划和控制时完全不需要它，极大地降低了计算开销。

上图展示了DINO-WM的整体流程。左侧是训练过程，模型学习从过去的观测(o_{t-k:t})和动作(a_{t:T-1})预测未来的DINOv2特征(o_{g})。通过在DINO-WM的想象中进行模型预测控制(MPC)来优化动作序列，以达到目标状态，即最小化预测的\hat{z}_{T}和实际的z_{g}的距离。

三、世界模型学习

与Dreamer在“想象”中学习策略不同，DINO-WM的学习过程完全在离线数据集上进行，目标是精准地预测预训练特征的动态变化。

训练数据: 来源于智能体与环境交互的离线轨迹数据集，只包含图像序列 (o_{1:T}) 和动作序列 (a_{1:T})。不需要奖励、折扣因子或终止信号。

损失函数: 模型训练的目标非常纯粹和直接：让预测的下一个状态的DINOv2特征与真实观测编码后的DINOv2特征尽可能接近。这是通过一个简单的均方误差损失（MSE）实现的：

其中，pθ 是转移模型，encθ 是冻结的DINOv2编码器，ϕ 是动作编码器。整个训练过程在潜在空间中进行。

与 Dreamer 的核心对比

特性	DINO-WM	Dreamer
视觉表征	使用预训练、冻结的DINOv2	从零开始学习离散的分类潜在变量
学习目标	预测未来的DINOv2特征	重建图像、预测奖励和折扣因子
训练范式	纯离线、任务无关	在线学习，需要与环境持续交互
数据需求	仅需 (观测, 动作) 序列	需要 (观测, 动作, 奖励, 折扣) 序列
计算开销	较低，无需图像重建	较高，需要训练图像解码器

四、行为学习与规划

DINO-WM的“行为学习”发生在测试阶段，是一种零样本规划 (Zero-shot Planning)。它不需要像Dreamer那样预先训练一个Actor-Critic策略网络。

问题设定: 给定当前观测图像 o_{0} 和一个目标状态的图像 o_{g}。
任务: 找到一个动作序列 a_{0:T-1}，使得智能体执行后，环境状态尽可能接近 o_{g}。
规划方法: 采用模型预测控制 (MPC)，并使用交叉熵方法 (CEM) 作为优化器。
1. 编码: 将当前观测 o_{0} 和目标图像 o_{g} 分别用DINOv2编码为潜在状态 z_{0} 和z_{g}。
2. 定义成本: 规划的成本函数是在潜在空间中定义的，即预测的最终状态 \hat{z}_T与目标状态 z_{g} 之间的均方误差：

C = ||\hat{z}_T-z_g||^2

优化: CEM算法通过迭代的方式，在每个时间步采样多组候选动作序列，用DINO-WM（转移模型）向前滚动，预测每组动作序列导致的最终状态，并计算成本。然后，选择成本最低的一批动作序列，更新采样分布，重复此过程，直到找到一个足够好的动作序列。
执行: 执行优化得到的动作序列中的第一个或前几个动作，然后回到第一步，用新的观测重新规划。