Policy and World Modeling Co-Training for Language Agents

1 Introduction

• Standard RL optimizes actions for reward without learning their consequences, while world modeling predicts next observations to help agents internalize environment dynamics.The paper notes that lacking consequence modeling can leave agents brittle to invalid operations, irreversible state changes, and delayed failures in long-horizon tasks.Source: §1 · p. 1
• Prior world-modeling approaches add separate simulators, extra training stages, or inference-time computation rather than jointly learning within standard RL training.The introduction describes simulator-based approaches (Gu et al., 2025; Xiao et al., 2026; Li et al., 2026; Liu et al., 2026a) and WM-first fine-tuning (Zhang et al., 2025; Yu et al., 2026).Source: §1 · p. 1, Figure 1 · p. 1
• PaW reuses on-policy rollouts as action-conditioned next-observation supervision, appending next-observation tokens and applying an auxiliary prediction loss to the same model.Causal attention keeps later observation tokens from affecting action logits, and inference requires no additional simulation steps.Source: §1 · pp. 1–2, Figure 1 · p. 1
• PaW consistently improves strong RL baselines across ALFWorld, WebShop, and search-augmented QA tasks, across models and RL algorithms.The paper also reports negligible training overhead.Source: §1 · p. 2
• PaW stabilizes auxiliary supervision with action-entropy-based transition selection, clipped MAE for noisy observations, and reward-adaptive balancing of WM and RL losses.These designs address uninformative transitions, unpredictable target tokens, and the need to balance the auxiliary WM loss with the RL loss.Source: §1 · p. 2

Problem setup

• A language agent pursues a user-specified goal by sampling actions from its policy based on the instruction and interaction history, then receiving rewards and next observations from the environment.These interactions form a trajectory of observations, actions, and rewards.Source: §2 · p. 2

On-policy Agentic RL

• On-policy agentic RL fine-tunes the policy to maximize expected trajectory return using sampled rollouts.Most on-policy algorithms are abstracted as minimizing an advantage-weighted action loss.Source: §2 · p. 2
• Algorithms such as GRPO (Shao et al., 2024) and GIGPO (Feng et al., 2025) differ mainly in advantage estimation and surrogate objectives, while PaW augments the base RL loss with world-modeling supervision.The method is orthogonal to these algorithmic choices.Source: §2 · p. 2

World-modeling for language agents

• For language agents, world modeling predicts the next textual observation autoregressively from the decision context and action.The likelihood is computed over observation tokens; learning this model enables agents to better understand action outcomes and make better decisions in long-horizon tasks (Zhang et al., 2025; Liu et al., 2026a).Source: §2 · p. 3

3.1 Constructing World Modeling Supervision from RL Rollouts

• PaW selects the top-α fraction of rollout transitions by action-token entropy for WM, while the RL loss still uses all generated actions.High-entropy actions indicate decision points with more diverse alternatives; the entropy is available from existing distributions, requiring no additional model forward.Source: §3.1 · p. 3

3.2 Clipped MAE Loss for Noisy Observation Prediction

• Clipped MAE reduces pressure to fit noisy observation tokens by bounding low-probability token gradients and masking tokens whose probability exceeds ρ.Unlike cross-entropy, MAE is less sensitive to unpredictable fragments; confidence clipping stops optimizing tokens already predicted with high confidence.Source: §3.2 · p. 4

3.3 Reward-Adaptive Loss Balancing

• PaW increases WM weight for rollout groups with low mean episode return and decreases it as mean return approaches the environment’s maximum attainable return.This reward-adaptive balance gives weaker-performing groups stronger next-observation supervision while allowing the policy objective to dominate for higher-performing groups.Source: §3.3 · p. 4

3.4 Training and Inference

• The final objective combines base on-policy RL loss with reward-weighted clipped-MAE WM loss on entropy-selected transitions, without changing inference or adding training stages or model forwards.PaW uses no additional environment interaction or separate model, and deployment retains the underlying RL agent’s interface and cost.Source: §3.4 · p. 5

4.1 Experiment Setup

• The setup covers interactive ALFWorld and WebShop tasks plus seven single- and multi-hop search-augmented QA benchmarks.Baselines include GRPO and GIGPO, each compared with its PaW-augmented variant.Source: §4.1 · pp. 5–6

4.3 Different Models and RL algorithms

• PaW gains generalize across RL algorithms, model families, and scales on WebShop.It improves PPO by +6.1 and RLOO by +4.5; on Qwen2.5-14B-Instruct the gain is +2.4, while Llama3.2-3B-Instruct rises from 4.0% to 62.2%.Source: Table 3 · p. 6, §4.3 · p. 7

4.2 Experimental Results

