Large language models (LLMs) have achieved strong performance in language-centric tasks. However, in agentic settings, LLMs often struggle to anticipate action consequences and adapt to environment dynamics, highlighting the need for world-modeling capabilities in LLM-based agents. We propose Reinforcement World Model Learning (RWML), a self-supervised method that learns action-conditioned world models for LLM-based agents on textual states using sim-to-real gap rewards. Our method aligns simulated next states produced by the model with realized next states observed from the environment, encouraging consistency between internal world simulations and actual environment dynamics in a pre-trained embedding space. Unlike next-state token prediction, which prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse, our method provides a more robust training signal and is empirically less susceptible to reward hacking than LLM-as-a-judge. We evaluate our method on ALFWorld and $τ^2$ Bench and observe significant gains over the base model, despite being entirely self-supervised. When combined with task-success rewards, our method outperforms direct task-success reward RL by 6.9 and 5.7 points on ALFWorld and $τ^2$ Bench respectively, while matching the performance of expert-data training.
核心贡献 · Key contributions
提出 RWML，一种利用仿真到真实差距奖励在基于 LLM 的智能体中学习动作条件世界模型的自监督方法。 Proposes RWML, a self-supervised method for learning action-conditioned world models in LLM-based agents using sim-to-real gap rewards.
在 ALFWorld 和τ² Bench 上评估 RWML，显示在无专家数据或任务奖励下相比基础模型有显著提升。 Evaluates RWML on ALFWorld and τ² Bench, showing significant gains over base model without expert data or task rewards.
将 RWML 与任务成功奖励结合，在两个基准上分别超越直接策略 RL 6.9 和 5.7 个百分点。 Combines RWML with task-success rewards, outperforming direct policy RL by 6.9 and 5.7 points on two benchmarks.
通过参数分析证明 RWML 比世界模型 SFT 更好地缓解灾难性遗忘。 Demonstrates RWML mitigates catastrophic forgetting better than world model SFT via parameter analysis.
展示 RWML 减少无效动作并提升智能体环境中的决策效率。 Shows RWML reduces invalid actions and improves decision-making efficiency in agentic environments.
通过消融研究确认基于嵌入的奖励和数据子采样对性能至关重要。 Provides ablation studies confirming embedding-based reward and data subsampling are crucial for performance.
局限 · Limitations
RWML 的有效性依赖于基础模型能力；较弱模型难以迁移世界知识。 RWML effectiveness depends on base model capability; weaker models struggle to transfer world knowledge.
方法仅在两个文本基准上评估；对其他环境的泛化性未验证。 Method evaluated only on two text-based benchmarks; generalization to other environments is unverified.
每个任务需要多次 rollout 收集数据，在复杂环境中可能成本高昂。 Requires multiple rollouts per task for data collection, which may be costly in complex environments.
基于嵌入的奖励依赖预训练嵌入质量；低质量嵌入可能降低性能。 Embedding-based reward relies on pre-trained embedding quality; poor embeddings may degrade performance.
数据子采样步骤引入超参数，可能需要针对不同任务进行调整。 Data subsampling step introduces hyperparameters that may need tuning for different tasks.
论文章节 · Sections（共 16）
摘要Abstract
1 引言1 Introduction
2.1 符号说明2.1 Notation
2.2 强化世界模型学习2.2 Reinforcement World Model Learning
3 实验3 Experiments
3.1 实验设置3.1 Experiment Setup
3.2 主要结果3.2 Main Results
3.3 RWML 遗忘更少3.3 RWML Forgets Less
3.4 消融研究3.4 Ablation Studies
4 讨论4 Discussion
4.1 RWML 对决策的影响4.1 Impact of RWML on Decision-Making