all lessons/ world_models/ 03 · learning a predictive latent statelesson 3 / 16

Learning a predictive latent state

Compressing history is not making a smaller photograph. It is deciding which pasts are allowed to blur together — and which differences you must never lose, because they change what happens next.

Where we are
Lesson 01 said a world model is state inference plus action-conditioned prediction. Lesson 02 said the ideal state is a full belief over hidden worlds — and that belief is far too big to store. So this lesson faces the forced problem: squeeze history into a compact learned state that still keeps what the future needs. We’ll derive the one rule that makes compression safe, watch different training losses keep different things, learn why representations quietly collapse, and end with a checklist for designing and testing a state.
Forced by 02The ideal state is a belief, too large to store. This stepCompress history into a predictive latent — sufficiency, not reconstruction. Forces 04A latent is inert until it can advance under an action.

1 · You can’t carry the whole movie

A single 256 × 256 color frame is 196,608 numbers. At 30 frames a second, one minute is over 350 million values — before you add depth, audio, touch, or actions. But to decide whether to push a box, you need almost none of that: where it is, how fast it’s going, what it’s made of, what’s touching it. Keeping the whole stream is not the same as understanding it. So compression isn’t optional; it’s the job.

Introduce an encoder E that folds the whole history Ht = (o1:t, a1:t−1) into a compact learned state xt, which the dynamics can then move forward:

xt = E(Ht)     and     x̂t+1 ∼ F(xt, at)

The word latent just means “not directly observed” — nothing more. It need not be mysterious, continuous, or neural: a latent can be a vector, a feature grid, discrete codes, object slots, a scene graph, particles, or a blend of named geometry and learned features. The only question that matters is what it keeps.

And here is the deep point. Compression is a many-to-one map: countless histories will land on the same latent. That is exactly what you want when their differences are noise, texture, or viewpoint that no query cares about. It is a disaster when their differences are velocity, contact, or intent that changes a future under some action. So learning a state is really learning which pasts are allowed to be treated as the same.

2 · The one rule that makes forgetting safe

When may two histories safely collapse to the same latent? Only when nothing you could do would ever tell them apart. Take histories H and H′ that the encoder merges, E(H)=E(H′). Downstream, the model can no longer know which one happened — so merging is safe only if every future you promised to predict is identical for both, under every action sequence in scope:

E(H)=E(H′) ⇒ p(yfuture | H, afuture, κ) = p(yfuture | H′, afuture, κ)

Here yfuture is whatever you predict (occupancy, reward, pixels, contact), afuture is a candidate action sequence, and κ names the questions you promise to answer. That implication is called predictive sufficiency relative to a query set — the whole lesson in one line.

It splits into two pressures pulling opposite ways:

You can have one without the other. A state can be sufficient but bloated (it tracks the object and memorizes the wallpaper). It can be minimal but broken (it stores the object’s class but forgets its velocity). The dream is a compact sufficient statistic; finite data and imperfect optimization make it a compromise.

Crucially, there is no universal “nuisance.” Paint color is nuisance for dodging a rigid box but state for quality inspection; fine texture is irrelevant to rigid motion but essential for estimating friction. So before you design a bottleneck, you must first name the questions, horizon, actions, and tolerable error. Nuisance is defined by the job.

3 · Make it concrete: one square, six factors, three jobs

Picture 64 × 64 videos of a small square sliding across a textured floor. Every frame is built from six ingredients:

  1. horizontal position x,
  2. horizontal velocity v,
  3. identity i — fragile glass or solid rubber,
  4. floor texture τ,
  5. brightness ℓ, and
  6. random camera noise n.

An action a ∈ {push left, no push, push right} nudges velocity by −1, 0, or +1 pixel/step, and position follows xt+1 = xt + vt+1. Hit the right wall and glass shatters while rubber bounces. Now watch how the same pixels serve three different jobs:

Run the arithmetic once. Position 55, velocity +2, wall at 63. With no push for three steps: 57, 59, 61. Push right first (velocity → +3): 58, 61, then 64 — the wall on step three. Push left first (velocity → +1): 56, 57, 58. If the square is glass, the right push is unsafe; if rubber, the hit might be fine or even useful. Same starting picture, opposite verdicts — because identity is a hidden factor the pixels barely show.

Now squeeze the encoder to just three scalars and train it to redraw the current frame (Q1). Since the textured background is 95% of the pixels, the loss happily spends capacity on τ and ℓ and stores a rough x — reconstructing beautifully while discarding velocity and identity. Two clips ending at x=55, one moving +2 and one moving −2, get nearly the same latent even though their futures diverge. Q1 is happy; Q2 and Q3 are doomed.

