25年12月来自哥伦比亚大学和NYU的论文"Closing the Train-Test Gap in World Models for Gradient-based Planning"。

世界模型与模型预测控制（MPC）相结合，可以在大规模专家轨迹数据集上进行离线训练，并在推理时实现对各种规划任务的泛化。与依赖缓慢搜索算法或迭代精确求解优化问题的传统MPC程序相比，基于梯度的规划提供一种计算效率更高的替代方案。然而，迄今为止，基于梯度的规划性能仍落后于其他方法。本文提出一种改进的世界模型训练方法，实现高效的基于梯度的规划。尽管世界模型是基于下一状态预测目标进行训练的，但在测试时，它却被用来估计一系列动作。该工作目标是弥合这种训练-测试差距。为此，提出一种训练-时数据合成技术，该技术能够利用现有的世界模型显著改进基于梯度的规划。在测试时，该方法在各种物体操作和导航任务中，以 10% 的时间预算，表现优于或与经典的无梯度交叉熵方法 (CEM) 相当。

---

在机器人任务中，预测智体的行为如何影响其环境状态，对于预测（Finn et al., 2016）和规划（Mohanan & Salgoankar, 2018; Kavraki et al., 2002）都至关重要。传统方法从基本原理出发，解析地推导出环境演化的模型，依赖于对环境、智体以及任何不确定性的先验知识（Goldstein et al., 1950; Siciliano et al., 2009; Spong et al., 2020）。相比之下，基于学习的方法直接从数据中推断此类模型，使其能够捕捉复杂的动态过程，从而提高模型的泛化能力和对不确定性的鲁棒性（Sutton et al., 1998; Schrittwieser et al., 2020; LeCun, 2022）。

世界模型（Ha & Schmidhuber，2018）尤其已成为一种强大的范式。给定当前状态和动作，世界模型可以预测由此产生的下一个状态。这些模型可以从精确的状态信息（Sutton，1991）或高维感觉输入（例如图像）中学习（Hafner et al.，2023）。后一种方法尤其引人注目，因为它能够利用预训练的视觉表征直接从原始图像中进行感知、预测和控制，并且无需测量实践中难以实现的精确环境状态（Assran et al.，2023；Bardes et al.，2024）。近年来，世界模型及其预测能力已被用于规划，使智体能够解决各种任务（Hafner et al., 2019a;b; Schrittwieser et al., 2020; Hafner et al., 2023; Zhou et al., 2025）。动态模型在离线状态下学习，而规划任务则在推理阶段被定义为一个约束优化问题：给定当前状态，找到一个动作序列，使目标状态尽可能接近目标状态。这种推理时优化为强化学习方法（Sutton et al., 1998）提供一种有效的替代方案，后者通常存在样本效率低下的问题。

世界模型与许多基于模型的规划算法兼容。传统方法，例如动态规划（DDP）（Mayne，1966）和迭代线性二次规划（iLQR）（Li & Todorov，2004），依赖于迭代求解由标称轨迹周围动力学的线性和二次近似导出的精确优化问题。虽然这些方法在低维情况下非常有效，但对于大规模世界模型而言，由于求解由此产生的优化问题在计算上难以处理，因此变得不切实际。作为一种替代方案，基于搜索的方法，例如交叉熵方法（CEM）（Rubinstein & Kroese，2004）和模型预测路径积分控制（MPPI）（Williams et al.，2017a），已被广泛采用，作为无梯度替代方案，并在实践中被证明是有效的。然而，这些方法计算量巨大，因为它们需要迭代地对候选解进行采样，并执行世界模型展开来评估每个解，而这一过程在高维空间中扩展性很差。相比之下，基于梯度的方法（SV，2023）通过直接利用世界模型的可微性来端到端地优化动作，从而避免采样的局限性。这些方法消除基于搜索的方法所需的高成本展开操作，因此在高维空间中扩展性更高。尽管如此，基于梯度的方法迄今为止在实证研究中取得的成功仍然有限。

