Atlasof World Models
Start

T-10 Province III · The Interior The Dreaming Fields 5 min 10 of 21 New

Learning in Imagination: Policy & Value

How an agent can learn what to do, and how good each situation is, by rehearsing entirely inside its own world model.

Learning in Imagination: Policy & Value: practise in the model, correct in reality Plain words: practise in your head, the real court and fixing the picture after a real try. Architect: actor pi(a | s) and critic V(s) trained on imagined rollouts of horizon 15 started from replayed real states, lambda-returns that bootstrap past the horizon and real episodes added to replay to correct the model. CIO: cheap practice, a costly real trial and the correction loop that is the true running cost, because a policy trained in imagination learns to please the model. T-10 · TWO LOOPS WM imagined : real 15 : 1 30 : 1 45 : 1 practise in your head the real court fix the picture how good it looks from here Practise many times in your head. Then try once for real. imagination, H = 15 actor pi(a|s) critic V(s) real episode, to replay model update lambda-return bootstraps past H Every training trajectory is imagined, started from real states. cheap practice costly trial correction loop value estimate may please the model The correction loop is the real running cost. Budget for it.
A cobalt ball resting on a bare observatory floor while a faint arc in the air traces where it could roll next.

In one sentence Let the agent practise thousands of times in its head, and only occasionally in the real world.

Illustration of Sonu, a ten-year-old girl with a dark ponytail in a yellow shirt, holding a small notepad. Sonu, 10

Think about a basketball player the night before a big game. She lies in bed with her eyes closed. She pictures the ball leaving her hands. She sees it hit the front of the rim and bounce out.

So in her head she pushes a little harder. This time it drops in.

Athletes really do practise like this. It helps them. They still go to the real court. The real court tells them where their pictures were wrong.

Some computers learn the same way. First they build a picture of how their world works. Then they practise inside that picture thousands of times.

They learn two things from all that practice. One is a habit, what to do next. The other is a feeling for how good things look from here.

A player at the free-throw line with a big lead feels calm. A player who is behind with two seconds left feels the pressure. The computer learns that feeling too, as a number.

Then the computer tries its habit for real. Real life shows where its pictures were wrong. It fixes its pictures and goes back to practising in its head.

The trap is the same one the player faces. If her picture of the ball is wrong, she practises the wrong shot all night. So the computer cannot only dream. It has to keep going back to the real court.

Try this at home. Put a box across the room as a target. Throw a soft ball five times and count the hits.

Now close your eyes and picture ten good throws. Then throw five more. Did the picturing help, or did you need the real throws?

Illustration of the Senior AI Architect, a person in glasses and a blue shirt holding a pen. Senior AI Architect

Dreamer (Hafner et al., 2019) trains an actor and a critic inside its latent model. Every training trajectory is imagined. Imagination starts from real states drawn from the replay data. The paper uses a 15-step horizon and tests horizons up to 40.

The actor is a policy π(a | s). The critic is a value V(s). The critic regresses toward λ-returns. These blend a few imagined rewards with bootstrapped values beyond them.

The value carries the agent past the end of its imagination. That keeps targets stable even when imagination is poor far ahead. Dreamer's own ablation shows it. The value model makes results much less sensitive to horizon length.

Gradients can flow back through the learned dynamics for continuous control. Discrete actions use REINFORCE-style estimators instead. DreamerV2 and DreamerV3 extended the same recipe to Atari and many other domains.

The results were strong for the time. On 20 visual control tasks, Dreamer averaged 823 after 5 million environment steps. PlaNet scored 332. The model-free D4PG reached 786 after 100 million steps.

The idea is older than Dreamer. Ha and Schmidhuber (2018) trained a controller entirely inside the dream of their world model. They then moved it back to the real game.

MuZero (Schrittwieser et al., 2019) takes a different cut. Its model never reconstructs observations. It predicts reward, value and policy at each step. That model plugs into Monte Carlo tree search at decision time.

MuZero matched AlphaZero on Go, chess and shogi without being given the rules. It also set a new state of the art on 57 Atari games.

The lesson across both lines is practical. A world model has to be accurate on the quantities the decision depends on. Faithful pixels are optional.

Take payment-plan offers in a bill-assistance program. A real trial takes a full billing cycle. A poor offer lands on a real household. A model of customer response lets a policy rehearse many offer sequences overnight.

The critic learns which account states look recoverable. The actor learns which offer to make next. Neither one ever touches a real account during training.

The danger is the familiar one. The policy will find offers the model loves and customers ignore. A held-out set of real billing cycles is the only honest score. Keep those cycles out of every training run.

Measure on the same start states:

  • Imagined return against real return. A widening gap means exploitation. - Critic error on real transitions, tracked over training. - How often the policy visits states that are rare in the replay data.

Illustration of My CIO, a grey-haired person in a purple jacket with arms folded. My CIO

This is where world models pay for themselves. An agent that rehearses in its own model needs far fewer real trials. That matters wherever real trials are slow, costly or risky.

Utility work is full of such trials. A new outreach rule for bill assistance takes a billing cycle to show results. A new dispatch rule for crews waits for the next storm. Nobody wants to run a thousand of those to train an agent.

Imagination changes that math. Compute is cheap next to a field trial. The agent can practise overnight on cases that would take years to collect for real. It can also rehearse rare cases that no pilot would ever see.

The caveat runs through this whole atlas. A policy trained in imagination will exploit the model's blind spots. It learns to please the model. On the cases the model gets wrong, it will be wrong with confidence.

The fix is a correction loop. Real outcomes flow back. The model is corrected and the policy retrains. That loop is the real cost of the program.

The loop needs data pipelines and a schedule. It also needs one person who owns it. Budget for it from the start. A program that funds the model and skips the loop has bought a confident guesser.

The decision is whether the loop is funded and owned. I would recommend to see a named owner first. A monthly schedule should exist before any policy goes near a customer.

Ask the team one question. How often will real outcomes flow back to correct the model?

Quick check

3 questions. Pick one answer each. Get them all right to mark this territory as checked.

1What two things does a Dreamer agent learn inside its imagination?

2Why does the critic help when the imagination horizon is short?

3What does MuZero's learned model predict?

Key terms

Further reading

  1. Dream to Control: Learning Behaviors by Latent Imagination Danijar Hafner, Timothy Lillicrap, Jimmy Ba, Mohammad Norouzi, 2019 Actor-critic inside a learned latent model, with a 15-step imagination horizon.
  2. Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model (MuZero) Julian Schrittwieser et al., 2019 A model that predicts only reward, value and policy, plus tree search.
  3. World Models David Ha, Jurgen Schmidhuber, 2018 A controller trained entirely inside its own dream, then moved back to the real environment.
  4. World Models · MBRL History Loqman Samani A readable history of model-based RL.

Search the atlas