长时间范围下的梯度规划世界模型
Berkeley AI Research前沿论文,探讨世界模型在更长时间维度上的规划方法。学术性强,多数程序员暂无直接应用场景。
Berkeley AI Research前沿论文,探讨世界模型在更长时间维度上的规划方法。学术性强,多数程序员暂无直接应用场景。
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(s_{t+1} \mid s_{t-h:t},; a_t)$。为简化讨论,在这篇博客文章中,我们将假设模型是马尔可夫模型 $P(s_{t+1} \mid s_{t-h:t},; a_t)$(这里的所有结果都可以推广到更一般的情形);当模型是确定性的时,它就简化为一个状态映射:
在实践中,状态 $s_t$ 通常是学习得到的潜在表示,例如由像素编码而来,因此模型运行在一个理论上紧凑且可微的空间中。关键在于,世界模型为你提供了一个可微模拟器:你可以在假设的动作序列下向前展开模型,并通过预测结果进行反向传播。
给定初始状态 $s_0$ 和目标 $g$,最简单的规划器会通过展开模型并最小化终点误差,选择动作序列 $\mathbf{a}=(a_0,\dots,a_{T-1})$:
这里,我们使用 $\mathcal{F}^T$ 作为穿过世界模型进行完整展开的简写,其中对模型参数 $\theta$ 的依赖是隐式的:
在短时域和低维系统中,这种方法的效果通常还不错。但随着时域增长,以及模型变得更大、表达能力更强,其弱点也会被进一步放大。
那么,为什么这种方法无法直接扩展到大规模场景?
对于更一般的世界模型,这里存在两个彼此独立的痛点,另外还有第三个特定于基于深度学习的学习型模型的问题。
熟悉随时间反向传播(BPTT)的人可能已经注意到,我们正在对一个反复作用于自身的模型求导,这会导致梯度爆炸或梯度消失问题。具体来说,如果我们对较早的动作(例如 $a_0$)求导(注意,我们是在对向量值函数求导,因此得到的是雅可比矩阵,这里记作 $D_x (\cdots)$):
可以看到,雅可比矩阵的条件性会随时间 $T$ 呈指数级变化:
从而导致梯度爆炸或梯度消失。
在短时域下,贪心解通常已经足够好,也就是每一步都直接朝目标移动。如果只需要向前规划几步,最优轨迹通常不会与每一步都“朝 $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)$ 优化以下惩罚函数:
在规划和机器人学文献中,这种方法有时也称为配点法。注意,提升后的形式与原始展开目标具有相同的全局极小值:只有当轨迹在动力学上可行时,两者才恰好为零。但两者的优化地形截然不同,而且我们会立即获得两个好处:
每次世界模型求值 $F_{\theta}(s_t,a_t)$ 都只依赖于局部变量,因此所有 $T$ 项都可以沿时间维度并行计算,从而在较长时域下获得巨大的速度提升;以及
为了获得学习信号,你不再需要穿过单个、由 $T$ 步复合而成的深层结构进行反向传播,因为之前的雅可比矩阵乘积现在被拆分成了求和,例如:
能够直接优化状态也有助于探索,因为我们可以暂时穿过不符合物理规律的区域,以找到最优规划:
然而,天下没有免费的午餐。事实上,尤其对于基于深度学习的世界模型而言,这里存在一个关键问题,使上述优化在实践中变得相当困难。
本节的简要结论是:通过基于深度学习的 $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)和强化学习(Gleave et al., 2019)中。
尽管存在训练更具对抗鲁棒性模型的方法,但模型性能与对抗鲁棒性之间存在已知的权衡(Tsipras et al., 2019):尤其是在存在大量弱相关变量的情况下,模型必须变得更加尖锐,才能获得更高的性能。事实上,无论是在计算机视觉领域还是 LLM 领域,大多数现代训练算法都不会通过训练消除对抗鲁棒性问题。因此,至少在深度学习发生重大范式转变之前,这是一个我们不得不面对的问题。
为什么对抗鲁棒性会成为世界模型规划中的问题?
考虑我们在升起的状态方法中优化的动力学损失的单个分量:
进一步关注基础状态:
由于世界模型通常是在状态/动作轨迹 $(s_1, a_1, s_2, a_2, \ldots)$ 上训练的,$F_{\theta}$ 的状态数据流形的维数受动作空间限制:
其中 $\mathcal{R}$ 是某个可选的增强空间(例如平移/旋转)。因此,我们通常可以预期 $\mathrm{dim}(\mathcal{M}_s)$ 远低于 $\mathrm{dim}(\mathcal{S})$,因此:很容易找到对抗性例子,将任何状态欺骗为任何其他所需的状态。
因此,动力学优化
感觉非常"粘滞",因为基础点 $s_t$ 可以轻松骗过 $F_{\theta}$,使其认为已经达成了局部目标。1
看看乘积的任何阶段是否有与数据流形正交的分量会发生什么。↩
这就是我们的新规划器 GRASP 出场的地方。主要观察是:虽然 $D_s F_{\theta}$ 不可信且具有对抗性,但动作空间通常是低维且经过充分训练的,所以 $D_a F_{\theta}$ 实际上是合理的优化对象,不会受到对抗鲁棒性问题的困扰!
GRASP 的核心是建立一个仅依赖于世界模型通过动作雅可比矩阵的一阶升起状态/配置法规划器。我们因此利用了学习世界模型 $F_{\theta}$ 的可微性,同时避免了陷入状态雅可比矩阵 $D_s F_{\theta}$ 的内在敏感性。
如前所述,我们从配置法规划目标开始,其中升起状态并将动力学松弛为一个惩罚:
然后我们做两个关键的补充。
要点 1:通过噪声注入进行探索
即使目标函数更平滑,规划也仍然是非凸的。我们通过在优化过程中将高斯噪声注入虚拟状态更新来引入探索。
动作仍然通过非随机下降更新:
状态噪声帮助你在升起的空间中"跳跃"不同的盆地,而动作由梯度引导。我们发现,特别是在这里噪声化状态(而不是动作)找到了探索和寻找尖锐极小值能力之间的良好平衡。2
要点 2:重塑梯度:停止脆弱的状态输入梯度,保留动作梯度
如前所述,脆弱的路径是流向世界模型状态输入的梯度 $D_s F_{\theta}$。最直接的初始做法是简单地停止流向 $F_{\theta}$ 的状态梯度:
令 $\bar{s}_t$ 是与 $s_t$ 相同的值,但梯度已停止。
定义停止梯度动力学损失:
这单独是不够的。注意现在状态只遵循前一个状态的步骤,没有任何东西强制基础状态追逐下一个状态。因此,存在微不足道的极小值,只是停在原点,然后仅最后的动作尝试在一步内到达目标。
我们可以将上述问题视为目标的信号从先前的状态中完全被切断。一种修复方法是在整个预测中简单地添加一个密集的目标项:
在正常情况下,这会过度偏向于直接追逐目标的贪心解决方案,但在我们的设置中,这通过停止梯度动力学损失对可行动力学的偏向得到平衡。最终目标如下:
其结果是一个不依赖于状态梯度的规划优化目标。
周期性"同步":短暂回到真实推演梯度
升起的停止梯度目标对于快速、引导性的探索很不错,但它仍然是原始串联推演目标的近似。
因此,每 $K_{\text{sync}}$ 次迭代,GRASP 进行一个短的精化阶段:
从 $s_0$ 使用当前动作 $\mathbf{a}$ 推演,并在原始串联损失上采取一些小的梯度步:
升起状态的优化仍然提供优化的核心,同时这个精化步骤增加了一些帮助,保持状态和动作朝向真实轨迹。这个精化步骤当然可以被你选择的串联规划器替换(例如 CEM);核心思想仍然是获得串联规划器的完整路径同步的一些好处,同时仍然主要使用升起状态规划的好处。
配置法规划器为长期规划提供了自然的解决方案,但由于现代世界模型中存在的对抗鲁棒性问题,这种优化相当困难。GRASP 为更平滑的配置法规划器提出了一个简单的解决方案,同时具有稳定的随机性用于探索。因此,较长地平线规划不仅最终成功次数更多,而且找到这样的成功也更快:
Push-T 结果。成功率 (%) / 成功的中位时间。粗体 = 行中最佳。注意中位成功时间会因较高的成功率而偏高;GRASP 设法在更高的成功率下仍然更快。
对于现代世界模型规划器,仍然有大量工作要做。我们希望利用学习世界模型的梯度结构,配置法(升起状态优化)对于长期规划是一个自然方法,但理解这里的典型梯度结构至关重要:平滑且信息丰富的动作梯度和脆弱的状态梯度。我们将 GRASP 视为这类规划器的初步迭代。
扩展到基于扩散的世界模型(更深的潜在时步可以被视为世界模型本身的平滑版本)、更复杂的优化器和噪声策略,以及将 GRASP 集成到闭环系统或用于适应性长期规划的强化学习策略学习中,都是自然而有趣的后续步骤。
我确实认为现在从事世界模型规划器的工作是一个激动人心的时刻。这是一个有趣的甜蜜点,背景文献(规划和控制总体)非常成熟和发达,但当前的设置(纯规划优化超过现代、大规模世界模型)仍然严重未被充分探索。但是,一旦我们理清所有正确的想法,世界模型规划器可能会变得像强化学习一样普遍。
更多详情,请阅读完整论文或访问项目网站。
@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}
}