Learning to Act with World Models by Valuing What Comes Next
Abstract
WM-PPO combines a world model with PPO’s critic to create auxiliary policy advantages and value targets from predicted successor states, improving task success without privileged information.
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 judgment by estimating the current policy's probability of success, but only for states the agent reaches. Our method, WM-PPO, elegantly brings the two together. For each action the policy samples, executed or not, the world model predicts the next observation and the critic values the predicted successor. In the policy channel, comparisons among these values give the policy auxiliary advantages; in the value channel, their mean gives the critic an auxiliary target. A shared gate weights both channels according to how reliably comparisons with similar margins predicted real outcomes earlier in training. The policy, critic and world model all learn from the agent's own interaction, with no teacher or privileged information, and only the policy is used at inference. With a comparable number of real environment steps, WM-PPO raises success from PPO's 63.5% to 87.5% on ALFWorld and from 39.3% to 52.9% on ScienceWorld. It matches the strongest PPO refinements we tested on ALFWorld and exceeds the best of them by 6.5 points on ScienceWorld. The trained policy more often completes prerequisite steps before placing an object or committing to an answer.
Community
TL;DR. A sparse terminal reward says little about individual steps, yet each state transition can change the chances of success long before the reward arrives. WM-PPO learns from this information: a world model trained on real transitions predicts what each sampled action would change, and PPO's own critic values the predicted successor by the current policy's probability of success.
Highlights
- Learning from actions that were never executed. For every action the policy samples, executed or not, the value of its predicted successor enters training twice: comparisons among the candidates give the policy auxiliary advantages (policy channel), and their mean gives the critic a one-step target (value channel).
- Grounded in real outcomes. A shared gate weights both channels by how reliably comparisons with similar value margins predicted real branch outcomes earlier in training.
- A minimal addition to PPO. The world model is the only learned component added. No teacher model or privileged information is used, and only the policy runs at inference.
- Results. With a comparable number of real environment steps, success rises from PPO's 63.5% to 87.5% on ALFWorld and from 39.3% to 52.9% on ScienceWorld. WM-PPO matches the strongest PPO refinements we tested on ALFWorld, exceeds the best of them by 6.5 points on ScienceWorld, and also surpasses the teacher-based methods we tested. On text Sokoban, it reaches 87.0% against 77.3% for ProAct's MC-PPO.
- Behavior. The trained policy more often completes prerequisite steps, such as heating, cooling or cleaning an object, before placing it or committing to an answer.
📄 Paper: https://arxiv.org/abs/2611.12200 · 💻 Code: https://github.com/nuskonghou/wm-ppo · 🤗 Checkpoints: https://huggingface.co/nuskonghou/wm-ppo
Get this paper in your agent:
hf papers read 2611.12200 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 1
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper