Learning latent dynamics

A compressed state that just sits there is a photograph. It becomes a world model only when you can poke it with an action and it tells you what happens next — and because the same scene can lead to different futures, that “next” has to be a spread, not a single arrow.

1 · Start from the contract, not the architecture

Forget architectures for a moment and ask what the transition is for. At time t the agent has seen o_≤t, taken a_<t, and is about to try a_t. A useful dynamics model answers one conditional question:

The physical state s is hidden, so we swap the whole history for an internal state and demand that, once that state is known, old frames add nothing decision-relevant:

with x_t the learned state, r_t reward or cost, and c_t continuation (does the episode keep going?). That single line is the whole contract: compress history, apply action, predict consequences. Every RNN, transformer, token model, or diffusion model in this lesson is just one way to honor it.

The simplest honoring is a deterministic step, x_t+1 = f_θ(x_t, a_t) — cheap and differentiable. But it has a fatal flaw: the same visible situation can lead to different futures, because of hidden intent, unresolved geometry, sensor noise, or plain randomness. A single output vector then has to average those futures (blur) or smuggle them into an unstable code. So we make the transition probabilistic. That one decision drives the rest of the lesson.

2 · Split the state into a story and a dice roll

A recurrent state-space model (RSSM) carries two kinds of state, and the split is the whole trick:

• Deterministic memory h_t — the running story. A recurrent update accumulates ordered context: which objects appeared, how they moved, what actions were applied. Fix its inputs and it’s fully determined.
• Stochastic state z_t — the dice roll for what’s unresolved. A sampled variable stands for what’s genuinely uncertain or newly revealed right now. Different draws describe different plausible local worlds.

The update runs in a deliberate order. First advance the story using the last dice roll and the last action:

In words: h_t−1 is the story so far, z_t−1 says which plausible world we were in, a_t−1 is the nudge we gave it, and f_θ (a GRU, state-space layer, or causal transformer block) rolls them forward. Crucially h_t is computed before seeing o_t — it’s a genuine prediction, not a summary that peeks at the current frame.

Now make two distributions over the same dice roll z_t — the guess before you look, and the correction after:

The prior guesses the next state from memory alone — it’s all you have when you imagine a future with no new camera frame. The posterior also gets to see the encoded real observation, so it corrects that guess with evidence — it’s what you use while filtering a real trajectory. (“Posterior” here isn’t the exact Bayesian posterior; it’s a learned approximation that infers the latent best explaining the current frame.)

Both routes must speak the same latent language. A posterior sample must be decodable and useful to the task heads; a prior sample must land in that same space. Otherwise you’ll reconstruct recorded frames beautifully while imagination wanders into alien states the instant observations stop.

3 · Build the objective one term at a time

Latents are never labeled, so we train them through the evidence they must explain. Fix the convention we’ll use all series: belief x_t picks action a_t; that action yields reward r_t, continuation c_t, and next observation o_t+1; the model represents those consequences in the destination state x_t+1=(h_t+1,z_t+1). The three heads factor cleanly:

For one transition, the standard loss is the negative evidence lower bound plus task weights:

Read it left to right and it’s just common sense. Draw the destination latent from the observation-corrected posterior. Ask that draw to explain the next frame and the consequences of the action that just happened — those log-likelihoods, made negative, are the prediction costs. The weights λ say which facts deserve capacity. And the last term, the KL, compares the corrected posterior with the imagination prior. (Some codebases index reward one step differently; harmless only if the loader, transition, heads, and return math all agree.)

Don’t dismiss the KL as “a regularizer that smooths latents.” Its job is precise: teach the prior — which must run blind, without future observations — to predict the latent distribution the posterior would infer after seeing them. In density form,

If the posterior puts mass where the prior thought nothing could happen, the ratio q/p blows up and the penalty climbs. Driving it down closes the gap between inference and imagination. But the two networks pull against each other. Too strong a KL forces q≈p before the latent learns anything from the frame — the decoder then leans on memory and the dice roll goes dead. Too weak a KL lets the posterior stuff in frame details the prior can’t foresee — reconstruction shines while imagined rollouts fall apart.

