PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement World Model Learning for LLM-based Agents

Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu

arXiv:2602.05842v2cs.CL

TL;DR

LLM-based agents often struggle to anticipate action consequences and adapt to changing environments. RWML learns action-conditioned world models by aligning simulated and observed next states, improving performance on ALFWorld and τ 2 Bench and outperforming direct task-success RL when combined with task-success rewards.

  • Problem

    LLM-based agents struggle to anticipate action consequences and adapt to environment dynamics, motivating world-modeling capabilities beyond standard language pretraining.

  • Method

    RWML is a self-supervised RL method that aligns simulated next states with observed environment states in a pretrained embedding space to learn action-conditioned world models.

  • Results

    RWML improved performance on ALFWorld and τ 2 Bench, while combined with task-success rewards it outperformed direct task-success RL by 6.9 points on ALFWorld and 5.7 on τ 2 Bench.

  • Takeaways & Limitations

    RWML provides a scalable self-supervised route to improve LLM-agent world modeling using interaction data and can match expert-data training when combined with task-success rewards.

  • Takeaways & Limitations

    The study trains and evaluates RWML exclusively on established, isolated benchmarks without real-world impact.

Abstract

from arXiv · show

Large language models (LLMs) have achieved strong performance in language-centric tasks. However, in agentic settings, LLMs often struggle to anticipate action consequences and adapt to environment dynamics, highlighting the need for world-modeling capabilities in LLM-based agents. We propose Reinforcement World Model Learning (RWML), a self-supervised method that learns action-conditioned world models for LLM-based agents on textual states using sim-to-real gap rewards. Our method aligns simulated next states produced by the model with realized next states observed from the environment, encouraging consistency between internal world simulations and actual environment dynamics in a pre-trained embedding space. Unlike next-state token prediction, which prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse, our method provides a more robust training signal and is empirically less susceptible to reward hacking than LLM-as-a-judge. We evaluate our method on ALFWorld and $τ^2$ Bench and observe significant gains over the base model, despite being entirely self-supervised. When combined with task-success rewards, our method outperforms direct task-success reward RL by 6.9 and 5.7 points on ALFWorld and $τ^2$ Bench respectively, while matching the performance of expert-data training.

1. Introduction

RWML addresses LLM-based agents’ difficulty anticipating action consequences and adapting to environment dynamics by learning action-conditioned world models through self-supervised sim-to-real alignment. Across ALFWorld and τ^2 Bench, it improves base performance without expert data and further outperforms direct task-success RL when combined with task-success rewards.

  • Motivation: LLM-based agents struggle to anticipate action consequences and adapt to environment dynamics despite strong linguistic and reasoning abilities.This reflects a gap between language competence and the demands of agentic interaction.
  • Motivation: Standard next-token pretraining emphasizes static-text language modeling, whereas successful long-horizon task completion requires reasoning about state evolution after actions.The mismatch motivates explicit world-modeling objectives for agentic use cases.
  • Method: RWML learns action-conditioned world models by minimizing the discrepancy between simulated and realized next states in a pre-trained embedding space.This sim-to-real alignment prioritizes semantic consistency rather than token-level fidelity.
  • Results: 19.6 and 6.9 points: RWML significantly improves base-model performance on the evaluated benchmarks without expert data, strong LLMs, or task-success rewards.The passage reports these gains across ALFWorld and τ^2 Bench but does not map each value to a specific benchmark.
  • Results: 6.9 and 5.7 points: with task-success rewards, RWML outperforms direct task-success reward RL on ALFWorld and τ^2 Bench, respectively.The method matches the performance of training with expert data.

2. Method

