GRASP：面向世界模型长时程规划的梯度规划器

GRASP 是一种新的基于梯度的规划器，用于学习到的动力学（即“世界模型”），它通过以下三点让长时程规划变得实用：(1) 将轨迹提升到虚拟状态，使优化在时间上可并行；(2) 直接在状态迭代上加入随机性以进行探索；(3) 重塑梯度，使动作获得清晰的信号，同时避免通过高维视觉模型产生脆弱的“状态-输入”梯度。

大型学习到的世界模型正变得越来越强大。它们能在高维视觉空间中预测长序列的未来观测，并以几年前难以想象的方式跨任务泛化。随着这些模型规模扩大，它们开始不再像特定任务的预测器，而更像通用模拟器。

但拥有一个强大的预测模型，并不等于能够有效地将其用于控制/学习/规划。在实践中，用现代世界模型进行长时程规划仍然很脆弱：优化会变得病态，非贪婪结构会产生糟糕的局部极小值，而高维潜空间会引入微妙的失败模式。

在这篇博文中，我描述了促使这个项目的问题以及我们解决这些问题的方法：为什么用现代世界模型进行规划会出奇地脆弱，为什么长时程才是真正的压力测试，以及我们做了哪些改变来让基于梯度的规划稳健得多。

这篇博文讨论了与 Mike Rabbat、Aditi Krishnapriyan、Yann LeCun 和 Amir Bar 合作完成的工作（* 表示同等指导），我们在其中提出了 GRASP。

什么是世界模型？

如今，“世界模型”这个术语被严重过载使用，根据语境，它既可以指显式的动力学模型，也可以指生成模型所依赖的某种隐式、可靠的内部状态（例如，当 LLM 生成国际象棋走法时，是否存在某种棋盘的内部表示）。我们在下面给出一个宽松的工作定义。

假设你采取动作 $a_t \in \mathcal{A}$ 并观测状态 $s_t \in \mathcal{S}$（图像、潜向量、本体感觉）。世界模型 是一种学习到的模型，给定当前状态和未来动作序列，它预测接下来会发生什么。形式上，它定义了在观测状态序列 $s_{t-h:t}$ 和当前动作 $a_t$ 上的预测分布：

\[P_\theta(s_{t+1} \mid s_{t-h:t},\; a_t)\]

它近似环境的真实条件分布 $P(s_{t+1} \mid s_{t-h:t},\; a_t)$。为简化起见，在这篇博文中，我们将假设一个马尔可夫模型 $P(s_{t+1} \mid s_{t-h:t},\; a_t)$（这里的所有结果都可以扩展到更一般的情况），当模型是确定性的时候，它简化为状态上的映射：

\[s_{t+1} = F_\theta(s_t, a_t).\]

在实践中，状态 $s_t$ 通常是学习到的潜表示（例如，从像素编码而来），因此模型在一个（理论上）紧凑、可微的空间中运行。关键点在于，世界模型为你提供了一个可微模拟器；你可以在假设的动作序列下将其向前展开，并通过预测进行反向传播。

规划：通过模型优化来选择动作

给定起点 $s_0$ 和目标 $g$，最简单的规划器通过展开模型并最小化终端误差来选择动作序列 $\mathbf{a}=(a_0,\dots,a_{T-1})$：

\[\min_{\mathbf{a}} \; \| s_T(\mathbf{a}) - g \|_2^2, \quad \text{其中 } s_T(\mathbf{a}) = \mathcal{F}_{\theta}^{T}(s_0,\mathbf{a}).\]

这里我们用 $\mathcal{F}^T$ 作为通过世界模型进行完整展开的简写（对模型参数 $\theta$ 的依赖是隐含的）：

\[\mathcal{F}_{\theta}^{T}(s_0, \mathbf{a}) = F_\theta(F_\theta(\cdots F_\theta(s_0, a_0), \cdots, a_{T-2}), a_{T-1}).\]

在短时域和低维系统中，这可以相当好地工作。但随着时域增长、模型变得更大且更具表现力，其弱点会被放大。

那么，为什么这在大规模下就不奏效呢？

为什么长时域规划很难（即使一切都是可微的）

对于更通用的世界模型，有两个独立的痛点，再加上第三个是学习型、基于深度学习的模型所特有的。

1）长时域展开会创建深层、病态的计算图

熟悉时间反向传播（BPTT）的人可能会注意到，我们正在通过一个反复应用于自身的模型进行微分，这会导致梯度爆炸/消失问题。也就是说，如果我们对较早的动作（例如 $a_0$）求导数（注意我们是在对向量值函数求导，得到雅可比矩阵，记为 $D_x (\cdots)$）：