• PaW consistently improves GRPO and GIGPO on ALFWorld and WebShop across model scales, and raises search-augmented QA average scores.On ALFWorld, success-rate gains are +7.9 and +3.0 for GRPO and +2.8 and +1.0 for GIGPO at 1.5B and 7B, respectively; on WebShop, gains are +8.0 and +4.0 for GRPO and +9.1 and +2.9 for GIGPO. For search-augmented QA, GRPO average scores rise from 43.9% to 44.8% at 3B and from 47.9% to 48.8% at 7B; GIGPO scores rise from 41.2% to 44.2% at 3B and from 45.8% to 47.5% at 7B.Source: §4.2 · pp. 6–7, Table 2 · p. 6

4.4 PaW Helps Where RL Fails

• In sparse-reward WebShop training, PaW's next-observation supervision helps Llama3.2-3B-Instruct escape reward-signal collapse and begin producing successful rollouts.Vanilla GRPO rarely obtains positive rewards, whereas PaW predicts next observations from collected state-action transitions, providing learning signals even when most trajectories receive zero task reward.Source: Figure 4 · p. 7, §4.4 · p. 7

4.5 Computational Overhead

• PaW adds minimal training overhead by reusing rollout data and the actor forward pass while adding the WM objective during the RL update.On ALFWorld with Qwen2.5-1.5B-Instruct, it adds 10.7 s per step (2.1% of the approximately 505 s GRPO step time); peak and average GPU memory increase by 2.4 GB (2.4%) and 2.2 GB (2.2%).Source: Figure 5 · p. 7, §4.5 · pp. 7–8

4.6 Ablation Study

• Both reward-adaptive loss balancing and the clipped MAE WM loss contribute to PaW's performance on ALFWorld and WebShop.Fixing λWM,g = 1 reduces scores from 77.9% to 75.5% on ALFWorld and from 68.6% to 67.0% on WebShop; replacing CMAE with cross-entropy yields 68.5% and 57.2%, respectively.Source: Table 4 · p. 8, §4.6 · p. 8

4.7 Hyperparameter Sensitivity

• PaW remains effective across a broad range of entropy-selection ratios α and clipping thresholds ρ, without delicate hyperparameter tuning.The strongest result uses α = 0.75 and ρ = 0.2; overly large clipping thresholds reduce performance.Source: §4.7 · p. 8

• LLM agents map instructions, interaction histories, and observations into executable actions for interactive tasks.Training beyond prompting includes supervised fine-tuning and reinforcement learning (Wang et al., 2023; Yao et al., 2023; Shinn et al., 2023; Deng et al., 2023; Zeng et al., 2024; Chen et al., 2023; Xi et al., 2024; Jin et al., 2025; Sun et al., 2025; Feng et al., 2025).Source: §5 · p. 8
• PaW reuses on-policy rollouts for dense next-observation supervision without changing the policy-gradient estimator.This makes PaW complementary to RL approaches that modify credit assignment or add auxiliary training signals.Source: §5 · p. 8
• Sparse and delayed agent rewards have motivated RL methods focused on credit assignment or auxiliary training signals.PaW complements these approaches rather than changing the policy-gradient estimator.Source: §5 · p. 8
• Prior language-agent world models support planning, verification, and policy learning, but often require separate models, simulator training, or inference-time planning.PaW trains the same policy model to predict next observations from its own rollouts, with no extra interactions or inference-time computation.Source: §5 · p. 8

6 Conclusion

• PaW reuses on-policy RL rollouts as action-conditioned next-observation supervision, optimizing an auxiliary world-modeling loss on the policy model.It avoids training a separate simulator or adding inference-time planning.Source: §6 · pp. 8–9
• PaW shows consistent gains on ALFWorld, WebShop, and search-augmented QA across RL algorithms and model scales, with minimal training overhead and no additional inference cost.The results suggest standard RL rollouts are a practical source of world-modeling supervision for language-agent training.Source: §6 · p. 9
• To make supervision effective in noisy agentic environments, PaW combines action-entropy-based transition selection, clipped MAE observation prediction, and reward-adaptive loss balancing.Source: §6 · p. 9

7 Limitations

• One-step next-observation supervision captures local dynamics but does not explicitly model longer-horizon dependencies or compounding prediction errors.Extending the co-training objective to multi-step world modeling is a proposed direction for future work.Source: §7 · p. 9
• Raw on-policy rollouts lack trajectory-level deduplication, so repeated trajectories may reduce supervision diversity and bias the auxiliary objective toward frequent patterns.Deduplication or diversity-aware sampling may improve world-modeling efficiency and effectiveness.Source: §7 · p. 9

A.1 Details of Training

• ALFWorld uses 128 rollout environments, 50-step episodes, and a 10-success/0-failure reward with a -0.1 invalid-action penalty.Group-based methods use groups of 8 and sample 16 groups per rollout; prompts are capped at 2048 tokens and responses at 512. Actor and PPO critic learning rates are 1e-6 and 1e-5, respectively; rollout and validation temperatures are 1.0 and 0.4, and the mini-batch size is 256.Source: §A.1 · p. 11
• WebShop uses 128 rollout environments and 15-step episodes, with a 10-success/0-failure reward and a -0.1 invalid-action penalty.Group-based methods use groups of 8 and sample 16 groups per rollout; prompts are capped at 4096 tokens and responses at 512. Actor and PPO critic learning rates are 1e-6 and 1e-5, respectively; rollout and validation temperatures are 1.0 and 0.4, and the mini-batch size is 64.Source: §A.1 · p. 11
• Search-augmented QA uses four turns, a 1-success/0-failure reward, and a -0.01 invalid-action penalty.Its prompt and response limits are 4096 and 512 tokens; training data size is 256, group size is 5, and the mini-batch size is 512. GIGPO uses ω=1, γ=0.95, and similarity threshold 0.9; rollout and validation temperatures are 1.0 and 0.0.Source: §A.1 · p. 12

A.2 Baseline Details

• The agent 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; Reflexion uses self-reflection and iterative improvement (Yao et al., 2023; Shinn et al., 2024). GPT-4o and Gemini-2.5-Pro are described as large-scale and comparable in scale and capability, respectively (Achiam et al., 2023; Team et al., 2023).Source: §A.2 · p. 12
• The RL baselines are PPO, RLOO, GRPO, and GiGPO, spanning classic policy optimization and group-based advantage estimation.The appendix describes RLOO as estimating advantages without value networks, GRPO as using trajectory-level advantages for multi-step tasks, and GiGPO as performing group-wise advantage estimation for LLM agents (Schulman et al., 2017; Ahmadian et al., 2024; Shao et al., 2024; Feng et al., 2025).Source: §A.2 · p. 12

A.3 Prompts

• Prompt templates use Python-style formatting, with placeholders such as {task_description}, {step_count}, and {current_observation} filled with runtime task context.ALFWorld and WebShop retain the two most recent history steps, while search-augmented QA uses the complete history.Source: §A.3 · pp. 12–13
• Prompts use <think> tags for step-by-step reasoning and <action> tags for the final action; search prompts also tag queries, retrieved evidence, and answers.The paper characterizes the reasoning format as chain-of-thought style (Wei et al., 2022); search queries use <search>, retrieved evidence uses <information>, and final answers use <anwser> tags.Source: §A.3 · p. 13

A.4 Noise-Gradient Analysis

• The WebShop diagnostic samples 100 search-result transitions from Qwen2.5-1.5B-Instruct rollouts and scores only next-observation tokens.It uses the world-model training serialization, with the decision context and action as prefix and the next observation as teacher-forced target.Source: §A.4 · p. 13
• The diagnostic compares CE loss −log p_t,i with MAE loss 1 − p_t,i, the latter matching the bounded per-token term in Clipped MAE before confidence masking.These correspond to the unbounded WM CE objective and the bounded term used by Clipped MAE, respectively.Source: §A.4 · p. 13
• To compare gradient allocation rather than loss scale, it normalizes summed target-logit gradient magnitudes by the total over scored next-observation tokens.The per-token magnitudes are 1 − p_t,i for CE and p_t,i(1 − p_t,i) for MAE.Source: §A.4 · p. 13
• Human-written rules label tokens overlapping ASIN-style product identifiers or brand-like uppercase spans of at least four letters as noisy; all other observation tokens are meaningful.Subword tokens inherit the noisy label when their character offsets overlap a noisy span.Source: §A.4 · p. 13
• The evaluation-only analysis updates no model parameters and isolates loss-function effects on gradient shares to test for over-allocation to noisy WebShop strings.It measures gradient allocation over the same sampled tokens.Source: §A.4 · p. 13

B.1 Full Results of ALFWorld and WebShop

• PaW consistently improves performance over the base RL methods on ALFWorld and WebShop; full results are averaged over three random seeds.The table reports ALFWorld subtask and overall success rates, and WebShop average score and success rate; no specific values are provided here.Source: §B.1 · p. 14, Table 5 · p. 14

B.2 Training Dynamics

• GRPO w/ PaW achieves better training reward than GRPO, while policy-gradient loss and clipped update ratio remain close to the baseline.The reported dynamics indicate that the reward improvement does not come from a large change in PPO-style update magnitude or clipping behavior.Source: §B.2 · p. 14
• The adaptive coefficient decreases during training because higher-reward rollouts receive less world-model supervision under reward-adaptive scaling.Source: §B.2 · p. 14
• World-model loss decreases as the shared policy learns next-observation prediction, while margin clipping bounds the auxiliary signal as predictions improve.The clipped token ratio tracks the number of observation tokens affected by margin clipping.Source: §B.2 · pp. 14–15