World models and embodied AI: Visual walkthrough

Everything below is a picture first, a symbol second. If you have never seen $ar g max$ , an integral, or a Gaussian, you will still be able to follow from line one.

Step 1 — A world is just "state now → state next"

WHAT. Draw the agent's whole life as a chain of snapshots. Each snapshot is called a state, written $s_{t}$ . The little $t$ is a clock reading: $s_{0}$ is "now", $s_{1}$ is "one tick later", and so on. Between two snapshots the agent does something — an action $a_{t}$ .

WHY. Before we can predict the future, we must agree what the future is made of. It is made of states, glued together by actions. That chain is the smallest honest picture of an environment.

PICTURE. Follow the beads on the string.

Nothing here is learned yet. We just named the beads.

Step 2 — The future is a fan, not a line

WHAT. From one state $s_{t}$ and one action $a_{t}$ , the next state is often not certain. Push a cup: it might slide, tip, or spin. So instead of a single arrow to one $s_{t + 1}$ , draw a fan of arrows, thicker where the outcome is more likely.

WHY. Real environments are noisy. If we pretend the future is a single line we will plan for a fantasy. The fan is honesty about uncertainty.

PICTURE. One action, a spray of possible next states.

The height of the fan at some point is a probability: a number between 0 and 1 saying how believable that outcome is. Tall = likely, short = rare, zero = impossible.

Step 3 — We don't have $P$ , so we learn a copy of it

WHAT. In the real world nobody hands us the true fan $P$ . The world-models idea: fit our own fan from watched experience and call it $\hat{P}_{θ}$ . The little hat $^$ means "our estimate, not the truth". The $θ$ is the bag of tunable numbers (weights of a neural net) we adjust until the copy matches reality.

WHY. Once we own a fan we can query it for free, forever, without touching the real (expensive, breakable) environment. This copy is exactly the forward model. It is built with the same machinery as a VAE (to compress observations) and a recurrent network (to carry memory through time).

PICTURE. Left: true fan (unknown). Right: our learned fan, adjusted to overlap it.

Step 4 — Why the copy must be a mixture, not one bump

WHAT. Suppose at a fork a car can go left or right. If our learned fan is a single hill (one Gaussian), its peak lands between the two roads — it predicts driving into the wall. Fix: build the fan out of several hills added together — a mixture.

WHY. The future is often multimodal: two or more genuinely different outcomes, with a gap between them. One bump can only ever say "somewhere in the middle." Several bumps can say "either here or there."

PICTURE. Blue = one lonely bump landing in the wall. Yellow = two bumps landing on the two roads.

Step 5 — Score a whole plan by imagining it forward

WHAT. A plan is a sequence of actions $a_{0}, a_{1}, \dots, a_{H - 1}$ over a horizon of $H$ steps. To score it, feed $a_{0}$ to the learned fan, sample an $s_{1}$ , feed $a_{1}$ , sample $s_{2}$ , … and add up the rewards. That is one imagined rollout.

WHY. We want the best plan, and "best" means "highest total reward across the whole horizon." We can only know a plan's total reward by playing it out — but now we play it out inside $\hat{P}_{θ}$ , for free.

PICTURE. One plan unrolled through the dream, rewards collected along the branch.

Step 6 — We can't sum infinite futures, so we sample them

WHAT. That average $E$ hides an impossible integral over every branch of every fan. Replace it with a cheap trick: roll out the plan $K$ times (each a different random branch) and average the $K$ totals.

WHY. A perfect average is uncomputable; a sampled average is one loop of code. The more samples $K$ , the closer the estimate. This is the same "just simulate it many times" spirit behind reinforcement learning returns.

PICTURE. $K$ imagined branches for the same plan; their total-rewards averaged into one score.

Step 7 — Pick the winning plan: Model Predictive Control

WHAT. Try many candidate plans, score each with Step 6, keep the best, execute only its first action, then re-plan from the new real state. This loop is Model Predictive Control (MPC).

WHY. Executing only the first action, then re-planning, keeps us honest: the learned fan drifts from reality over long horizons, so we never trust it far. Short, repeated look-aheads beat one long committed guess.

PICTURE. Several candidate plans fanned out; the green one wins; only its first step is taken.

If the fan is differentiable, we can skip blind trial and instead slide the actions uphill on the score using the chain rule — actions change future states which change future rewards — the gradient-based planning in the parent note. Blind search (CEM/CMA-ES) and gradient planning both maximise the same objective above.

Step 8 — Edge cases the fan must survive

WHAT & WHY. A plan is only as safe as the fan under stress. Four corners to check, each with what goes wrong and the fix.

PICTURE. The four failure corners on one board.

The one-picture summary

The whole page compressed: learn a fan ( $\hat{P}_{θ}$ ) → imagine plans through it (rollouts) → average many branches (Monte-Carlo return) → keep the best plan, act once, re-plan (MPC). The mixture keeps the fan honest at forks; short horizons keep it honest over time.

Recall Feynman retelling — say it back in plain words

Imagine a kid who wants to win a marble game. Instead of just flicking marbles and hoping, she first builds a little picture in her head of how marbles roll — she learned it by watching lots of flicks. That head-picture isn't one sure path; it's a spray of "maybe here, maybe there," and at a fork it keeps two possibilities instead of splitting the difference and aiming at the wall. Now, before any real flick, she plays whole games in imagination: try plan A, watch it bounce; try plan B; because each imagined game is a bit random she replays each one several times and averages the scores. She picks the plan that scored best, makes only the first flick for real, looks at where the marble actually landed, and re-imagines from there. She never trusts the head-picture too far, because tiny mistakes in imagining pile up the longer she dreams. That is a world model planning by imagination.

Recall Symbol check

What does $s_{t + 1} \sim P (\cdot ∣ s_{t}, a_{t})$ say? ::: the next state is drawn at random from the fan of outcomes opened by the current state and action. Why a mixture of Gaussians instead of one? ::: the future can be genuinely two-valued (left OR right); one bump predicts the average and crashes into the gap. What does $ar g max$ return? ::: the input that maximises the expression, not the maximum value itself. Why does MPC execute only the first action then re-plan? ::: because the learned model drifts from reality over long horizons, so we only trust its short-term prediction. What compounds over a long rollout? ::: per-step model error — it grows with rollout length, turning far-future states into fantasy.

Related vault notes: World models and embodied AI · Reinforcement-Learning-Basics · Variational-Autoencoders · Recurrent-Neural-Networks · Imitation-Learning · Sim-to-Real-Transfer · Vision-Transformers · Diffusion-Models.