# Policy and World Modeling Co-Training for Language Agents

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

## Key facts

- Cited by 15 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 46 of its 66 references that are in the map.
- First posted on arXiv 1 Jun 2026.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [SWE-bench: Can Language Models Resolve Real-World GitHub Issues?](https://mlsanity.com/paper/2310.06770/) (2023)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [A Survey on Large Language Model based Autonomous Agents](https://mlsanity.com/paper/2308.11432/) (2023)
- [WebArena: A Realistic Web Environment for Building Autonomous Agents](https://mlsanity.com/paper/2307.13854/) (2023)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [Mind2Web: Towards a Generalist Agent for the Web](https://mlsanity.com/paper/2306.06070/) (2023)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Measuring and Narrowing the Compositionality Gap in Language Models](https://mlsanity.com/paper/2210.03350/) (2022)
- [DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence](https://mlsanity.com/paper/2606.19348/) (2026)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [GLM-5: from Vibe Coding to Agentic Engineering](https://mlsanity.com/paper/2602.15763/) (2026)
- [Group-in-Group Policy Optimization for LLM Agent Training](https://mlsanity.com/paper/2505.10978/) (2025)
- [Robust Loss Functions under Label Noise for Deep Neural Networks](https://mlsanity.com/paper/1712.09482/) (2017)
- [AgentTuning: Enabling Generalized Agent Abilities for LLMs](https://mlsanity.com/paper/2310.12823/) (2023)
- [FireAct: Toward Language Agent Fine-tuning](https://mlsanity.com/paper/2310.05915/) (2023)
- [ZeroSearch: Incentivize the Search Capability of LLMs without Searching](https://mlsanity.com/paper/2505.04588/) (2025)
- [Is Your LLM Secretly a World Model of the Internet? Model-Based Planning for Web Agents](https://mlsanity.com/paper/2411.06559/) (2024)
- [Web Agents with World Models: Learning and Leveraging Environment Dynamics in Web Navigation](https://mlsanity.com/paper/2410.13232/) (2024)
- [AgentGym: Evolving Large Language Model-based Agents across Diverse Environments](https://mlsanity.com/paper/2406.04151/) (2024)

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

- [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)
- [EnvACE: Internalizing Environment Dynamics via World Rehearsal for Agentic Reinforcement Learning](https://mlsanity.com/paper/2608.06197/) (2026)
- [Co-Evolution in Agentic Systems: Toward Self-Directed Evolution Beyond Human Design](https://mlsanity.com/paper/2608.10299/) (2026)
- [Act More, Decide Less: Skill-Guided Adaptive Action Chunking for Long-Horizon LLM Agents](https://mlsanity.com/paper/2609.02042/) (2026)
- [Dual-Frontier: When Can an Agent Trust Its World Model?](https://mlsanity.com/paper/2609.26293/) (2026)
- [RLTL;DR: Self-improvement by Internalizing Self-generated Feedback](https://mlsanity.com/paper/2609.37633/) (2026)
- [TrajMind: Chaining Role-Specialized LoRAs for Fast-and-Slow Collective Trajectory Anomaly Diagnosis](https://mlsanity.com/paper/2609.02540/) (2026)
- [Don't Mask the Environment: Observation Supervision Changes How Agents Explore Under RL](https://mlsanity.com/paper/2609.20715/) (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)
- [Environmental Feedback Modeling Matters: Rethinking Feedback Treatment in Agentic Hindsight Self-Distillation](https://mlsanity.com/paper/2610.11384/) (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)
- [Model What Matters: Decision-Guided Selective World Modeling for Language Agents](https://mlsanity.com/paper/2610.19642/) (2026)
- [Co-Evolving Acting and Forecasting for Language Agents](https://mlsanity.com/paper/2610.20597/) (2026)

## Abstract (arXiv)

Reinforcement learning (RL) improves large language model (LLM) agents by teaching them which actions lead to high rewards, but provides little supervision on what those actions do to the environment. World modeling (WM) can fill this gap, yet existing approaches often require separate simulators, extra training stages, or additional inference-time computation. We observe that on-policy RL rollouts already contain the needed signal: each transition pairs an action with its resulting next observation. Based on this observation, we propose PaW, a Policy and World modeling co-training framework that adds auxiliary WM supervision to the same policy during RL, without changing the inference paradigm. To make auxiliary WM supervision informative and stable, PaW introduces three components: action-entropy-based WM data selection, noise-tolerant WM loss, and reward-adaptive loss balancing. Experiments on three agentic task benchmarks show consistent improvements over strong RL baselines across models and RL algorithms. These results suggest that standard RL rollouts are a practical source of WM supervision for language-agent training.