Reinforcement World Model Learning for LLM-based Agents

ArXiv: 2602.05842

🎯 Pitch

LLMs trained solely on static text fail catastrophically when trying to predict how their actions change the world—RWML fixes this by simply aligning a model's internal simulations with what actually happens, boosting downstream task success even without any task reward.

1. Executive Summary

This paper proposes Reinforcement World Model Learning (RWML), a self-supervised method that trains LLM-based agents to learn action-conditioned world models by aligning simulated next states with realized environment states in a pre-trained embedding space, improving downstream decision-making without requiring expert data or task-success rewards. Evaluated on ALFWorld and τ² Bench using Qwen models, RWML significantly boosts base model performance—by 19.6 and 6.9 points respectively—through world model learning alone, and when combined with task-success reward RL, outperforms direct policy RL by 6.9 and 5.7 points while matching expert-data training. The work establishes that RL-based world model learning serves as a scalable "mid-training" stage that improves agentic capability, but only when the base model possesses sufficient underlying capability to transfer world knowledge to decision-making.

2. Context and Motivation

The Core Problem: LLMs Lack Internal World Models for Agentic Settings

The fundamental problem this paper tackles is a capability gap that emerges when LLMs transition from standard language tasks to agentic environments. Standard LLM pretraining produces models with strong language competence — they can generate coherent text, answer questions, and follow instructions. But when these same models are deployed as autonomous agents operating in complex, long-horizon environments, they struggle with a qualitatively different demand: anticipating the consequences of their actions and adapting to environment dynamics (Section 1).

This distinction is not merely a matter of degree. The paper draws on neuroscience and cognitive science to argue that the ability to model potential future outcomes of one's actions — what they call world modeling — is central to biological intelligence (Craik, 1944; Tolman, 1948; Daw et al., 2005). Animals and humans use internal world models to reason, plan, explore, and learn from very few trials. The paper argues this capacity is "likewise essential for effective reasoning and planning in LLM-based agents" (Section 1, paragraph 3).

Concretely, in environments like ALFWorld (where an agent must navigate a household, find objects, and complete tasks using natural language) or τ² Bench (where a customer service agent uses tool calls while conversing with a simulated user), an LLM-based agent needs to understand questions like: If I go to the sidetable, what objects am I likely to find? If I issue this database lookup, what schema will the result have? If I tell the user their refund is being processed, how will they likely respond? These predictions about environment dynamics — the transition function $\mathcal{T}$ of the underlying MDP — are not reliably learned during standard language pretraining on static text corpora.

This gap is significant for several reasons the paper highlights implicitly throughout the introduction and method sections:

• Scalability of agent training: As tasks and environments become more complex, designing hand-crafted reward functions becomes increasingly difficult and domain-specific (Section 2.2). A method that learns from raw interaction data without task-success signals would scale more naturally.
• Data efficiency: Expert demonstrations are expensive to collect. Strong LLMs that could generate synthetic training data are themselves expensive to query at scale. A self-supervised approach that uses the agent's own interaction data addresses both bottlenecks.
• Complementarity with policy learning: The paper frames world model learning not as a replacement for task-success RL but as a preparatory "mid-training" stage that enables more effective downstream policy optimization (Figure 1). This two-stage pipeline — learn the environment dynamics first, then learn the task — mirrors classic model-based RL insights from the Dyna architecture (Sutton, 1991).

Conflicting Design Tensions in Prior Work

The paper is motivated by a set of tensions in how prior work approaches agent training, world model learning, and the choice of training objective (Section 1, Section 5).

Tension 1: SFT scales poorly for agent data. The dominant paradigm for training LLM agents has been supervised fine-tuning (SFT) on expert demonstrations or trajectories synthesized by stronger LLMs (Zeng et al., 2023; Chen et al., 2024; Su et al., 2025). While SFT provides dense supervision signals, it faces two scalability challenges that the paper explicitly identifies (Section 1, paragraph 4):

"they rely heavily on high-quality data from experts/strong LLMs; and they are based on SFT, which prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse."

The "token-level fidelity" critique is central to the paper's argument. When an SFT-trained model is asked to predict the next state $s_{t+1}$ given $(s_{\le t}, a_t)$, the training objective penalizes any deviation from the exact string that appeared in the training data — even if the model's prediction is semantically correct but uses different wording. This creates a mismatch: what the agent needs is semantic understanding of environment dynamics, but SFT optimizes for lexical reproduction. The paper argues this can lead to model collapse, where the model overfits to surface patterns rather than learning generalizable transition dynamics.

Tension 2: RL for agents requires expensive reward engineering. More recent work has explored using reinforcement learning (specifically PPO and GRPO) to train LLM agents directly from task-success rewards (Feng et al., 2025c; Tan et al., 2025; Luo et al., 2025a). While this bypasses the need for step-by-step demonstrations, it introduces a different bottleneck:

"designing task-success reward functions in complex environments still requires substantial human expertise" (Section 5, paragraph 1)

For benchmarks like ALFWorld, the task-success signal is naturally binary (task completed or not). But for more complex environments like τ² Bench — where success depends on nuanced criteria like whether the agent correctly followed policy, appropriately escalated issues, and maintained user satisfaction — the reward design becomes non-trivial and domain-specific (Barres et al., 2025). The paper argues this reliance on carefully designed terminal rewards "introduces scaling challenges" (Section 2.2, paragraph 1).

Tension 3: Existing world model approaches require separate models or external supervision. Prior work on world models for LLM agents has typically taken one of two approaches, both of which introduce overhead. Some methods train a separate world model to support inference-time planning algorithms like MCTS (Hao et al., 2023; Wu et al., 2025; Chae et al., 2025), requiring the overhead of training and serving multiple models. Others jointly learn world models and policies within a single model but rely on expert annotations or data synthesized by stronger LLMs during world model training (Zhang et al., 2025a; Yu et al., 2025a, c). The paper explicitly notes this dependency:

"these approaches either require expensive training/inference of multiple models, or rely on additional annotations from experts/strong LLMs during world model learning" (Section 5, paragraph 2)

Where Prior World Model SFT Approaches Fall Short

The paper positions itself most directly against prior work that trains LLMs to predict next states using SFT — which they call World Model SFT (WM SFT). The limitations they identify are both empirical and conceptual:

Empirical failure. In the paper's own experiments (Table 1), WM SFT alone actually degrades performance on ALFWorld compared to the base model (2.8% vs. 36.1% for the base ReACT model) and underperforms even the base model on τ² Bench (27.9% vs. 31.9%). This is a striking negative result: applying SFT to learn world dynamics from interaction data, without additional task-success RL, makes the agent worse at its primary job. The paper attributes this to the token-level fidelity problem — the SFT objective forces the model to memorize surface-level patterns rather than learning transferable world knowledge, and this overfitting interferes with the model's decision-making capabilities.

Forgetting. Table 3 quantifies the problem: WM SFT causes substantially more catastrophic forgetting of general capabilities (measured on MMLU-Redux, IFEval, MATH-500, GSM8k, GPQA-Diamond, and LiveCodeBench) compared to RWML. For example, Qwen2.5-7B drops from 77.26% to 67.16% on MMLU-Redux after WM SFT on ALFWorld (a -10.10 point decline), versus only to 74.88% after RWML (a -2.38 point decline). This is consistent with the "model collapse" concern: SFT aggressively rewrites model parameters to match the training distribution, damaging capabilities that were useful for the broader world.

Parameter disruption. The weight change analysis in Section 4.2 (Figure 3) reveals a mechanistic explanation: WM SFT induces "notably more parameter changes across layers compared to RWML" and subsequent Policy RL applied on top of WM SFT produces "substantially higher change ratios after policy optimization, reflecting stronger parametric interference." In plain terms: WM SFT scrambles the model's parameters in ways that both hurt existing capabilities and make subsequent policy learning harder.

The paper also notes that prior world model approaches like IWM (Implicit World Modeling) and SR (Self-Reflection) from Zhang et al. (2025a) and Yu et al. (2025c) — while effective — depend on expert rollouts and strong LLM-generated synthetic data. On ALFWorld, for example, SR requires GPT-4o to generate reflection data because the target model (Qwen2.5-7B) "cannot consistently generate coherent reflections" (Appendix B.3). This dependency on external supervision undermines the scalability argument.

How This Paper Positions Itself

The paper frames RWML as addressing a specific gap in the training pipeline for LLM agents (Figure 1): the space between pretraining and task-specific policy RL. Standard practice is either (a) go directly from pretrained model to policy RL with task-success rewards, or (b) insert an SFT stage using expert data or stronger LLMs. The paper argues for a third option: a self-supervised RL stage that teaches the model to understand environment dynamics using only its own interaction data, before task-specific policy learning begins.

The conceptual framework draws an explicit parallel to the classic Dyna architecture (Sutton, 1991), which used learned world models to augment model-free RL with simulated experience. But where Dyna trains a separate world model, RWML trains the policy model itself to also function as a world model — the same model that will later be optimized for task success first learns to anticipate environment transitions. This "single model" design choice is pragmatic: it avoids the overhead of training and serving multiple models, and the world model knowledge acquired during the mid-training stage transfers directly into improved decision-making during the downstream policy RL stage.

The paper's position relative to existing work can be summarized along three axes it explicitly distinguishes (Table A1):

