Policy and World Modeling Co-Training for Language Agents

2026-06-01

AI Summary

This work addresses the challenge of insufficient supervision in action–environment interactions during reinforcement learning (RL) for language agents, as well as the reliance of existing world models on external simulators or added inference overhead. The authors propose PaW, a novel framework that unifies policy training and world modeling within a single on-policy learning loop, leveraging state transitions from policy-sampled trajectories as supervision signals for the world model—without incurring additional computational costs. Key innovations include action-entropy-based data selection, a noise-robust loss function, and a reward-adaptive dynamic loss balancing mechanism. Experiments demonstrate that PaW significantly outperforms strong RL baselines across three agent benchmark tasks, confirming that standard RL trajectories can serve as an effective source of supervision for world modeling.

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.

Problem

• Reinforcement Learning
• World Modeling
• Language Agents
• Auxiliary Supervision
• On-policy Rollouts

Innovation

• world modeling
• reinforcement learning
• co-training
• language agents
• on-policy rollouts