Swap the objective to five-step position error under varied pushes, and now the state is forced to keep x and v. The third scalar stores identity only if the training data actually includes wall crashes and the loss distinguishes “break” from “bounce.” Otherwise identity predicts neither the short motion nor the objective, and it’s dropped. The lesson: the data’s coverage of actions and events decides whether a factor’s relevance is ever revealed.

So capacity alone never picks the right state. Three scalars can solve Q3 if they hold x, v, i; a 1,024-dim latent can still fail if the loss rewards texture shortcuts and the trajectories never reach a wall.

The bottleneck chooses what exists
Increase latent capacity. Too little loses dynamics; too much lets texture and sensor noise leak into state.
Interactive bottleneck diagram.
future score
—
reconstruction
—
nuisance
—
diagnosis
—

The widget shows an allocation pressure, not a magic “correct dimension.” Too few slots and decision-relevant factors collide; too many slots with a reconstruction-heavy objective and nuisance leaks in cheaply. In real systems the effective bottleneck isn’t just vector width — it’s noise, quantization, temporal stride, attention, pooling, and which gradients reach which tokens.

4 · Reconstruction learns to describe the picture

The obvious first idea: squash the image into a latent and try to paint it back. That’s an autoencoder, trained with squared pixel error:

Lrecon = average over pixels of (ot − ôt)²

Read it plainly: penalize each pixel’s squared miss and average. This has real virtues — dense supervision with no labels, pressure to keep broad scene information, and a decoder you can actually look at. It’s the right objective when generating the observation is the job, as in controllable video.

But notice what it optimizes: information gets valued by how many pixels it explains. A big textured floor outweighs a small moving square. Random details are costly to reproduce even though they can’t be predicted or controlled. And when several futures are possible, a squared-error decoder splits the difference and produces blur. A fancier stochastic decoder can paint multiple modes, but the state may still burn capacity on looks.

Predicting future frames is stronger than redrawing the present, because it forces the latent to carry dynamics — but it still tangles predictable meaning with rendering. This is exactly where the Computer Graphics track pays off: pixels come from geometry, materials, lighting, and camera, so a model asked to render must keep all of those, while a model asked only “will it collide?” need not.

5 · Contrastive learning: teach by comparison

What if you skip painting pixels entirely and just ask a comparison question: does this future belong to this past, more than to some random other clip? That is a contrastive objective — it scores the true future above negatives, so the signal is relational rather than pixel-perfect.

This can spend nothing on redrawing texture and instead sharpen high-level regularities. But the catch is subtle and important: the negatives define what counts as a useful distinction. Draw negatives from other videos with different floors, and the easy win is matching texture, not motion. Draw them from nearby times in the same video, and position and fine timing suddenly become valuable. And a “false negative” — a different observation that’s actually equivalent for your query — pushes the model apart when it should pull together. Contrastive loss doesn’t discover causal state; it makes the distinctions you selected cheap to learn.

6 · Predicting features — and the trap of collapse

Go one step further: don’t predict the future pixels, predict its features — another encoder’s representation of the future or a masked region:

Lembed = distance(P(xcontext, a), stopgrad(xtarget))

A predictor P guesses the target embedding; stopgrad means the target side is held fixed for this update. Because there’s no pixel decoder, the target is free to drop unpredictable texture and keep only stable structure.

But there’s a way to cheat, and it’s catastrophic: map every input to the same constant. Then the predictor matches perfectly while encoding nothing about the world. This is collapse, and every method here is really a scheme to forbid it:

And collapse hides. A representation can have perfectly healthy numerical variance yet only vary with brightness — loss stays low if brightness happens to predict the target. Catching this kind of “semantic collapse” needs controlled probes and interventions, not a glance at singular values or the loss curve.

7 · If you only make one kind of decision, keep only what changes it

When the model exists for a fixed control task, you can preserve the decisions — rewards, values, policies, search outcomes — instead of the whole observation. Two latents are value-equivalent when every action sequence earns the same expected return from both, even if their pixels differ wildly.

In the square world, “don’t break the glass” lets you throw away floor texture and brightness entirely; position, velocity, and identity stay because they change which push is safe. The result can be a wonderfully lean planning state.