RWML trains an LLM agent to predict action-conditioned next states from environment rollouts, using embedding-based rewards and GRPO before task-success reinforcement learning. Its self-supervised procedure emphasizes difficult examples and requires neither expert data nor task-success signals.

  • Motivation and method: RWML addresses sparse, expert-designed task-success rewards by learning world-model knowledge directly from environment dynamics before task-success RL.This staged approach is intended to improve scalability as tasks and environments become more complex.
  • World-model objective: The policy predicts a next state ˆst+1 conditioned on an action and interaction history, with reasoning tokens generated before the final prediction.The target is the realized next state observed from the environment.
  • World-model objective: RWML uses a binary reward based on the distance between predicted and ground-truth next states, primarily computed with off-the-shelf embedding cosine similarity.The reward is 1.0 when d(ˆst+1, st+1) < τd and 0.0 otherwise; binarization is empirically more robust and less susceptible to hacking.
  • Data collection and optimization: Training data come from target-model rollouts converted into ⟨s≤t, at, st+1⟩ triplets, with N > 1 rollouts per task to improve coverage and diversity.The resulting triplets support action-conditioned next-state learning from the environment rather than expert demonstrations.
  • Data collection and optimization: RWML optimizes the world-model reward with standard GRPO, subsampling easy examples and requiring no expert data, stronger LLMs, or task-success reward signals.Examples consistently achieving high reward across K = 10 attempts are subsampled, while easy examples are retained with probability p = 0.1.

3. Experiments

Experiments on ALFWorld and τ 2 Bench show that RWML improves agentic capability using only interaction data, and that its components contribute materially to performance. RWML also outperforms other training-based baselines when combined with task-success rewards.

  • Benchmarks: The evaluation uses ALFWorld for text-based household tasks and τ 2 Bench for interleaved customer-service tool use.Both benchmarks require agents to understand states, actions, and environment dynamics for task completion.
  • Baselines: Baselines span task-success reward learning, transition-function learning, and expert- or stronger-LLM-based training.Compared methods include RFT, Policy RL, WM SFT, IWM, Self-Reflection, Imitation Learning, and REACT-style prompting.
  • Main results: 19.6 points on ALFWorld and 7.9 points on τ 2 Bench are gained by RWML over the base model without expert data, stronger LLMs, or task-success signals.RWML is trained solely from interaction data and significantly improves agentic capability.
  • Main results: RWML combined with task-success reward outperforms all other training-based baselines.The paper reports this comparison in Table 1 after establishing RWML’s self-supervised gains.
  • Ablations: Ablations show that embedding-based rewards, easy-sample subsampling, and RWML training each contribute to performance.LLM-as-a-judge is unreliable and can sometimes be hacked during training, while weaker models are more sensitive to noisy reward functions.

4. Discussion

RWML improves decision-making accuracy and efficiency, reduces invalid or ineffective actions, and makes targeted parameter updates that preserve compatibility with later policy learning. Its transfer to decision-making on τ^2 Bench depends on base-model capability, with stronger models benefiting most.

  • Decision-making behavior: RWML-trained models make more accurate and efficient decisions by using improved environment knowledge.Examples include predicting a knife on the countertop in ALFWorld and considering airplane mode in τ^2 Bench.
  • Decision-making behavior: 24.90% to 8.84%: invalid tool calls per tool-call made decrease after RWML on τ^2 Bench.Invalid calls include made-up tool names and incorrect arguments.
  • Parameter updates: RWML induces fewer parameter changes across layers than WM SFT, encoding task-relevant information through smaller and more targeted updates.This pattern holds for Qwen3-8B on τ^2 Bench and Qwen2.5-7B-Instruct on ALFWorld.
  • Parameter updates: After Policy RL, RWML-initialized models retain weight-change ratios close to direct Policy RL, whereas WM SFT initialization causes substantially higher ratios and interference.The lower change ratio is consistent across both benchmarks and transformer components, aligning with better mitigation of catastrophic forgetting.
  • Base-model capability: RWML transfer on challenging τ^2 Bench depends on base-model capability: Qwen2.5-7B struggles, while Qwen3-8B and Qwen3-30B-A3B show substantial gains.The stronger models approach Qwen3-235B-A22B-Thinking-2507, indicating RWML is most effective for sufficiently strong base models.

5. Related Work

Related work spans training LLM-based decision-making agents with imitation learning and using world models to improve policy training and inference. Human- or LLM-generated demonstrations provide dense supervision but are costly to scale, motivating complementary world-model approaches.

  • Training Decision-Making Agents: LLM-based agents have been applied to interactive gaming, software engineering, and computer, phone, and browser use.
  • Training Decision-Making Agents: Many early language-agent methods use supervised fine-tuning on human demonstrations or trajectories synthesized from stronger LLMs with manually designed workflows or heuristics.These approaches rely on imitation learning to train decision-making agents.
  • Training Decision-Making Agents: High-quality supervised fine-tuning data provides dense supervision but is difficult to scale because collecting demonstrations is costly.
  • Training World Models: World-model research uses real-world interaction data beyond task-success rewards to support decision-making, including Dyna-style combinations of model-based and model-free learning.For LLM agents, recent work also trains separate world models for inference-time algorithms such as MCTS.

6. Conclusion

RWML is a scalable, self-supervised method that trains LLM-based agents as action-conditioned world models by aligning simulated and observed environment states. It improves environment understanding and decision-making before downstream task-success-reward reinforcement learning, and is evaluated on ALFWorld and τ 2 Bench.

  • 6. Conclusion: RWML enhances LLM-based agents’ environment understanding and decision-making before downstream reinforcement learning with task-success rewards.The method requires neither expert or strong-LLM annotations nor task-success signals during its self-supervised training.
  • 6. Conclusion: RWML trains an LLM as an action-conditioned world model by aligning simulated next states with observed environment states in a pre-trained embedding space.This alignment provides the method’s self-supervised training signal.
  • 6. Conclusion: RWML is evaluated on two long-horizon agent benchmarks: ALFWorld and τ 2 Bench.The supplied passage identifies these as the two evaluation benchmarks.

7. Impact Statements · B. More Details on ALFWorld · B.1. RWML Training Setup

The paper frames RWML as a scalable, self-supervised approach to improve world modeling using interaction data without expert annotations or stronger LLMs. Its ALFWorld setup uses challenging long-horizon tasks and filtered rollout triplets, with embedding-based training and distinct reasoning prompts for RWML and WM SFT.

  • 7. Impact Statements: RWML improves world modeling in LLM-based agents using interaction data without expert annotations or stronger LLMs.The method is trained exclusively on established, isolated benchmarks, while applications and data collection may still pose misuse risks.
  • B. More Details on ALFWorld: ALFWorld tasks can involve over 50 locations and require more than 50 expert-policy steps, challenging planning, subgoal tracking, and exploration.Agents must also identify likely household-item locations, such as desks, shelves, or dressers for desklamps.
  • B.1. RWML Training Setup: 2048 ALFWorld training tasks produced 21,011 training triplets and 2,288 validation triplets after rollout and minor invalid-action filtering.The target model was Qwen2.5-7B-Instruct, with N = 3 trajectories per task, temperature τ = 1.0, and a maximum of 30 steps.
  • B.1. RWML Training Setup: Training triplets use the form ⟨s≤t, at, st+1⟩, converting rollouts into state-history, action, and next-state examples.The triplets were collected for both RWML and WM SFT before postprocessing and dataset splitting.
  • B.1. RWML Training Setup: RWML uses Qwen3-Embedding-8B for rWM training with τd = 0.2 and prompts the model to generate reasoning before its final next-state prediction.WM SFT instead predicts the next state with empty reasoning tokens, a choice intended to support generalization and reasoning during second-stage Policy RL.
  • B.1. RWML Training Setup: WM SFT directly predicts next states with empty reasoning tokens because no reasoning data is available for the training triplets.The authors report that this training method can better enable generalization and reasoning during second-stage Policy RL training.

B.2. Policy RL Training Setup · B.3. Other Training Setup · C. More Details on τ 2 Bench