\[D_{a_0} \mathcal{F}_{\theta}^{T}(s_0, \mathbf{a}) = \Bigl(\prod_{t=1}^T D_s F_\theta(s_t, a_t)\Bigr) D_{a_0}F_\theta(s_0, a_0).\]

我们看到雅可比矩阵的条件数随时间 $T$ 呈指数级缩放：

\[\sigma_{\text{max/min}}(D_{a_0}\mathcal{F}_{\theta}^{T}) \sim \sigma_{\text{max/min}}(D_s F_\theta)^{T-1},\]

导致梯度爆炸或消失。

2）地形是非贪婪的，且充满陷阱

在短时域下，贪婪解——即每一步都径直朝目标移动——通常已经足够好。如果你只需要规划前方几步，最优轨迹通常与每一步“朝 $g$ 前进”相差不大。

随着时域增长，会发生两件事。首先，更长的任务更可能需要非贪婪行为：绕墙走、推动前重新定位、后退以走更好的路径。而且随着时域增长，通常需要更多这样的非贪婪步骤。其次，优化空间本身随时域缩放：$\mathrm{dim}(\mathcal{A} \times \cdots \times \mathcal{A}) = T\mathrm{dim}(\mathcal{A})$，进一步扩大了优化问题的局部极小值空间。

长时域的修复方法：解除动力学约束

假设我们将动力学约束 $s_{t+1} = F_{\theta}(s_t, a_t)$ 视为软约束，转而同时对动作 $(a_0,\ldots,a_{T-1})$ 和状态 $(s_0,\ldots,s_T)$ 优化以下惩罚函数：

\[\min_{\mathbf{s},\mathbf{a}} \mathcal{L}(\mathbf{s}, \mathbf{a}) = \sum_{t=0}^{T-1} \big\|F_\theta(s_t,a_t) - s_{t+1}\big\|_2^2, \quad \text{with } s_0 \text{ fixed and } s_T=g.\]

这在规划/机器人文献中有时也被称为配点法（collocation）。注意，提升后的形式与原始展开目标共享相同的全局极小值（当轨迹动力学可行时，两者都恰好为零）。但优化地形非常不同，我们立即获得两个好处：

• 每个世界模型评估 $F_{\theta}(s_t,a_t)$ 仅依赖于局部变量，因此所有 $T$ 项可以跨时间并行计算，从而在更长的时域下带来巨大的加速，并且
• 你不再通过单个深度 $T$ 步复合来反向传播以获取学习信号，因为之前的雅可比矩阵乘积现在分解为求和，例如：

\[D_{a_0} \mathcal{L} = 2(F_\theta(s_0, a_0) - s_1).\]

能够直接优化状态也有助于探索，因为我们可以暂时穿越非物理域来找到最优方案：

然而，天下没有免费的午餐。事实上，尤其是对于基于深度学习的世界模型，存在一个关键问题，使得上述优化在实践中相当困难。

基于深度学习的世界模型的一个问题：状态-输入梯度的敏感性

本节的tl;dr是：通过基于深度学习的 $F_{\theta}$ 直接优化状态极其脆弱，类似于对抗鲁棒性。即使你在低维状态空间中训练世界模型，世界模型的训练过程也会使未见过的状态景观变得非常尖锐，无论是未见过的状态本身，还是仅仅是与数据流形正交的方向。

对抗鲁棒性与“酒窝流形”模型

对抗鲁棒性最初研究分类模型 $f_\theta : \mathbb{R}^{w\times h \times c} \to \mathbb{R}^K$，并表明通过从基础图像 $x$（不属于类别 $k$）沿特定 logit 的梯度 $\nabla f_\theta^k$ 移动，你无需沿 $x’ = x + \epsilon\nabla f_\theta^k$ 移动很远，就能使 $f_\theta$ 将 $x’$ 分类为 $k$（Szegedy et al., 2014；Goodfellow et al., 2015）：

后来的工作为所发生的情况描绘了一幅几何图景：对于靠近低维流形 $\mathcal{M}$ 的数据，训练过程控制切向方向的行为，但不正则化正交方向的行为，从而导致敏感行为（Stutz et al., 2019）。另一种说法是：当仅考虑数据流形 $\mathcal{M}$ 的切向方向时，$f_\theta$ 具有合理的 Lipschitz 常数，但在法向方向上可能具有非常高的 Lipschitz 常数。事实上，在这些法向方向上更尖锐通常对模型有益，这样它可以更精确地拟合更复杂的函数。

