Building a First World Model › A Task to Think With

Abstract descriptions of "learning to predict the consequences of actions" are hard to reason about, so let us pin everything to one small environment for the whole chapter.

Imagine a puck sliding on a frictionful table. You control the puck by applying forces, and your goal is to bring it to a target location. Here is the twist that makes the problem interesting: you can only see where the puck is, never how fast it is moving. Velocity is invisible. Yet velocity determines the future, because momentum carries the puck after you stop pushing.

This is a deliberately engineered gap between what you observe and what you need to know. A controller that reacts only to the current position will overshoot the goal again and again, because it cannot anticipate that the puck will keep drifting. To succeed, the agent must somehow remember: recent positions imply recent motion, and recent motion implies future positions. Memory of the right kind is not a luxury here. It is the whole problem.

It is worth spelling out precisely why reactive control fails, because the failure mode is instructive. Suppose the puck sits directly on the goal but is moving quickly away from it. A position-only controller sees zero error and applies no force, or perhaps a tiny corrective nudge. Either way, momentum carries the puck off target, and only then does the controller react, chasing a puck that is already moving. The result is the familiar oscillating pursuit you may have seen in badly tuned thermostats and camera autofocus systems. The information needed to avoid this trap, namely velocity, exists in the world and even in the observation history, but not in any single observation. The agent's only route to success is to reconstruct the hidden variable from the sequence of things it has seen. This reconstruction-from-history requirement is the defining feature of partially observable problems, and our little puck is the cleanest possible instance of it.

The environment itself is simple enough to define in a few lines. The true physical state is a four-dimensional vector: position and velocity in the plane. Actions are two-dimensional forces, clipped to a bounded range. Observations are the two position coordinates plus a little sensor noise. Rewards encourage reaching the goal while penalizing large forces.

Each of these design choices earns its keep. The four-dimensional state is small enough that we can later inspect what the model has and has not learned, yet rich enough that one dimension of it is genuinely hidden. The bounded forces prevent the trivial degenerate solution of applying infinite force, which would make dynamics modeling unnecessary. Sensor noise matters more than its small magnitude suggests: it guarantees that the agent cannot recover velocity perfectly from two consecutive positions, so any useful velocity estimate must be an integrated, smoothed inference over many steps rather than a finite-difference calculation. Finally, evaluating the reward at the resulting state, a detail we return to during implementation, ensures the reward is predictable from quantities the model actually represents. Small design decisions like these determine whether a toy problem teaches the right lesson or an accidental one.