The paper specifies a shared Policy RL setup for ALFWorld, describes expert-data procedures for alternative training methods, and characterizes τ 2 Bench as a dynamic, tool-using customer-service environment. These details establish the training prompts, trajectory horizon, reward propagation, data sources, and evaluation setting.

  • B.2. Policy RL Training Setup: Policy RL follows prior setups and prompts from Feng et al. (2025c) and Yu et al. (2025a).The same setup is used across Policy RL experiments, including RWML+Policy RL and WM SFT.
  • B.2. Policy RL Training Setup: Training uses ALFWorld’s official training split, with reasoning tokens generated before each action.The prompting format requires <think>...</think> reasoning before action generation.
  • B.2. Policy RL Training Setup: 15 maximum steps and γ = 1.0 propagate terminal task-success rewards to every trajectory turn.These settings are used during Policy RL training.
  • B.3. Other Training Setup: Imitation Learning, IWM, and SR use expert data from the official ALFWorld dataset, following Zhang et al. (2025a).The setup distinguishes expert-data training from the reflection-data procedure used for SR.
  • B.3. Other Training Setup: SR reflection data uses a branching factor of 3 per expert action and is generated with Zhang et al. (2025a)’s prompt.The target model for reflection generation is Qwen2.5-7B-Instruct.
  • C. More Details on τ 2 Bench: τ 2 Bench evaluates customer-service ability in a text-based, long-horizon dual-control environment.Both the agent and user can use tool calls in a shared, dynamic environment.
  • C. More Details on τ 2 Bench: τ 2 Bench tasks require communication, information gathering, tool calls, and adaptation to evolving environment states.Users may also make tool calls, requiring the agent to adapt to those changes.

C.1. RWML Training Setup · C.2. Policy RL Training Setup

RWML training uses diverse multi-turn rollouts converted into state-action-next-state triplets, with safeguards against memorization and rewards tailored to user and tool responses. Policy RL uses official τ^2 Bench prompts with cost-saving simulator and terminal-success reward propagation.

  • C.1. RWML Training Setup: RWML data collection uses six trajectories per task from Qwen3-8B, combining GPT-4.1 and Qwen3-235B-A22B-Instruct user simulators on τ^2 Bench.The τ^2 Bench rollout allocation is N = 3 for each simulator, promoting diversity.
  • C.1. RWML Training Setup: Rollouts are converted into ⟨s_t, a_t, s_t+1⟩ triplets, while prompts withhold user intent and include basic user information to reduce memorization.The method also subsamples “too easy” triplets using τ_d = 0.15 and τ_easy = 0.0, yielding 5,578 triplets.
  • C.1. RWML Training Setup: The official τ^2 Bench evaluation setting uses GPT-4.1 as user simulator and a maximum step of 100.This passage specifies the evaluation configuration rather than a training hyperparameter.
  • C.1. RWML Training Setup: RWML uses Qwen3-Embedding-8B for rWM rewards and rouge-score to capture structured tool-use responses and missing keys or values.The passage states that rouge-score is more effective than the embedding-based measure for these structured outputs.
  • C.1. RWML Training Setup: The τ^2 Bench reward gives 1.0 to sufficiently close user responses, rounded ROUGE scores to tool-use responses, and 0.0 otherwise.The distance threshold is τ_d = 0.4 for user responses, while rounding to increments of 0.2 is intended to improve training stability.
  • C.1. RWML Training Setup: RWML and WM SFT each train for 2 epochs on 4xB200 GPUs, using learning rates of 1e-6 and 2e-6 respectively.RWML uses batch size 32 and group size 16; WM SFT uses an effective batch size of 32.
  • C.2. Policy RL Training Setup: Policy RL extends the codebase for multi-turn τ^2 Bench rollouts and uses official prompts for both training and evaluation.To reduce cost, Policy RL uses Qwen3-235B-A22B-Instruct instead of GPT-4.1 as the user simulator.
  • C.2. Policy RL Training Setup: Policy RL limits trajectories to 30 steps and uses γ = 1.0 to propagate terminal task-success rewards to every turn.The supplied passage ends before specifying the remaining Policy RL setup details.