因此，即使对于单个给定模型，此类对抗样本也极为常见。此外，这不仅仅是计算机视觉现象；对抗样本也出现在 LLM（Wallace et al., 2019）和 RL（Gleave et al., 2019）中。

虽然有方法可以训练出更具对抗鲁棒性的模型，但模型性能与对抗鲁棒性之间存在已知的权衡（Tsipras et al., 2019）：尤其是在存在许多弱相关变量的情况下，模型必须更尖锐才能实现更高性能。事实上，大多数现代训练算法，无论是在计算机视觉还是 LLM 中，都不会训练出对抗鲁棒性。因此，至少在深度学习发生重大范式转变之前，这是我们无法摆脱的问题。

为什么对抗鲁棒性对世界模型规划是个问题？

考虑我们在提升状态方法中优化的动力学损失的一个分量：

\[\min_{s_t, a_t, s_{t+1}} \|F_\theta(s_t, a_t) - s_{t+1}\|_2^2\]

让我们进一步只关注基础状态：

\[\min_{s_t} \|F_\theta(s_t, a_t) - s_{t+1}\|_2^2.\]

由于世界模型通常是在状态/动作轨迹 $(s_1, a_1, s_2, a_2, \ldots)$ 上训练的，$F_{\theta}$ 的状态数据流形维度受动作空间限制：

\[\mathrm{dim}(\mathcal{M}_s) \le \mathrm{dim}(\mathcal{A}) + 1 + \mathrm{dim}(\mathcal{R}),\]

其中 $\mathcal{R}$ 是某些可选的增强空间（例如平移/旋转）。因此，我们通常可以预期 $\mathrm{dim}(\mathcal{M}_s)$ 远低于 $\mathrm{dim}(\mathcal{S})$，从而：很容易找到对抗样本，将任何状态攻击到任何其他期望状态。

结果，动力学优化

\[\sum_{t=0}^{T-1} \big\|F_\theta(s_t,a_t) - s_{t+1}\big\|_2^2\]

感觉非常“粘滞”，因为基础点 $s_t$ 可以轻易欺骗 $F_{\theta}$ 认为它已经达到了局部目标。¹

1. 这种对抗鲁棒性问题，虽然对提升状态方法尤其严重，但并非它们独有。即使对于通过完整展开映射 $\mathcal{F}^T$ 进行优化的序列优化方法，也可能进入未见状态，在这些状态中，很容易将正常分量输入到 $D_s F_{\theta}$ 的敏感正常分量中。动作雅可比的链式法则展开为

\[\Bigl(\prod_{t=1}^T D_s F_\theta(s_t, a_t)\Bigr) D_{a_0}F_\theta(s_0, a_0).\]

看看如果乘积的任何阶段有任何分量垂直于数据流形会发生什么。↩

我们的修复

这就是我们新的规划器 GRASP 的用武之地。主要观察：虽然 $D_s F_{\theta}$ 不可信且具有对抗性，但动作空间通常是低维且经过充分训练的，因此 $D_a F_{\theta}$ 实际上可以通过优化进行合理利用，并且不会遭受对抗鲁棒性问题！

其核心是，GRASP 构建了一个一阶提升状态/配点规划器，仅通过世界模型依赖于动作雅可比。 因此，我们利用了学习到的世界模型 $F_{\theta}$ 的可微性，同时不成为状态雅可比 $D_s F_{\theta}$ 固有敏感性的受害者。

GRASP：梯度松弛随机规划器

如前所述，我们从配点规划目标开始，其中我们将状态提升并将动力学松弛为惩罚项：

\[\min_{\mathbf{s},\mathbf{a}} \mathcal{L}(\mathbf{s}, \mathbf{a}) = \sum_{t=0}^{T-1} \big\|F_\theta(s_t,a_t) - s_{t+1}\big\|_2^2, \quad \text{with } s_0 \text{ fixed and } s_T=g.\]

然后我们做了两个关键添加。

成分 1：通过对状态迭代加噪进行探索

即使目标更平滑，规划也是非凸的。我们通过在优化过程中向虚拟状态更新注入高斯噪声来引入探索。

一个简单版本：

\[s_t \leftarrow s_t - \eta_s \nabla_{s_t}\mathcal{L} + \sigma_{\text{state}} \xi, \qquad \xi\sim\mathcal{N}(0,I).\]

动作仍然通过非随机下降更新：

\[a_t \leftarrow a_t - \eta_a \nabla_{a_t}\mathcal{L}.\]

状态噪声帮助你在提升空间中在盆地之间“跳跃”，而动作仍然由梯度引导。我们发现，在这里专门对状态（而非动作）加噪声，能在探索和找到更尖锐极小值的能力之间取得良好平衡。²

