Learning to Act with World Models by Valuing What Comes Next

AI summary

Overview

Research area: Reinforcement learning for language-model agents in text-based environments, combining world models (learned environment dynamics) with actor-critic policy optimization (PPO).

Technical level: Advanced. The paper assumes familiarity with PPO, value functions, advantage estimation, and LoRA fine-tuning, though the central idea is intuitive.

Scope: The paper proposes WM-PPO, a method that uses a learned world model plus PPO's existing critic to value the predicted consequences of candidate actions, turning that valuation into auxiliary learning signals for both the policy and the critic, and evaluates it on ALFWorld and ScienceWorld.

What This Paper Is About

Agents trained on multi-step text tasks usually get only a single 0/1 reward at the very end, which says nothing about which intermediate steps helped or hurt. The authors' observation is that much of that missing information is already implicit in how the environment changes: finding an object or meeting a precondition shifts the odds of eventual success long before the reward appears. The goal is to extract that implicit signal by having a world model predict what each candidate action would change, and having PPO's critic judge whether the predicted outcome makes success more or less likely.

Key Contributions

1. Learning from dynamics by valuing predicted consequences. The authors show that applying PPO's critic to world-model-predicted successor states ("transition values") converts information implicit in environment dynamics into feedback that supplements a sparse terminal reward, and they verify with diagnostics that these values support both action selection and actual policy improvement.

2. Calibrating value-based feedback against real outcomes. Because critic estimates can be noisy and biased, a shared reliability gate weights both learning channels according to how often comparisons with similar value margins agreed with real branch outcomes observed earlier in training.

3. WM-PPO: two complementary learning channels. Comparisons among transition values become auxiliary policy advantages (policy channel), while their mean becomes an auxiliary one-step critic target (value channel); the design adds only one learned component to PPO (the world model), uses no teacher model or privileged information, and requires only the policy at inference.

4. Empirical validation across two benchmarks. On ALFWorld, success rises from PPO's 63.5% to 87.5%; on ScienceWorld, from 39.3% to 52.9%, at a comparable number of real environment steps.

Main Findings

• Valuing real successors improves action selection. Holding the policy fixed on ALFWorld (PPO policy and critic after 24 iterations, 53.6% success on 274 evaluation games), selecting among four sampled candidates by the critic's value of the real next state raises success from 56.9% (majority vote) to 70.4%, a 13.5-point gain (SE 2.8, p < 10⁻⁴). A single sampled candidate yields 54.4%, and a control that only exploits immediate task completion reaches 59.5%, indicating the gain comes from valuing intermediate transitions.

• Predicted successors retain most of that gain. Replacing one-step execution with the world model's greedy prediction reaches 66.4%, 9.5 points above majority vote (SE 2.5, p < 0.001) and close to the 70.4% obtained with true consequences. An untrained base model used as the world model reaches 61.7%, so training the world model adds 4.7 points.

• Reliability depends on the value margin. Of 195 candidate pairs with differing estimated success probabilities, the critic orders 57% correctly and the median value gap is 0.008; ranking accuracy rises from 53% in the lower half of gaps to 71% in the top quarter, and among the 12 pairs where one candidate succeeded in all three continuations and the other in none, 11 were ranked correctly. This motivates the margin-binned gate.

• Predicted transitions are exact for familiar observations, uncertain for novel ones. Predictions match exactly for 99.7% of held-out observations that follow from what the agent has already seen, but for only 26% of those that reveal objects it has not seen, where they instead supply a prior over what will be found.

• The comparisons alone can improve the policy. One offline PPO-clip epoch on 8,042 decision points from 500 new training games, using consequence-based labels and no task reward, raises policy-only success from 57.1% to 65.3% (+8.2 points, SE 2.0, p = 0.0003). Majority-vote labels on the same data and update count give 56.9%, essentially no change, even though the two label sets agree at 44% of decision points.

• Online training beats PPO and tested refinements. WM-PPO raises ALFWorld success from 63.5% to 87.5% and ScienceWorld from 39.3% to 52.9%; it matches the strongest PPO refinements tested on ALFWorld and exceeds the best of them by 6.5 points on ScienceWorld. The main table compares 17 methods and variants, and the interaction-efficiency analysis places WM-PPO among 23 other methods and settings on ALFWorld and 20 on ScienceWorld.

• Behavior changes qualitatively. The trained policy more often completes prerequisite steps before placing an object or committing to an answer, matching the motivating example of cleaning a plate before storing it.

• Branch comparisons are informative but not always available. The policy proposed a different action within eight resamples in 62–67% of branches on ScienceWorld and 40–61% on ALFWorld; branches where all eight samples repeat the chain action yield no comparison.

Methodology in Plain English

The setup is a text game where the agent writes a short action each turn and receives a text observation back. With a terminal reward of 1 or 0 and a discount factor of 1, the critic's value for a state is literally the probability that the current policy will eventually succeed from there.

The method works as follows. At each training state, the policy samples four candidate actions but executes only the first, so the behavior policy is unchanged. For every distinct candidate (executed or not), the world model—a separate LoRA adapter on the same base model—greedily predicts the next observation; a deterministic prompt constructor appends that prediction to the history to form a predicted successor state. The same critic that PPO already uses then values that predicted successor, giving each candidate a "transition value." Predicted endings are valued directly as 1 or 0.

Those transition values feed two channels. In the policy channel, the candidate with the highest transition value is treated as preferred, and the four sampled responses get auxiliary advantages of the group-relative form (a 1 for the preferred candidate minus its mean), which in effect push up the log-probability of preferred responses relative to the others at the same state. In the value channel, the mean of the four transition values becomes an auxiliary one-step regression target for the critic at the current state, added to the critic's main regression on real returns.

Because critic-based comparisons can be noisy, both channels are scaled by a shared weight per margin bin. Each iteration plays 64 tasks, each with one main "chain" episode plus three "branches" that replay the chain to a random step and resample until a different action is taken, then continue to the end—giving 256 episodes per iteration and real paired outcomes for different actions from the same state. Earlier branches are compared against the transition-value ordering, and margin bins that historically agreed more often with real outcomes receive higher weight. Real returns remain the critic's main supervision, the world model is trained by next-observation cross-entropy on executed transitions only, and after training only the policy is deployed. Models are LoRA adapters on Qwen3-4B-Instruct-2507, with four candidates, an auxiliary strength of 1, clipping bounds of 0.2 and 0.28, a value-channel coefficient of 0.5, and GAE λ of 0.95.

Why This Matters

Impact on research. The paper reframes world models in RL for language agents: instead of keeping the world model in the deployed loop for planning (as in prior WMA, ITP, and CoMAP-style systems) or using it to generate foresight traces from a teacher, it uses predictions purely as an internal training signal grounded by the critic already present in PPO. It also offers a concrete alternative to distillation and privileged-information approaches, since no teacher, hindsight observations, reference solutions, or simulator access is required—and it addresses the known problem of noisy LLM value estimates with an outcome-calibrated reliability gate rather than a fixed heuristic.

Real-world applications:

• Training dialogue and tool-use agents where success is only known at the end of a multi-turn session (booking flows, form completion, troubleshooting assistants).
• Household and embodied instruction-following agents, the setting of ALFWorld, where ordering constraints (clean before storing) matter more than the final disposition of an object.
• Scientific or procedural assistants, the setting of ScienceWorld, where intermediate conditions such as heating or mixing must be satisfied before answering.
• Any sparse-reward sequential decision problem with an available text or structured environment where rollouts are expensive and simulation-based planning at inference is impractical.

Industry relevance. The method is deliberately cheap to adopt: it adds one learned component to an existing PPO pipeline, reuses the transitions already being collected, requires no additional teacher model, and adds no inference-time cost since only the policy is deployed. That profile fits teams already running PPO-style RL fine-tuning on open-weight models and looking for denser learning signals without buying more environment interaction or inference latency.

Future Directions

• World-model fidelity for unseen content. Predictions are exact for 99.7% of observations implied by what the agent has seen but only 26% for those revealing unseen objects; improving prediction of genuinely novel observations could close the remaining gap between predicted-consequence selection (66.4%) and true-consequence selection (70.4%), which is not resolved in the paper.

• Extending to environments where endings are not reported. Only ScienceWorld's observations report episode termination, so in ALFWorld every transition value is a critic value; the question of how to detect predicted task completion in environments without explicit end markers is left open.

• Widening the reliability calibration. The gate currently bins by value margin and uses branches from preceding iterations; whether richer calibration (for example conditioning on state, task type, or prediction confidence) improves the gate is not explored.

• Transfer beyond text environments and this model family. All models here are LoRA adapters on Qwen3-4B-Instruct-2507, and both benchmarks are text environments with 40-step limits; generalization to other model sizes and to continuous-control or multimodal domains is not reported.

• Cost accounting. The paper reports comparable real environment steps and an interaction-efficiency comparison against 23 other methods/settings on ALFWorld and 20 on ScienceWorld, but a full breakdown of the extra compute from training and running a second model is not given in the provided content.

Target Audience

Researchers and practitioners in reinforcement learning and language-agent training who already know PPO and value-based methods and want a principled way to exploit environment dynamics for credit assignment under sparse terminal rewards. It is most useful to engineers building multi-turn LLM agents who can afford PPO-style RL fine-tuning but want denser learning signals without a teacher model, privileged simulator information, or inference-time planning. Readers new to actor-critic methods will find the diagnostics (Sections 2.1–2.3) the most accessible entry point, since they isolate each claim with controlled experiments before the full method is introduced.

Authors’ abstract

A sparse terminal reward tells an agent little about its steps. Much of the missing information is implicit in the environment's dynamics, because each state transition can change the conditions for success long before the reward arrives. A world model trained on real transitions can learn to predict what each action changes, but a predicted observation does not by itself show whether success has become more likely. With a binary terminal reward, PPO's critic makes this…