通过扩散策略优化扩展世界模型强化学习

# 通过扩散策略优化实现世界模型强化学习的规模化 来源：https://arxiv.org/html/2605.26282

Xiaoyuan Cheng\(^1\), Wenxuan Yuan\(^{2,*}\), Zhancun Mu\(^3\), Yuanzhao Zhang\(^4\), Yiming Yang\(^1\), Hai Wang\(^1\), Zhuo Sun\(^{5,\dagger}\), Che Liu\(^{6,\dagger}\)

\(^1\)Dynamic Systems Lab, University College London
\(^2\)College of Computing and Data Science, Nanyang Technological University
\(^3\)School of Intelligence Science and Technology, Peking University
\(^4\)Santa Fe Institute
\(^5\)School of Statistics and Data Science, Shanghai University of Finance and Economics
\(^6\)Department of Computing, Imperial College London

代码：https://github.com/Edmond1Cheng/MBDPO
Hugging Face：https://huggingface.co/BruceYuan/MBDPO

核心贡献者：Xiaoyuan Cheng: [email protected] (https://arxiv.org/html/2605.26282v1/[email protected]) 和 Wenxuan Yuan: [email protected] (https://arxiv.org/html/2605.26282v1/[email protected])。

†通讯作者：Zhuo Sun: [email protected] (https://arxiv.org/html/2605.26282v1/[email protected]), Che Liu: [email protected] (https://arxiv.org/html/2605.26282v1/[email protected])。

###### 摘要

基于模型的强化学习（RL）可以通过使用世界模型在大规模场景下得到有效支持。然而，在实践中，这种方法的规模化仍然受到根本限制。一个普遍公认的挑战是模型偏差和误差累积，这会降低长期预测的质量。除了这些问题之外，我们识别出一个更关键但尚未充分探索的瓶颈：现有世界模型方法中搜索与价值学习之间的结构性错位。具体而言，策略改进通常依赖于由独立的非搜索策略推导出的价值函数，导致训练不一致，最终产生次优学习。为了解决这一局限性，我们提出了基于模型的扩散策略优化（MBDPO）在世界模型中，这是一个通过扩散策略表示统一搜索和策略优化的框架，从而释放世界模型在可扩展策略学习中的潜力。我们不构建基于学习到世界模型的显式规划器，而是将策略优化重新表述为潜在世界模型中搜索轨迹上的扩散过程。在此观点下，我们从收集的数据集中提取一个隐式能量函数，该函数锚定策略，使MBDPO能够为策略优化修正分数场，同时减轻错位问题。我们在广泛的环境下评估MBDPO，包括多任务离线预训练、在线学习以及离线到在线微调。在离线场景中，我们进一步通过在大规模数据集上进行预训练来研究其规模化行为，观察到随着模型容量增加，性能持续单调提升。

参见图注

图1：离线与在线性能概览。（左）我们的方法（MBDPO）在多任务离线预训练中显著优于TD-MPC2（Hansen et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib11)），当模型参数从1.7M增加到340M时，表现出清晰的单调规模化行为。（右）在从头开始的在线设置中，MBDPO在4个基准测试、121个任务上持续获得优于或具有竞争力的结果。

## 1 引言

基于模型的强化学习（RL）长期以来被认为是通用情况下进行样本高效规划的有前途范式。通过学习一个显式动力学系统，智能体可以在不与真实环境进行多次交互的情况下进行“想象”和规划（Sutton, 1996 (https://arxiv.org/html/2605.26282#bib.bib1); Sutton et al., 1998 (https://arxiv.org/html/2605.26282#bib.bib2)）。最新进展进一步扩展了这一范式（Garrido et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib3); Zhou et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib4); Hafner et al., 2025a (https://arxiv.org/html/2605.26282#bib.bib5); Psenka et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib40)），使得大模型能够捕捉跨多样任务的丰富环境动力学。现代世界模型不再仅仅是局部帧预测器（S V et al., 2023 (https://arxiv.org/html/2605.26282#bib.bib6); Thrun et al., 1990 (https://arxiv.org/html/2605.26282#bib.bib7); Guan et al., 2023 (https://arxiv.org/html/2605.26282#bib.bib8)），而是旨在学习可泛化的环境动力学表示，这些表示可以跨下游控制问题重用。同时，大规模数据集的出现为世界模型的大规模预训练创造了新的机会（Wu et al., 2023 (https://arxiv.org/html/2605.26282#bib.bib9); Hafner et al., 2025b (https://arxiv.org/html/2605.26282#bib.bib10)）。类比于视觉和语言领域的基础模型，预训练世界模型有潜力释放多任务中可扩展的基于模型的RL（Hansen et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib11)）。

尽管有最新进展，但基于世界模型的基于模型的RL方法仍然难以显著超越其无模型对应方法（Van Hasselt et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib12); Fujimoto et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib13); Chang et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib14)）。已经提出了两种解释。最常见的解释将这种不足归因于有限的模型精度和累积预测误差（M. Moerland et al., 2023 (https://arxiv.org/html/2605.26282#bib.bib17); Lambert et al., 2022 (https://arxiv.org/html/2605.26282#bib.bib15)）。当使用学习到的模型模拟长程轨迹时，即使单步预测的小误差也会随时间累积，导致越来越不可靠的状态估计。基于这一观察，先前的工作集中于减轻不完美模型动力学导致的策略退化。代表性方法包括提高模型精度（Nagabandi et al., 2018 (https://arxiv.org/html/2605.26282#bib.bib18); Hafner et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib19)）、融入不确定性感知建模（Deisenroth and Rasmussen, 2011 (https://arxiv.org/html/2605.26282#bib.bib20); Curi et al., 2020 (https://arxiv.org/html/2605.26282#bib.bib21); Fu et al., 2022 (https://arxiv.org/html/2605.26282#bib.bib22)）、采用截断或分支的展开（Janner et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib16); Park et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib23)）以及增强长期预测（Ma et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib24); Farebrother et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib25)）。这些努力背后有一个隐含假设：足够精确的世界模型自然会转化为更好的规划，并最终带来更好的控制性能。然而，最近的工作挑战了这一观点（Chang et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib14)），认为模型精度本身并不能决定规划有效性。即使是非常精确的模型也可能无法产生强策略，导致世界模型内的RL在性能和可扩展性上都有困难。核心障碍在于策略搜索与价值学习之间的错位。由于搜索的计算成本高，价值函数通常使用由学习到的非搜索策略网络生成的展开轨迹进行训练（Hansen et al., 2022 (https://arxiv.org/html/2605.26282#bib.bib26), 2024 (https://arxiv.org/html/2605.26282#bib.bib11), 2025 (https://arxiv.org/html/2605.26282#bib.bib27); Wang et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib67); Zhan et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib28); Shimizu and Tomizuka, 2024 (https://arxiv.org/html/2605.26282#bib.bib53)）。相比之下，策略改进依赖于通过在学习到的世界模型上执行搜索过程（即模型预测路径积分（Williams et al., 2015 (https://arxiv.org/html/2605.26282#bib.bib54)））产生的轨迹。这种不匹配引入了分布差异：价值函数在由非搜索策略生成的“状态-动作”对上训练，而策略执行和改进则由通过世界模型搜索产生的轨迹主导。结果，价值函数常常在分布外的“状态-动作”对上被查询，此时其估计不可靠且通常过于乐观；搜索过程会利用这些误差，导致系统性高估偏差，最终损害性能和可扩展性（Fujimoto et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib29)）。

参见图注

图2：MBDPO的核心框架。目标策略分布通过一系列逐步转移核逐渐塑造，从πφ^N到πφ^0，通过多步精炼将高斯先验N(0,I)转换为最优吉布斯策略πφ^*。关键是，每个转移核πφ^(τ-1)(a_{t:t+H}^{τ-1}|a_{t:t+H}^{τ},z_t)由在学习的world model内完全估计的分数函数φ^控制：蒙特卡洛动作序列样本通过潜在动力学F展开，根据其能量正则化累积回报G̃进行评估，并重新加权以恢复指向更高回报动作的分数。通过这种方式，世界模型重新定义并修正了分数场，将标准生成式去噪过程转变为基于模型的策略优化。此外，隐式能量E∝-logβ将策略锚定在KL信任区域内的非搜索策略分布上，从而统一了搜索与策略优化。

受先前方法瓶颈的启发，我们研究如何克服这两个限制，释放基于模型的RL在世界模型中的可扩展性潜力。在这项工作中，我们引入了一个基于扩散模型的新视角，并重新思考如何使世界模型内的基于模型的RL具有可扩展性。从这个视角出发，我们证明世界模型能够实现直接策略优化，其中分数匹配通过想象的轨迹实现。我们的贡献有三点：（1）我们重新审视了策略优化的基础，并从值迭代和策略改进两个角度表明，搜索与价值学习之间的错位可能导致现有世界模型方法的次优性。（2）我们提出了基于模型的扩散策略优化（MBDPO），这是一个统一的框架，将搜索和策略优化表述为世界模型想象轨迹上的扩散过程，从而释放世界模型在可扩展策略学习中的潜力（见图2 (https://arxiv.org/html/2605.26282#S1.F2)）。在此公式中，缓冲区数据集诱导一个隐式能量函数，将策略更新锚定到行为分布。这种机制使得分数修正能够超越先前方法的次优性。通过这样做，MBDPO统一了搜索和策略优化，消除了它们的结构性错位。（3）我们开发了一个可扩展的扩散RL算法，与世界模型联合训练，并在多任务离线预训练、在线学习和离线到在线微调等多种设置下，展示了随着模型容量增加而持续、单调的性能改进。除了这些定量结果，我们还通过可视化受控的潜在轨迹来评估底层表示，发现学习到的策略始终表现出与物理直觉一致的结构化行为。

相关工作。我们回顾两个领域的前期工作：用于RL的世界模型和扩散策略。

世界模型。世界模型的概念源于Jay W. Forrester，他将其用于构建复杂系统的心理表征（Forrester, 1971 (https://arxiv.org/html/2605.26282#bib.bib51)）。后来被引入机器学习，使智能体能够在学习到的内部环境中学习和行动（Ha and Schmidhuber, 2018 (https://arxiv.org/html/2605.26282#bib.bib31)）。随后，Dreamer系列（Hafner et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib19), 2020 (https://arxiv.org/html/2605.26282#bib.bib33), 2025a (https://arxiv.org/html/2605.26282#bib.bib5), 2025b (https://arxiv.org/html/2605.26282#bib.bib10)）将潜在世界模型从概念验证转变为可扩展且稳健的RL范式。然而，由于模型偏差，完全在学习到的世界模型内训练策略可能导致分布偏移（Janner et al., 2019 (https://arxiv.org/html/2605.26282#bib.bib16)）。另一条研究路线，TD-MPC（Hansen et al., 2022 (https://arxiv.org/html/2605.26282#bib.bib26), 2024 (https://arxiv.org/html/2605.26282#bib.bib11), 2025 (https://arxiv.org/html/2605.26282#bib.bib27)），结合了基于模型和无模型RL。然而，搜索与价值之间的错位阻碍了其可扩展性和性能。最新方法，包括带有世界模型的视觉-语言-动作模型（Zhen et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib35); Cen et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib37); Team et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib39); Intelligence et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib62)）以及世界动作模型（WAMs）（Han et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib38); Ye et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib34); Li et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib36)），已被提出用于生成动作块和状态表示，这得益于大型预训练视觉模型的快速发展。与RL相比，WAM系列不执行显式策略优化，这限制了其实现长期控制和稳健泛化的能力。

扩散策略。扩散策略可以大致分为无模型和基于模型的方法。无模型扩散策略不依赖于学习到的动力学；相反，它们直接从固定数据集（Janner et al., 2022 (https://arxiv.org/html/2605.26282#bib.bib41); Chi et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib42)）或学习到的Q函数（Psenka et al., 2023 (https://arxiv.org/html/2605.26282#bib.bib43); Dong et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib44); Cheng et al., 2026 (https://arxiv.org/html/2605.26282#bib.bib45); Li and Levine, 2026 (https://arxiv.org/html/2605.26282#bib.bib52)）学习策略分布。相比之下，基于模型的扩散技术利用环境动力学来估计分数函数，而无需依赖离线数据（Pan et al., 2024 (https://arxiv.org/html/2605.26282#bib.bib46); Xue et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib47); Cheng et al., 2025 (https://arxiv.org/html/2605.26282#bib.bib48)）。关键的是，计算这个分数函数需要访问转移动力学，这使得世界模型成为理想的基座；它们提供了必要的模拟环境。然而，在许多实际场景中，真实动力学不可用，限制了它们的直接适用性。这个瓶颈反过来凸显了学习到的世界模型的关键必要性。通过近似多样化的动力学谱，世界模型是释放基于模型的扩散用于可扩展策略学习全部潜力的关键。尽管有巨大前景，但在学习到的世界模型内有效的基于模型的扩散策略仍然未被充分探索。

## 2 预备知识

基于模型的RL。我们考虑一个由元组 M = (S, A, T, r, ρ, γ)