2. 因为我们只对状态（而非动作）加噪声，相应的动力学并非真正的朗之万动力学。↩

要素 2：重塑梯度：停止脆弱的状态输入梯度，保留动作梯度

如前所述，脆弱的路径是流入世界模型状态输入的梯度，即 \(D_s F_{\theta}\)。最初最直接的做法就是直接停止流入 \(F_{\theta}\) 的状态梯度：

• 设 $\bar{s}_t$ 与 $s_t$ 取值相同，但梯度被停止。

定义停止梯度动力学损失：

\[\mathcal{L}_{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) = \sum_{t=0}^{T-1} \big\|F_\theta(\bar{s}_t, a_t) - s_{t+1}\big\|_2^2.\]

仅此还不够。注意现在状态只跟随前一个状态的步进，没有任何东西迫使基础状态去追赶下一个状态。结果就是存在平凡的极小值：直接停在原点，然后只让最后一个动作试图一步到达目标。

稠密目标塑形

我们可以把上述问题看作目标的信号被完全从先前状态中切断。一种修复方法是简单地在整个预测过程中加入稠密目标项：

\[\mathcal{L}_{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a}) = \sum_{t=0}^{T-1} \big\|F_\theta(\bar{s}_t, a_t) - g\big\|_2^2.\]

在常规设置中，这会过度偏向于直接追逐目标的贪心解，但在我们的设置中，这被停止梯度动力学损失对可行动力学的偏好所平衡。最终目标如下：

\[\mathcal{L}(\mathbf{s},\mathbf{a}) = \mathcal{L}_{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) + \gamma \, \mathcal{L}_{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a}).\]

结果是一个不依赖状态梯度的规划优化目标。

周期性“同步”：短暂回到真实展开梯度

提升的停止梯度目标非常适合快速、有引导的探索，但它仍然是原始串行展开目标的一个近似。

因此，每隔 $K_{\text{sync}}$ 次迭代，GRASP 会进行一个短暂的细化阶段：

1. 使用当前动作 $\mathbf{a}$ 从 $s_0$ 展开，并在原始串行损失上执行几次小的梯度步：

\[\mathbf{a} \leftarrow \mathbf{a} - \eta_{\text{sync}}\,\nabla_{\mathbf{a}}\,\|s_T(\mathbf{a})-g\|_2^2.\]

提升状态优化仍然提供优化的核心，而这个细化步骤提供一些辅助，使状态和动作保持贴近真实轨迹。这个细化步骤当然可以替换为你选择的串行规划器（例如 CEM）；核心思想是仍然获得串行规划器全路径同步的部分好处，同时仍然主要利用提升状态规划的优势。

GRASP 如何解决长程规划

基于配点法的规划器为长时程规划提供了一种自然的解决方案，但通过现代世界模型进行这种优化相当困难，原因在于对抗鲁棒性问题。GRASP 提出了一种简单的解决方案，以实现更平滑的基于配点法的规划器，同时为探索提供稳定的随机性。因此，长时程规划最终不仅能取得更多成功，还能更快地找到这些成功：

Push-T 结果。成功率（%）/ 成功所需时间中位数。粗体 = 行内最佳。注意，成功率越高，成功时间中位数会偏高；尽管成功率更高，GRASP 仍能更快完成。

下一步是什么？

现代世界模型规划器仍有大量工作要做。我们希望利用学习到的世界模型的梯度结构，而配点法（提升状态优化）是长时程规划的一种自然方法，但理解这里典型的梯度结构至关重要：平滑且信息丰富的动作梯度和脆弱的状态梯度。我们将 GRASP 视为此类规划器的一个初步迭代。

扩展到基于扩散的世界模型（更深的潜在时间步可视为世界模型本身的平滑版本）、更复杂的优化器和加噪策略，以及将 GRASP 集成到闭环系统或 RL 策略学习中以实现自适应长时程规划，都是自然且有趣的下一步。

我确实认为，现在正是研究世界模型规划器的激动人心的时刻。这是一个有趣的甜蜜点：背景文献（整体规划与控制）已经极其成熟和完善，但当前设定（对现代大规模世界模型进行纯规划优化）仍然严重缺乏探索。但是，一旦我们弄清了所有正确的思路，世界模型规划器很可能会像 RL 一样普及。

更多详情，请阅读完整论文或访问项目网站。

引用

@article{psenka2026grasp,
  title={Parallel Stochastic Gradient-Based Planning for World Models},
  author={Michael Psenka and Michael Rabbat and Aditi Krishnapriyan and Yann LeCun and Amir Bar},
  year={2026},
  eprint={2602.00475},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2602.00475}
}