Real systems manage that tension with a KL weight β, “free nats” (a small information budget that isn’t penalized), or split gradients (one term pulls the prior toward a stopped-gradient posterior, another nudges the posterior toward a stopped-gradient prior). These are optimization knobs, not changes to the contract.

4 · Do the KL by hand: what is it actually measuring?

Numbers make the KL concrete. Take a one-dimensional latent for a car’s lateral position. Before the next frame, the prior predicts p(z_t+1|h_t+1) = N(2.0, 0.5²). After the camera catches the lane markings, the posterior infers q(z_t+1|h_t+1,o_t+1) = N(2.4, 0.4²). For two univariate Gaussians,

Plug in: log(0.5/0.4)=0.223; posterior variance 0.16; squared mean-surprise (2.4−2.0)²=0.16; and 2σ_p²=0.50. So KL=0.223+(0.32/0.50)−0.5=0.363 nats.

That one number bundles two disagreements. The mean term says the evidence shifted the car 0.4 latent units from where memory predicted. The variance terms say seeing the lane made the position more certain. Training pushes the prior toward mean ≈ 2.4 and a tighter spread when past context makes that correction predictable, and simultaneously discourages the posterior from encoding things that can’t be anticipated from history. But if the frame reveals a genuine surprise — a pedestrian stepping out — the posterior should move; forcing zero KL would erase the surprise instead of modeling it.

Once z_t+1 is in hand, a decoder predicts o_t+1 and the reward/continuation heads predict the action’s cost and whether the episode ends. The latent coordinate needn’t literally be meters — it only has to support these predictions and stay reachable by both prior and posterior.

5 · What the objective does and does not pin down

It’s tempting to read each latent coordinate as a physical variable, but the loss doesn’t require that. Apply any invertible remap g to every latent, let the transition and decoders learn the matching inverse, and the observable predictions are unchanged — a rotated latent basis models the same world as an axis-aligned one. So what’s actually pinned down is a predictive equivalence class: states are “the same” when they imply the same future consequences under every relevant action.

This decides whether you should even demand “disentanglement.” If a planner only needs accurate reward and constraint queries, an entangled-but-stable state is fine. If engineers need to edit object pose, transfer across robots, or read failures, structure becomes a product requirement — and you get it by injecting bias: object slots, coordinate frames, equivariant layers, bottlenecks, or supervised probes. The variational objective alone will never spontaneously label “position,” “velocity,” and “intent.”

The Markov property, likewise, lives in the learned state, not in a single image. We want the story-plus-dice to be enough:

Recurrent memory is the mechanism that makes that approximation hold. If velocity needs three frames but you reset memory after two, no transition head can conjure the missing information. Test for leftover history-dependence by handing a probe extra past context: if that sharply improves next-event prediction beyond (h,z), the state isn’t sufficient for that query.

Finally, the shape of the dice matters. A diagonal Gaussian treats nearby alternatives as smooth coordinate shifts and enables the reparameterization z=μ+σ⊙ε. A categorical latent with probabilities π_1:K treats alternatives as separated symbols, with KL

Discrete states keep contact modes or object identities separate instead of averaging them, at the cost of within-mode detail and trickier gradients. Many RSSMs use several small categoricals so their combinations give lots of states while each softmax stays manageable. Choose by the topology of your uncertainty, not by a belief that one latent family is universally richer.

6 · The architecture follows the unit you predict

All four families obey the same contract; they differ in what one prediction is. RSSMs compress time into recurrent memory and predict a small continuous or categorical dice roll. Per-step cost is tiny, so millions of imagined control steps are practical — but recurrence can bottleneck information, and a plain Gaussian latent can under-represent separated futures.

Autoregressive token models quantize images or patches, then predict the future as a sequence of categorical picks. Cross-entropy is stable and expressive and transformers bring their whole toolbox — but token order is a modeling choice, errors compound within a frame as well as across time, and fine quantization stretches context length.

Diffusion or flow transitions start from noise and iteratively shape the next latent or trajectory, capturing rich multimodal densities with no fixed mixture count. The price is many network evaluations per prediction — often unaffordable inside a planner scoring thousands of action sequences.

Predictive-embedding models predict a target encoder’s features instead of pixels, sparing capacity from texture and emphasizing stable semantics. Their danger is collapse: if both encoders emit a constant, prediction is trivial and useless — so stop-gradient targets, variance constraints, or task heads must keep the features informative.

None of these wins on a universal axis. Pick the predicted unit from the downstream query: a control agent may want a compact RSSM with reward and continuation heads; a video editor may pay diffusion’s compute for detail; a search agent may want discrete tokens for branching and caching. The architecture is “right” only relative to latency, horizon, uncertainty, and output fidelity.

7 · The heads decide what the state is allowed to forget

Say the scene has a red emergency switch, a moving crate, and an approaching wall. The latent can’t keep every photon. A tempting plan says “learn a general representation first, attach prediction heads later.” That’s false during training: every head sends gradients into the latent, so every head votes on what the latent must remember. The representation isn’t learned first and queried second — it’s shaped by the queries as it forms. Walk it head by head.

1. An observation head asks the state to predict pixels, depth, or proprioception. These targets are dense, so almost every step gives a signal, and the latent keeps appearance and geometry. The trap is equating dense supervision with relevance: unpredictable wallpaper can eat capacity while the tiny switch contributes few pixels. So ground the state in observations, but don’t let observation fidelity be the only definition of useful. In practice: keep the head, lower its weight or predict stable features when nuisance dominates, and diagnose by swapping texture while holding geometry fixed.
2. A reward or cost head asks for a task scalar. Its gradient keeps variables whose changes alter consequences — distance to the wall, whether the switch is live — which is efficient because it targets decisions directly. But reward is often sparse: two very different situations both read zero until one turns dangerous. So pair task pressure with denser targets, and diagnose reward-only compression by changing geometry where immediate reward stays flat and checking whether the latent notices before the delayed consequence lands.
3. A continuation head predicts whether the trajectory survives or hits an absorbing terminal. Without it, an imagined rollout keeps collecting value past the point where the real episode would have ended. The subtlety: continuation labels can reflect data-collection rules (a time limit) rather than physics. So distinguish true terminals from administrative cutoffs, and diagnose by inspecting predicted return around every terminal type — mass must not leak past absorbing states.
4. Geometry or object heads (occupancy, pose, tracks, contact) force the state to keep objects and relations across time — easier to inspect and exactly what collision/manipulation planners need. Hoping pixel reconstruction will spontaneously invent stable identity usually fails under occlusion. The cost is labels or reliable self-supervision; the diagnostic is object permanence: hide an object, reveal it, and check identity, pose, and contacts survived.
5. Value or policy heads keep controllability and value-equivalent distinctions and can make planning blazing fast — but only for the current task. Change the goal and two states that shared an optimal action may now need to differ, and a value-shaped latent may already have merged them. So decide up front: narrow task controller, or reusable environment model? Keep broader heads if goals may change, and diagnose overspecialization by swapping rewards while freezing the model and seeing what can no longer be recovered.

A robust system usually blends dense grounding with task heads, then watches whether one loss hijacks the shared encoder. The real test isn’t “did total loss fall?” It’s “which counterfactual changes move the latent — and are those exactly the changes that alter the future?”

8 · A step-by-step training procedure

1. Collect sequences, not shuffled frames. Store observations, actions, rewards/costs, terminal flags, timestamps. Shuffle away the time order and you destroy the very thing dynamics is made of.
2. Encode each observation: e_t=E_ϕ(o_t). Normalize modalities, preserve action timing.
3. Initialize memory with zeros, a learned state, or an encoder over a context window — and reset it at real episode boundaries.
4. Advance the deterministic path: h_t=f(h_t−1,z_t−1,a_t−1). This ordering keeps the current frame out of the prior.
5. Parameterize both distributions: the prior reads h_t; the posterior reads h_t,e_t. Output means and scales or categorical logits; keep scales away from zero.
6. Sample differentiably: Gaussian z=μ+σ⊙ε, ε∼N(0,I); for discrete latents use a straight-through or relaxed sample.
7. Predict evidence and consequences from (h_t,z_t), using a likelihood suited to each variable rather than forcing unit-variance MSE on everything.
8. Optimize reconstruction, task, and KL terms — and log each component separately, because a healthy total loss can hide a collapsed latent or a neglected reward head.
9. Validate the prior on its own: infer a start state with the posterior, then remove observations and roll the prior under held-out actions. Compare predicted events, rewards, geometry, and uncertainty against the real continuation.

A short burn-in lets memory fill before you score losses. Truncated backprop caps memory cost — but truncate shorter than the real dependency and the model can never learn it. Batch construction, reset masks, and action–observation alignment matter as much as the neural block.

9 · Separate the routes and the failures become obvious

The prior/posterior split is also a debugging instrument. Instead of staring at one total loss, climb three increasingly harsh tests: posterior filtering (every real frame can correct belief), one-step prior prediction (dynamics must cross one gap with no correction), and free-running prior rollout (the model eats its own states). The first test that breaks names the missing mechanism.

1. Sharp reconstruction, useless imagination. Recorded frames reconstruct cleanly, but imagined futures turn to mush once observations stop. The trap — “the decoder works, so the model works” — confuses posterior inference with prior prediction. The likely cause: the posterior carries frame info the prior can’t predict. Confirm by measuring the gap between posterior-conditioned and prior-only metrics from the same start, then strengthen dynamics, tighten KL alignment, and add multi-step prior losses. Success shows up first at one step, then across the rollout — if only reconstruction changes, you fixed the wrong thing.
2. The latent ignores the observation. Filtering and imagination look identical even when the frame carries news. An over-strong KL or too-powerful memory has collapsed the posterior onto the prior. Look for KL near zero, then mask or swap the current frame — if the state barely moves, evidence isn’t getting in. Reopen the channel with KL warm-up, free nats, a weaker decoder, or better gradient balance; verify with controlled frames that reveal hidden state, where the posterior should move exactly when warranted.
3. Prior variance explodes. Everything becomes uncertain as the horizon grows, even in easy regimes. Calling that “honest uncertainty” is too generous — an underfit transition or unstable scale can inflate variance to hide systematic error. Plot predicted scale and entropy by horizon and regime against real residuals; bound log-scales, supply missing state/action inputs, add data where error is high. A good fix makes uncertainty track difficulty, not rise everywhere.
4. The latent memorizes texture. Background changes move the state more than reward-critical geometry, because dense reconstruction drowns the task signal. Run a nuisance counterfactual: change wallpaper/lighting while keeping objects and dynamics, and watch the latent and task outputs. If they lurch, predict stable features or objects, reweight task heads, augment nuisance — then re-run; invariance should improve without erasing real geometry.
5. State jumps or averages at contact. Smooth transitions work until impact, where the world branches into stick, slide, bounce, or separate. One smooth latent averages those regimes and follows none. Stratify errors just before and after contact; add stochastic or discrete mode variables or a contact head, and check each sampled branch stays physically coherent rather than just adding pixel variety.
6. Great offline loss, poor control. The data may only cover the behavior policy’s actions, or the latent may omit variables that matter only under intervention. Passive next-frame accuracy can’t prove action sensitivity. Evaluate counterfactual and planner-selected actions, not just recorded ones; if error rises off-policy, collect exploratory data and add action-sensitive predictions. The final check is closed-loop: change an action from a fixed start belief and the predicted consequence should move the way the real system does.