从感官数据中学习世界模型。基于学习的动力学模型已成为控制和决策的核心，为依赖第一性原理建模的传统方法提供一种数据驱动的替代方案（Goldstein et al., 1950; Schmidt & Lipson, 2009; Macchelli et al., 2009）。早期研究侧重于在低维状态空间中对动力学进行建模（Deisenroth & Rasmussen, 2011; Lenz et al., 2015; Henaff et al., 2017; Sharma et al., 2019），而最近的方法则直接从图像等高维感官输入中学习。像素空间预测方法（Finn，2016；Kaiser，2019）已在人体运动预测（Finn，2016）、机器人操作（Finn&Levine，2016；Agrawal，2016；Zhang，2019）以及Atari游戏求解（Kaiser，2019）等应用中取得成功，但由于图像重建的成本，其计算量仍然很大。为了解决这个问题，其他方法学习一种紧凑的潜表示，其中动态特性被建模（Karl，2016；Hafner，2019b；Shi，2022；Karypidis，2024）。这些模型通常采用两种监督方式：一是解码潜预测以匹配真实观测值（Edwards et al., 2018; Zhang et al., 2021; Bounou et al., 2021; Hu et al., 2022; Akan & Güney, 2022; Hafner et al., 2019b）；二是使用直接在潜空间中运行的预测目标，例如联合嵌入预测架构（JEPA）中的预测目标（LeCun, 2022; Bardes et al., 2024; Drozdov et al., 2024; Guan et al., 2024; Zhou et al., 2025）。

使用世界模型进行规划。由于目标函数的非线性和非凸性，使用世界模型进行规划极具挑战性。诸如 CEM（Rubinstein & Kroese，2004）和 MPPI（Williams，2017a）等基于搜索的方法被广泛应用于此领域（Williams，2017b；Nagabandi，2019；Hafner，2019b；Zhan，2021；Zhou，2025）。这些方法能够有效地探索动作空间，帮助跳出局部最小值，但由于其基于采样的特性，通常在高维环境下扩展性较差。相比之下，基于梯度的方法利用世界模型的可微性，通过反向传播直接优化动作，从而提供一种更具可扩展性的替代方案。尽管这些方法效率很高，但在高度非平滑的损失函数地形中，它们容易陷入局部最小值（Bharadhwaj et al., 2020a; Xu et al., 2022; Chen et al., 2022; Wang et al., 2023），并且梯度优化可能会诱发利用模型误差的对抗性动作序列（Schiewer et al., 2024; Jackson et al., 2024）。Zhou et al. (2025) 观察到，当使用基于预训练视觉嵌入（例如 DINOv2）构建的世界模型时，梯度提升算法（GBP）尤其脆弱（Oquab et al., 2024），其性能通常不如 CEM。为了应对这些挑战，人们提出了几种稳定化技术。例如，随机采样射击通过在动作序列中注入噪声并在轨迹优化过程中探索更广泛的动作集来帮助缓解对抗性轨迹（Nagabandi，2018）；Zhang（2025）通过在推理时选择性地扰动状态输入，对学习的策略引入对抗性攻击，使其对环境扰动具有鲁棒性。Florence（2022）在训练隐式策略函数时添加梯度惩罚以提高其平滑度并稳定优化过程，但他们的方法不涉及训练或使用世界模型。其他方法旨在采用混合方法，结合搜索和梯度步骤来平衡全局探索和局部优化（Bharadhwaj，2020a）。

世界模型中的训练-测试差距。使用学习的世界模型进行规划时，一个关键挑战是训练目标和规划目标之间的不匹配（Lambert，2020）。事实上，在训练阶段，世界模型通常被优化以最小化从专家演示或行为策略中收集的轨迹的单步预测或重建误差。然而，在测试阶段，相同的模型被用于规划器中以优化多步动作序列。因此，训练和测试阶段的目标本质上是不同的，导致训练期间观察的轨迹与规划期间遇到的轨迹之间存在分布偏移。这种不匹配会导致规划器将模型驱动到状态空间的分布外区域，在这些区域，预测误差会随时间的推移而累积，模型对于长期优化变得不可靠（Ajay et al., 2018; Ke et al., 2019; Zhu et al., 2023）。解决训练集与测试集差距的常用策略是数据集聚合（Ross et al., 2011），它通过扩展规划算法生成的动作轨迹并将其添加到训练集中来扩大训练分布（Talvitie, 2014; Nagabandi et al., 2018）。

世界模型通常使用专家轨迹数据集，以下一状态预测为目标进行训练。然而，在测试时，它们却被用于优化一系列动作的规划目标。这种不匹配是导致基于梯度的规划（GBP）实证性能不佳的根本原因，并提出两个假设来解释其原因。 （1）在规划过程中，梯度下降法探索的中间动作序列会将世界模型驱动到训练过程中未曾遇到的状态。在这些分布外的状态下，模型误差会累积，使得世界模型作为优化代理变得不可靠。（2）世界模型诱导的动作级优化空间可能难以遍历，包含许多较差的局部最小值或平坦区域，这阻碍有效的基于梯度优化。

本文针对上述两个挑战提出两种算法：在线世界建模和对抗世界建模。这两种算法都通过不断向数据集中添加新的轨迹并基于这些轨迹对世界模型进行微调来扩展已知潜状态的区域。为了管理离线专家轨迹和规划预测轨迹之间的分布偏移，在线世界建模使用环境模拟器来修正GBP生成轨迹上的状态。基于这些修正后的轨迹进行微调，可以确保当GBP进入专家轨迹分布之外潜状态空间区域时，世界模型仍能表现良好。为了克服GBP过程中在非平滑损失曲面上进行优化的困难，对抗世界建模会沿最大化世界模型损失的方向扰动专家轨迹。对抗微调平滑由此产生的动作损失曲面，使其更容易通过基于梯度的规划进行优化。如图展示这两种方法的示意图：

---

问题描述

世界模型通过预测在当前状态下采取某个动作后的状态来学习环境动态。然后在测试时，学习的世界模型可以通过模拟未来轨迹并指导动作优化来进行规划。形式上，世界模型近似于（可能未知的）动态函数 h : S × A → S，其中 S 表示状态空间，A 表示动作空间。环境根据以下公式演化：

s_t+1 = h(s_t, a_t)，对于所有 t，(1)

其中 s_t ∈ S，a_t ∈ A 分别表示 t 时刻的状态和动作。

潜世界模型。在实践中，通常无法获取环境的确切状态；相反，只能获得部分观测数据，例如图像。为了使世界模型能够在高维观测空间 O 中高效学习，采用嵌入函数 Φ_μ : O → Z 将观测值映射到低维潜空间 Z。给定嵌入函数 Φ_μ，目标是学习一个潜世界模型 f_θ : Z × A → Z，使得

z_t = Φ_μ(o_t), z_t+1 = f_θ(z_t, a_t), 对所有 t 成立。(2)

Φ_μ 的选择直接影响潜世界模型的表达能力。本文使用一个预先通过自监督学习训练的固定编码器，该编码器能够开箱即用地提供丰富的特征表示。

训练。为了训练潜世界模型，从离线轨迹数据集 T 中采样形如 (o_t, a_t, o_t+1) 的三元组，并最小化真实下一个潜状态 z_t+1 = Φ_μ (o_t+1) 与预测的下一个潜状态 zˆ_t+1 之间的 l2 距离。此过程由以下教师强制目标函数表示：

min_θ E_(o_t,a_t,o_t+1)∼T ||f_θ (Φ_μ (o_t), a_t)−Φ_μ (o_t+1)||²_2. (3)

值得注意的是，仅针对世界模型的参数 θ 最小化此目标函数，而不是针对可能很大的嵌入函数参数。

规划。在测试阶段，我用已学习的世界模型来优化到达目标状态的候选动作序列。通过递归地将世界模型应用于从初始潜状态开始的动作序列，可以获得预测的潜目标状态，进而得到潜空间中与真实目标状态的距离。这能够找到最优动作序列：

{aˆ∗_t} = argmin_{aˆ_t} ||zˆ_H+1 −z_goal||²_2. (4)

其中 zˆ_H+1 由递归过程生成：

zˆ_2 = f_θ(z_1, aˆ_1), zˆ_t+1 = f_θ(zˆ_t,aˆ_t)，t>1 (5)

用函数 rollout_f 来表示此递归过程。

基于梯度的规划 (GBP) 通过梯度下降法求解规划目标 (4)。关键在于，由于世界模型是可微分的，因此 ∇_{aˆ_t} zˆ_H+1 = ∇_ {aˆt} rollout_f (z_1, {aˆ_t})_H+1 是定义良好的。

相比之下，基于搜索的 CEM 算法无需梯度，但需要评估更多的动作序列。在算法 1 中详细介绍 GBP 算法。

如下算法 4 总结CEM：

由于误差可能在较长的时间范围内传播，模型预测控制 (MPC) 通常用于通过优化 H 步动作序列来反复重新规划，但每次只执行前 K ≤ H 个动作，然后根据更新后的状态重规划。

由于规划目标完全由世界模型决定，GBP 的成功取决于：(1) 模型能够准确预测任何候选动作序列下的未来状态；(2) 该可微优化的稳定性。下面提出两种旨在改进这两方面的微调方法。

在线世界建模

在基于梯度的规划过程中，待优化的动作序列并不受限于训练过程中观察的行为分布。世界模型通常在固定的专家轨迹数据集上进行训练，而基于梯度的规划（GBP）选择动作仅仅是为了提升规划目标，并不考虑这些动作是否与专家行为相似。因此，优化过程经常会提出超出分布范围的动作序列。已知在这种情况下通过已学习的模型进行优化会引入对抗性输入（Szegedy et al., 2013; Goodfellow et al., 2014）。在设置中，这些对抗性动作序列会将世界模型驱动到潜状态空间中训练过程中很少或从未观察的区域，从而导致较大的预测误差。即使误差最初很小，也会随着规划器向前滚动模型而累积，最终降低长期规划的性能。

为了解决这个问题，提出在线世界建模方法，该方法迭代地修正GBP生成的轨迹，并基于由此产生的演示结果对世界模型进行微调。并非仅基于专家演示进行训练，而是反复地将规划器自身产生的轨迹纳入模型，从而扩展世界模型能够可靠预测的潜状态区域。

首先，用专家轨迹 τ 的初始和目标潜状态进行梯度下降 (GBP) 训练，得到预测动作序列 {aˆ_t}。这些动作可能会将世界模型引入训练分布之外的潜空间区域。为了解决这个问题，获取一个修正后的轨迹：即在真实动力学模拟器 h 中执行动作序列 {aˆ_t} 后产生的实际状态序列。将修正后的轨迹，

τ′ = (z_1, aˆ_1 , z′_2, aˆ_2, …, z′_H+1), (6)

添加到世界模型每次更新时使用的训练数据集中。基于这些修正后的轨迹进行重训练，可以扩展训练分布，覆盖由梯度下降规划引起的潜空间区域，从而减轻规划过程中预测误差的累积。在算法 2 中提供更多细节，并在上图中展示了该方法。

此过程类似于 DAgger（数据集聚合）（Ross，2011），这是一种在线模仿学习方法，其中基础策略网络在其自身的展开（rollouts）上进行迭代训练，并将动作预测替换为专家策略的预测。类似地，引入真实世界模拟器作为模仿的专家世界模型。

对抗性世界建模

由于世界模型仅针对下一状态预测目标进行训练，因此它们的输入梯度无需表现良好。对抗性训练已被证明可以产生表现更好的输入梯度（Mejia，2019），从而平滑输入损失曲面。基于此观察，提出一种对抗性训练目标，该目标明确地针对状态-动作空间中世界模型预期表现较差的区域。这些对抗样本可能位于专家轨迹分布之外，这使得模型能够精确地接触到对动作优化至关重要的区域。这种被称为对抗世界建模的方法确实能够平滑规划目标的损失曲面（如图所示），从而提高动作序列优化的稳定性。

对抗训练通过优化最坏情况下的扰动来提高模型的鲁棒性（Madry，2018）。对抗样本的生成方法是对输入施加扰动 δ，使模型的损失最大化。为了在对抗样本上训练世界模型，使用目标函数如下：

