Policy and World Modeling Co-Training for Language Agents

Researchers propose PaW, a co-training framework that enhances language model agents by simultaneously optimizing reinforcement learning policies and world models using data from standard RL rollouts. The approach eliminates the need for separate simulators or training stages while demonstrating consistent improvements across multiple benchmarks.

PaW addresses a fundamental limitation in current RL-based LLM agent training: while RL teaches agents which actions yield rewards, it provides minimal supervision about environmental consequences. This gap has traditionally required separate world modeling infrastructure, increasing computational overhead and training complexity. The researchers' key insight is elegantly simple—RL rollouts already contain the necessary signal in the form of action-observation pairs that naturally encode causal relationships. By leveraging this existing data through auxiliary supervision, PaW achieves efficiency gains without architectural modifications.

The framework builds on established trends in agent training where environmental understanding increasingly matters as tasks become complex and multi-step. Prior work acknowledged world modeling's value but struggled with practical implementation, often requiring external simulators or multi-stage training pipelines. PaW's contribution lies in demonstrating that co-training during standard RL is both feasible and beneficial, with three technical innovations—action-entropy-based data selection, noise-tolerant loss functions, and reward-adaptive balancing—ensuring supervision remains informative across diverse training scenarios.

For the broader AI development ecosystem, this research reduces barriers to building capable language agents by streamlining the training pipeline. Developers can now improve agent performance without substantial additional computational investment, making advanced agent capabilities more accessible to organizations with limited resources. The consistent improvements across multiple benchmarks and RL algorithms suggest the approach generalizes well rather than optimizing for specific scenarios.

• →PaW co-trains policies and world models simultaneously using RL rollout data without requiring separate simulators or inference-time overhead.
• →Three technical components—action-entropy selection, noise-tolerant loss, and adaptive balancing—stabilize auxiliary supervision during training.
• →Experiments demonstrate consistent performance gains across multiple benchmarks, models, and RL algorithms.
• →Standard RL rollouts contain sufficient causal information for effective world modeling supervision.
• →The approach reduces training complexity and computational requirements for developing capable language agents.