PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2606.02388v1cs.LGcs.AI

TL;DR

RL improves language-agent actions but gives limited supervision about their environmental consequences. PaW co-trains the same policy with world-model supervision from on-policy rollouts, yielding consistent gains across tasks, model scales, and RL algorithms without extra inference cost.

  • Problem

    Standard RL optimizes language-agent actions for reward without explicitly learning their consequences, limiting supervision for environmental dynamics.

  • Method

    PaW reuses on-policy rollouts to co-train the policy with auxiliary next-observation prediction using selected transitions, clipped MAE loss, and adaptive loss balancing.

  • Results

    PaW consistently improves performance across ALFWorld, WebShop, and search-augmented QA, RL algorithms, and model scales without additional inference cost.

  • Takeaways & Limitations

    Standard on-policy RL rollouts are a simple, practical source of world-modeling supervision for language-agent training.

  • Takeaways & Limitations

    PaW uses one-step next-observation supervision, capturing local dynamics without explicitly modeling longer-horizon dependencies or compounding prediction errors.

Abstract

from arXiv · show

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.

1 Introduction

Standard RL improves reward-seeking actions but does not teach agents their environmental consequences, while existing world-modeling methods add separate training, simulators, or inference-time planning. PaW instead reuses next observations from on-policy rollouts for stable auxiliary world-model supervision during unchanged policy learning, with three targeted design components.

  • Motivation: RL optimizes actions for reward maximization without learning their consequences, leaving agents vulnerable to invalid operations, irreversible state changes, and delayed failures.World modeling addresses this gap by predicting next observations from environment dynamics.
  • Prior limitations: Existing language-agent world-modeling methods typically rely on separate simulators, additional training stages, or inference-time planning outside standard RL.These approaches may train simulators to generate imagined trajectories or first instill world-modeling ability before fine-tuning.
  • Key observation: On-policy RL rollouts provide both action-based policy supervision and resulting next observations as dense action-conditioned world-model supervision, without additional rollouts.Each interaction step pairs an action with the next observation revealing its consequence.
  • PaW framework: PaW appends next-observation tokens to reused RL rollouts and trains the same model with an auxiliary next-token-prediction loss while leaving policy learning and inference unchanged.Causal attention prevents later next-observation tokens from affecting action logits, so inference remains standard policy execution.
  • PaW designs: PaW combines action-entropy-based data selection, clipped MAE for noisy observations, and reward-adaptive loss balancing to make auxiliary world-model supervision informative and stable.Rollout observations can be uninformative, unpredictable, or noisy, requiring the auxiliary loss to balance with the RL loss.
  • Contributions: PaW yields consistent improvements over strong RL baselines across three agentic tasks, multiple models, and RL algorithms.The paper presents this as an evaluation result and contribution of the framework.

2 Preliminaries

Language-agent tasks model multi-turn interaction as trajectories in which a policy maps instruction-and-history contexts to actions, while the environment returns rewards and next observations. On-policy RL optimizes expected return, and world modeling predicts action-conditioned next observations to support decisions in long-horizon tasks.

  • Problem setup: A policy πθ solves a user-specified goal through multi-turn interaction, sampling textual or serialized actions from decision contexts built from instructions and history.Each transition records the action, reward, and resulting next observation within a trajectory.
  • On-policy Agentic RL: On-policy agentic RL fine-tunes πθ to maximize expected trajectory return using sampled rollouts and an advantage-weighted action loss.The advantage is reward-derived, while action likelihood sums token-level log-likelihoods.
  • On-policy Agentic RL: GRPO and GIGPO mainly differ in advantage estimation and surrogate-objective instantiation, while PaW is orthogonal to these choices and augments the base RL loss with world-modeling supervision.This positions the proposed supervision as compatible with multiple on-policy RL algorithms.
  • World-modeling for language agents: For language agents, world modeling predicts the textual next observation conditioned on the decision context and action.The objective is autoregressive and computes likelihood over observation tokens.
  • World-modeling for language agents: Learning world modeling helps agents understand action outcomes and make better decisions in long-horizon tasks.This motivation follows from modeling action-conditioned environment dynamics.

