Learning to Act with World Models by Valuing What Comes Next

Learning to Act with World Models by Valuing What Comes Next. By Chenghao Gao and four others, from the National University of Singapore, UCLA, and one other institution. Published on the ninth of November.

A sparse reward at the end of an episode tells an agent very little about the individual steps that led there. Much of the missing information is already present, implicitly, in the environment's dynamics, because every state transition can change the conditions for success long before any reward arrives. Two existing tools each hold half of what is needed. A world model trained on real transitions can learn to predict what an action changes, yet a predicted observation does not by itself say whether success has become more likely. And with a binary terminal reward, the critic in PPO does make that judgment, by estimating the current policy's probability of success, but only for states the agent actually reaches.

WM-PPO joins the two. For every action the policy samples, whether or not it is executed, the world model predicts the next observation, and the critic then values that predicted successor. The critic's judgment is thereby extended to states the agent never visited, one imagined step ahead of each candidate action. Each partner covers the other's gap. The world model supplies successors the agent did not reach, and the critic supplies the judgment the world model cannot make: whether that successor brings success closer.

Those values are put to work in two channels. In the policy channel, comparisons among the values of the sampled actions give the policy auxiliary advantages. In the value channel, their mean gives the critic an auxiliary target. A single shared gate weights both channels, and it sets that weight according to how reliably comparisons with similar margins predicted real outcomes earlier in training. In effect, the auxiliary signals are trusted in proportion to their track record: where similar comparisons used to line up with what really happened, they count for more.

Everything is learned from the agent's own interaction. The policy, the critic and the world model are trained without a teacher and without privileged information, and at inference time only the policy is used. The world model and the critic are scaffolding for training, and are set aside once the policy is learned.

With a comparable number of real environment steps, WM-PPO lifts success on ALFWorld from the sixty-three point five percent of plain PPO to eighty-seven point five percent. On ScienceWorld it goes from thirty-nine point three percent to fifty-two point nine percent.

Against the strongest PPO refinements the authors tested, it matches them on ALFWorld and exceeds the best of them by six and a half points on ScienceWorld. The behaviour changes too: the trained policy more often completes prerequisite steps before placing an object or committing to an answer.

A reward that comes only at the end leaves each step unexplained. Here the explanation is built from parts the agent already trains: a world model that says what an action would do, and a critic that says how much that outcome is worth for success. Pairing them, and trusting the pairing only as far as it has proved reliable, gives every sampled action a signal, including the ones the agent never took.