What Is a World Model? The Birth of Imagination in Machines

What Is a World Model? The Birth of Imagination in Machines

1.1 Can a Machine Imagine?

Imagine you are standing in your kitchen, holding a full glass of water. You know, without any conscious calculation, that if you tilt your hand too quickly the water will slosh over the rim. You know that the cat will dart away the moment you reach for it. You know that the next step down the staircase is exactly where your foot will land, because you have walked these stairs a thousand times. None of this knowledge came from a rulebook. It came from something quieter: an internal model of how the world behaves, built gradually from experience, tested silently every time you act.

Now ask a harder question. Can a machine possess the same kind of knowledge? Can it hold, inside its circuits, a representation of cause and effect, of gravity and friction, of what happens when it pushes versus when it pulls? Can it "imagine" the consequences of an action it has never actually performed, and choose well on the basis of that imagination?

This book is built on the conviction that the answer is yes, and that this capacity is among the most consequential ideas in modern artificial intelligence. The technical name for this capacity is a world model: an internal system that an agent uses to predict how its environment will respond to its actions. A world model is, in a precise sense, a machine's imagination. It is the thing that allows an artificial agent to look before it leaps, to test ideas without paying the price of failure, and to compress a lifetime of raw experience into a structured understanding of how things work.

The goal of this chapter is to give you a rigorous yet intuitive grasp of what a world model is, why the concept has such deep roots in both cognitive science and engineering, and why it has become the central design idea behind a new generation of intelligent systems. By the end of the chapter you should be able to do three things: write down the formal mathematical definition of a world model, explain in plain language why separating prediction from action is a powerful design principle, and walk through a concrete example in which an agent learns to navigate a small grid environment using nothing but a learned model of transitions. The chapter closes with a preview of how this abstract idea becomes the engine behind self-driving cars, humanoid robots, and game-playing AIs that rival the best human players in history.

To begin, we need to separate two questions that are often muddled together. The first question is whether an agent can act intelligently in an environment. The second question is whether the agent must do so by predicting that environment. Model-free approaches to intelligence try to answer the first question directly, mapping observations to actions through sheer learned association. Model-based approaches add the second layer: before acting, the agent consults a learned predictor of the world. A world model is the formal realization of this second layer. Understanding it well is the foundation for almost everything that follows.

1.2 A Formal Definition: Prediction as the Foundation

Let us begin with the precise definition, because the rest of the chapter depends on it. A world model is a function that, given the current state of the world and an action the agent is considering, predicts what the next state will be and what reward will be received. In symbols:

$$f: \mathcal{S} \times \mathcal{A} \rightarrow \mathcal{S} \times \mathbb{R}$$

Here $\mathcal{S}$ is the set of all possible states the agent might be in, and $\mathcal{A}$ is the set of all actions it might take. The function $f$ takes a state-action pair $(s_t, a_t)$ at time $t$ and produces a prediction of the next state $s_{t+1}$ and the immediate reward $r_t$. In a stochastic world model, the prediction is not a single next state but a probability distribution over possible next states:

$$f(s_t, a_t) \sim P(s_{t+1}, r_t \mid s_t, a_t)$$

This formulation is so compact that it can be easy to miss its power. Three things deserve close attention.

First, the state $s_t$ is a summary. It is not the raw sensory flood of pixels, sounds, and forces that an agent receives at each moment. It is a compressed, sufficient description of everything about the past that matters for predicting the future. In a chess game, $s_t$ is the configuration of pieces on the board; in a self-driving car, it is a vector that encodes the positions and velocities of nearby vehicles, the state of traffic lights, and the curvature of the road ahead. The world model operates on these summaries, not on raw sensory input. We will return to the question of how those summaries are built in Section 1.5, where the compression property of world models is treated as one of their three defining benefits.

Second, the action $a_t$ is whatever the agent can do. In a grid-world it might be the choice to move north, south, east, or west. In a robot arm it is a vector of joint torques. In a language model it might be the next token to emit. The world model does not need to know what kind of action it is handling; it simply learns the mapping from any state-action pair to its consequences. This generality is part of what makes the framework so widely applicable.

Third, the reward $r_t$ is the scalar signal that defines the agent's goals. In reinforcement learning the reward is the formal encoding of "what we want"; in a navigation task it might be $+1$ for reaching a goal and $-1$ for falling into a pit. Many world models include reward prediction as part of their output, because rewards are part of how the world responds to action. Other world models are reward-agnostic, predicting only the next state and leaving the reward structure to a separate evaluator. Either choice is valid; both are common.

Putting these pieces together, a learned world model is any parametric function—typically a neural network—trained to approximate the true transition dynamics of an environment. If we denote the parameters of the network by $\theta$, then the trained world model is

$$\hat{f}_\theta(s_t, a_t) \approx P_{\text{env}}(s_{t+1}, r_t \mid s_t, a_t)$$

The hat over $f$ reminds us that the model is an approximation. No model is perfect. The art of building a good world model lies in choosing the right architecture, the right state representation, and the right training data so that this approximation is accurate enough to be useful.

To see why this matters, consider what happens when an agent has such a model. Suppose it finds itself in state $s_t$ and must decide among several candidate actions. It can simply roll out each candidate: feed $(s_t, a_1)$ into the model, read out the predicted $s_{t+1}$, feed that into the model again with the next candidate action, and so on. After several simulated steps, the agent has a forecast of where each action will lead and what rewards will accumulate. It can then pick the action whose forecast is best. This is the essence of planning, and it is impossible without a model. Without one, the agent is reduced to trial and error in the real world, which is often costly, irreversible, or fatal.

This is also why a world model is fundamentally different from a policy. A policy is a function $\pi: \mathcal{S} \rightarrow \mathcal{A}$ that maps states directly to actions. A world model is a function $f: \mathcal{S} \times \mathcal{A} \rightarrow \mathcal{S} \times \mathbb{R}$ that maps state-action pairs to predicted consequences. The policy says *what to do*; the world model says *what would happen if you did*. They are complementary. The most powerful agents use both: a world model for look-ahead reasoning, and a policy for fast action when look-ahead is unnecessary.

For the mathematically inclined reader, this is also the point where the Markov property becomes important. A state $s_t$ is said to be Markov if it contains all inf