在这些受到对抗扰动的轨迹上进行训练，为在线世界建模提供一种替代方法，用于揭示规划过程中可能遇到的状态，而无需依赖 GBP 展开。这在仿真成本高昂或不可行的情况下是一个显著的优势。

用快速梯度符号法 (FGSM) (Goodfellow，2014) 生成对抗性潜状态，该方法能够有效地逼近最大化预测误差的最坏情况扰动 (Wong，2020)。虽然可以使用更强的迭代攻击方法，例如投影梯度下降 (PGD)，但 FGSM 在 GBP 性能方面能够取得相当的提升，同时计算效率显著更高。这样能够在整个大规模离线模仿学习数据集上生成对抗样本。

对于给定小批量中的每个状态-动作对，寻找能够最大程度增加世界模型预测误差的潜状态或动作的微小变化。令 ε_a 和 ε_z 分别表示对动作 {a_t} 和潜在状态 {z_t} 的扰动半径。计算关于扰动的梯度 ∇_δ_a,δ_z ||f_θ (z_t + δ_z, a_t + δ_a) − z_t+1||²_2，并进行有符号梯度上升（即沿降低预测精度的方向），步长为 α_a = 1.25 ε_a, α_z = 1.25 ε_z。对结果进行裁剪，使扰动的每个元素都保持在半径范围内。此过程对应于 PGD 式攻击的单步，产生的扰动位于允许区域的边缘，在该区域内对模型构成最大挑战。详细说明参见算法 3的总结。

为了初始化扰动半径 ε_a 和 ε_z，用缩放因子 λ_a 和 λ_z，对抗世界建模在 0 ≤ λ_a ≤ 1 和 0 ≤ λ_z ≤ 0.5 的范围内具有鲁棒性。此外，将 ε_a 和 ε_z 固定为初始小批量数据的标准差在所有实验中均保持稳定。如算法 3 中那样对每个批次更新此估计值并不能持续改善最终规划性能。

---

通过在三个任务（PushT、PointMaze 和 Wall）上微调预训练的世界模型来评估方法，这些模型使用下一状态预测目标。对于每个任务，测量从初始配置 o_1 到达目标配置 o_goal 的成功率。在开环模式下，从 o_1 开始运行一次算法 1，并评估预测的动作序列。在 MPC 模式下，对每个 MPC 步骤运行一次算法 1（使用 Φ_μ(o′_1) 作为第一个 MPC 步骤的初始潜状态），在环境模拟器中展开预测的动作 {aˆ_t} 以到达潜状态 zˆ_H+1，并将 zˆ_1 = zˆ_H+1 设置为下一次 MPC 迭代。

用 DINO-WM（Zhou，2025）作为初始世界模型，因为它在选择的任务中与 CEM 算法配合使用时表现出色。嵌入函数 Φ_μ 采用预训练的 DINOv2 编码器（Oquab，2024），并在微调转移模型 f_θ 时保持不变。f_θ 使用 ViT 架构（Dosovitskiy，2021）实现。此外，还训练一个 VQVAE 解码器（van den Oord，2018）来可视化潜状态，尽管它在规划过程中不起作用。为了验证方法的广泛适用性，还研究了 IRIS（Micheli，2023）世界模型架构的应用。

为了初始化用于规划优化的动作序列，评估从标准正态分布中随机抽样和使用初始化网络两种方法。初始化网络 g_θ 的训练条件为 g_θ (z_1, z_g) = {aˆ_t}。实验中，随机初始化往往优于初始化网络。

在 GBP 过程中，将算法 1 中的 L_goal 设置为加权目标损失，以便从每个预测状态而非仅仅是最后一个状态获得梯度。通过实验发现，这种任务假设可以推广到导航任务（例如 PointMaze 和 Wall）和非导航任务（例如 PushT）；也就是说，无论任务是否具有子目标可分解性，该目标都能改进或达到最终状态损失的性能。此外，还评估在GBP中使用Adam优化器（Kingma & Ba，2014）的性能。虽然在实验中，对于所有世界模型，使用Adam优化器相比GD优化器都能显著提高性能，但单独使用Adam优化器无法将性能扩展到与CEM算法相当或更优的水平。