Reinforcement World Model Learning for LLM-based Agents¶
会议: NeurIPS 2026
arXiv: 2602.05842
OpenReview: https://openreview.net/forum?id=VUh75nf7CQ
论文: https://arxiv.org/abs/2602.05842
领域: LLM Agent
关键词: LLM智能体, 世界模型, 强化学习, GRPO, 模拟到真实对齐
一句话总结¶
本文提出强化世界模型学习（RWML），无需专家轨迹或任务胜负奖励，直接利用文本嵌入空间的模拟-真实差距构建轻量化奖励，通过 GRPO 赋予单模型大语言智能体自监督动作条件环境动态建模与推理能力。
研究背景与动机¶
大语言模型在问答、代码编写及复杂多步推理等语言为中心的任务中已展现出卓越能力，推动了将其作为自主智能体部署于长程真实交互环境（如具身模拟 ALFWorld、客户服务与工具调用 τ²-Bench）的浪潮。然而，现有 LLM 智能体在需要前瞻性预测动作后果并灵活适应环境动力学的场景下常常频繁受挫。这种语言能力与智能体智能之间的割裂，根源在于标准自回归下一个词预测（Next-token Prediction）主要关注静态语料库上的表面语言统计规律，而长程环境决策要求模型不仅理解当前状态，还必须在采取特定动作前推演环境潜在的转移后果。
早期尝试赋予语言智能体世界模型能力的研究主要依赖监督微调（SFT），让模型直接拟合专家轨迹或由更强 LLM 合成的转移数据。这类方法面临两大严重瓶颈：一是极其依赖昂贵的专家演示或强教师模型合成，在动作空间广阔的多步场景下难以大规模扩展；二是 SFT 的监督损失过于执着于词表级的表面逐字对齐（Token-level Fidelity），而非状态语义层面的动态等价，极易因强制记忆环境描述的特定措辞而诱发灾难性遗忘与模型崩溃。另一方面，若直接使用最终任务成功信号进行强化学习（Policy RL），环境反馈不仅极度稀疏、奖励工程繁复脆弱，而且在复杂任务中存在极大的探索盲目性。
本文的切入角度是打破传统的 SFT 模仿与端到端任务奖励依赖，将世界模型建模视为一种无需外源监督的“中期训练”（Mid-training）任务。核心 idea：让智能体直接在环境探索生成的动作-状态三元组上运行自我监督，利用预训练文本嵌入空间的余弦距离构建 Sim-to-Real 差距奖励，在 GRPO 框架下引导智能体先生成推理链再预测下一状态，既规避了逐词过拟合与奖励黑客攻击，又保留了参数可塑性并显著增强后续决策能力。
方法详解¶
整体框架¶
RWML 旨在不依赖专家数据、更强教师模型或外源任务成功标签的前提下，直接通过环境自主交互让单一 LLM 策略 \(\pi_\theta\) 学会环境动力学 \(\mathcal{T}\)。算法分为三个阶段：首先利用当前策略 \(\pi_\theta\) 在环境中收集若干多步轨迹，切分为 \(\langle s_{\le t}, a_t, s_{t+1} \rangle\) 三元组并滤除平庸样本；随后利用预训练嵌入模型的相似度构建轻量化二值奖励，通过群组相对策略优化（GRPO）训练智能体生成带思考链的下一状态预测；最后将具备内部世界模型的模型权重作为良好先验，衔接下游面向任务成功奖励的策略强化学习（Policy RL）。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["环境交互采样<br/>模型自主 Rollout 获取转移三元组"] --> B["样本难度过滤：剔除高频平庸样本"]
B --> C["Sim-to-Real 嵌入奖励计算<br/>语义距离与阈值二值化"]
C --> D["群组相对策略优化 GRPO<br/>思考链与下一状态端到端更新"]
D --> E["下游任务策略强化学习 Policy RL<br/>结合稀疏成功奖励微调决策"]关键设计¶
1. 样本难度过滤：剔除高频平庸样本，聚焦非平凡环境转移
如果直接在智能体随机或贪心探索收集到的所有连续帧三元组上进行强化学习，大量由于简单格式重述或微小移动产生的“显而易见”的状态转移将淹没具有高信息量的关键交互（例如打开容器、调用特定 API、寻找特定隐藏物品）。为此，RWML 引入了难度加权子采样机制。系统先在 10% 的初始三元组子集上监督微调一个轻量预测模型 \(\pi'_\theta\)，随后让 \(\pi'_\theta\) 对剩余 90% 的样本进行 \(K=10\) 次预测采样。若样本在 \(K\) 次采样中几乎全都能达到预测阈值： $\(\frac{1}{K} \sum_{k=1}^K r^{\mathrm{WM}}(\hat{s}_{t+1}^{(k)}, s_{t+1}) \ge \tau_{\mathrm{easy}}\)$ 则判定该样本为简单样本。对于此类过于平庸的状态转移，仅以保留概率 \(p=0.1\) 加入训练集，使得最终训练集保留约 70% 具有中高认知难度的转移，促使智能体集中计算资源学习非平凡的环境因果逻辑。
2. 嵌入语义差距奖励：在表征空间对齐模拟与真实转移
传统世界模型 SFT 强制模型逐字匹配观测文本 \(s_{t+1}\)，然而在自然语言环境与复杂交互中，同一个环境状态存在大量的语义等价表达（例如语序颠倒、同义词替换或字典键顺序差异），严苛的 Cross-Entropy 损失会迫使模型参数剧烈振荡。RWML 使用轻量、确定性的预训练文本嵌入模型 \(E(\cdot)\)（如 Qwen3-Embedding-8B）度量生成状态 \(\hat{s}_{t+1}\) 与真实环境反馈 \(s_{t+1}\) 之间的余弦距离： $\(d(\hat{s}_{t+1}, s_{t+1}) = 1 - \cos\left(E(\hat{s}_{t+1}), E(s_{t+1})\right)\)$ 为了进一步提升强化学习过程的稳健性并彻底切断奖励模型被“作弊”（Reward Hacking）的漏洞，系统通过超参数阈值 \(\tau_d\) 将距离映射为二值化奖励： $\(r^{\mathrm{WM}}(\hat{s}_{t+1}, s_{t+1}) = \begin{cases} 1.0, & \text{若 } d(\hat{s}_{t+1}, s_{t+1}) < \tau_d \\ 0.0, & \text{否则} \end{cases}\)$ 实验证实，这种二值化嵌入奖励在防止模型生成冗长废话迎合 LLM-as-a-judge 评分方面展现了极其稳固的鲁棒性。
3. GRPO 驱动的思考-预测协同学习：保持参数空间平滑与可塑性
模型接收历史观测序列与候选动作，被提示首先在 <think> ... </think> 标签内输出显式的多步推理链（总结环境关键线索、分析前序动作影响与空间位置规律），随后在 <next_state> ... </next_state> 中生成最终的环境预测。整个序列通过标准群组相对策略优化（GRPO）进行更新：
$\(\mathbb{E}_{\pi_{\theta_{\mathrm{old}}}}\left[\min\left(\rho_\theta A, \text{clip}(\rho_\theta, 1 \pm \epsilon) A\right) - \beta D_{\mathrm{KL}}(\pi_\theta \parallel \pi_{\theta_{\mathrm{ref}}})\right]\)$
其中重要性权重 \(\rho_\theta = \pi_\theta(y \mid x) / \pi_{\theta_{\mathrm{old}}}(y \mid x)\)，优势函数 \(A\) 基于同组采样的二值奖励标准化计算得出。与剧烈改变底层权重的 SFT 相比，基于在线 RL 的 RWML 能够在仅微幅调整注意力与 MLP 投影权重的条件下赋予模型因果感知能力，有效遏制对预训练通用知识的灾难性遗忘，并为后续下游 Policy RL 的平稳衔接创造了理想的参数流形。
损失函数 / 训练策略¶
RWML 在训练中设置每任务收集 \(N\) 条轨迹（ALFWorld \(N=3\)，\(\tau^2\)-Bench \(N=6\)），温度 \(\tau=1.0\)。数据过滤阶段 ALFWorld 设置 \(\tau_d=0.1\)、\(\tau^2\)-Bench 设置 \(\tau_d=0.15\)，保留概率 \(p=0.1\)；正式 RL 阶段 ALFWorld 奖励阈值设为 \(\tau_d=0.2\)，\(\tau^2\)-Bench 设为 \(\tau_d=0.4\)。在 2 台 B200 GPU 上采用学习率 \(1\times 10^{-6}\)、批大小 32、组大小 8 训练 2 个 epoch。完成 RWML 后，第二阶段 Policy RL 直接在该检查点上运行任务成功回报的 GRPO 优化，折现因子 \(\gamma=1.0\)，最大训练步数设为 15（评估上限 30 步）。
实验关键数据¶
主实验¶
评估在体现具身探索的 ALFWorld（测试集包含分布内 ID 与分布外 OOD 场景）与涉及复杂工具调用和模拟用户对话的 \(\tau^2\)-Bench（涵盖 Retail、Telecom、Airline 三大垂直行业）上展开，基座模型分别为 Qwen2.5-7B-Instruct 与 Qwen3-8B。
| 方法 | 额外监督/数据类型 | ALFWorld (ID) | ALFWorld (OOD) | ALFWorld (AVG) | \(\tau^2\)-Bench (Retail) | \(\tau^2\)-Bench (Telecom) | \(\tau^2\)-Bench (Airline) | \(\tau^2\)-Bench (AVG) |
|---|---|---|---|---|---|---|---|---|
| ReACT (Qwen2.5-7B / Qwen3-8B) | 纯提示词工程 | 16.2 ± 1.0 | 6.8 ± 2.0 | 13.0 ± 1.3 | 37.7 ± 4.9 | 31.2 ± 4.2 | 21.6 ± 5.2 | 31.9 ± 2.9 |
| ReACT (GPT-5) | 闭源大模型 | 51.6 ± 1.3 | 44.8 ± 0.7 | 49.3 ± 0.9 | 55.8 ± 7.2 | 65.0 ± 5.4 | 55.0 ± 4.1 | 59.3 ± 0.5 |
| RFT | 任务成功拒绝采样 SFT | 34.4 ± 3.8 | 34.4 ± 3.4 | 34.4 ± 2.6 | 43.3 ± 3.1 | 33.3 ± 3.1 | 13.3 ± 2.4 | 33.3 ± 1.7 |
| Policy RL | 任务成功奖励在线 RL | 82.1 ± 3.6 | 79.2 ± 2.0 | 81.0 ± 1.6 | 40.8 ± 1.2 | 39.2 ± 1.2 | 30.0 ± 8.2 | 38.0 ± 1.6 |
| WM SFT | 自监督转移逐字 SFT | 3.1 ± 0.0 | 2.1 ± 0.7 | 2.8 ± 0.3 | 32.3 ± 5.3 | 24.1 ± 6.1 | 26.9 ± 6.6 | 27.9 ± 3.1 |
| RWML (本文) | 无任务标签 / 自监督 RL | 34.4 ± 0.6 | 29.2 ± 7.5 | 32.6 ± 2.1 | 40.8 ± 4.0 | 40.5 ± 4.9 | 31.3 ± 6.7 | 38.8 ± 2.5 |
| WM SFT + Policy RL | 转移 SFT + 任务 RL | 76.2 ± 3.4 | 82.3 ± 0.7 | 80.4 ± 1.5 | 40.8 ± 4.2 | 45.0 ± 5.4 | 30.0 ± 7.1 | 40.3 ± 3.9 |
| RWML + Policy RL (本文) | 自监督 RL + 任务 RL | 86.7 ± 2.8 | 90.1 ± 0.7 | 87.9 ± 1.6 | 44.2 ± 2.1 | 45.8 ± 2.4 | 38.3 ± 2.4 | 43.7 ± 2.1 |
| 专家模仿学习 (Imitation) | 人类/强模型专家轨迹 | 84.9 ± 1.9 | 77.6 ± 3.2 | 82.5 ± 2.3 | 48.3 ± 1.2 | 41.7 ± 3.1 | 38.3 ± 2.4 | 43.7 ± 1.3 |
| IWM (隐式世界模型) | 专家轨迹 + 转移数据 | 85.6 ± 1.6 | 78.1 ± 1.6 | 83.1 ± 1.0 | 40.8 ± 4.3 | 44.2 ± 5.1 | 46.7 ± 2.4 | 43.3 ± 2.6 |
| SR (自我反思) | 专家轨迹 + 强模型推理 | 83.9 ± 1.0 | 82.3 ± 0.7 | 83.3 ± 0.4 | 45.0 ± 3.5 | 45.8 ± 8.3 | 43.3 ± 2.4 | 45.0 ± 3.6 |
消融实验¶
| 配置 | 核心改动 | ALFWorld (AVG) | \(\tau^2\)-Bench (AVG) | 说明 |
|---|---|---|---|---|
| RWML 完整配置 | 嵌入余弦二值奖励 + 难样本过滤 | 32.6 ± 2.1 | 38.8 ± 2.5 | 完整自监督世界模型能力 |
| - 替换为二值化 LLM-as-a-judge | 用 Qwen3-235B 二值判别 | 14.5 ± 1.3 | 31.3 ± 2.9 | ALFWorld 性能腰斩，判定准则噪声大 |
| - 替换为连续值 LLM-as-a-judge | 用 Qwen3-235B 打分 \(r \in [0,1]\) | 3.6 ± 1.3 | 33.7 ± 3.9 | 发生严重奖励作弊，模型生成无意义模式套分 |
| - 去除难样本子采样 (w/o subsample) | 保留全部平庸易预测三元组 | 2.9 ± 1.0 | 36.3 ± 2.5 | 简单样本淹没梯度更新，导致自监督学习崩溃 |
| - 移除训练 (基座零样本) | 原始 Qwen 基座模型直接决策 | 13.0 ± 1.3 | 31.9 ± 2.9 | 基座未掌握环境转移因果规律 |
关键发现¶
- 纯自监督即产生实质决策增益：在完全不接触终局胜负奖励、无专家标注的情况下，RWML 单独使 Qwen2.5-7B 在 ALFWorld 上的平均得分从 13.0 跃升至 32.6（提升 19.6 点），使 Qwen3-8B 在 \(\tau^2\)-Bench 从 31.9 提升至 38.8（提升 6.9 点）。
- 与下游任务 RL 高度互补：在两阶段训练中，RWML + Policy RL 达到了 87.9（ALFWorld）和 43.7（\(\tau^2\)-Bench），相比直接 Policy RL 分别提高 6.9 和 5.7 点，不仅全面击败了 WM SFT + Policy RL，甚至赶平或超越了依赖高昂专家演示的 Imitation Learning（82.5/43.7）和 IWM（83.1/43.3）。
- 对通用能力的灾难性遗忘显著更少：灾难性遗忘评测显示，WM SFT 导致 Qwen2.5-7B 在 MMLU-Redux 下滑 10.10 点、GPQA 下滑 7.58 点；而 RWML 仅微跌 2.38 点和 4.05 点，在 MATH-500 和 LiveCodeBench 上几乎维持原性能。权重变化率分析显示 RWML 的点对点参数更新幅度极其保守，在各层显著低于 SFT。
亮点与洞察¶
- 自监督 Sim-to-Real 嵌入奖励设计：利用通用文本嵌入模型的余弦相似度并实施距离阈值二值化，既化解了 SFT 逐词对齐引发的“语义等价却扣分”的窘境，又消除了大模型裁判（LLM-as-a-judge）易被文字游戏黑客攻击的漏洞。
- “RL 中期训练 + RL 后训练”范式的优越性：传统认知常采用 SFT 注入知识后再用 RL 对齐行为，而本文发现用在线 RL 进行世界模型建模能以最精简的参数更新完成环境动力学内化，为后续策略优化保留极佳的优化曲率与参数相容性。
- 可复用的智能体预探索方案：通过对探索数据进行在线预测验证并剔除高频平庸转移，使得自监督数据无需人工筛选即能自然聚焦于长程任务的关键状态转移，这一采样思路可无缝迁移至具身控制与自动化 GUI 等各类交互系统。
局限与展望¶
- 基座模型能力门槛依赖：实验表明在极其严苛的工具调用任务（\(\tau^2\)-Bench）上，参数较小或基础推理能力较弱的模型（如 Qwen2.5-7B）难以顺利将习得的世界模型知识迁移转化至下游直接决策，该方法的转移增益在 Qwen3-8B 与 30B 模型上才充分展现。
- 离散文本抽象对物理连续环境的覆盖限制：当前实验集中于文本具身（ALFWorld）与工具/API 对话（\(\tau^2\)-Bench），状态均以文本描述表示，尚未向高维视觉多模态连续观测空间扩展。
- 未来方向：探究多模态具身环境下如何构建基于视觉潜空间的自监督强化世界模型，并进一步探索在推理阶段显式结合世界模型进行蒙特卡洛树搜索（MCTS）的双重增益。
相关工作与启发¶
- vs World Model SFT: 传统方法使用交叉熵强制匹配下一个观测的所有 token，导致逐字模式过拟合与严重的参数灾难性遗忘（在 ALFWorld 上引发退化至 2.8 分）；RWML 采用表征空间嵌入距离强化学习，实现了平滑泛化。
- vs IWM & Self-Reflection (Zhang et al. / Yu et al.): 现有前沿世界模型方法高度依赖人类专家轨迹或极具深度的强模型离线反思合成；RWML 证明纯靠智能体自身在环境中的随机/弱策略试错与自监督奖励，即可达到甚至超越专家数据的微调水平。
- vs Policy RL (直接以任务成功为奖励): 纯端到端 Policy RL 面对长程复杂任务探索效率低且极易陷入局部最优；RWML 作为自监督阶段先行为智能体构筑对环境转移的因果认知，极大地提升了后续策略优化的探索效率与最终上限。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出将环境转移的嵌入相似度作为自监督二值奖励，在不依赖专家与外源标签的前提下通过 GRPO 实现智能体内部世界模型构建。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 ALFWorld 与 \(\tau^2\)-Bench，设计了包括基线对比、专家对标、遗忘度基准、权重层级变化量化和奖励黑客漏洞等多维度的详实评测。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰，逻辑严密，从动机、数学建模、伪代码流程到实证机理解释层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为自主大语言智能体开辟了一条完全自监督、易于大规模扩展且不破坏通用预训练底座的世界模型后训练新路径。