Policy and World Modeling Co-Training for Language Agents

Ning Lu, Baijiong Lin, Shengcai Liu, Jiahao Wu, Haoze Lv, Yanbin Wei, Lingting Zhu, Shengju Qian, Xin Wang, Ying-Cong Chen, Qi Wang, Ke Tang


Abstract
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.
Anthology ID:
2026.emnlp-main.1885
Volume:
Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
EMNLP
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
42962–42979
Language:
URL:
https://aclanthology.org/2026.emnlp-main.1885/
DOI:
Bibkey:
Cite (ACL):
Ning Lu, Baijiong Lin, Shengcai Liu, Jiahao Wu, Haoze Lv, Yanbin Wei, Lingting Zhu, Shengju Qian, Xin Wang, Ying-Cong Chen, Qi Wang, and Ke Tang. 2026. Policy and World Modeling Co-Training for Language Agents. In Proceedings of the 2026 Conference on Empirical Methods in Natural Language Processing, pages 42962–42979, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
Policy and World Modeling Co-Training for Language Agents (Lu et al., EMNLP 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.emnlp-main.1885.pdf