Building a First World Model › The Shape of a Minimal World Model

This section covers the architecture of our model: what it stores, how it updates that storage when new observations arrive, and how it predicts forward when they do not. We build intuition first and write down the mathematics alongside.

The environment's true state has four numbers, but the agent never sees two of them. So the agent must construct its own state, and we call this constructed state the latent state. It is "latent" because it is never observed directly. It exists only inside the model, inferred from the stream of observations and actions.

Why not skip latents entirely and train a network to map raw observation histories directly to predictions? You could, for short histories. But a fixed-size latent state that is recursively updated gives you something a growing raw-history window does not: a bounded, persistent belief that can be maintained indefinitely with constant computation per step. When we formalize this in Part II: Inference, Decisions, and Control, this idea becomes the belief state of a partially observable decision process. Here we simply build one and watch it work.

There is a second, equally important reason for latents. The quantities we ultimately care about, like "where will the puck be in ten steps" and "how much reward will I collect," are functions of the physical state, not of any single observation. A latent state gives the network a place to accumulate exactly the information those predictions require, and nothing forces it to waste capacity replaying irrelevant pixel-level detail.

A third reason deserves mention even at this early stage, because it shapes everything downstream. A fixed-size latent is a bottleneck, and bottlenecks force compression, and compression forces the network to decide what matters. Given unlimited capacity, a model can hedge by memorizing everything about its inputs; given a tight latent, it must prioritize the features that actually reduce prediction error. In our puck environment this prioritization is easy to see: the single most predictive unobserved quantity is velocity, so pressure from the loss pushes the latent toward representing it. In richer environments the same pressure operates on messier trade-offs, deciding whether surface texture or object identity better serves future prediction, and those trade-offs become a central design question rather than an accident. The bottleneck is thus not merely an efficiency measure. It is the mechanism by which the objective's preferences get written into the representation.

Our minimal world model has four neural networks, each with one job:

• Encoder $e_\phi$: compresses a raw observation $o_t$ into an embedding, the sensory contribution to the latent state.
• Transition model $w_\theta$: the heart of the system. Given the current latent state and a candidate action, it predicts the next latent state. This is the "physics engine" of the model, learned rather than programmed.
• Decoder $d_\psi$: maps a latent state back to observation space, so the model can be trained by reconstructing what actually happened.
• Reward head $g_\omega$: maps a latent state and action to a predicted reward, connecting the model to the task rather than only to the sensations.

The subscripts $\phi, \theta, \psi, \omega$ denote the parameters of each network. All four train jointly, end to end, from nothing but recorded interaction data.

It helps to think of these components as answering four distinct questions. The encoder answers "what did I just see?", translating raw sensation into a form the rest of the system can consume. The transition model answers "what happens next if I act?", and it is the only component consulted during imagination. The decoder answers "what would I see if my belief were true?", providing the training signal that anchors latents to reality. The reward head answers "how good is this situation for my task?", tying the whole apparatus to the objective the agent actually cares about. Notice the asymmetry between them: the encoder, decoder, and reward head are ordinary feedforward networks, while the transition model is the piece whose quality under repeated self-application determines whether planning works. Most of the failure modes we discuss later live in that one component.

One might ask why the decoder is needed at all if the goal is control. The answer lies in how gradients reach the latent state. Without a decoder, the only pressure shaping the latent comes from the reward head, and reward signals are sparse and weak in most environments. Reconstruction provides a dense, per-step learning signal that forces the latent to retain enough information to describe the world, and the reward head then builds task structure on top of that foundation. The division of labor is deliberate: reconstruction keeps the latent honest, reward prediction keeps it relevant.

Here is the conceptual core of the chapter, and the distinction most often muddled in casual discussions. The transition machinery plays two different roles depending on whether a fresh observation is available.

When the agent acts and then sees what happened, it performs a filtering update (also called the posterior step). The latent state is revised using both the action just taken and the observation just received:

$$h_t = \mathrm{GRU}\big(h_{t-1},\; [\,e_\phi(o_t)\,;\, a_{t-1}\,]\big)$$

where $[\cdot\,;\cdot]$ denotes concatenation and the GRU cell is a small recurrent network. Because the update sees the newest observation, errors in $h_{t-1}$ get corrected. This is the model operating online, in contact with reality.

When the agent wants to imagine, no future observation exists. It applies the prior transition (the imagination step), which relies on the action alone:

$$\bar{h}_{t+1} = w_\theta(h_t, a_t)$$

From this predicted latent, the decoder produces a predicted observation $\hat{o}_{t+1} = d_\psi(\bar{h}_{t+1})$ and the reward head produces a predicted reward $\hat{r}_t = g_\omega([\bar{h}_{t+1}; a_t])$. Crucially, the output of the prior step can be fed back into itself: $\bar{h}_{t+2} = w_\theta(\bar{h}_{t+1}, a_{t+1})$, and so on for as many steps as we like. This recursion is what turns a one-step predictor into a simulator.