3 Methodology

PaW co-trains policy learning and next-observation modeling in the same policy during on-policy RL. It selects informative transitions, uses clipped MAE world-model supervision, and adapts its weight to rollout rewards without changing deployment cost.

  • Core framework: PaW augments the base on-policy RL objective with an auxiliary next-observation prediction objective that updates the same parameters θ.The RL loss improves action selection, while the world-model loss predicts next observations from rollout transitions.
  • Transition selection: Action-entropy selection retains the top-α fraction of transitions, using high-entropy actions as more informative targets for action-conditioned transition learning.The RL loss uses all generated actions, whereas the world-model loss uses only selected transitions.
  • Robust world modeling: Clipped MAE replaces cross-entropy to bound low-probability-token gradients and masks tokens whose prediction confidence exceeds threshold ρ.This focuses supervision on insufficiently predicted observation tokens while reducing pressure to fit noisy or non-action-relevant details.
  • Adaptive balancing: Reward-adaptive weighting assigns larger λWM,g to low-return rollout groups and decreases it as mean return approaches Rmax, allowing stronger dynamics supervision when performance is poor.This prevents dense auxiliary gradients from dominating policy learning while reducing auxiliary pressure for high-performing groups.
  • Training and inference: PaW preserves the base RL loss while adding reward-adaptively weighted CMAE supervision on entropy-selected transitions, requiring no additional environment interaction or model forward.Deployment keeps the same interface and cost as the underlying RL-trained agent.

4 Experiments

Experiments across interactive decision-making and search-augmented QA show that PaW consistently improves GRPO and GIGPO across model scales and task types. Additional analyses indicate that PaW supports sparse-reward learning, generalizes across algorithms and backbones, and adds minimal training overhead.

  • Interactive decision-making: PaW improves GRPO and GIGPO success rates on ALFWorld and WebShop across Qwen2.5-1.5B and 7B models.On ALFWorld, gains reach +7.9 for GRPO and +2.8 for GIGPO at 1.5B; on WebShop, gains reach +8.0 and +9.1, respectively.
  • Search-augmented QA: For search-augmented QA, PaW raises GRPO averages from 43.9% to 44.8% at 3B and from 47.9% to 48.8% at 7B.GIGPO averages similarly increase from 41.2% to 44.2% at 3B and from 45.8% to 47.5% at 7B.
  • Generality: PaW generalizes beyond GRPO and GIGPO by combining with PPO and RLOO and by applying GRPO with PaW to Qwen3-1.7B, Llama3.2-3B-Instruct, and Qwen2.5-14B-Instruct.These evaluations test generality across both RL algorithms and model backbones.
  • Sparse-reward learning: In sparse-reward WebShop training, vanilla GRPO rarely obtains positive rewards, whereas PaW supplies dense world-model supervision from state-action transitions.The auxiliary next-observation prediction objective prevents the training signal from collapsing in this setting.
  • Computational overhead: PaW adds only 10.7 s per step, or 2.1% of the approximately 505 s GRPO step time, while peak and average GPU memory increase by 2.4 GB (2.4%) and 2.2%.The method reuses rollout data and the actor forward pass, adding only the WM objective during the RL update.
  • Ablations: Adaptive WM loss balancing and clipped MAE outperform fixed weighting and standard cross-entropy, with ALFWorld falling from 77.9% to 75.5% and then 68.5%.On WebShop, performance falls from 68.6% to 67.0% with fixed weighting and to 57.2% with cross-entropy.

5 Related Work

Related work trains LLM agents with supervised fine-tuning or reinforcement learning, while world-modeling methods predict future states, rewards, or transitions for planning and policy learning. PaW complements existing approaches that address sparse, delayed rewards or use future-observation modeling.

  • Training LLM agents: LLM agents map instructions, interaction histories, and observations into executable actions across web, tool-use, embodied, and other interactive tasks.Recent work trains these agents with supervised fine-tuning or reinforcement learning.
  • Training LLM agents: Because agent rewards are often sparse and delayed, existing reinforcement-learning methods improve credit assignment or add auxiliary training signals.The passage positions PaW as complementary to these approaches.
  • World modeling for language agents: World models learn environment dynamics by predicting future states or rewards, and language-agent methods use LLMs as world models, simulators, or transition predictors.These methods support planning, verification, and policy learning through future-observation modeling.

