Internalizing World Models via Self-Play Finetuning for Agentic RL

TL;DR

The paper tackles the generalization gap of Large Language Model agents in out-of-distribution reinforcement learning by proposing SPA, a simple framework that internalizes a world model through self-play finetuning. It decomposes the world model into state estimation and transition modeling, training them via supervised fine-tuning on self-generated trajectories, and then bootstraps policy optimization with PPO using the learned dynamics priors. Empirical results on Sokoban, FrozenLake, and Sudoku show substantial improvements in Pass@k, outperforming online world-modeling baselines and reward-grounding methods, with strong easy-to-hard transfer within the same environment family. Overall, SPA provides a minimal, reproducible world-modeling backbone that enhances exploration, credit assignment, and generalization for agentic RL in interactive, OOD settings.

Abstract

Large Language Models (LLMs) as agents often struggle in out-of-distribution (OOD) scenarios. Real-world environments are complex and dynamic, governed by task-specific rules and stochasticity, which makes it difficult for LLMs to ground their internal knowledge in those dynamics. Under such OOD conditions, vanilla RL training often fails to scale; we observe Pass@k--the probability that at least one of (k) sampled trajectories succeeds--drops markedly across training steps, indicating brittle exploration and limited generalization. Inspired by model-based reinforcement learning, we hypothesize that equipping LLM agents with an internal world model can better align reasoning with environmental dynamics and improve decision-making. We show how to encode this world model by decomposing it into two components: state representation and transition modeling. Building on this, we introduce SPA, a simple reinforcement learning framework that cold-starts the policy via a Self-Play supervised finetuning (SFT) stage to learn the world model by interacting with the environment, then uses it to simulate future states prior to policy optimization. This simple initialization outperforms the online world-modeling baseline and greatly boosts the RL-based agent training performance. Experiments across diverse environments like Sokoban, FrozenLake, and Sudoku show that our approach significantly improves performance. For example, SPA boosts the Sokoban success rate from 25.6% to 59.8% and raises the FrozenLake score from 22.1% to 70.9% for the Qwen2.5-1.5B-Instruct model.

Figures (10)

• Figure 1: Validation score on five environments. Sokoban, FrozenLake and Sudoku are OOD for the models, where Pass@k ($k$=8) drops with training. However, in ALFWorld and WebShop, Pass@k increases greatly with training.
• Figure 2: Validation performance on Sokoban and Frozen Lake. Red: Vanilla RL(StarPO baseline); Yellow: state-estimation RL; Blue: SPA (world-modeling SFT as a cold start, then state-estimation RL). Left: Pass@k; Right: Pass@1.
• Figure 3: SPA equips an LLM agent with an inner world model. The world model is decomposed into two parts: (i) State Estimation, which converts raw observations into structured, informative natural-language descriptions of the world state; and (ii) Transition Modeling, trained in a self-play Supervised Fine-Tuning setting to predict next-state representations from current states.
• Figure 4: RL performance after world-modeling SFT under three settings: State Estimation RL baseline, No masking vs. Masking current states with next states for SPA.
• Figure 5: Pass@8 and Pass@1 comparison by using correct abstraction and incorrect abstraction on Sokoban during RL training. Left: Abstraction PPO with correct coordinates. Right: randomized coordinates (mis-specified abstraction).