The order is the point: can evidence be encoded → can one transition predict it → does that transition survive self-composition → can a planner use it under new actions. Also plot per-dimension KL, not just its sum — a few active dimensions can mask widespread collapse. This turns a vague “bad world model” into a pinpointed inference, dynamics, stability, or coverage failure.

10 · System-design implications

Make action semantics explicit. The transition must know whether a_t is a requested command, an actuator setpoint, or a measured executed action. With delay, a command at t may only bite at t+2; a low-level controller may reshape steering before the wheels feel it. Store only requested actions and the model will read execution variation as randomness — so put actuator state and delay in the latent or transition. A quick alignment test shifts the action stream forward and backward and checks which offset best predicts the next posterior.

Define the imagination interface. A planner should copy a latent, batch candidate actions, sample transition branches, and query reward, continuation, constraints, and uncertainty without decoding full pixels — rendering is optional and usually the latency hog.

Keep inference and simulation consistent. Live, you filter through the posterior; the planner forks that state and rolls the prior. Version encoder, transition, and heads together — a latent has no stable meaning across independently changed parts.

Budget uncertainty deliberately. More stochastic samples improve coverage but multiply planning cost. Pick sample count from decision sensitivity, reuse common random numbers when comparing actions, and protect rare safety outcomes instead of pruning by average probability.

Instrument reset and timing. A one-frame action offset can masquerade as irreducible noise. Episode boundaries, dropped frames, actuator delay, and sensor latency belong in the data model; hide them and the RSSM wastes stochastic capacity explaining pipeline bugs.

Match capacity to queries. More dimensions can cut reconstruction loss without helping control. Size the state by held-out prior rollouts and downstream decisions, not compression ratio — and if goals may change, keep general predictive heads instead of shaping everything around one reward.

Separate model state from application state. A latent checkpoint is an implementation detail, not a database schema. Persist raw observations, actions, timestamps, and outcomes so beliefs can be recomputed after retraining; when a service hands state between machines, version the whole bundle and refuse mixed versions rather than trusting two same-sized vectors to mean the same thing.

11 · Takeaway — and why a transition distribution is only the start

We now have a working loop: memory advances under action; a prior proposes the next latent blind; a posterior corrects it with evidence; heads ground the latent; and the KL trains those two routes to share one usable language. But writing p(z_t+1|h_t+1) doesn’t guarantee the distribution captures the right uncertainty. A Gaussian can still smear mass between two physically separate futures, and “diverse” samples can carry meaningless or miscalibrated probabilities.

So Lesson 05 asks the next forced question: when several futures are genuinely valid, how do you represent their separate modes, tell aleatoric ambiguity from model ignorance, and test that the probabilities are honest? We move from having a stochastic transition to making its uncertainty worth acting on.

Interview prompts

• Why does an RSSM contain both h_t and z_t? h deterministically compresses ordered history; z represents locally uncertain or newly revealed state, allowing multiple plausible futures.
• What is the difference between the prior and posterior? The prior predicts latent state from past memory and is available during imagination. The posterior also sees the current observation and corrects that prediction during filtering and training.
• What does KL(q||p) do operationally? It penalizes posterior states that the imagination prior assigns low probability, teaching the prior to anticipate observation-informed inference.
• What happens when the KL weight is too high or too low? Too high can collapse the posterior so it ignores evidence; too low lets the posterior create states the prior cannot reach, producing good reconstruction but bad rollouts.
• Why must the current observation not enter the prior? A planner will not possess future observations. Letting them leak into the prior creates a training path unavailable at deployment.
• How would you detect a weak transition hidden by a strong decoder? Compare posterior-conditioned prediction with prior-only one-step and open-loop rollouts; the gap exposes dependence on correction.
• When would tokens be preferable to a continuous RSSM? When discrete branching, transformer scaling, or stable categorical likelihoods matter more than compact, very fast imagined control steps.
• How do auxiliary heads alter the state? Their gradients decide which information survives compression; reward heads preserve consequences, geometry heads preserve spatial structure, and pixel heads preserve appearance.