6 Conclusion

PaW co-trains policy and world modeling by reusing on-policy RL rollouts as action-conditioned next-observation supervision, avoiding separate simulators and inference-time planning. Its transition selection, clipped MAE prediction, and reward-adaptive balancing yield consistent gains with minimal training overhead and no additional inference cost.

  • Conclusion: PaW reuses on-policy RL rollouts as action-conditioned next-observation supervision for an auxiliary world-modeling loss on the same policy model.This avoids training a separate simulator or adding inference-time planning.
  • Conclusion: PaW combines action-entropy-based transition selection, clipped MAE observation prediction, and reward-adaptive loss balancing for noisy agentic environments.These components make the auxiliary supervision more effective and stable.
  • Conclusion: Experiments on ALFWorld, WebShop, and search-augmented QA show consistent gains across RL algorithms and model scales.The reported improvements span three agentic task benchmarks.
  • Conclusion: PaW achieves these gains with minimal training overhead and no additional inference cost.The framework preserves the inference paradigm while adding auxiliary world-modeling supervision during RL.

7 Limitations · A Implementation Details

PaW improves RL performance but remains limited by one-step world-model supervision and duplicated raw on-policy trajectories. The paper identifies multi-step modeling and trajectory-level deduplication as directions for addressing these limitations.

  • 7 Limitations: PaW consistently improves RL performance, but the framework has two main limitations.
  • 7 Limitations: One-step next-observation supervision captures local dynamics without explicitly modeling longer-horizon dependencies.
  • 7 Limitations: One-step modeling can incur compounding prediction errors over extended horizons.
  • 7 Limitations: Extending the co-training objective to multi-step world modeling is proposed as future work.
  • 7 Limitations: WM supervision is constructed from raw on-policy rollouts without trajectory-level deduplication.
  • 7 Limitations: Repeated trajectories may reduce supervision diversity and bias the auxiliary objective.

A.1 Details of Training · A.2 Baseline Details · A.3 Prompts

The appendix specifies task-specific training configurations, baseline agents and RL algorithms, and prompt templates for ALFWorld, WebShop, and search-augmented QA. It also defines shared rollout settings, runtime prompt slots, interaction-history limits, and structured reasoning/action or search-answer tags.

  • A.1 Details of Training: ALFWorld uses 2048-token prompts, 512-token responses, 50 environment steps per episode, actor learning rate 1e-6, and critic learning rate 1e-5 in PPO.Success receives reward 10, failure receives 0, and invalid actions incur a -0.1 penalty.
  • A.1 Details of Training: WebShop uses 4096-token prompts, 512-token responses, 15 environment steps per episode, actor learning rate 1e-6, and critic learning rate 1e-5 in PPO.Its rule-based reward is 10 for success and 0 for failure; invalid-action penalties are also applied.
  • A.1 Details of Training: Group-based RL methods use groups of 8 and sample 16 groups per rollout, while PPO uses 128 distinct rollout environments; GIGPO uses ω = 1 and γ = 0.95.The rollout temperature is 1.0, validation temperature 0.4, and mini-batch size 64.
  • A.1 Details of Training: Search-augmented QA uses 4096-token prompts, 512-token responses, four turns, actor learning rate 1e-6, reward 1 for success, reward 0 for failure, and invalid-action penalty -0.01.Training uses 256 examples and group size 5.
  • A.2 Baseline Details: The baselines include GPT-4o and Gemini-2.5-Pro as closed-source LLMs, plus ReAct and Reflexion as prompting-based agents.ReAct interleaves reasoning and acting, whereas Reflexion uses self-reflection and iterative improvement.
  • A.2 Baseline Details: The evaluated RL baselines are PPO, RLOO, GRPO, and GiGPO, representing classic, group-based, trajectory-level, and hierarchical policy-learning approaches.The passage identifies RLOO as avoiding value networks and GRPO as targeting multi-step tasks.
  • A.3 Prompts: Prompt templates use Python-style runtime formatting for semantic slots such as {task_description}, {step_count}, and {current_observation}, with recent or complete interaction history supplied by task.ALFWorld and WebShop retain two recent history steps, while search QA uses complete history; tags structure reasoning, actions, queries, answers, and retrieved information.