The word "filtering" is borrowed from estimation theory, where a filter maintains a running estimate of a hidden state by combining predictions with measurements. Our GRU plays exactly that role, though it learns the combination rule rather than deriving it from assumptions about noise statistics. The intuition worth carrying forward is that filtering is a correction operation: whatever nonsense accumulated in the latent last step, the fresh observation offers a chance to repair it. Imagination, by contrast, is a commitment operation: the model stakes its entire picture of the future on the accuracy of its current belief and its transition function, with no external check until reality eventually weighs in. Every design tension in this book, from multi-step losses to uncertainty penalties to replanning schedules, traces back to managing the gap between these two modes of operation.

Notice what the latent state must contain for either step to work. Since observations reveal only position, the encoder output at a single time step says nothing about velocity. But the recurrent filtering update receives a sequence of embeddings, and a sufficiently trained GRU can integrate them into an implicit velocity estimate. We will verify empirically that this happens, by comparing against a model stripped of its recurrence.

This verification plan is worth pausing on, because it illustrates a general experimental strategy. We cannot directly inspect the latent and read off "velocity lives here," since distributed representations do not label their dimensions. What we can do is construct a controlled comparison: two models identical in every respect except the presence of memory, trained identically, evaluated identically. If the recurrent variant systematically outperforms the memoryless one on precisely the tasks that require velocity knowledge, we have strong indirect evidence that the recurrence is being used to encode velocity. Controlled ablation of this kind is often more informative than inspection, and we will rely on it repeatedly.

Training data comes from an exploratory policy, here just random forces, collected as episodes of transitions $(o_t, a_t, r_t, o_{t+1})$. We cut these episodes into overlapping windows of length $L$ and train on the following objective. Starting from the first observation of a window, alternate prior predictions and filtering updates:

$$\mathcal{L} \;=\; \frac{1}{L}\sum_{t=0}^{L-1}\Big[\underbrace{\big\| o_{t+1} - d_\psi(w_\theta(h_t, a_t)) \big\|_2^2}_{\text{observation prediction}} \;+\; \underbrace{\big(r_t - g_\omega([\,w_\theta(h_t, a_t);\, a_t\,])\big)^2}_{\text{reward prediction}}\Big]$$

Read the summand carefully, because every symbol earns its place. At step $t$ of the window we hold the filtered latent $h_t$, which has digested observations up to and including $o_t$. We apply the prior transition with action $a_t$, obtaining a predicted next latent. The decoder must reconstruct the actually observed $o_{t+1}$ from that predicted latent, and the reward head must predict the actually received $r_t$. Only then do we perform the filtering update with the real $o_{t+1}$ and move on. The gradients therefore flow through chains of prior transitions, teaching the transition model to be accurate over the exact kind of recursive application it will face at planning time.

That last sentence deserves emphasis, because it is easy to misread the objective as pure one-step training. Within a window of length $L$, the loss at step $t$ depends on the latent $h_t$, which was itself shaped by the filtering updates of earlier steps, and gradient descent adjusts earlier transitions based on how well later predictions turn out. The alternation of prior steps and filtering steps means the model is always asked to predict from beliefs that were built up through its own prior machinery, punctuated by corrections. This is a middle ground between pure one-step training and full open-loop rollouts, and the window length $L$ controls where that middle ground sits: short windows approximate one-step training, long windows approach open-loop training with all of its stability challenges. Choosing $L$ is therefore not an arbitrary hyperparameter but a statement about how far into the future you want the model to be accountable.

A few explicit assumptions are baked into this setup, and naming them now will save confusion later:

• Markov property in the latent: we assume the filtered latent $h_t$ summarizes everything relevant for predicting the future. If important information falls outside the GRU's capacity, predictions degrade silently.
• Approximate determinism: our transition model outputs a single predicted latent, not a distribution. This is adequate when the environment is nearly deterministic apart from small sensor noise, and inadequate when outcomes are genuinely multimodal. Probabilistic transitions arrive in Part V: World-Model Architectures.
• Stationarity: the environment's dynamics do not change over time, so one model fit on old data remains valid.
• Off-policy data: training uses random actions, while evaluation uses planned ones. The model is never explicitly warned about this mismatch, a theme that will recur throughout the book.

Each assumption hides a potential failure. The Markov assumption fails gracefully here, since a modest GRU easily integrates enough history to estimate velocity, but in environments with long-range dependencies, say a door opened twenty minutes ago that still affects airflow, a fixed-size latent may simply be too small, and the degradation appears only as mysteriously poor predictions in rare situations. Approximate determinism fails whenever the same latent and action legitimately lead to different outcomes: a point-estimate model averages those outcomes, producing a prediction of a future that never actually occurs, and a planner acting on that average can be systematically misled. Stationarity fails in any adaptive environment, including environments containing other learning agents. And the off-policy mismatch is subtler than it looks: random actions visit a broad but shallow coverage of the state space, while planned actions drive the system along narrow corridors of high-value states, and the model's accuracy along those corridors is largely untested at training time. None of these assumptions is wrong for our toy problem, but each defines a boundary beyond which the simple recipe breaks, and the parts of this handbook that follow are largely a tour of what lies beyond those boundaries.