Policy and World Modeling Co-Training for Language Agents

• URL: http://arxiv.org/abs/2606.02388v1
• Date: Mon, 01 Jun 2026 15:35:40 GMT
• Title: Policy and World Modeling Co-Training for Language Agents
• Abstract summary: 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.<n>World modeling (WM) can fill this gap, yet existing approaches often require separate simulators, extra training stages, or additional inference-time computation.<n>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.
• Score: 44.550078924046375
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Scaling Automatic Research Agents via World Models [42.9808933592206]
  Recent research (AutoResearch) agents show the capability to independently implement solutions and learn from the execution.<n>Behind these gains, post-training (especially RL) plays a central role.<n>World Model RL (WMRL) replaces environment execution with a world model to remove this bottleneck.
  arXiv  Detail & Related papers  (2026-08-12T20:11:25Z)
• TAPO: Transition-Aware Policy Optimization for LLM Agents [27.95488836373643]
  TAPO is a unified training framework that alternates between policy optimization and transition supervision.<n>It serves as a computationally lightweight, plug-and-play enhancement module for existing agent RL algorithms.<n>We conduct systematic experiments on WebShop and ALFWorld, integrating foundation models of various scales with different policy optimization algorithms.
  arXiv  Detail & Related papers  (2026-07-30T10:17:55Z)
• Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL [0.8546394944138487]
  We formalize text-based world modeling as a steerable transition-dynamics problem decomposed into initial state.<n>We curate 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families.<n>We introduce a plug-and-play GRPO training framework with deterministic state checks, and perform zero-shot transfer ablations on three OOD environments.
  arXiv  Detail & Related papers  (2026-05-07T00:40:32Z)
• Agent2World: Learning to Generate Symbolic World Models via Adaptive Multi-Agent Feedback [51.22403664895878]
  Agent2World is a tool-augmented multi-agent framework that achieves strong inference-time world-model generation.<n>It also serves as a data engine for supervised fine-tuning, by grounding generation in multi-agent feedback.
  arXiv  Detail & Related papers  (2025-12-26T18:54:14Z)
• Ring-lite: Scalable Reasoning via C3PO-Stabilized Reinforcement Learning for LLMs [51.21041884010009]
  Ring-lite is a Mixture-of-Experts (MoE)-based large language model optimized via reinforcement learning (RL)<n>Our approach matches the performance of state-of-the-art (SOTA) small-scale reasoning models on challenging benchmarks.
  arXiv  Detail & Related papers  (2025-06-17T17:12:34Z)
• RAGEN: Understanding Self-Evolution in LLM Agents via Multi-Turn Reinforcement Learning [125.96848846966087]
  Training large language models (LLMs) as interactive agents presents unique challenges.<n>While reinforcement learning has enabled progress in static tasks, multi-turn agent RL training remains underexplored.<n>We propose StarPO, a general framework for trajectory-level agent RL, and introduce RAGEN, a modular system for training and evaluating LLM agents.
  arXiv  Detail & Related papers  (2025-04-24T17:57:08Z)
• Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining [74.83412846804977]
  Reinforcement learning (RL)-based fine-tuning has become a crucial step in post-training language models.<n>We present a systematic end-to-end study of RL fine-tuning for mathematical reasoning by training models entirely from scratch.
  arXiv  Detail & Related papers  (2025-04-10T17:15:53Z)
• Mastering the Unsupervised Reinforcement Learning Benchmark from Pixels [112.63440666617494]
  Reinforcement learning algorithms can succeed but require large amounts of interactions between the agent and the environment. We propose a new method to solve it, using unsupervised model-based RL, for pre-training the agent. We show robust performance on the Real-Word RL benchmark, hinting at resiliency to environment perturbations during adaptation.
  arXiv  Detail & Related papers  (2022-09-24T14:22:29Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.