PYTHON · NUMPY

Learn in imagination

Learn a model of the world's hidden dynamics from random play, then plan a policy that acts correctly in reality -- without ever training in the real world.

01Challenge

The generalization wall had a second answer. Instead of only gathering more data, learn a MODEL of the world and act inside it. The Bench drops you in a world with a hidden twist -- it silently rotates every action 120 degrees -- and lets an agent play at random. From that play you fit a world model, a 2x2 matrix that captures the twist. Then the payoff: a naive policy that ignores the model drives the wrong way and never arrives; a policy that plans with the learned model undoes the twist and reaches every goal. This is the world-action-model idea -- Dreamer, NVIDIA Cosmos, LeRobot's 2026 world-model policies -- at its smallest.

Predict firstThe world secretly rotates every action 120 degrees. A naive policy that ignores this -- does it reach the goal?

02Model

A world model predicts what your actions do. Once you have it, you can plan or train a policy inside it -- in imagination -- without touching the real world. Here the model is a linear map from action to motion, learned by least squares from random play; inverting it gives the action that produces the motion you want. That a policy planned entirely in the learned model works in reality is the whole thesis of learning to act by imagining.

03Guided practice

Step A (worked): random play gives (action, motion) pairs; least squares recovers the world's 120-degree twist as the matrix Ghat -- run it and read the recovered angle. Step B: the policy uses K = identity, ignoring the model, so it drives 120 degrees off and reaches almost nothing -- the FAIL. Step C (the fix): K = np.linalg.inv(Ghat), so the policy uses the learned model to undo the twist. Step D: it now reaches every goal in the real world, planned entirely from the model.

04Feedback

PASS WHEN PASS when a policy that plans with the learned model reaches at least 36 of 40 real goals.

If your run fails, check:

Read the bench’s printed result — it passes when a policy that plans with the learned model reaches at least 36 of 40 real goals. Run the cells top to bottom; if one errors, the FAIL line names the step to fix.
The policy still ignores the model -- set K = np.linalg.inv(Ghat) to undo the world's twist.

05Retrieve & space

What does a world model predict?
Why does the naive policy fail where the model-based one succeeds?
What does it mean to train or plan a policy 'in imagination'?

06Mastery & project

Pass the lab: a policy that plans inside a learned world model reaches every goal, having never trained in the real world. That is the world-action-model idea -- the fourth verb, imagine.

The live 'Imagine, then act' lab lets you watch the whole arc: play, learn the model, dream a policy inside it, and deploy it for real.

GO FURTHER · THE LIVE PLAYGROUND

Do it live: Imagine, then act ↗ — play, learn a world model, train a policy in the dream, then act for real.
World models, in depth ↗ — the full course on learned simulators you can imagine in.

← The generalization wall, and what a VLA adds Learn the reward →

PYTHON · NUMPY

Learn in imagination

Learn a model of the world's hidden dynamics from random play, then plan a policy that acts correctly in reality -- without ever training in the real world.

You read

the arrays and values already in scope

You change

the code you write in each cell

Fixed

the dataset and the checks that grade you