Reinforcement World Model Learning for LLM-based Agents

Abstract: 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.

• The paper introduces RWML, framing world model learning as an RL problem that uses embedding-based rewards to align predicted and observed states semantically.
• RWML’s methodology significantly improves performance on long-horizon, text-based tasks by reducing invalid actions and mitigating catastrophic forgetting without expert data.
• Experimental results show that RWML combined with policy optimization rivals expert-data methods, offering a scalable and efficient approach for training adaptive LLM agents.

Reinforcement World Model Learning (RWML) for LLM-Based Agents

Motivation and Problem Setting

The increasing utilization of LLMs as autonomous agents in complex, real-world environments exposes a key limitation: LLMs, typically pretrained via next-token prediction objectives, lack the inductive bias and explicit training to model environmental dynamics and anticipate action consequences. This deficit impedes their ability to make informed decisions and adapt effectively to environment transitions, particularly in long-horizon, agentic tasks where reasoning about the consequences of actions is essential.

Previous approaches attempting to equip LLM-based agents with world models have largely relied on supervised fine-tuning (SFT) to predict next textual states, often using data from expert policies or stronger LLMs. However, these approaches face two major issues: (1) scaling limitations due to dependence on high-quality annotated or synthesized data, and (2) a modeling bias that emphasizes token-level similarity rather than semantic or task-relevant world state equivalence, leading to model collapse and insufficient generalization.

The RWML Approach

The paper introduces Reinforcement World Model Learning (RWML), a scalable, self-supervised framework for LLM-based agents. RWML formulates world model learning as a reinforcement learning (RL) problem over environment trajectories, in which the agent is trained to produce action-conditioned next-state predictions that align semantically with actual observed next states. Instead of optimizing for token-level accuracy, RWML employs a reward function based on the distance (typically cosine distance) between the predicted and observed states in a pretrained embedding space, thereby incentivizing semantic consistency over surface-form matching.

Key algorithmic features:

• Self-supervised Data Collection: Rollouts are generated by the target model itself, requiring only environment interaction, not expert demonstrations.
• Reward Computation: Simulated and real next states are compared using pretrained embeddings; rewards are binary or scalar thresholds on embedding similarity to improve robustness and minimize reward hacking.
• Training Algorithm: Grouped Relative Policy Optimization (GRPO), an RL variant, is used to optimize the world model in a manner that maintains strong generalization and avoids catastrophic forgetting.
• Data Subsampling: To focus the agent on non-trivial transitions, RWML subsamples ‘easy’ instances—those trivially solved by a lightweight SFT world model—prioritizing challenging, informative cases.

In contrast to LLM-as-a-judge reward schemes, which can be gamed by degenerate outputs, or direct SFT on next state tokens, RWML’s embedding-based, RL-driven loss provides a more reliable and semantically meaningful training signal.

Experimental Results and Analysis

RWML is evaluated on ALFWorld and Tau2-Bench, two challenging, long-horizon, text-based agentic environments requiring nuanced world model utilization for task success. The following points succinctly capture the main empirical findings:

• Self-supervised world model learning (RWML alone) yields drastic improvements over untuned base LLM agents: Performance increased by 19.6 points (ALFWorld) and 7.9 points (Tau2-Bench) without the use of expert data, step rewards, or strong LLMs.
• RWML + task-success reward RL consistently outperforms plain Policy RL: Combining RWML’s world model pretraining with downstream policy optimization achieves a 6.9 point (ALFWorld) and 5.7 point (Tau2-Bench) boost over direct reward-based RL.
• Equivalence to expert-data regimes: In several domains, RWML plus Policy RL matches or exceeds the performance of methods relying on expert trajectories and high-quality synthetic labels.
• Catastrophic forgetting is notably reduced with RWML: Compared to SFT-based world model learning, models tuned with RWML undergo fewer destructive parameter updates, retain prior knowledge on general, math, and coding benchmarks more robustly, and demonstrate lower layer-wise and module-wise parameter change ratios during further policy optimization.
• Invalid action rates and inefficient decisions decrease: RWML-trained models exhibit a marked reduction in invalid or low-quality actions, passing more consistently through long-horizon tasks.
• Ablation studies: All major RWML design choices are justified empirically. Reward ablation reveals embedding-based scores outperform LLM-as-a-judge and token-level metrics. Removing data subsampling or world model RL significantly degrades performance, especially for smaller LLM backbones.

Theoretical and Practical Implications

RWML demonstrates that RL-based, self-supervised world model learning can provide a scalable and annotation-efficient path towards robust, adaptive LLM-based agents. By learning to semantically align internal simulations with real environmental transitions, LLMs can acquire more generalizable, task-relevant knowledge about domains beyond the surface form of language.

On the practical side, the decoupling from expert data and task-specific reward shaping enables large-scale deployment and adaptation in heterogeneous, evolving environments where collecting high-quality demonstrations or engineering terminal rewards is costly or infeasible. The compactness and parameter efficiency of the resulting models also suggest better resilience under continual or transfer learning regimes, with less interference during downstream policy updates.

One limitation identified is the dependence on the base model's capacity: weaker LLMs, even when world-model-trained with RWML, may struggle in highly complex environments compared to stronger base models. This motivates future lines of research in architecture or training procedures facilitating better transfer from world modeling to policy optimization in resource-constrained agents.

Impact on Future Developments

The RWML paradigm hints at a modular, staged training pipeline where self-supervised world model learning via RL serves as a robust bridge between foundation model pretraining and downstream RL for agentic tasks. This architecture could foster both sample efficiency and safer, more predictable adaptation as LLM-agents are further integrated into high-stakes, interactive, or continually evolving environments.

Future research directions include hybrid model-based/model-free agent training that exploits RWML-internal simulations for planning and exploration, sophisticated reward inference circumventing explicit reward signals, and rigorous theoretical grounding in optimization landscapes induced by RL-based mid-training. Further mechanistic studies on knowledge retention and parameter update patterns will be crucial for scalable, lifelong agent learning.

Conclusion

Reinforcement World Model Learning (RWML) establishes an effective, scalable methodology for equipping LLM-based agents with robust, semantically grounded world models. By bridging the sim-to-real modeling gap with RL in the embedding space, RWML advances the state of LLM agents both in sample complexity and in their adaptability to unfamiliar, long-horizon decision-making tasks. Its approach sets a foundation for future mid-training protocols and agent architectures emphasizing generalizable world knowledge and efficient downstream adaptation (2602.05842).