基于世界模型的Q学习

# 使用世界模型的Q学习
来源：https://arxiv.org/html/2608.17163  
作者：贾悦如（单位：北京大学）  
切尔西·芬恩（单位：斯坦福大学）  
多尔萨·萨迪格（单位：斯坦福大学）

#### 摘要  
离策略强化学习（RL）的样本效率已显著提升，使得诸如将视觉-语言-动作模型通过RL微调为可靠高效策略等应用成为可能。世界模型进一步提升了样本效率，因其能预测状态变化而非仅预测动作，但其成功主要局限于监督式策略学习。传统的基于模型的RL方法常直接在想象的轨迹上优化策略或价值函数，这容易导致偏差累积，并且难以扩展到大规模高维问题（如真实世界机器人学），随着任务时间跨度和视觉复杂度增加，这一问题会加剧。在本文中，我们研究能否直接在标准Q学习之上利用世界模型来提升性能，同时确保训练和推理基于真实在线环境。我们提出**QWM**框架，该框架在Q学习基础上利用世界模型，在想象轨迹上进行测试时搜索，以在在线推演和评估阶段选择高价值动作。由于策略和价值函数仅基于真实状态转移进行训练，QWM避免了模型偏差累积，同时获得了预测性搜索带来的样本效率增益。在具有挑战性的操作基准测试Robomimic和LIBERO上，QWM在样本效率和性能上均显著优于现有的先进方法。

**脚注**：贡献相等。  
**††脚注**：通讯作者：[email protected]  

