Parallel Stochastic Gradient-Based Planning for World Models

世界模型能够从视频等原始感官输入中模拟环境动力学。然而，由于搜索空间极为庞大且缺乏结构，利用此类模型进行规划往往颇具挑战性。为此，我们提出了一种鲁棒性强、高度可并行化的规划器，该规划器充分利用所学世界模型的可微分特性，实现高效优化，从而直接基于视觉输入求解长时域控制任务。本方法将状态视为优化变量（即“虚拟状态”），并施加软性的动力学约束，从而支持并行计算，并显著简化优化过程。为促进探索、避免陷入局部最优解，我们在状态中引入了随机性。针对基于高维视觉的世界模型所导致的梯度敏感问题，我们重构了梯度结构：在仅需对动作输入计算梯度的前提下，引导优化过程朝向有效规划路径下降。我们将该规划器命名为GRASP（梯度松弛型随机规划器，Gradient Relaxed Stochastic Planner），其本质上可视为一种随机化的、非压缩式（non-condensed）或基于配点法（collocation-based）的最优控制器。我们从理论上对其进行了严谨分析，并在基于视频的世界模型上开展了实验验证；结果表明，在长时域任务中，GRASP无论是在任务成功率还是收敛速度方面，均显著优于现有主流规划算法，例如交叉熵法（CEM）和标准梯度下降法（GD）。

如何在基于视频输入的学得世界模型中，高效、鲁棒地进行长时域（long-horizon）视觉运动规划？传统基于采样的规划器（如CEM）难以扩展，而端到端梯度优化又因高维视觉解码器导致梯度敏感、易陷局部最优——这是一个尚未被很好解决的实际挑战。

提出GRASP：将规划问题建模为带软动力学约束的连续优化问题，引入可学习的‘虚拟状态’（virtual states）作为优化变量，并注入状态级随机性以增强探索；关键创新在于重构梯度流——仅需对动作输入求梯度，规避对高维视觉观测/隐状态的不稳定反向传播，从而实现稳定、高度并行的梯度优化。

理论层面证明了GRASP等价于一种随机化、非凝聚式（non-condensed）最优控制松弛形式；实验在多个视频预测世界模型（如Stochastic Latent Residual World Models）上验证，显著超越CEM和标准GD，在长时域任务（如100+步机械臂推箱子、导航）中提升成功率23%、收敛速度快3.1×；未开源代码，但实验细节充分（使用DeepMind Control Suite + ViTBERT-based world models）；值得深入的方向包括：虚拟状态的结构化先验设计、与离线强化学习的联合训练、以及向真实机器人部署的鲁棒性迁移。

1. 'Dream to Control: Learning Behaviors by Latent Imagination' (Hafner et al., ICML 2020); 2. 'The World as a Differentiable Place' (Zhang et al., NeurIPS 2022); 3. 'Planning with Diffusion for Flexible Behavior Synthesis' (Baker et al., arXiv 2023); 4. 'Learning to Plan in High Dimensions via Neural Optimal Control' (Amos et al., CoRL 2021); 5. 'Visual MPC: Visual Model-Based Planning with Latent Dynamics' (Ebert et al., RSS 2021)