[Paper] 缩小基于梯度规划的世界模型的训练-测试差距

Source: arXiv - 2512.09929v1

概览

基于世界模型的强化学习承诺在大规模专家数据集上离线训练单一动力学预测器，然后在许多下游规划任务中重复使用它。虽然基于梯度的规划器速度快，但历史上它们的表现不及经典的无梯度方法，如交叉熵方法（CEM）。本文指出了原因——模型的训练方式（下一状态预测）与测试时的使用方式（动作序列优化）之间存在不匹配——并提出了具体的修正措施，弥合这一差距，实现了既更快又在精度上具有竞争力的基于梯度的规划。

关键贡献

• 训练‑测试差距分析： 形式化了下一状态预测训练目标与推理时动作优化使用场景之间的差异。
• 数据合成技巧： 引入了简单、模型无关的增强（例如想象滚动、动作扰动采样），让世界模型接触到规划过程中会看到的轨迹类型。
• 改进的基于梯度的规划器： 通过合成数据，标准的动作序列梯度下降能够匹配或超越 CEM 的性能，同时仅使用约 10 % 的计算预算。
• 广泛的实证验证： 在多样的操作（如块堆叠）和导航（如迷宫）环境上进行基准测试，展示了跨任务的一致提升。
• 开源实现： 提供代码和预训练模型，降低实践者采用该技术的门槛。

方法论

1. 基线世界模型： 在大量专家轨迹上训练神经动力学模型 (f_\theta(s_t, a_t) \rightarrow s_{t+1})，使用常规的均方误差进行下一状态预测。
2. 识别差距： 推理时，规划器将模型视为可微分模拟器，并对定义在未来奖励上的损失进行反向传播，以更新候选动作序列 (\mathbf{a}_{0:H})。然而模型从未接触过由其自身不完美预测产生的状态分布。
3. 训练时数据合成：
   • 想象滚动： 从真实状态出发，使用随机采样的动作让当前模型前向滚动，生成合成轨迹。
   • 动作扰动回放： 对专家动作加入噪声并重新模拟，促使模型对离策略动作具有鲁棒性。
   • 奖励感知采样： 根据估计的未来奖励对合成样本加权，使模型倾向于后续规划时会探索的区域。
4. 联合训练： 将真实专家数据与合成样本混合继续训练动力学模型。无需额外的损失项；对两类数据均使用相同的下一状态预测目标。
5. 基于梯度的规划： 推理时，随机初始化一个动作序列，使用训练好的世界模型计算预测轨迹，评估任务特定奖励，并将奖励梯度反向传播以细化动作（例如使用 Adam）。

结果与发现

• 通过合成数据训练的模型弥合了性能差距：基于梯度的规划器现在在所有测试任务上匹配或略微超越 CEM。
• 计算节省显著——梯度下降在约 10 % 的迭代次数内收敛，相比 CEM 大幅降低延迟和能耗。
• 消融实验表明，每个合成组件（想象滚动、扰动、奖励感知采样）均有正向贡献；去除任意一项都会导致成功率和速度下降。

实际意义

• 更快的在线规划： 机器人或智能体可以在毫秒级别重新规划，而不是秒级，从而在操作（如移动输送带上的抓取‑放置）和自主导航（如无人机避障）中实现实时响应。
• 硬件需求降低： 基于梯度的规划器仅依赖简单的反向传播，可在普通 GPU 或甚至端侧加速器上高效运行，区别于 CEM 的大规模并行采样。
• 简化流水线： 开发者可以保持单一的世界模型训练循环，并在许多下游任务中直接复用同一模型，无需为每个任务重新训练策略。
• 可扩展到大规模数据集： 该方法仅增加了低成本的合成滚动，因而能够很好地适应大规模离线数据（例如自动驾驶车队的日志）。
• 混合系统的潜力： 该方法可与模型自由的微调结合，形成“取长补短”的系统——世界模型提供强先验，基于梯度的规划实现快速适应。

局限性与未来工作

• 模型偏差仍然存在： 该方法缓解但未消除长时滚动中的误差累积；极深的规划视野仍可能受影响。
• 任务特定奖励设计： 基于梯度的规划仍需可微分的奖励信号；为复杂稀疏任务设计此类奖励并非易事。
• 仅限确定性动力学： 当前形式假设世界模型是确定性的；将其扩展到随机或部分可观测环境仍是开放挑战。
• 未来方向：
  • 引入不确定性估计（如集成）以指导合成过程。
  • 探索逐步增加滚动长度的课程式合成。
  • 在更高维、感知丰富的领域（如视觉驱动的操作）进行测试，因状态估计会带来额外难度。

结论： 通过使训练数据分布与世界模型在推理时的实际使用方式保持一致，本文实现了基于梯度的规划在速度优势上的释放，同时不牺牲性能——这是迈向更敏捷、数据高效自主系统的重要一步。

作者

• Arjun Parthasarathy
• Nimit Kalra
• Rohun Agrawal
• Yann LeCun
• Oumayma Bounou
• Pavel Izmailov
• Micah Goldblum

论文信息

• arXiv ID: 2512.09929v1
• 分类: cs.LG, cs.RO
• 发表时间: 2025 年 12 月 10 日
• PDF: Download PDF