C.3. Additional Evaluation Results · C.4. Other Training Setup

Additional evaluation uses the official GPT-4.1 user-simulator setting with a 100-step maximum and finds RWML surpasses methods without experts or strong LLMs while remaining competitive with methods that use them. For τ^2 Bench, expert rollouts are constructed through rejection sampling with a strong LLM because expert trajectory annotations are unavailable.

  • C.3. Additional Evaluation Results: RWML is evaluated under the official setting with GPT-4.1 as the user simulator and a maximum of 100 steps.These results are reported in Table A3.
  • C.3. Additional Evaluation Results: The additional evaluation covers representative models from each category in Tables 1 and 2 because GPT-4.1 is expensive.
  • C.3. Additional Evaluation Results: RWML surpasses all training methods that do not use experts or strong LLMs.
  • C.3. Additional Evaluation Results: RWML is competitive with methods that use experts or strong LLMs.
  • C.4. Other Training Setup: τ^2 Bench lacks expert trajectory annotations, so expert rollouts are collected using rejection sampling with Qwen3-235B-A22B-Thinking-2507.
  • C.4. Other Training Setup: The rollout collection uses N_total = 6, split between N = 3 with GPT-4.1 and N = 3 with Qwen3-235B-A22B-Instruct as user simulators.
  • C.4. Other Training Setup: The Qwen3-235B-A22B-Instruct user simulator matches the model used in Section C.1.
  • C.4. Other Training Setup: Only rollouts that successfully solve the tasks are retained as expert rollouts.

D. More Details on Ablation Studies

The ablation details highlight that LLM-as-a-judge can be unreliable for evaluating learned environment dynamics, including rewarding hacked predictions that lack genuine understanding.

  • LLM-as-a-judge: LLM-as-a-judge can award high scores to predictions that do not demonstrate genuine understanding of task-relevant environment dynamics.The authors characterize this evaluator as generally unreliable.
  • LLM-as-a-judge: Table A8 provides an example prompt and hacked response from ALFWorld illustrating this evaluation failure.The example accompanies the LLM-as-a-judge results in Table 4.
  • LLM-as-a-judge: The trained model is Qwen2.5-7B-Instruct, while the judge model is Qwen3-235B-A22B-Instruct.The passage distinguishes the model being trained from the model performing evaluation.

E. More Details on Parameter Change Analysis

Parameter-change analysis finds that RWML produces the most stable updates across layers and modules, while WM SFT causes broader changes. Policy RL preserves similar update profiles after RWML initialization, supporting more coherent training dynamics but offering descriptive rather than mechanistic evidence.

  • Layer-wise analysis: RWML yields the lowest proportion of substantially updated parameters across transformer layers on both benchmarks, whereas WM SFT produces broader modifications.These trends remain highly consistent throughout network depth.
  • Policy RL initialization: Policy RL has similar parameter-change profiles after RWML or base-model initialization, while WM SFT initialization substantially elevates update ratios.The RWML and base-model profiles closely resemble each other, unlike the WM SFT profile.
  • Module-wise analysis: Module-wise changes across attention projections and MLP parameters mirror the layer-wise trends, indicating RWML-induced stability is not localized to specific components.The analysis covers Q/K/V/O attention projections and MLP projection parameters.
  • Overall implication: The results provide additional evidence that using RL during both mid-training and post-training yields more coherent and stable updates than the conventional SFT-then-RL pipeline.This conclusion synthesizes the layer-wise, module-wise, and initialization analyses.
  • Limitations: The analyses are empirical and descriptive, leaving a systematic mechanistic account of RL-based training dynamics for future work.The paper specifically identifies optimization dynamics and mechanistic interpretability as directions for deeper study.
Loading 2602.05842v2…