# Reinforcement World Model Learning for LLM-based Agents

arXiv: https://arxiv.org/abs/2602.05842 · HTML: https://mlsanity.com/paper/2602.05842/

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 35 of its 63 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- LLM agents and tool use / [General LLM agent frameworks and benchmarks](https://mlsanity.com/agents/general-llm-agent-frameworks-and-benchmarks/)

## Builds on (most cited first, 30 of 35)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Voyager: An Open-Ended Embodied Agent with Large Language Models](https://mlsanity.com/paper/2305.16291/) (2023)
- [WebArena: A Realistic Web Environment for Building Autonomous Agents](https://mlsanity.com/paper/2307.13854/) (2023)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models](https://mlsanity.com/paper/2506.05176/) (2025)
- [Mind2Web: Towards a Generalist Agent for the Web](https://mlsanity.com/paper/2306.06070/) (2023)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [Cross-lingual Language Model Pretraining](https://mlsanity.com/paper/1901.07291/) (2019)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning](https://mlsanity.com/paper/2308.08747/) (2023)
- [τ2-Bench: Evaluating Conversational Agents in a Dual-Control Environment](https://mlsanity.com/paper/2506.07982/) (2025)
- [AgentTuning: Enabling Generalized Agent Abilities for LLMs](https://mlsanity.com/paper/2310.12823/) (2023)
- [Training Software Engineering Agents and Verifiers with SWE-Gym](https://mlsanity.com/paper/2412.21139/) (2024)
- [GUI-R1 : A Generalist R1-Style Vision-Language Action Model For GUI Agents](https://mlsanity.com/paper/2504.10458/) (2025)
- [Is Your LLM Secretly a World Model of the Internet? Model-Based Planning for Web Agents](https://mlsanity.com/paper/2411.06559/) (2024)
- [GUICourse: From General Vision Language Models to Versatile GUI Agents](https://mlsanity.com/paper/2406.11317/) (2024)
- [Learn-by-interact: A Data-Centric Framework for Self-Adaptive Agents in Realistic Environments](https://mlsanity.com/paper/2501.10893/) (2025)

## Built on by (most cited first, 24 of 24)

- [World Action Models: The Next Frontier in Embodied AI](https://mlsanity.com/paper/2605.12090/) (2026)
- [Code as Agent Harness](https://mlsanity.com/paper/2605.18747/) (2026)
- [Qwen-AgentWorld: Language World Models for General Agents](https://mlsanity.com/paper/2606.24597/) (2026)
- [Agentic World Modeling: Foundations, Capabilities, Laws, and Beyond](https://mlsanity.com/paper/2604.22748/) (2026)
- [Progress-conditioned Group Policy Optimization for Long-Horizon Agentic Tasks](https://mlsanity.com/paper/2607.22724/) (2026)
- [Bridging the Agent-World Gap: Text World Models for LLM-based Agents](https://mlsanity.com/paper/2606.09032/) (2026)
- [The Dark Room in the Reward Channel: Dense Prediction Rewards Collapse GRPO-Trained LLM Agents -- and The Channel, Not the Content, Decides What Works](https://mlsanity.com/paper/2607.21273/) (2026)
- [TAPO: Transition-Aware Policy Optimization for LLM Agents](https://mlsanity.com/paper/2607.27973/) (2026)
- [State2State: Environment-Derived Mid-Training for LLM Agents](https://mlsanity.com/paper/2608.04934/) (2026)
- [EnvACE: Internalizing Environment Dynamics via World Rehearsal for Agentic Reinforcement Learning](https://mlsanity.com/paper/2608.06197/) (2026)
- [PriorZero: Bridging Language Priors and World Models for Decision Making](https://mlsanity.com/paper/2605.12289/) (2026)
- [Learning from Environmental Feedback: Credit Assignment across Multiple Timescales for Agentic Reinforcement Learning](https://mlsanity.com/paper/2608.08255/) (2026)
- [Towards a Belief-Based World Model for LLM Agents](https://mlsanity.com/paper/2609.00455/) (2026)
- [The Hitchhiker's Guide to Agentic AI: From Foundations to Systems](https://mlsanity.com/paper/2606.24937/) (2026)
- [An Irreducible Quantum Advantage in Aligning World Models with Reality](https://mlsanity.com/paper/2608.19779/) (2026)
- [How do World Models and Policies Compose in LLM Agents? A Joint Spectral and Behavioral Account](https://mlsanity.com/paper/2608.30067/) (2026)
- [Don't Mask the Environment: Observation Supervision Changes How Agents Explore Under RL](https://mlsanity.com/paper/2609.20715/) (2026)
- [Generative Embodied Multiple Behavior Control Systems for Human-like Agents](https://mlsanity.com/paper/2609.22691/) (2026)
- [Agent-Editing World Model: Rethinking World Modeling for LLM Agents](https://mlsanity.com/paper/2609.28416/) (2026)
- [Does Learning to Predict the World Help Agents Act? Auditing World-Model Post-Training](https://mlsanity.com/paper/2609.33335/) (2026)
- [EVOKE: Eliciting World Knowledge in Agents for Transferable Decision-Making](https://mlsanity.com/paper/2609.38334/) (2026)
- [ObsSpec: Accelerating Agentic Reinforcement Learning with Observation Speculation](https://mlsanity.com/paper/2610.15246/) (2026)
- [ANTE: Anticipate, Then Execute -- Learning Internal World Models for Coding Agents from Verified Beliefs](https://mlsanity.com/paper/2610.15368/) (2026)
- [Co-Evolving Acting and Forecasting for Language Agents](https://mlsanity.com/paper/2610.20597/) (2026)

## Abstract (arXiv)

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.