The key innovation is not the idea of world models for agents — that has a long history in both classical RL and recent LLM agent work. Rather, it is the combination of three design choices that prior work had not simultaneously realized: (1) self-supervised reward (no expert data or task labels), (2) RL training objective (avoiding SFT's token-fidelity pathologies), and (3) single-model architecture (the world model is the policy model).

The paper also positions RWML as complementary to, rather than competitive with, task-success RL. The central claim is that RWML serves as an effective "mid-training" stage that improves the starting point for Policy RL, and the combination (RWML + Policy RL) consistently outperforms Policy RL alone across both benchmarks (Table 1). This two-stage approach — learn the world first, then learn the task — is presented as a general recipe that can sit alongside existing post-training methods.

The Motivation for Embedding-Space Rewards vs. Token-Level Objectives

A subtle but important motivation that runs through the paper is the choice of reward signal. Why use embedding cosine similarity rather than next-token prediction? The paper's argument (Section 2.2, Section 3.4) is multi-layered:

1. Semantic equivalence over lexical fidelity: In agentic environments, the same environment state can be described in many lexically different but semantically equivalent ways. Consider: "You arrive at sidetable 1. On the sidetable you see a mug, a pepper shaker, and a tomato" vs. "You are at sidetable 1. There is a mug, pepper shaker, and tomato on it." These are semantically identical, but SFT would penalize the model for any deviation. The embedding-space reward, by using cosine similarity, recognizes semantic equivalence.

2. Resistance to reward hacking: The ablation studies (Section 3.4, Table 4) compare embedding-based rewards with LLM-as-a-judge rewards (where a stronger LLM evaluates prediction quality). The paper finds that LLM-as-a-judge "is unreliable and can sometimes be hacked during training" — the model learns to produce outputs that superficially satisfy the judge LLM without demonstrating genuine environment understanding (see the example in Appendix D, Table A8, where a vague prediction receives a score of 1.0). The embedding-based reward, being a simpler, non-parametric signal, is "less susceptible to reward hacking" (Section 2.2, footnote 2).

3. Binarization for robustness: The paper empirically finds that binarizing the continuous cosine similarity into a 0/1 reward (with threshold $\tau_d$) improves robustness compared to using raw similarity scores (footnote 2). This is a practical design insight: a crisp "good enough" signal is more stable for RL training than a continuous, noisy reward.

These motivations are not merely theoretical — they are validated in the ablation that shows removing the embedding-based reward (substituting LLM-as-a-judge) substantially degrades performance (Table 4).

3. Technical Approach

3.1 Reader orientation

This paper proposes a training method — not a new model architecture or inference algorithm — that teaches an LLM-based agent to understand environment dynamics before learning to solve tasks. The core idea is to use reinforcement learning with a self-supervised reward (embedding cosine similarity between predicted and actual next states) to train the model as an action-conditioned world model, avoiding the need for expert demonstrations, stronger LLMs, or manually designed task-success rewards during this "mid-training" stage. The method solves the problem that LLMs, despite strong language capabilities, lack internal models of how environments respond to actions — and prior approaches to teaching world models through supervised fine-tuning (SFT) fail because they optimize token-level fidelity rather than semantic understanding.

3.2 Big-picture architecture (diagram in words)

The RWML system has four major components arranged in a pipeline:

1. Rollout Collector — uses the target model $\pi_\theta$ to interact with the environment, collecting trajectories of (state, action, next_state) triplets. This requires no expert data or external supervision; the model explores using its own (initially suboptimal) policy.

2. Data Filter — subsamples the collected triplets to remove "too easy" examples using a separately fine-tuned SFT model, keeping a training set biased toward medium-to-hard samples that require genuine world model learning.

3. World Model RL Trainer — trains $\pi_\theta$ using GRPO to predict next states $\hat{s}_{t+1}$ given $(s_{\le t}, a_t)$, with rewards computed by embedding cosine similarity between predicted and actual next states. This component teaches the model to reason about environment dynamics.

4. Downstream Policy RL (optional second stage) — takes the RWML-trained model and further fine-tunes it with GRPO using standard task-success rewards, benefiting from the world model knowledge acquired in the previous stage.

The information flow is: the agent interacts with the environment using its current policy (component 1) → collected trajectories are converted to world model training triplets and filtered (component 2) → the model is trained via RL to predict next states with embedding-space rewards (component 3) → optionally, the resulting model is further trained with task-success RL (component 4). Steps 1–3 are entirely self-supervised — no ground-truth task labels or expert data are needed.

3.3 Roadmap for the deep dive

• First, the mathematical framework (MDP notation, state/action definitions, what the world model learning objective looks like formally), because everything else builds on this abstraction and the distinction between real states $s_t$ and predicted states $\hat{s}_t$.
• Second, the reward function — the embedding-space cosine similarity reward and why it is chosen over token-level SFT or LLM-as-a-judge — since the reward signal is the core innovation that makes self-supervised RL viable for world model learning.
• Third, the data collection and subsampling pipeline, explaining how training triplets are gathered from the model's own rollouts and how "too easy" samples are filtered out to focus learning on non-trivial world dynamics.
• Fourth, the GRPO training process — the RL algorithm, group-relative advantages, KL regularization, and how the model is prompted to generate reasoning before predicting next states.
• Fifth, the downstream integration with Policy RL, showing how the RWML-trained model serves as an improved initialization for task-success RL and why this two-stage approach outperforms both standalone methods.

3.4 Detailed, sentence-based technical breakdown

This is primarily a methods paper proposing a new training algorithm. The core idea is that an LLM-based agent can improve its decision-making by first learning to function as an action-conditioned world model through RL with embedding-space rewards, entirely from its own interaction data without external supervision.

MDP Formulation and World Model Objective

The paper formalizes the agentic task as a Markov Decision Process (Section 2.1) with the standard tuple $\langle \mathcal{S}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma \rangle$, where $\mathcal{S}$ is the state space (textual observations from the environment), $\mathcal{A}$ is the action space (natural language commands or tool calls), $\mathcal{T}$ is the environment transition function (how the environment responds to actions), $\mathcal{R}$ is the reward function (task-success signal, only available at episode termination and only during the optional second stage), and $\gamma \in (0, 1]$ is the discount factor.

The agent $\pi_\theta$ operates as follows: at time step $t$, it receives the current state $s_t$ (a textual observation) and the interaction history of up to $H$ previous turns $s_{\le t} \equiv \langle s_{t-H}, a_{t-H}, ..., s_t \rangle$, generates an action $a_t \sim \pi_\theta(\cdot | s_{\le t})$, and the environment returns a new state $s_{t+1} \sim \mathcal{T}(\cdot | s_t, a_t)$. The paper explicitly notes that $s$ is technically an observation (as in a POMDP), but simplifies notation by treating it as the state since all textual information received from the environment is included.

The critical notational distinction introduced for world model learning is the use of $\hat{s}_{t+1}$ to denote simulated/predicted next states produced by the model, versus $s_{t+1}$ for realized environment states. This distinction is essential because the entire training signal comes from comparing these two.

The world model learning objective is introduced in Section 2.2. Given a history $s_{\le t}$ and a proposed action $a_t$, the model is trained to produce:

$(\text{reason}, \hat{s}_{t+1}) \sim \pi_\theta(\cdot | s_{\le t}, a_t)$

where $\text{reason}$ denotes reasoning tokens generated before the final prediction $\hat{s}_{t+1}$, and $s_{\le t} \equiv \langle s_{t-H}, a_{t-H}, ..., s_t \rangle$ is the $H$-step interaction history.

What it computes: the model autoregressively generates a sequence of tokens — first a chain-of-thought reasoning about the current situation, then a prediction of what the environment will return as the next state after executing action $a_t$. The reasoning tokens are free-form text that the model uses to articulate its understanding of the environment before committing to a prediction.

Why this form: the inclusion of reasoning tokens before the prediction is standard in LLM training recipes (used in R1-style reasoning, GRPO, and ReACT) because it gives the model space to perform intermediate computation — essentially acting as a scratchpad for world model simulation. The paper contrasts this with WM SFT, where the model is trained to directly output $\hat{s}_{t+1}$ without reasoning, noting that "since there is no reasoning data available for the triplets, we find this training method for WM SFT can better enable generalization/reasoning during the second stage Policy RL training" (Appendix B.1). In other words, forcing the SFT model to predict without reasoning is actually beneficial for WM SFT because it prevents the model from learning to rely on a reasoning scaffold that may not transfer to the policy learning stage, whereas RWML uses RL to teach the model to generate its own reasoning.

The World Model Reward Function

The core innovation of RWML — and what distinguishes it from prior work — is the choice of reward signal for training the world model. The paper defines a binary reward function (Section 2.2):

$r^{\text{WM}}(\hat{s}_{t+1}, s_{t+1}) = \begin{cases} 1.0, & \text{if } d(\hat{s}_{t+1}, s_{t+1}) < \tau_d, \\ 0.0, & \text{otherwise}. \end{cases}$

where $\tau_d$ is a distance threshold hyperparameter (set to 0.2 for ALFWorld, 0.4 for τ² Bench), and $d(\cdot, \cdot)$ is a distance function measuring the discrepancy between predicted and actual next states.

What it computes: a binary 0/1 reward — the model receives a reward of 1.0 if its predicted next state is "close enough" to the actual next state (below the threshold $\tau_d$), and 0.0 otherwise. "Close enough" is measured in an embedding space, not in token space.

Why this form: the binary reward is chosen empirically over continuous rewards, with the paper noting in a footnote that "binarized rewards are more robust and less susceptible to hacking" (Section 2.2, footnote 2). This is a practical insight from the RL literature: binary rewards create a clear success/failure signal that is harder for the model to exploit through degenerate solutions compared to continuous rewards, where the model might learn to produce outputs that receive moderate-but-meaningless scores. The threshold $\tau_d$ controls how strict the "good enough" criterion is — a lower threshold demands more precise predictions but may make the reward too sparse.

The distance function $d$ is implemented using an off-the-shelf embedding model $E(\cdot)$ with cosine similarity:

$d(\hat{s}_{t+1}, s_{t+1}) = 1 - \cos(E(\hat{s}_{t+1}), E(s_{t+1}))$

where $E(\cdot)$ is a pre-trained text embedding model (Qwen3-Embedding-8B, Zhang et al., 2025b), and $\cos(\mathbf{u}, \mathbf{v}) = \frac{\mathbf{u} \cdot \mathbf{v}}{\|\mathbf{u}\| \|\mathbf{v}\|}$ is cosine similarity.

What it computes: the embedding model maps both the predicted and actual next-state strings to dense vector representations. Cosine similarity between these vectors ranges from -1 (opposite directions) to 1 (identical direction). The distance $1 - \cos(\cdot, \cdot)$ ranges from 0 (identical embeddings) to 2 (opposite embeddings). When this distance is below $\tau_d$, the prediction is considered semantically close enough to the ground truth.

Why this form: this is the crucial design choice that separates RWML from SFT-based world model learning. Cosine similarity in embedding space measures semantic equivalence rather than token-level fidelity. Two different wordings of the same environment observation — "You arrive at sidetable 1. On the sidetable you see a mug" versus "You are at sidetable 1. There is a mug on it" — will have high cosine similarity because embedding models are trained to map semantically similar texts to nearby vectors, even when the surface forms differ. In contrast, SFT would penalize the model for any token deviation from the exact training string. The choice of Qwen3-Embedding-8B specifically matters because it is "a fast, strong, open-source LLM that can be hosted locally" (Section 3.4) — it provides a frozen, non-trainable judge that doesn't introduce training instability.

The paper also introduces a specialized variant of the reward for τ² Bench (Appendix C.1), because this environment includes both natural language user responses and structured tool-use responses (often JSON). The modified reward is:

$r^{\text{WM}}(\hat{s}_{t+1}, s_{t+1}) = \begin{cases} 1.0, & \text{if } d(\hat{s}_{t+1}, s_{t+1}) < \tau_d \text{ and } s_{t+1} \text{ is a user response}, \\ \text{round(rouge}(\hat{s}_{t+1}, s_{t+1}), 0.2), & \text{if } s_{t+1} \text{ is a tool-use response}, \\ 0.0, & \text{otherwise}. \end{cases}$

where $\text{rouge}(\cdot, \cdot)$ computes ROUGE scores (Lin, 2004) measuring n-gram overlap between predicted and actual tool responses, and $\text{round}(\cdot, 0.2)$ rounds to the nearest 0.2 increment.

What it computes: for user responses (natural language), the same embedding similarity reward with threshold $\tau_d = 0.4$ is used (the threshold is higher than ALFWorld's 0.2 because "user responses are highly non-deterministic" — a user simulator might respond in many valid ways to the same agent action). For tool responses (structured JSON), ROUGE score captures structural fidelity — whether the predicted response has the right keys, value types, and overall format — and is binned into discrete reward levels (0.0, 0.2, 0.4, 0.6, 0.8, 1.0) for stability.

Why this hybrid form: the paper notes that "tool-use responses are generally structured outputs, [and] using rouge-score is more effective at capturing these structures and any missing keys/values" (Appendix C.1). Embedding similarity might give high scores to JSON that is semantically related but structurally wrong (e.g., {"name": "John"} vs. {"full_name": "John", "customer_id": "abc123"}), whereas ROUGE penalizes missing keys and schema mismatches. The round-to-0.2 discretization promotes "better training stability" by reducing noise in the continuous reward signal during RL.

Data Collection Pipeline

The training data for RWML is collected entirely from the target model's own interactions with the environment — no external data, expert demonstrations, or stronger LLMs are needed (Section 2.2, paragraph 4). The process follows a clear sequence:

Step 1: Trajectory rollout. For each training task, the target model $\pi_\theta$ (in its current, pre-training state) generates $N > 1$ complete trajectories $(s_0, a_0, s_1, a_1, ... , s_T)$ by interacting with the environment with temperature $\tau = 1.0$. The paper sets $N = 3$ for ALFWorld (2048 training tasks) and $N_{\text{total}} = 6$ for τ² Bench (178 training tasks). The choice of $N > 1$ is motivated to "improve coverage and diversity" — by generating multiple rollouts per task, the dataset captures different possible trajectories and environment responses, reducing the risk that the world model only learns a narrow distribution of transitions. For τ² Bench specifically, the 6 rollouts per task are split equally: $N = 3$ using GPT-4.1 as the user simulator and $N = 3$ using Qwen3-235B-A22B-Instruct as the user simulator, again "to promote diversity" (Appendix C.1).

Step 2: Triplet extraction. Each rollout is converted into training triplets by iterating over all time steps $t$: for each step, the history up to that point $s_{\le t}$, the action taken $a_t$, and the resulting next state $s_{t+1}$ form a training sample $\langle s_{\le t}, a_t, s_{t+1} \rangle$. This means every intermediate step in every trajectory becomes a separate world model training example — the model learns to predict next states not just at the beginning of episodes but at every decision point. After minor postprocessing (removing samples with invalid actions), this yields 21,011 training triplets and 2,288 validation triplets for ALFWorld, and 5,578 training triplets for τ² Bench after subsampling (see next step). For τ² Bench, an additional preprocessing step is applied: database values in tool responses are masked by replacing them with their OpenAPI schema definitions. For example, {"customer_id": "abc123", "full_name": "John Doe"} becomes {"type": "object", "properties": {"customer_id": {"type": "string"}, "full_name": {"type": "string"}}}. This prevents the model from memorizing specific database entries rather than learning the general structure of tool responses.

Step 3: "Too easy" sample subsampling. A key practical component of the data pipeline is filtering out training examples that are trivially easy for the model to predict. The motivation, drawn from Snell et al. (2024) and Sun et al. (2025), is that "easy" samples don't contribute to learning meaningful world dynamics and may dilute the training signal. The procedure is:

1. Train a separate SFT model $\pi'_{\theta}$ on 10% of the dataset (the validation split) to predict $\hat{s}_{t+1}$ from $\langle s_{\le t}, a_t \rangle$. This model serves as a weak preliminary world model.
2. Use $\pi'_{\theta}$ to generate $K = 10$ predictions per sample on the remaining 90% of the dataset (the training split).
3. Compute the average reward across these $K$ attempts:

$\frac{1}{K} \sum_{K} r^{\text{WM}}(\hat{s}_{t+1}, s_{t+1}) \geq \tau_{\text{easy}}$

where $\tau_{\text{easy}} = 0.0$ for both benchmarks (meaning any sample where the weak SFT model achieves any reward at least once in 10 attempts is considered "easy").

4. Samples that meet this threshold (consistently high reward) are retained in the final training set with probability $p = 0.1$. Samples below the threshold are always retained.

What it computes: this procedure identifies training examples that even a simple SFT model (trained on only 10% of the data) can already predict correctly. These "easy" samples are heavily downsampled (retained only 10% of the time), while "hard" samples — those the SFT model consistently fails on — are always kept. The result is a final dataset biased toward medium-to-hard examples while preserving some easy samples for diversity.

Why this form: the paper notes that this value is "chosen heuristically without tuning" and is fixed across both benchmarks (Appendix B.1, C.1). The $\tau_{\text{easy}} = 0.0$ threshold means a sample is "easy" if the weak model can generate any prediction that receives a non-zero reward in 10 attempts — an intentionally low bar that identifies genuinely trivial cases. The 10% retention probability ($p = 0.1$) ensures the dataset doesn't collapse to only hard samples, which could cause the RL training to fail from excessive difficulty. The specific percentages are not ablated, but the ablation study (Table 4) shows that removing the subsampling step entirely ("w/o subsample") degrades performance on τ² Bench from 38.8 to 35.8, confirming that focusing on non-trivial samples is beneficial.

The resulting filtered datasets contain 15,813 triplets for ALFWorld and 5,578 triplets for τ² Bench, with approximately 60% of τ² Bench next-states being tool-use responses and 40% being user responses (Appendix C.1).

GRPO Training for World Model Learning

Once the filtered training dataset is constructed, RWML uses Group Relative Policy Optimization (GRPO) (Shao et al., 2024; DeepSeek-AI et al., 2025) to train the model $\pi_\theta$ to predict next states (Section 2.2). The GRPO objective is:

$\mathbb{E}_{\pi_{\theta_{\text{old}}}} \left[ \min\left(\rho_\theta A, \text{clip}(\rho_\theta, 1 \pm \epsilon) A\right) - \beta D_{\text{KL}}(\pi_\theta || \pi_{\theta_{\text{ref}}}) \right]$

where $\rho_\theta = \pi_\theta(y|x) / \pi_{\theta_{\text{ref}}}(y|x)$ is the importance sampling ratio between the current policy and a reference policy, $A = [r^{\text{WM}} - \text{mean}(r^{\text{WM}})] / \text{std}(r^{\text{WM}})$ is the group-relative advantage computed by normalizing rewards within each training batch, $\epsilon$ is the clipping parameter (standard PPO-style clipping to prevent too-large policy updates), and $\beta$ is the KL regularization coefficient penalizing deviation from the reference policy $\pi_{\theta_{\text{ref}}}$ (the frozen pre-training checkpoint).

What it computes: for each training batch (called a "group" in GRPO), the model generates $G$ predictions per input (where the group size $G = 8$ for ALFWorld and $G = 16$ for τ² Bench). Each prediction receives a world model reward $r^{\text{WM}}$. These rewards are normalized within the group to zero mean and unit variance to produce advantages $A$ — predictions that are better than the group average get positive advantages, worse ones get negative advantages. The PPO-style clipped surrogate objective then updates the policy to increase the probability of predictions that received positive advantages while staying close to the reference policy via the KL penalty.

Why this form: GRPO is chosen over standard PPO because it uses group-relative advantages rather than requiring a learned value function (critic). This matters practically for LLM training because it avoids the memory and compute overhead of training a separate value head. The group-relative normalization also provides a natural curriculum: as the model improves and all predictions become better, the advantage normalization automatically re-centers the signal around the current performance level rather than an absolute scale. The KL penalty with coefficient $\beta$ prevents the model from diverging too far from its pretrained distribution, which is crucial for preserving general capabilities and avoiding catastrophic forgetting (consistent with the findings in Section 3.3).

Prompting format for world model prediction. The model is prompted to generate reasoning before predicting the next state, producing output in the format:

<thinking> [reasoning tokens] </thinking>
<next_state> [predicted next state] </next_state>

The specific prompts differ between benchmarks but follow the same structure (Appendix Tables A4 and A6). For ALFWorld (Table A4):

"Your task is to predict the immediate next observation after taking the potential action above. You should first briefly reason step-by-step about the previous steps and current situation — summarize key information you've learned about the environment that is relevant to the task. This reflection and reasoning process must be enclosed within <thinking> </thinking> tags. Once you've finished your reasoning, you should describe the next observation (use the past and current observations as examples!) and present them within <next_state></next_state> tags."

For τ² Bench (Table A6), the prompt additionally provides available user information (e.g., "I am John Doe. My phone number is 123-456-7890") and instructs the model to predict the OpenAPI schema rather than actual values for JSON tool responses.

Contrast with WM SFT prompting. The WM SFT baseline (Tables A5 and A7) uses a deliberately different prompt that instructs the model to NOT generate reasoning:

"Directly present your final prediction of the next observation ... within <next_state></next_state> tags. DO NOT generate anything else."

and the ground-truth label includes empty reasoning tags: Label: <thinking> </thinking> <next_state> {next_state} </next_state>. The paper's reasoning (Appendix B.1) is that since no reasoning data is available for the triplets, forcing WM SFT to predict without reasoning during training "can better enable generalization/reasoning during the second stage Policy RL training" — the SFT-learned world model knowledge is encoded directly in the model weights without being conditioned on a particular reasoning style, making it more transferable to the policy learning stage.

Training hyperparameters. For RWML on ALFWorld: learning rate $1 \times 10^{-6}$, batch size 32, group size 8, 2 epochs, trained on 2×B200 GPUs. For RWML on τ² Bench: learning rate $1 \times 10^{-6}$, batch size 32, group size 16, 2 epochs, trained on 4×B200 GPUs. For WM SFT on both benchmarks: learning rate $2 \times 10^{-6}$, effective batch size 32, 2 epochs. The SFT learning rate is higher than the RL learning rate — a common practice because SFT updates are typically more stable and can tolerate larger steps, while RL updates need more conservative rates to prevent policy collapse.

Integration with Downstream Policy RL

The RWML-trained model serves as an improved initialization for a subsequent Policy RL stage that optimizes for task-success rewards (Section 3.1, baselines). This second stage is conceptually simple but practically important: it uses the same GRPO algorithm but with a completely different reward signal — the binary task-success signal provided by the environment at episode termination, propagated to all intermediate steps with $\gamma = 1.0$.

The Policy RL setup (Appendix B.2 for ALFWorld, C.2 for τ² Bench) follows the same prompting format as standard agent interaction: the model is prompted to generate reasoning (using <thinking> </thinking> tags) before producing an action. The task-success reward is binary: 1.0 if the task is completed successfully, 0.0 otherwise. On ALFWorld, a maximum of 15 steps per task is allowed during training (though 30 steps during evaluation), and GRPO is run for 300 steps with group size 8 on 2×B200 GPUs (average training time: 28 hours). On τ² Bench, a maximum of 30 steps is allowed, and GRPO is run for 200 steps with group size 8 on 8×B200 GPUs (average training time: 5 days). The longer training time for τ² Bench reflects both the larger model (8B vs. 7B parameters) and the more complex environment dynamics (tool calls, user interactions, policy compliance).

Why this two-stage approach works. The paper's central empirical claim (Table 1) is that RWML + Policy RL significantly outperforms Policy RL alone — by 6.9 points on ALFWorld (87.9% vs. 81.0%) and 5.7 points on τ² Bench (43.7% vs. 38.0%). The mechanistic explanation (Section 4.2) is that RWML learns a "parameter landscape more compatible with policy learning, reducing conflict and redundancy during post-training." In other words, the world model knowledge acquired during RWML creates better internal representations that make the subsequent task-learning RL more efficient and effective. The weight change analysis (Figure 3) supports this: when Policy RL is applied on top of an RWML-trained model, the parameter change ratios are similar to applying Policy RL directly on the base model — suggesting that RWML doesn't interfere with or disrupt the representations that Policy RL needs. In contrast, applying Policy RL on top of WM SFT produces substantially larger parameter changes, indicating that the SFT-trained world model has disrupted the model's representations in ways that require more aggressive re-adaptation during policy learning.

The RL-only variant of RWML. In the main experiments (Table 1), the "RWML (ours)" row shows the performance of the model after only the world model RL stage — before any Policy RL. On ALFWorld, this base + RWML model achieves 32.6% (vs. 36.1% for the base ReACT model and 2.8% for WM SFT alone). On τ² Bench, it achieves 38.8% (vs. 31.9% for the base model and 27.9% for WM SFT). Critically, RWML alone improves over the base model on τ² Bench (31.9% → 38.8%, a +6.9 point gain) while WM SFT alone degrades performance (31.9% → 27.9%). This demonstrates that the world model knowledge learned through RL genuinely transfers to improved decision-making even without explicit task-success training — the model leverages its improved understanding of environment dynamics to make better decisions during simple ReACT-style inference.

4. Key Insights and Innovations

Innovation 1: World Model Learning as a Mid-Training Stage, Not an Inference-Time Aid

The dominant framing in prior work on world models for LLM agents treats world model learning as an inference-time capability: the model learns to simulate environment transitions so that it can plan, look ahead, or evaluate candidate actions during decision-making. This is the pattern in Hao et al. (2023), where a separate world model is queried during MCTS-based planning; in Chae et al. (2025) and Gu et al. (2025), where learned world models augment web navigation at test time; and even in Zhang et al. (2025a), where the world model and policy are jointly trained but the world model serves an explicit simulation role during inference.

RWML reframes world model learning as a purely training-time intervention — a "mid-training" stage whose purpose is to reshape the model's internal representations before task-specific policy learning begins, without being deployed as a simulator during inference. The world model predictions are never used to make decisions at test time. Instead, the act of learning to predict next states — through RL with embedding-space rewards — produces parameter updates that make the model fundamentally better at understanding environment dynamics, and this improved understanding manifests as better action selection during standard ReACT-style inference.

This is a conceptual shift from "world models as tools" to "world models as representation learning." It echoes the classic insight from model-based RL that learning a transition model can accelerate policy learning (Sutton's Dyna, 1991), but adapts it to the modern LLM fine-tuning paradigm where the "world model" and the "policy" are the same neural network, and the benefit comes from shared representations rather than simulated rollouts. The evidence that this reframing is empirically meaningful — not just a philosophical distinction — is that RWML alone (without any task-success training) improves decision-making on τ² Bench by +6.9 points over the base model (Table 1), even though the model never uses its world-modeling output during inference. The world model knowledge transfers implicitly through weight sharing.

The significance of this reframing extends beyond the specific method. It suggests a general principle for LLM agent training: auxiliary tasks that teach the model to understand environment structure can serve as effective pretraining for downstream policy learning, even when the auxiliary task's outputs are discarded at deployment. This opens a broad design space — what other auxiliary objectives (predicting user intent, anticipating tool failures, forecasting task duration) might similarly serve as representation-learning pretraining for agent policies?

Innovation 2: Embedding-Space Rewards as a Scalable Alternative to Token-Level Supervision

The paper's most operationally distinctive contribution is the use of a frozen embedding model's cosine similarity as a reward signal for world model training, in deliberate contrast to the two dominant alternatives in prior work: SFT with token-level cross-entropy loss (used in WM SFT, IWM, and most prior world model work) and LLM-as-a-judge evaluation (used in various self-improvement and RL-from-feedback pipelines).

The standard assumption in LLM fine-tuning — inherited from pretraining — is that token-level prediction accuracy is the correct training objective. If you want a model to predict next states, you train it with next-token prediction on the target string. This paper provides direct empirical evidence that this assumption is actively harmful for world model learning. WM SFT, which uses token-level SFT on the identical data as RWML, produces a model that is worse than the untrained base model on ALFWorld (2.8% vs. 36.1%, Table 1), exhibits severe catastrophic forgetting (Table 3, up to −10.1 points on MMLU-Redux), and creates parameter disruptions that interfere with subsequent policy RL (Figure 3). These are not marginal underperformance — they are qualitative failures of the training objective.

The conceptual contribution is the diagnosis of why token-level SFT fails: it "prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse" (Section 1). In environments where the same state can be described in many lexically different but semantically equivalent ways, optimizing for exact string reproduction causes the model to overfit to surface patterns rather than learning generalizable transition dynamics. The embedding-space reward, by using cosine similarity in a pre-trained embedding space, provides a training signal that rewards being right about what matters rather than being identical to a reference string.

The paper also provides evidence against the seemingly more sophisticated alternative of LLM-as-a-judge rewards. The ablation study (Table 4) shows that substituting the embedding reward with a strong LLM judge (Qwen3-235B-A22B-Instruct) substantially degrades performance, and the qualitative example in Appendix D (Table A8) demonstrates why: the judge LLM can be "hacked" — the model learns to produce vague, non-committal predictions that the judge scores highly without demonstrating genuine environment understanding. The embedding reward, by being a simpler, non-trainable, non-parametric signal, is more resistant to this form of reward hacking.

This finding is significant beyond the specific application because it identifies a general tension in RL-based LLM training: more powerful reward models are not necessarily more robust, and simpler, embedding-based similarity metrics may provide more reliable training signals for tasks where semantic equivalence matters more than surface form. This has implications for the broader RLHF and constitutional AI literature, where the sophistication of the reward model is often assumed to be the primary bottleneck.

Innovation 3: Diagnostic Evidence That RL (Not Just SFT) Enables Transfer from World Modeling to Decision-Making

Prior work on world model learning for LLM agents (Zhang et al., 2025a; Yu et al., 2025c; Feng et al., 2025b) has uniformly used SFT as the training objective, and has demonstrated that SFT-trained world models can improve agent performance — but only when combined with expert data, strong LLM-generated synthetic data, or explicit world-model-guided planning at inference time. What was unknown before this paper is whether the SFT objective itself was the limiting factor, or whether world model knowledge simply requires external supervision to be useful.

RWML provides a clean controlled experiment to answer this question. The RWML and WM SFT baselines use identical training data (the same triplets from the model's own rollouts, with the same subsampling), but differ only in the training objective: RL with embedding rewards vs. SFT with token-level cross-entropy. The results are stark:

• WM SFT alone is harmful: it degrades ALFWorld performance from 36.1% (base model) to 2.8% (Table 1) — a catastrophic collapse. On τ² Bench, it degrades from 31.9% to 27.9%.
• RWML alone is beneficial: it improves ALFWorld performance to 32.6% (still below the base model's 36.1%, but vastly better than WM SFT's 2.8%) and improves τ² Bench from 31.9% to 38.8%.
• When combined with Policy RL, the gap persists: RWML + Policy RL achieves 87.9% on ALFWorld vs. 80.4% for WM SFT + Policy RL (Table 1), and 43.7% vs. 40.3% on τ² Bench.

These results demonstrate that the training objective, not just the data, determines whether world model knowledge transfers to improved decision-making. SFT on world model data teaches the model to reproduce strings, but this skill doesn't transfer to the structurally different task of selecting actions — and may actively interfere with it. RL on world model data, by contrast, teaches the model to understand environment dynamics in a way that generalizes to action selection, even though the model is never explicitly trained to connect its world model predictions to its action choices.

The weight change analysis (Section 4.2, Figure 3) provides a mechanistic hypothesis for why: RWML induces "notably fewer parameter changes across layers compared to WM SFT" and produces a "parameter landscape more compatible with policy learning, reducing conflict and redundancy during post-training." In other words, RL-based world model learning makes targeted, conservative updates to the model's representations that preserve existing capabilities while adding environment understanding, whereas SFT makes broad, disruptive updates that overwrite useful knowledge.

This is a fundamental finding, not an incremental one, because it challenges the default assumption in LLM fine-tuning pipelines that SFT is the appropriate tool for teaching any new capability. The paper demonstrates a case where SFT is not just suboptimal but actively destructive, and RL provides a qualitatively different learning dynamic that enables knowledge transfer where SFT cannot. This has implications beyond world model learning — it suggests that for any auxiliary task where the goal is to improve the model's understanding (rather than its output format), RL with appropriate reward design may be preferable to SFT.

Innovation 4: The Difficulty-Conditioned Data Subsampling Strategy

While data subsampling is not the paper's central contribution, the specific mechanism introduced — using a weak SFT model to identify and downsample "too easy" training examples — represents a practical innovation with broader applicability. The insight is that for world model learning, not all environment transitions are equally informative. Transitions that are trivially predictable (e.g., "after saying hello, the user says hello back") don't teach the model anything useful about environment dynamics, and including too many of them in the training data may dilute the signal from genuinely challenging transitions.

The specific procedure (train a weak SFT model on 10% of the data, use it to score the remaining 90%, downsample examples it consistently predicts correctly) is computationally lightweight and requires no external supervision. Its effectiveness is validated in the ablation (Table 4): removing subsampling degrades RWML performance on τ² Bench from 38.8 to 35.8.

The broader significance is methodological: this is a general-purpose technique for curating RL training data based on model-specific difficulty estimates that could apply to any self-supervised RL task where some training examples are trivially easy and others are genuinely challenging. It echoes the difficulty-conditioned ideas from the test-time compute scaling literature (Snell et al., 2024, which the paper explicitly cites) but applies them to training data construction rather than inference-time strategy selection.

This innovation is incremental rather than fundamental — the core ideas come from prior work, and the specific implementation (10% split, K=10 attempts, p=0.1 retention) is described as "chosen heuristically without tuning" (Appendix B.1) — but it addresses a genuine practical challenge that would otherwise limit the method's effectiveness.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on two long-horizon agent benchmarks: ALFWorld (Shridhar et al., 2021), a text-based embodied environment with household tasks using the official training (2048 tasks) and test splits; and τ² Bench (Barres et al., 2025), a customer-service environment with interleaved tool-use and user conversation spanning retail, telecom, and airline domains, using its 178-task training split and official test splits. These benchmarks are chosen because they require agents to reason about environment dynamics — ALFWorld demands knowing object locations ("where is the knife likely to be?"), while τ² Bench requires understanding tool behaviors and user response patterns.

• Base model(s). On ALFWorld, all training-based methods use Qwen2.5-7B-Instruct (Qwen et al., 2025). On τ² Bench, the base model is Qwen3-8B (Yang et al., 2025), chosen because of "the difficulty of the benchmark and the enhanced tool-use capabilities from Qwen3 models" (Section 3.1). For Policy RL training on τ² Bench, Qwen3-235B-A22B-Instruct serves as the user simulator to reduce costs compared to the official GPT-4.1 user. Additional reference points use ReACT-prompted Qwen2.5-7B, Qwen3-8B, Qwen3-235B-A22B, GPT-4.1, and GPT-5.

• Metrics. The primary metric is task success rate (%), measured as the fraction of test tasks completed successfully within a maximum of 30 steps. For ALFWorld, the paper also reports separate ID (in-distribution task types seen during training) and OOD (out-of-distribution, novel task types) success rates. For τ² Bench, results are additionally broken out by domain (Retail, Telecom, Airline). All results are averaged over 3 independent evaluation runs with standard deviations reported.

• Baselines. The paper organizes baselines into three categories (Section 3.1). (1) Learning from task-success reward: Reinforced Finetuning (RFT) uses rejection sampling — SFT on trajectories that successfully solved the task — and Policy RL uses GRPO to directly optimize for task-success rewards from online rollouts (Feng et al., 2025c; Yu et al., 2025a). (2) Learning from interaction/transition function: World Model SFT (WM SFT) trains on identical data as RWML but uses SFT to predict $s_{t+1}$ directly from $\langle s_{\le t}, a_t \rangle$, with no reasoning tokens. (3) Learning from expert/strong LLMs: Imitation Learning (SFT on expert trajectories), Implicit World Modeling (IWM) and Self-Reflection (SR) from Zhang et al. (2025a) and Yu et al. (2025c) — these augment expert rollouts with model-generated alternative actions, synthesize world model or reflection data using stronger LLMs (GPT-4o for ALFWorld, GPT-4.1 for τ² Bench), and SFT on the combined dataset. The paper also includes ReACT-prompted (Yao et al., 2023) performance on Qwen and GPT models as non-training references.

• Generation budget / compute accounting. All training-based methods are compared at the same task level using comparable numbers of training trajectories. RWML uses $N=3$ rollouts per training task on ALFWorld and $N_{\text{total}}=6$ on τ² Bench; RFT and Policy RL also use multiple rollouts per task. The paper reports training hardware (B200 GPUs) and wall-clock time (28 hours for ALFWorld Policy RL, 5 days for τ² Bench Policy RL), providing a rough practical cost comparison. Testing uses a fixed maximum of 30 steps per task across all methods. No FLOPs-matched analysis between training methods is performed.

• Cross-validation / statistical protocol. All evaluation results are averaged over 3 independent runs with standard deviations reported in Table 1. The paper does not describe cross-validation for hyperparameter selection; key hyperparameters (e.g., $p=0.1$ for subsampling, $\tau_{\text{easy}}=0.0$, $\tau_d$ values) are described as "chosen heuristically without tuning" and fixed across benchmarks. For τ² Bench, additional evaluation is performed under the official setting (GPT-4.1 as user simulator, max 100 steps) on representative models to validate that training results transfer.

Main Quantitative Results

RWML as a Standalone Self-Supervised Method (Table 1, "Self-Supervised" rows)

The paper's most foundational empirical claim is that RWML — trained entirely without expert data, strong LLM supervision, or task-success rewards — improves agent performance over the untrained base model. On τ² Bench, this claim holds cleanly: the base Qwen3-8B model (ReACT-prompted) achieves 31.9% average across domains, while RWML alone reaches 38.8% — a +6.9 point gain. This is a meaningful improvement from purely self-supervised world model learning, and it is larger than the policy RL baseline's gain (Policy RL alone reaches only 38.0%, just 0.8 points below RWML alone). On ALFWorld, the result is more nuanced: the base Qwen2.5-7B-Instruct model achieves 37.7% (ReACT-prompted), while RWML alone reaches 32.6% — actually below the base model, though substantially above WM SFT's catastrophic 2.8%. The paper characterizes RWML's ALFWorld performance as advancing "19.6 and 7.9 points on ALFWorld and τ² Bench, respectively" compared to the base model (Section 3.2), but Table 1 shows the ALFWorld base model for the 7B model is 37.7%, making the comparison baseline ambiguous — the 19.6-point figure appears to reference an implicit lower-performing version of the base model used during data collection, not the best ReACT-prompted baseline.

The critical comparison within the self-supervised category is RWML vs. WM SFT. These two methods use identical training data but different objectives: RL with embedding rewards vs. SFT with token-level cross-entropy. On ALFWorld, WM SFT alone achieves 2.8% — a catastrophic collapse from any reasonable baseline. On τ² Bench, WM SFT achieves 27.9%, which is actually below the base model's 31.9% but far above ALFWorld-level collapse. The gap between RWML and WM SFT is +29.8 points on ALFWorld (32.6% vs. 2.8%) and +10.9 points on τ² Bench (38.8% vs. 27.9%). This is the paper's cleanest demonstration that the training objective, not the data, determines whether world model knowledge transfers to improved decision-making. SFT on world model triplets is actively destructive on ALFWorld; RL on the same data is neutral-to-positive.

RWML Combined with Policy RL (Table 1, "Self-Supervised + Policy RL" rows)

The paper's central practical claim is that RWML serves as an effective mid-training stage before Policy RL, and that the combination outperforms Policy RL from scratch. The headline numbers: RWML + Policy RL achieves 87.9% on ALFWorld (averaged across ID: 86.7% and OOD: 90.1%), compared to 81.0% for Policy RL alone — a +6.9 point gain. On τ² Bench, RWML + Policy RL achieves 43.7% vs. 38.0% for Policy RL alone — a +5.7 point gain. Both improvements are statistically meaningful relative to the reported standard deviations (e.g., Policy RL on ALFWorld: 81.0 ± 1.6; RWML + Policy RL: 87.9 ± 1.6).

The WM SFT + Policy RL baseline provides a critical control: does any world model training before Policy RL help, or is RL-based world model training specifically important? WM SFT + Policy RL achieves 80.4% on ALFWorld — nearly identical to Policy RL alone (81.0%), and substantially below RWML + Policy RL (87.9%). On τ² Bench, WM SFT + Policy RL achieves 40.3%, which is between Policy RL alone (38.0%) and RWML + Policy RL (43.7%). The interpretation: WM SFT provides some benefit on τ² Bench (+2.3 over Policy RL alone) but essentially none on ALFWorld (-0.6), while RWML provides substantial gains on both. This suggests that the benefit of mid-training world model learning depends on the training objective, with RL-based learning transferring more effectively to downstream policy optimization.

Comparison Against Methods Using Expert Data / Strong LLMs (Table 2)

The paper makes the bold claim that RWML + Policy RL matches or exceeds methods that use expert demonstrations or strong LLM-generated synthetic data, despite being self-supervised. On ALFWorld, this claim is strongly supported: RWML + Policy RL achieves 91.0% on ID tasks and 90.1% on OOD tasks (average 90.6%) — the highest scores in Table 2. Imitation Learning achieves 78.6% (ID) and 75.0% (OOD); IWM achieves 77.9% (ID) and 82.3% (OOD); SR achieves 75.6% (ID) and 76.0% (OOD). The paper notes that Zhang et al. (2025a) originally reported 78.1/82.8/82.0 for ID and 64.1/70.3/71.1 for OOD on these methods — the reproduced numbers differ, possibly due to model choice (Qwen2.5-7B vs. the original paper's base model) or implementation details.

On τ² Bench, the picture is more qualified. RWML + Policy RL (43.7%) achieves the second-highest overall score, behind only Imitation Learning (44.2%) and closely competitive with IWM (42.2%) and SR (42.2%). The differences are small and within overlapping standard deviations (RWML + Policy RL: 43.7 ± 2.1; Imitation Learning: 44.2 ± 2.1). The fair characterization is that RWML + Policy RL is competitive with but does not clearly outperform expert-data methods on τ² Bench, and the paper correctly frames this as "the second best overall score, despite not accessing any expert data/strong LLMs" (Section 3.2).

A notable detail: the Imitation Learning baseline on τ² Bench (44.2%) outperforms RWML + Policy RL (43.7%) by a narrow margin. This baseline uses expert rollouts collected via rejection sampling from Qwen3-235B-A22B-Thinking-2507 — a model that is substantially larger and more capable than the 8B model being trained. The fact that RWML + Policy RL nearly matches the performance of SFT on a 235B model's successful trajectories, using only self-generated interaction data, is a strong result, but it also reveals that expert data (when available) provides a slightly superior signal on this benchmark.

Catastrophic Forgetting Analysis (Table 3)

Section 3.3 evaluates how much general capability the models retain after world model training. The comparison is between WM SFT and RWML, evaluated on six benchmarks spanning general knowledge (MMLU-Redux), instruction-following (IFEval), math (MATH-500, GSM8k), STEM reasoning (GPQA-Diamond), and coding (LiveCodeBench). The headline finding: RWML consistently forgets less than WM SFT across nearly all benchmarks.

On ALFWorld (Qwen2.5-7B base): WM SFT causes a -10.10 point decline on MMLU-Redux (77.26% → 67.16%) while RWML causes only -2.38 points (77.26% → 74.88%). On GPQA-Diamond, WM SFT drops by -7.58 points (32.83% → 25.25%) vs. RWML's -4.05 points (32.83% → 28.79%). On LiveCodeBench, WM SFT drops by -3.85 points (19.23% → 15.38%) vs. RWML's -2.75 points (19.23% → 16.48%). On τ² Bench (Qwen3-8B base): WM SFT causes a -2.39 point decline on IFEval (84.46% → 82.07%) while RWML causes only -1.10 points (84.46% → 83.36%). On some benchmarks (MATH-500, LiveCodeBench on τ² Bench), RWML shows zero degradation.

One counterexample: on ALFWorld's MATH-500, RWML shows zero change (75.40% → 75.40%) while WM SFT drops by -3.80 points. This is consistent with the overall pattern but highlights that MATH-500 is already near the base model's capability ceiling and neither training method substantially affects it.

The paper attributes this difference to the on-policy nature of RL: "online RL preserves prior knowledge and capabilities significantly better than SFT due to its on-policy nature" (Section 3.3), citing Shenfeld et al. (2025) and Chen et al. (2025a). The weight change analysis in Section 4.2 (Figure 3) corroborates this: WM SFT induces substantially more parameter changes across layers, while RWML makes smaller, more targeted updates.

Scaling with Base Model Capability (Figure 4)

Section 4.3 examines how the effectiveness of RWML depends on the underlying model's capability, testing three base models on τ² Bench: Qwen2.5-7B, Qwen3-8B, and Qwen3-30B-A3B-Thinking-2507. The finding is that stronger base models benefit more from RWML. Qwen2.5-7B shows minimal improvement from RWML on this challenging benchmark (the paper characterizes it as "struggling to transfer world knowledge to decision-making"). Qwen3-8B shows substantial gains, reaching approximately 38.8% from a base of 31.9%. Qwen3-30B-A3B-Thinking-2507 shows gains but with less room for improvement since it starts from a higher baseline (closer to 50% in base ReACT performance). The paper notes that this model is "an enhanced version of Qwen3-30B-A3B post-trained with additional reasoning and agent data, leaving less room for further improvement."

This result has important practical implications: RWML is not a universal cure for weak models. If the base model lacks sufficient capability to even begin understanding environment dynamics from its own interaction data, self-supervised world model learning provides limited benefit. The paper frames this as a boundary condition: "RWML is most effective for (sufficiently) strong base models" (Section 4.3).

Ablation Studies and Robustness Checks

Embedding-based reward vs. LLM-as-a-judge (Table 4, columns 2-4). Replacing the embedding cosine similarity reward with an LLM-as-a-judge reward significantly degrades performance on both benchmarks. On ALFWorld, the LLM-as-a-judge variant (binary) achieves only 28.9% vs. RWML's 32.6% — a -3.7 point decline. The continuous LLM-as-a-judge variant (real-valued reward in [0, 1]) is even worse at 23.9%. On τ² Bench, the binary LLM-as-a-judge variant achieves 35.8% vs. RWML's 38.8% (-3.0 points), while the continuous variant achieves 37.8% (-1.0 point). The paper notes that the Qwen2.5-7B model on ALFWorld is "more susceptible to data quality/noisy reward functions" than the stronger Qwen3-8B on τ² Bench, and provides a qualitative example in Appendix D (Table A8) showing the model generating a vague prediction ("The immediate next observation after going to countertop 1 will describe the contents of countertop 1 or indicate if there is anything special about the countertop") that receives a perfect score of 1.0 from the LLM judge despite containing no actual environment knowledge. This is direct evidence of reward hacking against the LLM judge.

Data subsampling (Table 4, "w/o subsample" column). Removing the "too easy" sample subsampling step degrades RWML performance on both benchmarks. On ALFWorld, performance drops from 32.6% to 30.2% (-2.4 points). On τ² Bench, the drop is from 38.8% to 35.8% (-3.0 points). Both degradations are substantial relative to the overall RWML improvements, confirming that focusing training on medium-to-hard examples is beneficial. Interestingly, on τ² Bench, removing subsampling actually makes performance worse than the base model (35.8% vs. 38.8%), suggesting that the subsampling is particularly important for more complex environments where many transitions may be trivially predictable (e.g., standard tool call responses that don't require reasoning).

Removing training entirely (Table 4, "w/o training"). This corresponds to the untrained base model and serves as a sanity check: 36.1% on ALFWorld and 31.9% on τ² Bench (the ReACT baselines from Table 1). The fact that RWML without subsampling on τ² Bench (35.8%) nearly matches the untrained baseline (31.9%) — representing only a +3.9 point gain — demonstrates that the subsampling step is responsible for a meaningful fraction of RWML's total benefit. On ALFWorld, even RWML without subsampling (30.2%) outperforms the closest SFT baseline (WM SFT: 2.8%), confirming that the RL objective alone provides substantial benefit even without data curation.

WM SFT as an alternative mid-training stage (Table 1, "WM SFT + Policy RL" rows). While not presented as an ablation, the WM SFT + Policy RL baseline serves as one: it tests whether any world model pretraining helps Policy RL, or whether RL-based pretraining is specifically important. On ALFWorld, WM SFT + Policy RL (80.4%) is nearly identical to Policy RL alone (81.0%) and substantially below RWML + Policy RL (87.9%). This demonstrates that the benefit of mid-training is not a generic "pretraining helps" effect — SFT-based world model learning, despite being trained on the same data, provides essentially no downstream benefit on ALFWorld. On τ² Bench, WM SFT + Policy RL (40.3%) provides a modest +2.3 point gain over Policy RL alone (38.0%), but still substantially underperforms RWML + Policy RL (43.7%). This asymmetry between benchmarks is interesting: on τ² Bench, even SFT-based world model learning transfers somewhat usefully to policy optimization, while on ALFWorld it transfers nothing.

Reasoning tokens in RWML vs. WM SFT (training setup, Appendices B.1 and C.1). The WM SFT baseline deliberately excludes reasoning tokens during training (Tables A5 and A7 instruct the model "DO NOT perform any thinking"), while RWML includes reasoning generation (Tables A4 and A6). This is not varied independently, so the ablation is confounded — WM SFT differs from RWML in both training objective (SFT vs. RL) and reasoning inclusion (no reasoning vs. reasoning). The paper justifies this design choice by noting that "since there is no reasoning data available for the triplets, we find this training method for WM SFT can better enable generalization/reasoning during the second stage Policy RL training" (Appendix B.1). In other words, they deliberately gave WM SFT what they believed to be its best configuration, and it still underperformed. However, the lack of a WM SFT variant that includes reasoning tokens means we cannot fully isolate whether the RL objective or the reasoning scaffold is responsible for the performance gap.

Alternative reward formulations for τ² Bench (Appendix C.1). On τ² Bench, the paper uses a hybrid reward (embedding similarity for user responses, ROUGE for tool responses) rather than pure embedding similarity. This is a domain adaptation, not an ablation with a controlled comparison. The paper does not report results using pure embedding similarity on τ² Bench, so we cannot assess how much the hybrid formulation matters. The justification — that "tool-use responses are generally structured outputs" and ROUGE better captures missing keys/values — is plausible but untested.

Impact of temperature and rollout count (not ablated). The paper uses temperature $\tau = 1.0$ for rollout collection and $N = 3$ (ALFWorld) or $N = 6$ (τ² Bench) rollouts per task. Neither the temperature nor the number of rollouts is ablated. Higher temperature would produce more diverse (but potentially lower-quality) training data; more rollouts would provide better coverage of possible trajectories. These are practical hyperparameters whose sensitivity is unknown, and the paper's decision to "choose heuristically without tuning" applies here as well.

Critical Assessment

Does RWML + Policy RL genuinely outperform Policy RL from scratch?

Yes, with strong evidence. The +6.9 point gain on ALFWorld (87.9% vs. 81.0%) and +5.7 point gain on τ² Bench (43.7% vs. 38.0%) are statistically significant relative to the reported standard deviations (Table 1). The result is replicated across two different benchmarks with different base models (7B and 8B) and different environment dynamics. However, it is worth noting that the Policy RL baseline on ALFWorld (81.0%) may itself be improvable — the paper follows training setups from Feng et al. (2025c) and Yu et al. (2025a), but does not demonstrate that Policy RL has converged or that hyperparameters are optimized for the standalone Policy RL case. If Policy RL were given more training steps, a different learning rate, or a different KL penalty, it might close some of the gap with RWML + Policy RL. The paper does not ablate Policy RL hyperparameters independently for the two initialization conditions (base model vs. RWML-pretrained).

Does the self-supervised RWML stage improve over the base model without task-success rewards?

On τ² Bench: yes (+6.9 points). On ALFWorld: the evidence is ambiguous. The ReACT baseline for Qwen2.5-7B-Instruct on ALFWorld is 37.7% (Table 1), while RWML alone achieves 32.6% — a decline of -5.1 points. The paper's headline of "19.6 points" improvement appears to compare against a different (lower-performing) base model configuration, possibly the model used for data collection rather than the best ReACT-prompted version. This discrepancy is not explained in the paper. The fair characterization is: on τ² Bench, RWML alone is clearly beneficial; on ALFWorld, RWML alone is potentially neutral to slightly negative compared to the best base model configuration, though it is dramatically better than WM SFT (32.6% vs. 2.8%) and provides a substantially better initialization for downstream Policy RL (87.9% vs. 81.0%). The transfer of world model knowledge to decision-making appears to be more direct on τ² Bench and more dependent on the Policy RL bridge on ALFWorld.

Does RWML match the performance of expert-data training?

On ALFWorld: yes, and exceeds it. RWML + Policy RL (90.6% average of ID/OOD) substantially outperforms all expert-data methods in Table 2 (Imitation Learning: 76.8%, IWM: 80.1%, SR: 75.8%). On τ² Bench: competitive but does not clearly exceed it. RWML + Policy RL (43.7%) is marginally below Imitation Learning (44.2%) and statistically indistinguishable from IWM (42.2%) and SR (42.2%). The claim of "matching the performance of expert-data training" (Abstract) is more precisely "matching on τ² Bench, exceeding on ALFWorld." The paper's reproduction of expert-data methods on τ² Bench uses rejection sampling from Qwen3-235B-A22B-Thinking-2507 (a 235B model) to collect "expert" rollouts, which is a form of strong LLM distillation rather than genuine expert annotation. This makes the comparison somewhat circular: the expert-data methods use a 235B model to train an 8B model, while RWML uses only the 8B model's own interactions.

Do the ablation studies convincingly isolate the contribution of each component?

Partially. The reward function ablation (embedding vs. LLM-as-a-judge) and the subsampling ablation are clean and informative. However, three important components are confounded and not independently ablated: (1) the RL objective vs. the reasoning token scaffold — WM SFT differs from RWML in both respects, so we cannot attribute the performance gap solely to RL vs. SFT; (2) the hybrid reward on τ² Bench (embedding + ROUGE) is not compared against pure embedding similarity; (3) the number of rollouts ($N$), temperature ($\tau$), and other data collection hyperparameters are fixed rather than swept. The paper would be stronger with a WM SFT variant that includes reasoning tokens (to isolate the training objective effect) and an RWML variant on τ² Bench that uses only embedding similarity (to isolate the hybrid reward effect).

Are the forgetting results (Table 3) a fair comparison?

Yes, with the caveat that "forgetting" is measured as relative decline from the base model rather than absolute post-training capability. On MMLU-Redux, for example, RWML on ALFWorld achieves 74.88% while WM SFT achieves 67.16% — both are below the base model's 77.26%, and RWML simply forgets less. The paper frames this as RWML "better mitigates catastrophic forgetting," which is accurate. However, the absolute capability loss from RWML (2-4 points on most benchmarks) is small but non-zero, and on some benchmarks (GPQA-Diamond on ALFWorld: -4.05 points), it is non-trivial. The practical implication is that RWML still causes some degradation of general capabilities, just substantially less than SFT. The paper does not investigate whether this residual forgetting can be eliminated through KL regularization tuning or other mitigation strategies.

Missing experiments that would strengthen the paper

1. No test of RWML on a third benchmark or model family. All results are on two benchmarks (ALFWorld and τ² Bench) with Qwen models. Testing on WebArena, SWE-bench, or another agent benchmark — or with a different model family (e.g., Llama) — would substantially strengthen the generality claims.

2. No comparison against alternative mid-training objectives. The paper compares RWML against WM SFT (next-state prediction with SFT), but not against other auxiliary tasks that might similarly serve as mid-training: predicting task success from intermediate states, predicting action outcomes, or learning a value function. This leaves open the question of whether next-state prediction specifically, or RL-based auxiliary learning generally, is the key factor.

3. No ablation of rollout count ($N$) or temperature ($\tau$). The quality and diversity of the self-generated training data almost certainly depends on these parameters. A model with temperature 0.1 would produce narrow, high-confidence trajectories; a model with temperature 1.0 and more rollouts would see more diverse (and potentially erroneous) transitions. Understanding this tradeoff would inform practical deployment.

4. No test of RWML at larger model scales. All experiments use 7B-8B models. The scaling experiment in Figure 4 (30B model) only tests RWML alone, not RWML + Policy RL. Whether the +5-7 point gain from RWML + Policy RL over Policy RL alone persists at larger scales (where Policy RL itself may be more effective) is unknown.

5. No analysis of computational cost vs. benefit. The paper reports hardware and wall-clock times but does not quantify the total FLOPs or GPU-hours of RWML training vs. Policy RL training, or how much additional cost RWML adds to the total training pipeline. For practitioners weighing whether to adopt RWML, this cost-benefit analysis is essential: if RWML training takes as long as Policy RL and provides a +5 point gain, the cost doubles for a modest improvement; if it's cheap relative to Policy RL, the tradeoff is more favorable.

6. Limitations and Trade-offs

Difficulty Estimation Cost Is Not Accounted For

The assumption or constraint. The entire compute-optimal framework from the reference paper depends on estimating prompt difficulty before allocating the inference budget. The paper's method for doing so is extraordinarily expensive:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity" (Section 3.2)

The procedure generates 2048 samples per question and scores them with the PRM to assign a difficulty bin. In a realistic deployment, this cost would dominate the test-time budget being optimized—generating 2048 samples far exceeds the largest budgets studied (256–512 generations). The paper acknowledges this gap but defers it to future work, noting that "training models to directly predict difficulty of a question" is a key direction (Section 8).

The consequence. The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. If difficulty estimation costs are included, the total compute budget would be difficulty estimation cost + strategy execution cost, and the 4× advantage could shrink dramatically or even reverse. In practical deployments—especially those with small per-query budgets—the overhead of generating 2048 samples would make the approach infeasible. The paper's claim of 4× efficiency should therefore be understood as an upper bound on achievable efficiency rather than a realized deployment gain.

What evidence exists in the paper. The paper explicitly acknowledges this limitation in Section 3.2. The predicted-vs-oracle difficulty comparison (Figures 4 and 8) shows that predicted bins perform nearly as well as oracle bins, confirming that ground-truth labels are not the bottleneck—but the computational cost of generating 2048 samples per question remains regardless of whether the correctness check uses ground-truth or PRM scores.

Mitigation status. Not addressed. The paper frames this as "a key avenue for future work" (Section 3.2) and suggests training a lightweight difficulty predictor or using adaptive estimation, but neither is implemented or evaluated. Future work on cheap difficulty estimation is explicitly called for in Section 8.

Verifier Over-Optimization Is Documented But Not Solved

The assumption or constraint. The paper's central finding on search (Section 5.3) is that PRM-guided search suffers from verifier over-optimization: beam search degrades performance on easy problems at high budgets (Figure 3, right), lookahead search paradoxically performs worst overall despite being the most powerful optimizer (Figure 3, left), and qualitative examples show degenerate outputs that score highly under the PRM but are incorrect (repetitive low-information steps, overly short 1–2 step solutions, Appendix M). The reason is that the PRM is an imperfect proxy for actual correctness, and aggressive optimization against it eventually exploits its errors rather than finding genuinely better solutions.

The consequence. The compute-optimal policy mitigates this by routing easy problems away from aggressive search (using best-of-N instead of beam search on bins 1–2), but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed (bins 3–4), over-optimization still limits the scaling ceiling—beam search curves in Figure 3 flatten and sometimes decline at high budgets, meaning there is a hard ceiling on what test-time compute can achieve regardless of budget. This ceiling is determined by verifier quality, not by the search algorithm or the base model's capability. The practical implication is that scaling test-time compute indefinitely will hit diminishing returns bounded by how well the verifier can distinguish correct and incorrect solutions.

What evidence exists in the paper. The evidence is extensive and consistent across multiple experiments:

• Figure 3 (right): Beam search accuracy on bin 1 decreases from ~78% to ~77% as budget goes from 4 to 256 generations.
• Figure 3 (left): Lookahead search underperforms all methods at the same generation budget.
• Appendix M (qualitative): Examples of search producing degenerate, high-scoring outputs (repetitive steps, overly short solutions).
• Section 5.3: The paper explicitly identifies over-optimization as the cause: "search finds solutions that score highly under the PRM but are actually incorrect."

Mitigation status. The paper mitigates this symptomatically (by routing easy problems to best-of-N in the compute-optimal policy) but does not address the root cause. Improving verifier robustness—through adversarial training, ensemble methods, or constrained search with KL penalties—is mentioned as future work (Section 8). The current approach therefore operates under the existing verifier's reliability ceiling, and the gains reported are specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D.

Sequential Revisions Introduce Unresolved Latency vs. Parallelism Tradeoffs

The assumption or constraint. The paper measures test-time compute in "generations" (number of complete solutions sampled), which is a reasonable proxy for total FLOPs but ignores the wall-clock time implications of sequential vs. parallel computation. The compute-optimal policy for revisions (Section 6) frequently favors sequential or sequential-heavy strategies—for example, on easy problems (bin 1–2), fully sequential revisions are optimal, and even on medium-hard problems, the optimal ratio involves substantial sequential depth (Figure 7, right). However, sequential revisions are inherently serial: each revision depends on the previous one, meaning a chain of 64 sequential revisions takes roughly 64× longer wall-clock time than 64 parallel samples executed simultaneously on sufficient hardware.

The consequence. For latency-sensitive applications—interactive assistants, real-time decision-making, customer-facing deployments—the sequential-heavy strategies favored by the compute-optimal policy on easy problems may be practically unusable regardless of their accuracy advantages. A strategy that requires 64 sequential forward passes adds unacceptable latency even if the total FLOPs are equivalent to a parallel strategy. The paper does not discuss this tradeoff, meaning a practitioner following the compute-optimal policy might inadvertently choose a strategy that is optimal in FLOPs but catastrophic for user experience. The reverse is also true: for throughput-oriented batch processing where latency doesn't matter, the sequential strategies' serial dependency could still bottleneck hardware utilization compared to parallel strategies that fully saturate compute.

What evidence exists in the paper. There is no explicit discussion of latency or wall-clock time in the paper. The implication can be inferred from Figure 5 (right panel), which shows the architecture of sequential vs. parallel revision sampling—the sequential chain has a clear serial dependency (each output feeds into the next), while the parallel sampling is independent. The paper reports that "sequential revision marginally outperforms parallel sampling under both verifier-based and majority-based answer selection" (Section 6.1, Figure 6 right), but this comparison is in FLOPs-equivalent terms, not wall-clock-equivalent terms.

Mitigation status. Not addressed. The paper does not mention latency, throughput, or wall-clock time as considerations. For practitioners, this means the efficiency numbers must be interpreted in the context of their deployment's latency requirements. A hybrid approach—running some parallel chains with moderate sequential depth to balance latency and accuracy—is the natural compromise, but the paper provides no guidance on this axis.

Difficulty Bins Are Static, Coarse, and Require Manual Pre-Computation

The assumption or constraint. The compute-optimal policy operates by discretizing questions into five fixed difficulty quintiles (Section 3.2), pre-computing the best strategy for each bin using two-fold cross-validation on a held-out set, and then using a lookup table at test time. This requires: (1) choosing the number of bins (5) a priori, (2) having a large enough test set to support bin-level cross-validation (the 500-question MATH test set, split into ~100 questions per quintile, then further split by two-fold cross-validation means strategy selection is based on ~50 questions per fold per bin), and (3) assuming that the optimal strategy for a bin is stable and that the bin boundaries are meaningful.

The consequence. Several problems arise. First, the coarseness: within a single bin, there may be substantial heterogeneity—a question at the easy end of bin 3 and one at the hard end of bin 3 receive the identical strategy, even though the optimal strategy might differ. A finer-grained or continuous difficulty estimate could improve allocation but would require more data to estimate the optimal policy per bin, which the small test set (500 questions) cannot support. Second, the static nature: difficulty bins are computed once and treated as fixed. There is no mechanism for dynamically adjusting the strategy mid-computation—for instance, starting with a few parallel samples, assessing the score distribution, and reallocating the budget if the initial assessment suggests the problem is easier or harder than expected. Third, the data requirements for cross-validation: the paper uses two-fold cross-validation on the 500-question test set, meaning the policy for each bin is selected based on ~50 questions. This is a small sample, and the selected strategies may not be robust—the paper does not report confidence intervals on the compute-optimal scaling curves, making it difficult to assess whether the observed gains are statistically reliable at this sample size.

What evidence exists in the paper. The binning structure is described in Section 3.2, and the cross-validation protocol is explained there. The test set size of 500 questions is stated in Section 4. Figure 4 and Figure 8 show the compute-optimal scaling curves but without error bars or confidence intervals. The paper does not ablate the number of bins or test sensitivity to bin count.

Mitigation status. Partially addressed through the predicted-vs-oracle difficulty comparison (Figures 4, 8), which shows that the difficulty estimation method (oracle vs. predicted) does not substantially affect results. However, the issues of bin coarseness, static allocation, and small sample sizes for strategy selection are not addressed or acknowledged as limitations. Future work on continuous difficulty estimates or dynamic allocation is not explicitly suggested, though the paper does flag the cost of difficulty estimation as future work (Section 8).

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate

The assumption or constraint. The revision model is trained exclusively on trajectories where all in-context answers are incorrect followed by a correct target answer (Section 6.1). This is a deliberate training data construction choice—the model learns to take incorrect answers as input and produce improved answers as output. However, this creates a fundamental asymmetry: during inference, the revision chain may produce a correct answer at step k, but at step k+1, the model conditions on this correct answer and, having never seen correct answers in its training context, may incorrectly "revise" it into a wrong answer.

The consequence. The paper reports that approximately 38% of correct answers get converted back to incorrect ones during revision chains when using a naive approach (Section 6.1). This means that longer revision chains are not monotonically improving—they can oscillate between correct and incorrect, and taking the last revision in the chain is unreliable. The paper mitigates this with a selection mechanism (majority voting or verifier-based selection across the entire chain, picking the best answer from any point), but this is an imperfect patch: the model is still wasting computation generating revisions that degrade quality, and the selection mechanism adds its own potential for error (the verifier might select the wrong answer, or majority voting might favor an incorrect answer that appears frequently in the chain).

What evidence exists in the paper. The 38% reversion rate is explicitly stated in Section 6.1 (paragraph on "the correct-to-incorrect reversion problem"). The mitigation (majority voting or verifier-based selection across the chain) is described and evaluated in Figure 6, which shows that sequential + best-of-N weighted and sequential + majority both outperform parallel-only approaches, but neither eliminates the underlying issue. The fact that revision chain accuracy plateaus rather than continuously improving in Figure 6 (left)—improving from ~18.2% at step 1 to ~24–25% by steps 15–20 and remaining flat out to step 64—is consistent with the reversion problem capping the benefits of additional sequential depth.

Mitigation status. The paper mitigates the symptoms (by selecting the best answer across the chain rather than taking the final revision) but does not address the cause (the training data asymmetry). A more principled solution—such as including correct-to-correct trajectories in the training data, teaching the model to recognize when no revision is needed, or using a stopping criterion based on verifier confidence—is not explored. The ReST$^{\text{EM}}$ experiment (Appendix K, Figure 16) further highlights the fragility of revision training: attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, suggesting that the revision approach is sensitive to training methodology in ways that are not fully understood.

Results Are Specific to a Single Benchmark and Single Model Family

The assumption or constraint. All experiments in the paper use the MATH benchmark (Hendrycks et al., 2021) exclusively—a dataset of high-school competition-level math problems—and the PaLM 2-S* model family (Anil et al., 2023). The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" and that MATH is chosen because "test-time compute is expected to help most when the model already possesses the necessary knowledge and the challenge is drawing complex inferences" (Section 4). This deliberately scopes the findings to symbolic reasoning tasks where correctness is well-defined and the model's capability is partially but not fully mature.

The consequence. Several aspects of the findings could be model-specific or benchmark-specific and may not generalize. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s particular output distribution and calibration properties—a model with different error patterns or different calibration might exhibit different difficulty-dependent scaling curves. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. MATH consists exclusively of competition-level math problems requiring symbolic reasoning with unambiguous ground-truth answers—it is unclear whether the difficulty-dependent patterns (beam search hurting easy problems, revisions helping easy problems) generalize to other reasoning domains (code generation, logical reasoning, scientific QA) or to tasks requiring factual knowledge or open-ended generation rather than inference. Most critically, many real-world applications lack clean correctness signals for both difficulty estimation (via pass@1) and PRM training (via Monte Carlo rollout correctness), making the entire pipeline potentially inapplicable outside of closed-form answer problems.

What evidence exists in the paper. The paper acknowledges the scope limitation implicitly by using only MATH, but does not claim broader generality. The model choice is justified in Section 4. No results on other benchmarks or model families are reported. The paper does not discuss whether the method would transfer to open-ended generation, dialogue, or planning tasks where correctness is ambiguous or multi-dimensional.

Mitigation status. Not addressed. The authors acknowledge in Section 8 that future work should investigate "extension to other domains and modalities," but no experiments or analysis are provided. For practitioners considering deploying RWML, this means the method's effectiveness is only validated for mathematical reasoning with PaLM 2 models, and adaptation to other tasks or model families would require independent verification.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a conceptual reframing of how world model learning integrates into the LLM agent training pipeline. The dominant paradigm in prior work treats world models either as inference-time planning tools (separate simulators queried during MCTS or lookahead search) or as capabilities learned through supervised fine-tuning on expert-generated or strong-LLM-synthesized transition data. RWML challenges both assumptions simultaneously by demonstrating that world model learning through reinforcement learning with self-supervised embedding rewards can serve as an effective mid-training stage that reshapes internal representations before task-specific policy learning begins, without requiring world model outputs at inference time.

This is not a paradigm shift in the sense of introducing an entirely new problem formulation or training objective — the MDP framework, the GRPO algorithm, and the concept of learning transition dynamics are all well-established. Rather, the paper's contribution is a diagnostic reframing that changes which research directions appear promising and which appear limited. Specifically:

The SFT-for-world-models assumption is empirically undermined. Prior work on teaching LLMs world model capabilities through next-state prediction (Zhang et al., 2025a; Yu et al., 2025c) uniformly used SFT as the training objective, with the implicit assumption that token-level prediction accuracy is the appropriate learning signal. RWML provides direct counterevidence: on identical training data, SFT-based world model learning (WM SFT) produces catastrophic performance degradation on ALFWorld (2.8% vs. 36.1% base) and underperforms the base model on τ² Bench (27.9% vs. 31.9%), while RL-based world model learning (RWML) either preserves or improves performance (32.6% on ALFWorld, 38.8% on τ² Bench). The weight change analysis (Figure 3) provides mechanistic evidence that SFT induces broader, more disruptive parameter updates that interfere with existing capabilities and hinder subsequent policy learning. This finding challenges the default assumption in LLM fine-tuning pipelines that SFT is the appropriate tool for teaching any new capability — RWML demonstrates a case where SFT is actively destructive and RL provides qualitatively different learning dynamics that enable knowledge transfer.

RL emerges as a practical tool for auxiliary representation learning, not just policy optimization. The paper's most distinctive design choice — using a frozen embedding model's cosine similarity as a reward for training a world model — represents the introduction of a new reward design pattern that sits between two well-explored extremes. On one side are token-level SFT objectives, which the paper shows fail for tasks requiring semantic understanding rather than surface-form reproduction. On the other side are LLM-as-a-judge rewards, which the ablation (Table 4) shows are susceptible to reward hacking (the vague prediction in Appendix D receiving a perfect score). The embedding-space reward offers a middle ground: it is simple, non-trainable, and resistant to hacking, while still capturing semantic equivalence rather than exact string matching. This pattern is not specific to world model learning — it could apply to any auxiliary task where the goal is to improve the model's semantic understanding of a domain, such as predicting user intent, anticipating tool failures, or forecasting task outcomes.

The training objective determines whether world model knowledge transfers to decision-making. Prior work had not cleanly separated the effect of the training objective from the effect of the training data in world model learning. Methods like IWM and SR from Zhang et al. (2025a) combine expert-data SFT with world model SFT, making it impossible to isolate which component drives performance. RWML provides a controlled experiment: RWML and WM SFT use identical data (same triplets, same subsampling) but different objectives (RL vs. SFT), and the performance gap is both large and consistent (29.8 points on ALFWorld, 10.9 points on τ² Bench). This demonstrates that the training objective is not an implementation detail but a first-order determinant of whether world model knowledge improves agent performance.

This work reconciles a contradiction between the Dyna intuition and SFT-based world model implementations. The Dyna architecture (Sutton, 1991) established that learned transition models can accelerate policy learning — an insight that should naturally extend to LLM agents. But prior attempts to implement this insight with SFT-based world models (Zhang et al., 2025a; Yu et al., 2025c) showed mixed results, often requiring expert data or strong LLM supervision to be effective. This paper suggests that the failure was not in the Dyna intuition but in the training objective used to implement it. SFT's token-level fidelity requirement is poorly matched to the semantic nature of environment dynamics; RL with embedding-space rewards better captures the structure that matters for downstream decision-making. This resolution shifts research attention from "what data should we train world models on?" toward "what objective should we use to train world models?"

Two research directions become more attractive based on these results:

1. RL-based auxiliary task pretraining for agents. The paper establishes that RL with appropriately designed self-supervised rewards can improve agent performance without task-specific labels. This opens a broad design space: what other auxiliary tasks (predicting user satisfaction, anticipating environment surprises, forecasting remaining steps) could similarly serve as representation-learning pretraining for agent policies, and what reward designs make them effective?

2. Conservative parameter update methods for continual agent learning. The weight change analysis (Figure 3) and forgetting results (Table 3) provide evidence that RL-based training induces smaller, more targeted parameter updates than SFT, better preserving existing capabilities. This suggests that RL — beyond its role in policy optimization — may be a general-purpose tool for safe fine-tuning that minimizes catastrophic interference, a finding that extends well beyond the agent domain.

One research direction becomes less attractive based on these results:

3. SFT on interaction data as a standalone agent training method. The paper provides strong evidence that SFT on raw interaction data (without expert filtering or strong LLM curation) can be actively destructive to agent performance. The WM SFT baseline's 2.8% on ALFWorld — a catastrophic collapse from any reasonable baseline — suggests that the common practice of simply fine-tuning on the model's own trajectories should be approached with caution, particularly in environments where the model's initial policy is weak and the training data contains many erroneous transitions.

Follow-Up Research This Work Enables

1. Train a lightweight difficulty predictor to replace the expensive 2048-sample estimation and measure the amortized efficiency gain. The paper's difficulty estimation procedure (Section 3.2) generates 2048 samples per question and scores them with the PRM to assign difficulty bins. This cost is not accounted for in the reported 4× efficiency gains, making the practical deployment advantage uncertain. A direct follow-up would train a classifier — possibly a small model or even a linear probe on top of the base model's embeddings — to predict difficulty bins from the question text alone, using the PRM-based difficulty labels as training targets. The key metric is not just classification accuracy but amortized efficiency: does the total cost (difficulty prediction + strategy execution) still beat best-of-N with the same total budget? A strong result would show that a cheap predictor (e.g., a single forward pass of a small classifier) achieves bin assignment accuracy comparable to the PRM-based method, recovering most of the compute-optimal gains without the estimation overhead. A negative result — the classifier cannot reliably distinguish difficulty levels from text alone, or the accuracy loss eliminates the efficiency advantage — would identify a fundamental limitation of text-based difficulty estimation and motivate adaptive methods.

2. Combine PRM-guided beam search with the revision model as the proposal distribution, measuring whether the complementary strengths compound or interfere. The paper studies PRM search and iterative revisions as independent mechanisms (Sections 5 and 6), explicitly noting they were never combined (Section 8). The two mechanisms have complementary difficulty-dependent strengths: revisions substantially help on easy problems (bin 1–2, Figure 7 right) by refining roughly-correct answers, while beam search helps on medium problems (bins 3–4, Figure 3 right) by exploring diverse solution strategies. The natural experiment is to use the revision model as the proposal distribution within beam search: at each expansion step of the search tree, condition the model on previous (possibly incorrect) branches in context and generate the next steps using the revision-trained model rather than the base model. Would this combination achieve additive gains (revision improvement + search improvement) or do the mechanisms interfere (e.g., the revision model's training makes it less diverse, reducing the exploration benefit of beam search)? The cleanest test would compare four conditions at matched generation budgets on MATH, broken out by difficulty bin: (a) base model + best-of-N, (b) revision model + best-of-N, (c) base model + beam search, (d) revision model + beam search. Condition (d) would test whether the combination exceeds each individual mechanism.

3. Ablate whether the RWML benefit comes from the RL objective or the reasoning token scaffold by adding a "WM SFT with reasoning" baseline. The current WM SFT baseline deliberately excludes reasoning tokens during training (Appendix B.1: "DO NOT perform any thinking"), while RWML trains the model to generate reasoning before predicting next states. This confounds the training objective (SFT vs. RL) with the reasoning scaffold (absent vs. present). The cleanest ablation would construct a WM SFT variant that includes reasoning tokens: for each triplet ⟨s_≤t, a_t, s_{t+1}⟩, use a strong LLM or the RWML-trained model itself to generate a reasoning trace, then train the SFT model to predict both the reasoning and the next state. If WM SFT + reasoning matches RWML performance, then the reasoning scaffold — not the RL objective — is the active ingredient, and the paper's central claim about RL's advantage over SFT would need revision. If WM SFT + reasoning still substantially underperforms RWML, the paper's conclusion that RL provides a qualitatively different learning dynamic is strengthened. A third condition — RWML without reasoning tokens (the model predicts next states directly, with the embedding reward) — would further isolate the effect of reasoning by testing whether RL benefits from the reasoning scaffold or whether the RL objective alone is sufficient.

4. Replicate the compute-optimal scaling framework on a non-symbolic benchmark (e.g., code generation with HumanEval or MBPP) to test domain generality. All experiments use MATH, which consists of competition-level math problems with unambiguous correct answers and a well-defined symbolic reasoning structure. The difficulty-dependent patterns — beam search hurting easy problems, revisions helping easy problems, neither helping on hard problems — may be specific to mathematical reasoning. Code generation shares some properties with math (well-defined correctness via unit tests, multi-step reasoning, clean verifier training signal) but differs in output structure (code is more constrained than mathematical prose, error patterns are different). A replication on MBPP or HumanEval would test whether the same difficulty bins emerge, whether the same strategy assignments (best-of-N for easy, beam search for medium) are optimal, and whether the 4× efficiency gain over best-of-N transfers. A negative result — the patterns differ substantially or the efficiency gain disappears — would scope the paper's contribution to symbolic math reasoning specifically. A positive result — the patterns replicate — would suggest the framework generalizes to a broader class of well-defined reasoning tasks.

5. Train a verifier adversarially on search-generated solutions rather than i.i.d. samples, and measure whether over-optimization resistance improves. The paper documents verifier over-optimization as the primary bottleneck for test-time compute scaling (Section 5.3): beam search degrades easy-problem performance because the PRM's errors get exploited under aggressive optimization. The current PRM is trained on i.i.d. sampled solutions using Monte Carlo rollouts (Section 5.1). A natural improvement is adversarial training: after the initial PRM training, run beam search against the PRM on training questions, collect the solutions that receive high PRM scores but are actually incorrect (false positives), and fine-tune the PRM on these adversarial examples with corrected labels (low score for incorrect solutions). The key metric is whether the beam search degradation on easy problems (bin 1–2, Figure 3 right) is reduced or eliminated after adversarial retraining. A positive result — the degradation curve flattens or reverses, allowing beam search to be safely deployed on easy problems — would substantially increase the efficiency gains from compute-optimal scaling by expanding the set of problems where aggressive search is beneficial. A negative result — adversarial training helps marginally but the over-optimization problem persists — would suggest that the PRM capacity or training procedure has fundamental limitations, motivating ensemble or architectural approaches.

6. Measure whether the forgetting advantage of RWML over WM SFT (Table 3) persists when the KL regularization coefficient β is tuned identically for both methods. The paper shows that RWML causes less catastrophic forgetting than WM SFT across nearly all benchmarks (Table 3), attributing this to the "on-policy nature" of RL. However, the GRPO objective for RWML includes an explicit KL penalty term (β D_KL) that regularizes toward the reference policy, while the WM SFT baseline presumably uses standard SFT without such regularization (the paper does not specify whether any regularization beyond standard weight decay is applied). The forgetting advantage could therefore be partly or entirely due to the KL regularization rather than the RL objective per se. A clean comparison would train an SFT variant with an equivalent KL regularization (e.g., by adding a KL penalty term to the SFT loss or using a constrained optimization approach) and measure whether the forgetting gap persists. If the gap largely disappears, the practical implication shifts from "use RL instead of SFT" to "use KL regularization during fine-tuning" — which is simpler and more broadly applicable. If the gap persists with matched regularization, the paper's claim about RL's inherent forgetting advantage is strengthened.

Practical Applications and Downstream Use Cases

1. Mid-training stage for agent models in production pipelines. Organizations training LLM-based agents for customer service, technical support, or internal tool-use applications — where the agent needs to understand how backend systems, databases, and users respond to actions — can insert RWML as a self-supervised mid-training stage between pretraining and task-specific RL. On τ² Bench (a customer-service environment), RWML + Policy RL achieved 43.7% vs. 38.0% for Policy RL alone (Table 1) — a +5.7 point gain — without requiring any expert annotations or strong LLM data. For a production deployment processing millions of customer interactions, a 5-percentage-point improvement in task success rate translates to substantial cost savings in reduced escalations to human agents. The self-supervised nature of RWML means it can be run on each new model version or each new environment without waiting for expert annotation pipelines.

2. Data-efficient fine-tuning of open-source models for proprietary environments. Companies deploying open-source LLMs (Qwen, Llama) in proprietary environments — where expert demonstrations are unavailable and strong LLM APIs are either too expensive or restricted by data privacy policies — can use RWML to improve agent performance using only the model's own interaction data. The entire RWML pipeline requires zero external supervision: the model explores the environment, collects its own trajectories, and learns world dynamics from embedding-space rewards. On τ² Bench, RWML alone (no task-success rewards, no expert data) improved the base Qwen3-8B from 31.9% to 38.8% (Table 1) — a meaningful gain from a purely self-supervised stage. For environments where building a reward function is itself expensive (the paper notes this as a key motivation in Section 2.2), starting with RWML provides immediate improvement before investing in reward engineering.

3. Reducing catastrophic forgetting during domain-specific fine-tuning. The forgetting analysis (Table 3) shows that RWML causes substantially less degradation of general capabilities than SFT — a -2.38 point decline on MMLU-Redux vs. -10.10 for WM SFT on ALFWorld, and consistently smaller declines across all six benchmarks tested. For organizations that need their fine-tuned agent models to retain general reasoning, instruction-following, and coding capabilities (for example, an agent that handles customer service but must also answer general questions, follow complex policies, and occasionally debug issues), RWML provides a more capability-preserving training method than SFT on interaction data. The weight change analysis (Section 4.2, Figure 3) provides mechanistic evidence that RWML makes smaller, more targeted parameter updates, which directly translates to better capability retention.

4. Bootstrapping agent performance in new environments where reward functions don't yet exist. The paper's core motivation (Section 2.2) is that "designing task-success reward functions in complex environments still requires substantial human expertise, limiting scalability." When deploying an LLM agent in a new domain — a new software tool, a new database system, a new user interaction protocol — the environment dynamics are initially unknown to the model, and designing reliable reward functions requires understanding those dynamics (a chicken-and-egg problem). RWML breaks this circularity: the model can first learn environment dynamics through self-supervised interaction (the world model stage), then the improved model can be used either for direct deployment (RWML alone on τ² Bench: 38.8%) or as a better starting point for reward function design and policy RL. The τ² Bench results are particularly relevant here because τ² Bench involves complex, interleaved tool-use and user interactions — exactly the kind of environment where reward function design is non-trivial (the paper notes that τ² Bench success depends on "whether the agent correctly followed policy, appropriately escalated issues, and maintained user satisfaction").

When to Prefer This Method

The paper explicitly positions RWML against three categories of alternatives (Section 3.1, Table A1): methods that learn from task-success rewards (Policy RL, RFT), methods that learn from SFT on interaction data (WM SFT), and methods that learn from expert/strong LLM data (Imitation Learning, IWM, SR). The following decision rule emerges from the paper's empirical results:

Prefer RWML (with or without downstream Policy RL) when:

• Expert demonstrations or strong LLM-generated data are unavailable, too expensive, or restricted by privacy concerns (RWML uses only the target model's own interaction data)
• The base model has sufficient underlying capability that its own rollouts contain meaningful information about environment dynamics — Section 4.3 shows that Qwen2.5-7B struggles to transfer world knowledge on τ² Bench while Qwen3-8B shows substantial gains
• Catastrophic forgetting of general capabilities is a concern — Table 3 shows RWML causes substantially less degradation than WM SFT on all six general benchmarks tested
• The environment involves structured tool-use responses where ROUGE-based rewards can complement embedding similarity (τ² Bench: 38.8% with hybrid reward vs. unknown performance with pure embedding reward)
• A two-stage training pipeline is acceptable (RWML mid-training followed by Policy RL) — the combination (87.9% on ALFWorld, 43.7% on τ² Bench) consistently outperforms Policy RL alone (81.0%, 38.0%)

Prefer direct Policy RL (without RWML) when:

• The base model is weak and its self-generated trajectories are too noisy or erroneous to provide useful world model training signal (Section 4.3: Qwen2.5-7B on τ² Bench shows minimal transfer)
• Task-success rewards are easy to design and provide dense, reliable signal throughout training (the paper's rationale for Policy RL as a strong baseline)
• The computational budget for mid-training is constrained — RWML adds an entire training stage (2 epochs of GRPO on the collected triplets) before Policy RL begins

Prefer expert-data methods (Imitation Learning, IWM, SR) when:

• High-quality expert demonstrations or strong LLM-generated trajectories are available at scale — on τ² Bench, Imitation Learning (44.2%) marginally outperforms RWML + Policy RL (43.7%), suggesting expert data provides a small additional benefit when available
• The base model is too weak to generate useful self-training data (the paper used GPT-4o and GPT-4.1 to generate reflection data for the SR baseline on ALFWorld and τ² Bench because the target models "cannot consistently generate coherent reflections")

Prefer WM SFT only when it serves as a stepping stone to Policy RL and only on benchmarks where the model already shows some competence — on τ² Bench, WM SFT + Policy RL (40.3%) outperforms Policy RL alone (38.0%) by +2.3 points, suggesting SFT-based world model learning provides modest benefit on more complex environments. However, on ALFWorld, WM SFT + Policy RL (80.4%) is nearly identical to Policy RL alone (81.0%), meaning the SFT world model stage adds cost without benefit. And WM SFT alone is never preferable — it degrades performance below the base model on both benchmarks.