## 1 引言  
离策略强化学习（RL）的最新进展显著提高了其样本效率，使得诸如将视觉-语言-动作（VLA）模型通过RL微调为高度可靠、高性能的策略等应用成为可能（28 (https://arxiv.org/html/2608.17163#bib.bib49);5 (https://arxiv.org/html/2608.17163#bib.bib58)）。随着RL应用于日益复杂的场景，进一步提高样本效率变得极其宝贵，甚至必不可少。世界模型提供了一种弥合这一差距的自然杠杆，因为它们能够预测状态的变化，而不仅仅是动作。然而到目前为止，它们的成功主要局限于监督式策略学习。这就引出一个自然的问题：是否存在一种简单的方法来利用世界模型提升在线RL的性能，而无需采用传统基于模型的RL方法，后者容易将世界模型的偏差累积到训练过程中？  

先前关于基于模型的RL的工作探索了学习动力学模型，并直接在其中优化策略或价值函数（24 (https://arxiv.org/html/2608.17163#bib.bib66);22 (https://arxiv.org/html/2608.17163#bib.bib9)），但将这种方法扩展到大规模高维问题（如真实世界机器人学）仍然困难：主要在不完美世界模型内推演训练的策略会继承并累积其偏差，且随着任务时间跨度和视觉复杂度增加，这一问题会加剧。我们并非将世界模型作为真实交互的替代品，而是探究它是否能反过来提升一个仍然基于真实在线环境训练和校准的RL智能体。在本文中，我们研究结合世界模型的Q学习，其中模型并非环境交互的替代品，而是用于提升性能的一种方式。我们的核心洞察是，学习到的世界模型可以应用于标准Q学习之上，通过想象未来轨迹实现测试时动作搜索，并且这种改进的动作选择本身就能直接转化为Q学习方法性能的提升。  

我们提出**基于世界模型的Q学习（QWM）**，这是一个利用世界模型在标准Q学习之上直接进行测试时搜索的框架。与基于模型的RL中通常使用状态价值函数进行搜索不同，学习到的Q函数使得更强大的搜索成为可能，可以在候选动作执行前评估它们并选择价值最高的一个。这种测试时的动作搜索既应用于在线推演（以获取更好的在线RL数据），也应用于评估阶段（使学习到的策略行动更有效）。QWM利用世界模型的想象力为下游性能选择最佳动作，而不是从策略中采样单个动作，或采样多个动作并选择Q值最高的一个，同时避免了在学习模型内训练策略所带来的偏差累积和不稳定性。  

我们的主要贡献是**QWM**，一个简单通用的框架，利用学习到的世界模型通过测试时搜索来改进Q学习。由于底层策略和价值函数仍然是针对真实环境转移在线训练的，QWM避免了主要依赖想象推演的方法所面临的偏差累积问题，同时仍能在决策时获得此类预测模型带来的样本效率增益。我们在具有挑战性的操作基准测试Robomimic和LIBERO上评估QWM，涵盖基于状态和视觉观察的设置，发现它在样本效率和性能上均显著优于现有强方法。

见标题  
图1：QWM概览。左：世界模型树搜索。从当前状态出发，策略提出候选动作，世界模型想象未来状态。每个初始动作的值通过聚合中间节点和叶节点的价值来计算。右：在线动作选择。在每个环境步骤中，QWM使用树搜索分数从策略提出的动作中进行选择，然后在环境中执行选中的动作。树搜索既用于在线采样也用于评估，而策略和评论家仅基于真实环境转移进行训练。

## 2 相关工作  
**使用先验数据的强化学习**。为了提高在线RL的性能和样本效率，先前工作研究了使用离线数据集加速在线学习的问题。流行方法涉及平衡探索与利用（61 (https://arxiv.org/html/2608.17163#bib.bib33);64 (https://arxiv.org/html/2608.17163#bib.bib16);42 (https://arxiv.org/html/2608.17163#bib.bib32)），校准价值估计（45 (https://arxiv.org/html/2608.17163#bib.bib14);10 (https://arxiv.org/html/2608.17163#bib.bib37)），或在微调过程中同时保留离线数据与新收集的在线数据（56 (https://arxiv.org/html/2608.17163#bib.bib12);44 (https://arxiv.org/html/2608.17163#bib.bib8);1 (https://arxiv.org/html/2608.17163#bib.bib2);7 (https://arxiv.org/html/2608.17163#bib.bib3);8 (https://arxiv.org/html/2608.17163#bib.bib11)）。离线RL算法也已被直接应用于在线微调（16 (https://arxiv.org/html/2608.17163#bib.bib21);15 (https://arxiv.org/html/2608.17163#bib.bib17);26 (https://arxiv.org/html/2608.17163#bib.bib5);46 (https://arxiv.org/html/2608.17163#bib.bib18);9 (https://arxiv.org/html/2608.17163#bib.bib31);6 (https://arxiv.org/html/2608.17163#bib.bib30);12 (https://arxiv.org/html/2608.17163#bib.bib10)），即先用离线RL训练先验数据，然后再在线获取新样本。我们的工作与这些研究方向的不同之处在于，我们专注于利用世界模型来改进在线RL，而不是研究先验数据如何纳入RL目标本身；这使我们的方法适用于任何RL微调算法。

**基于模型的RL**。基于模型的RL方法学习环境的动力学模型，并将其用于策略或价值学习。一类方法通过向前滚动学习模型生成额外的合成转移，并将其视为真实经验（54 (https://arxiv.org/html/2608.17163#bib.bib19);17 (https://arxiv.org/html/2608.17163#bib.bib55);34 (https://arxiv.org/html/2608.17163#bib.bib65);36 (https://arxiv.org/html/2608.17163#bib.bib61);33 (https://arxiv.org/html/2608.17163#bib.bib60)），或直接基于想象的轨迹优化策略（65 (https://arxiv.org/html/2608.17163#bib.bib52);20 (https://arxiv.org/html/2608.17163#bib.bib22);21 (https://arxiv.org/html/2608.17163#bib.bib23);22 (https://arxiv.org/html/2608.17163#bib.bib9);23 (https://arxiv.org/html/2608.17163#bib.bib29)）或通过与真实数据混合的短分支推演（14 (https://arxiv.org/html/2608.17163#bib.bib59);30 (https://arxiv.org/html/2608.17163#bib.bib26)）。第二类方法则仅使用学习模型在想象推演上进行规划，使用轨迹优化器如交叉熵方法或MPPI（43 (https://arxiv.org/html/2608.17163#bib.bib25);13 (https://arxiv.org/html/2608.17163#bib.bib38);4 (https://arxiv.org/html/2608.17163#bib.bib28);25 (https://arxiv.org/html/2608.17163#bib.bib53);24 (https://arxiv.org/html/2608.17163#bib.bib66)）；其中一些方法进一步学习一个终止价值函数来界定规划时间跨度（25 (https://arxiv.org/html/2608.17163#bib.bib53);24 (https://arxiv.org/html/2608.17163#bib.bib66);35 (https://arxiv.org/html/2608.17163#bib.bib24)）。第三类方法将学习到的（或已知的）模型与蒙特卡洛树搜索相结合，并由学习到的价值函数引导，主要用于离散动作、高密度模拟领域（52 (https://arxiv.org/html/2608.17163#bib.bib51);49 (https://arxiv.org/html/2608.17163#bib.bib27);63 (https://arxiv.org/html/2608.17163#bib.bib62);50 (https://arxiv.org/html/2608.17163#bib.bib64)），近期工作将其扩展到连续控制（27 (https://arxiv.org/html/2608.17163#bib.bib63);59 (https://arxiv.org/html/2608.17163#bib.bib56)）。这些方法依赖于某种组合，例如已知动力学模型、离散动作，或从基于模型想象推演的搜索统计量中引导策略和价值目标。此外，许多工作依赖于丰富的模拟数据，其中模型误差和样本预算问题远不如真实世界稀疏奖励机器人操作那么突出。在所有这些方法中，主要基于模型生成推演训练的策略或价值函数，会继承并将模型自身的偏差直接带入训练，这一问题随着任务时间跨度和视觉复杂度增加而加剧（30 (https://arxiv.org/html/2608.17163#bib.bib26);4 (https://arxiv.org/html/2608.17163#bib.bib28);33 (https://arxiv.org/html/2608.17163#bib.bib60);20 (https://arxiv.org/html/2608.17163#bib.bib22)）。**QWM**通过仅在测试时利用模型而非将其作为额外训练数据源来规避此问题，并在真实数据上训练策略和评论家。

#### 用于VLA的世界模型。  
一个密切相关的工作方向是将世界模型应用于VLA策略。VLAW（18 (https://arxiv.org/html/2608.17163#bib.bib48)）和World-VLA-Loop（39 (https://arxiv.org/html/2608.17163#bib.bib47)）在真实推演上微调世界模型，并用其生成合成训练数据用于策略学习之间交替进行。WMPO（66 (https://arxiv.org/html/2608.17163#bib.bib44)）和World4RL（31 (https://arxiv.org/html/2608.17163#bib.bib20)）直接在想象推演上运行在策略RL，而WoVR（32 (https://arxiv.org/html/2608.17163#bib.bib15)）则针对由此产生的奖励幻觉和分布外漂移，采用关键帧初始化的推演和策略对齐的协同进化。RISE（62 (https://arxiv.org/html/2608.17163#bib.bib41)）、World-Gymnast（51 (https://arxiv.org/html/2608.17163#bib.bib40)）和GigaBrain-0.5M（55 (https://arxiv.org/html/2608.17163#bib.bib39)）使用学习到的进度/价值模型或VLM对想象推演进行评分，并将得到的奖励用于策略梯度更新。所有这些方法都将世界模型主要用作训练数据生成器，这容易继承累积偏差。SAILOR（29 (https://arxiv.org/html/2608.17163#bib.bib46)）学习一个世界模型和奖励模型来指导模仿学习并从分布外状态恢复。**QWM**可应用于任何使用Q函数的VLA RL微调方法，我们不是使用世界模型生成数据或指导模仿学习，而是利用世界模型在测试时改进Q学习。

#### 测试时缩放与最佳采样推理。  
大量工作表明，测试时缩放（例如通过最佳采样）可以显著提高生成模型在各领域的性能，包括自回归语言模型（60 (https://arxiv.org/html/2608.17163#bib.bib4);53 (https://arxiv.org/html/2608.17163#bib.bib6);3 (https://arxiv.org/html/2608.17163#bib.bib7)）、扩散模型（40 (https://arxiv.org/html/2608.17163#bib.bib1)）和策略学习（8 (https://arxiv.org/html/2608.17163#bib.bib11);26 (https://arxiv.org/html/2608.17163#bib.bib5);11 (https://arxiv.org/html/2608.17163#bib.bib45);2 (https://arxiv.org/html/2608.17163#bib.bib13)）。在大多数设置中，候选方案由验证器或价值函数独立评分，仅保留最佳样本，而不建模候选方案近期后果如何展开。**QWM**则可被视为一种测试时缩放形式，其中搜索根据候选方案的预测下游价值进行评估，从而更准确地判断哪些动作是最佳的。

## 3 预备知识  
我们考虑一个由元组 \(\{\mathcal{S},\mathcal{A},\rho,r,\gamma,T\}\) 指定的马尔可夫决策过程（MDP），其中 \(\mathcal{S}\) 表示状态空间，\(\mathcal{A}\) 表示动作空间，\(\rho(s)\) 是初始状态的分布，\(r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}\) 是奖励函数，\(\gamma\in[0,1)\) 是折扣因子，\(T(s'|s,a)\) 是支配环境的状态转移概率。RL的目标是找到最大化期望折扣回报 \(\mathbb{E}_{\pi}\left[\sum_{t=0}^{T}\gamma^{t}r(s_{t},a_{t})\right]\) 的策略 \(\pi\)。  

在本文中，我们研究如何利用学习到的世界模型在决策时改进在线RL微调。具体来说，我们假设可以访问一个环境转移数据集 \(\mathcal{D}_{\text{offline}}=\{(s,a,r,s')\}\)，这些数据可以是离线收集的或在线累积的，用于预训练一个世界模型 \(M_{\psi}(s'|s,a)\)，该模型预测给定当前状态和动作的下一个状态。世界模型在决策时使用。在线训练期间，智能体通过环境交互收集元组 \((s_t,a_t,r_t,s_{t+1})\)；这些元组被添加到回放缓冲区 \(\mathcal{D}\) 中，并用于更新策略和Q函数以获得更高的回报。我们的核心问题是如何利用世界模型来改进Q学习的性能。我们专注于离策略RL，其中Q函数估计给定状态和动作下策略的折扣回报，并通过TD学习进行训练：  

\[
\mathcal{L}(\phi)=\mathbb{E}_{(s_t,a_t,s_{t+1})\sim\mathcal{D}}\left[\left(r_t+\gamma Q_{\phi'}(s_{t+1},\tilde{a}^*_{t+1})-Q_{\phi}(s_t,a_t)\right)^{2}\right], \tag{1}
\]  

其中 \(Q_{\phi'}\) 是目标网络，\(\tilde{a}^*_{t+1}\) 是由RL策略选择的下一个动作。我们基于两个最先进的方法实现了**QWM**……