The price is task blindness. Ask that same robot later to inspect the floor’s texture and the information may be gone for good. Worse, reward can be sparse: until a trajectory actually reaches the wall, identity never affects observed return, so the encoder gets no reason to keep it. Multi-task targets, auxiliary predictive losses, deliberate coverage of consequential events, or simply a bigger general-purpose state all hedge against tomorrow’s task.

8 · Mixing objectives: give each loss one job

A loss is not a label you stick on afterward — it is the pressure that decides which mistakes are allowed to consume capacity. To reason about a hybrid objective, imagine feeding the same history in and asking, one loss at a time, exactly what gradient follows:

  1. Current reconstruction asks “can you recover what the sensor sees now?” Every region gives dense supervision, so the state keeps broad detail and you get a visual diagnostic. But nothing here demands motion — if texture dominates, velocity can vanish. Catch it with identical last frames reached by opposite motions.
  2. Future reconstruction adds “can you render what it will see later?” Now representation and dynamics must cooperate — stronger than copying the present. The trap is equating pretty output with useful state: appearance can still dominate, multiple futures blur, and full-frame generation is costly for planning. Check structured outcomes apart from pixels.
  3. Contrastive prediction asks “which candidate future belongs here?” Dropping the decoder buys strong discriminative features — but the negative sampler is a hidden curriculum. Easy negatives (texture, timestamp, camera) reward shortcuts; false negatives punish equivalent futures. Change the negatives and see if the state changes for the right reason.
  4. Predictive embedding asks “can you match an abstract target feature?” The target can shed unpredictable detail and focus on structure — but the shortcut (map everything to one feature) is fatal. Stop-gradient, a teacher, or variance control blocks full collapse; partial collapse and teacher bias remain. Probe decisions, not feature variance.
  5. Inverse/action prediction asks “which action connected these two states?” To answer, the state must highlight controllable change — a great fix when passive appearance dominates. But it can ignore passive motion and may just read a visible actuator trace. Remove the trace and test motion that happens without the agent.
  6. Reward/value/task heads ask “which information changes the decision?” They yield a lean state for a known goal — and that leanness is the boundary: anything irrelevant to the training reward is dropped, even if a future task needs it. Test transfer under changed rewards.
  7. Geometry/flow/depth heads ask for structured spatial quantities. They ground motion and viewpoint in variables many queries reuse — but need labels or assumptions that can be wrong. Validate under camera motion and calibration shift rather than trusting the auxiliary target.

Read the list as a map of complementary blind spots: current reconstruction gives breadth but not motion; future prediction adds time but keeps a rendering burden; action prediction emphasizes control but misses passive change; task heads sharpen decisions but narrow transfer; geometry heads add structure but import assumptions. Hybrid training is popular because one pressure patches another’s hole — a light decoder prevents information loss, a future-feature loss adds predictable semantics, an inverse head adds controllability, reward heads add task relevance.

But the naive move — add every scalar loss together — does not guarantee each role survives. Their scales, gradient directions, data frequency, and head capacity decide what the shared encoder actually becomes.

So treat loss weights as resource allocation. Watch per-objective gradient size and downstream probes, not just the total number. A strong private head can ace its target without improving shared state; a dominant pixel loss can drown out a rare collision signal.

Audit each loss against a pair
For every loss, name the physical distinction it’s supposed to keep, build two examples differing only in that distinction, and check the shared latent actually changes — then build a nuisance-only pair that should stay equivalent. If the head succeeds while the shared state fails these pairs, the objective is being solved by a shortcut or a private head, not by building the world state you wanted.

9 · The bottleneck has many dials, not one

Latent width is just the most visible knob. The effective bottleneck also includes:

So an architecture is really a prior over which facts are cheap to encode. A convolutional grid makes locality cheap. Object slots make instance persistence and interaction cheap, but discovery and identity hard. A transformer makes flexible relations cheap but exact geometry and long memory expensive. A voxel or point cloud fits geometric queries but needs pose and calibration. Choose the biases that match the invariances your task actually has.

10 · Should the state ignore a change, or move with it?

It’s tempting to demand invariance to every nuisance — rotate the camera, keep the same state. But some changes should move the state in a predictable way, not vanish. That is equivariance. Slide an object two meters and its spatial representation should slide two meters; the collision map must not stay put.

Different coordinate choices make this concrete:

The 3D Vision lessons cover projection, pose, and multiview geometry — exactly the structure that helps when a world model must reason across camera motion. Purely learned features can work, but built-in geometric equivariance usually buys data efficiency and interpretability.

