RWML: Reinforcement World Model Learning for LLM-based Agents

Paper: 2602.05842 Authors: Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu Categories: cs.CL

Abstract

RWML is a self-supervised approach enabling language models to develop action-conditioned world models for textual environments. The method uses sim-to-real gap rewards to align predicted states with actual environment outcomes. Unlike token-level prediction, this technique emphasizes semantic consistency rather than exact wording, reducing model collapse risks.

Key Contributions

• Self-supervised world model learning: No need for expert demonstrations or environment-specific reward engineering
• Sim-to-real gap rewards: Uses pre-trained embedding space for semantic alignment between predicted and actual states
• Robust training signal: More stable than “LLM-as-a-judge” approaches

The Agent World Model Problem

LLM-based agents operating in interactive environments need to predict the consequences of their actions. Current approaches face several challenges:

• Token-level prediction: Requiring exact text reproduction is brittle and prone to model collapse
• LLM-as-a-judge: Using another LLM to evaluate predictions is expensive and inconsistent
• Expert demonstrations: Requiring human-labeled trajectories limits scalability

RWML: Self-Supervised World Models

RWML takes a different approach by training the LLM to predict semantic outcomes rather than exact text:

Sim-to-Real Gap Rewards

Instead of comparing predicted and actual states at the token level, RWML computes similarity in a pre-trained embedding space. This means the model is rewarded for capturing the meaning of state transitions, not their exact wording.

Self-Supervised Training

The agent generates predictions about what will happen after taking an action, then receives reward based on how semantically close its prediction was to the actual outcome. No expert labels needed.

Results

On ALFWorld and τ² Bench benchmarks:

• Outperformed direct reward RL by 6.9 and 5.7 points respectively
• Achieved performance comparable to expert-trained models
• More robust training dynamics than token-level prediction approaches

Why This Matters

1. Scalable agent training: Self-supervised approach eliminates the need for expert demonstrations
2. Semantic robustness: Embedding-space rewards are more stable than token-level or LLM-judge approaches
3. World model as auxiliary task: Learning to predict outcomes improves the agent’s planning ability
4. Generalizable framework: Applicable to any text-based interactive environment