A.4 Noise-Gradient Analysis · B Additional Experimental Results

The noise-gradient analysis evaluates how CE and MAE allocate world-model gradient budget across noisy and meaningful WebShop observation tokens. Additional results report full benchmark performance and training dynamics, showing PaW’s effects on reward, losses, and update behavior.

  • A.4 Noise-Gradient Analysis: The diagnostic samples 100 WebShop search-result transitions from Qwen2.5-1.5B-Instruct rollout traces.It uses the same serialization format as world-model training.
  • A.4 Noise-Gradient Analysis: It scores only next-observation tokens, using the decision context and generated action as prefix and the next observation as the teacherforced target.The analysis follows the world-model training serialization format.
  • A.4 Noise-Gradient Analysis: CE assigns target-logit gradient magnitude 1 − pt,i, whereas MAE assigns pt,i(1 −pt,i), enabling comparison of normalized gradient allocation.The diagnostic compares token-level losses while controlling for raw loss scale.
  • A.4 Noise-Gradient Analysis: Noisy tokens are WebShop spans overlapping random product identifiers or brand strings, while all other next-observation tokens are meaningful.Random identifiers follow B[0-9A-Z]{9}; brand-like uppercase spans require at least four letters and may include digits.
  • A.4 Noise-Gradient Analysis: The evaluation-only analysis isolates loss-function effects on gradient allocation without updating model parameters.Normalized shares indicate whether a loss over-allocates gradient to noisy WebShop surface strings.
  • B Additional Experimental Results: Full benchmark results average three random seeds and report ALFWorld subtask and overall success rates plus WebShop score and success rate.The table covers ALFWorld and WebShop performance with statistical statistics.
  • B Additional Experimental Results: On WebShop, PaW improves training reward over GRPO while policy-gradient loss and clipped update ratio remain broadly comparable.This suggests auxiliary world modeling improves learning without substantially changing the main policy-optimization dynamics.
  • B Additional Experimental Results: As training reward improves, PaW’s reward-adaptive coefficient decreases, world modeling loss trends downward, and clipped token ratio tracks clipped observation-token supervision.The clipped token ratio reports how often the margin objective clips observation-token supervision.

B.1 Full Results of ALFWorld and WebShop

Table 5 reports full ALFWorld and WebShop results, including statistical statistics, averaged over three random seeds. PaW consistently improves performance over the base RL methods.

  • Table 5 presents full ALFWorld and WebShop results, including statistical statistics.Results are averaged over 3 random seeds.
  • PaW consistently improves performance over the base RL methods.

B.2 Training Dynamics

PaW improves policy-side training reward without materially changing policy-gradient loss or clipped update ratio. Its auxiliary world-model objective learns from RL rollouts, with reward-adaptive supervision decreasing as rewards rise and clipping bounded as prediction improves.

  • B.2 Training Dynamics: GRPO w/ PaW achieves better training reward than the GRPO baseline, indicating that auxiliary world-model supervision improves the learned policy during RL.The policy-gradient loss and clipped update ratio remain close to GRPO, suggesting the gain does not arise from substantially larger PPO-style updates or clipping changes.
  • B.2 Training Dynamics: The adaptive coefficient decreases during training because higher-reward rollouts receive less world-model supervision under reward-adaptive scaling.This reflects the framework’s reward-adaptive loss balancing during RL training.
  • B.2 Training Dynamics: The world-model loss decreases, showing that the shared policy model learns next-observation prediction targets from RL rollouts.The auxiliary objective therefore receives usable supervision from on-policy transitions.
  • B.2 Training Dynamics: The clipped token ratio tracks observation tokens affected by margin clipping, which keeps the auxiliary signal bounded as world-model prediction improves.The ratio measures the extent of clipping applied to observation-token updates.
Loading 2606.02388v1…