11 · Keep the belief, not just the best guess

Lesson 02 was blunt: the ideal state is a distribution. A deterministic latent vector can still carry distribution parameters — means, variances, mixture weights — but only if the training targets demand them. Reward a decoder for one best reconstruction and the encoder will quietly drop the losing hypotheses.

Three common ways to keep the doubt:

  1. Parametric uncertainty: output a Gaussian, categorical, or mixture. Cheap, but only as expressive as the family.
  2. Stochastic latent state: factor state as x=(h,z), sampling z from a learned posterior in training and from the transition prior in imagination. This carries branching futures — unless a too-strong decoder makes the model ignore z.
  3. Ensembles or particles: keep several hypotheses and read off their disagreement — extra compute for honest multimodality.

A healthy state should widen its doubt when it’s flying blind and narrow it when good evidence lands. If every sampled future differs only in texture while sharing the same wrong trajectory, that’s visual variety, not decision-relevant uncertainty.

12 · Test the state, not just the loss

A low training loss only proves the encoder-plus-head optimized its objective on its data. It does not prove x holds the state you wanted. Use direct tests:

  1. Frozen probes: freeze the encoder and read off position, velocity, identity, contact, reward, and uncertainty with deliberately simple heads. Easy readout means accessible organization (though a failed probe might just be weak).
  2. Interventions: change background/lighting with dynamics fixed, then change velocity/action with appearance fixed. The latent should track the causal change, not the nuisance.
  3. Counterfactual pairs: identical last frame, opposite hidden velocity. If their latents merge, memory is insufficient.
  4. Rollout tests: probe imagined latents at growing horizons — measure drift, permanence, constraint violations, calibration.
  5. Nearest neighbors: see which clips land close in latent space. Texture-based neighbors expose shortcuts averages hide.
  6. Ablation sweeps: vary dimension, tokens, losses, history — look for the smallest state that keeps downstream performance, not the prettiest reconstruction.
  7. Closed-loop: let a planner use the state and compare reward, safety, and regret against model-free, oracle-state, and observation-only baselines.
A probe is evidence, not a verdict
A big nonlinear probe can dig out information the dynamics or planner could never use; a tiny probe can miss information stored in a distributed form. Match the probe’s power to the intended consumer, and always pair probes with causal interventions and closed-loop outcomes.

13 · Tempting beliefs that are wrong

14 · How to reason about it out loud

Given “design the latent state for a robot manipulation world model,” walk this order:

  1. Query contract: predict contact, object pose, success, and uncertainty for candidate end-effector commands over two seconds.
  2. Required factors: gripper pose and velocity, object poses and shapes, articulation, support/contact, rough material behavior, and belief about occluded objects.
  3. Intended nuisances: sensor grain and lighting that don’t affect material inference — stated as caveats, not a blanket “appearance is irrelevant.”
  4. Structure: an ego-centric 3D feature field for free space, object slots for persistent entities, and a global recurrent token for unmodeled context.
  5. Objective roles: future-feature prediction for dynamics, geometry/pose heads for grounding, contact/reward heads for decisions, and a modest decoder for diagnostics and breadth.
  6. Anti-collapse: teacher/stop-gradient or contrastive separation, balanced slot use, rollout consistency, action variety.
  7. Uncertainty: multimodal object states or stochastic transitions plus ensemble disagreement on unfamiliar interactions.
  8. Targeted tests: same-image/opposite-velocity, lighting swaps, long occlusion, novel pushes, contact events, calibration, and closed-loop planning.

A strong answer names the tradeoffs. Metric 3D state helps spatial planning but leans on localization and calibration; object slots help persistence but struggle with amorphous material; a big transformer state is flexible but expensive to roll thousands of times; a compact task latent plans fast but transfers poorly. The right architecture is the one whose biases, compute, data, and outputs meet the stated contract.

15 · Takeaway

Takeaway
A predictive latent is a lossy, query-relative contract. Two histories may share a latent only when their relevant futures agree under the candidate actions in scope. Bottleneck, architecture, data coverage, and objective jointly decide which distinctions survive — and reconstruction, contrastive, predictive-embedding, geometry, and value losses each keep a different world. Validate with probes, interventions, rollout calibration, and closed-loop decisions, never the training loss alone.

We’ve decided what the latent should keep — but we treated its transition F as a black box. Lesson 04 opens that box: how deterministic memory, stochastic latents, priors, posteriors, and sequence models actually learn to move a state forward under an action.

Interview prompts