Learning Latent Dynamics

The agent squeezes camera frames into a small code and imagines futures in that code instead of predicting pixels, so planning runs thousands of rollouts per second.

RSSM splits state into deterministic history and stochastic belief so imagination rolls forward with KL 0.1681 and reward loss 0.02 and no pixel decoding

Why Does This Exist?

A visual agent sees $64 \times 64 \times 3$ frames, about 12,000 numbers per step. Predicting the next full frame with a deconvolutional decoder at every planning step is slow, and most pixels are distractors like clouds or reflections. Small blur compounds too, so a 50 step pixel rollout melts into gray static.

Model-based RL plans with a learned model instead of the real world. This page covers the compact version of that idea: learn transitions between small latent codes, train with a variational bottleneck, then imagine in latent space only. Dreamer is the full agent that uses this trick to learn behavior inside its own imagination. Planning with learned models shows how those imagined rollouts turn into actions.

The one idea you need first

A latent state is a compressed code $z$ that keeps what matters for control and drops the rest. An encoder maps observation $x$ to code, a dynamics model maps $(z, a)$ to next code $z'$ , and a small head predicts reward $r$ . The full treatment of planning with such a model is in model-based RL.

Think of It Like This

Chess in your head, not photos of the board

You play a blindfold chess game. You do not picture wood grain, lamp glare, or your opponent hands. You track a short list: piece squares, whose turn, king safety, open files. You try a line five moves deep inside that list and pick the move that looks best.

Latent dynamics works the same way. The frame is the photo. The latent code is the piece list. The prior is your guess of the next position without looking. The posterior is the updated list after you glance at the board. Imagination is playing lines without touching pieces.

Where it stops: you arrive with a lifetime of knowing what a knight is. The agent starts with zero priors, so a naive compressor can throw away a tiny but vital cue before it learns its value.

How It Actually Works

PlaNet and Dreamer package this as the Recurrent State Space Model, or RSSM. Call $h_t$ the deterministic history kept by a recurrent cell, $s_t$ the stochastic belief about the present, $x_t$ the frame, $a_t$ the action, $e_t$ the encoded frame. The posterior $q$ sees the frame. The prior $p$ guesses without it.

$h_t = f(h_{t-1}, s_{t-1}, a_{t-1})$

$q(s_t \mid h_t, x_t) = \mathcal{N}(\mu_q, \sigma_q^2)$

$p(s_t \mid h_t) = \mathcal{N}(\mu_p, \sigma_p^2)$

In words: $h_t$ carries memory, $s_t$ carries uncertainty, $q$ is belief with evidence, $p$ is belief from memory alone. Training pushes $p$ toward $q$ so imagination matches reality.

The training objective

Parameters fit the variational lower bound over horizon $T$ :

$\mathcal{L} = \mathbb{E}\left[\sum_{t=1}^T \left(\log p(x_t \mid h_t, s_t) + \log p(r_t \mid h_t, s_t) - D_{KL}(q \parallel p)\right)\right]$

The first two terms force the code to reconstruct frames and rewards. The $D_{KL}$ term is the bottleneck that keeps imagination honest. Once trained, the policy $\pi(a \mid h, s)$ rolls out with $p$ only, predicts rewards $\hat{r}$ and values $\hat{V}$ , and never decodes a pixel.

Worked example

Take $h_t = 0.8$ . Prior without observation:

$p(s_t \mid h_t) = \mathcal{N}(1.0, 1.0)$

A frame arrives and the posterior becomes:

$q(s_t \mid h_t, x_t) = \mathcal{N}(1.5, 0.64)$

Reward head predicts $\hat{r}_t = 1.0 \times 0.8 + 0.8 \times 1.5 = 2.0$ against true $r_t = 1.8$ .

KL between univariate Gaussians:

$D_{KL}(q \parallel p) = \frac{1}{2}\left[\frac{0.64}{1.0} + \frac{(0.5)^2}{1.0} - 1 + \ln\left(\frac{1.0}{0.64}\right)\right]$

That is $0.5 \times (0.64 + 0.25 - 1.0 + 0.4463) = 0.5 \times 0.3363 = 0.1681$ .

Reward loss with unit variance:

$\mathcal{L}_{reward} = \frac{1}{2}(2.0 - 1.8)^2 = 0.0200$

Total to minimize is $0.0200 + 0.1681 = 0.1881$ , so ELBO is $-0.1881$ . The diagram shows the same split: posterior $1.5$ , prior $1.0$ , KL $0.1681$ , reward $2.0$ versus $1.8$ .

Code

This check computes the exact KL and reward terms from the worked example. It uses only the standard library.

1import math2
3
4def kl_1d(mu_q: float, var_q: float, mu_p: float, var_p: float) -> float:5    """Closed form KL between univariate Gaussians q and p."""6    return 0.5 * (7        var_q / var_p8        + ((mu_q - mu_p) ** 2) / var_p9        - 1.010        + math.log(var_p / var_q)11    )12
13
14def reward_pred(h: float, s: float) -> float:15    """Reward head used in the worked example."""16    return 1.0 * h + 0.8 * s17
18
19if __name__ == "__main__":20    kl = kl_1d(1.5, 0.64, 1.0, 1.0)21    rhat = reward_pred(0.8, 1.5)22    rloss = 0.5 * (rhat - 1.8) ** 223    print(f"KL = {kl:.4f}")  # -> KL = 0.168124    print(f"rhat = {rhat:.2f}")  # -> rhat = 2.0025    print(f"reward loss = {rloss:.4f}")  # -> reward loss = 0.020026    print(f"ELBO = {-(rloss + kl):.4f}")  # -> ELBO = -0.1881

Watch Out For

Compressor drops the cue that matters later

Training on reward alone teaches the encoder to keep only what predicts imminent reward. In sparse tasks a distant hazard gives zero reward for hundreds of steps, so it gets filed as background noise and imagined paths drive straight through it. Symptom: great training curves, collisions in deployment.

Keep an auxiliary image or contrastive loss so the code preserves geometry even when reward is silent. Probe reconstructions for small hazards before you trust the model. If a brake light or pit vanishes in reconstruction, your bottleneck is too tight.

Prior never catches the posterior

If the KL weight is too high early, the posterior collapses to a dull prior and stochastic states carry no information. If it is too low, imagination drifts because the prior never learned to match informed belief. Symptom: KL near zero with blurry predictions, or small training KL but huge imagination error.

Use KL balancing with free bits, meaning a minimum KL per dimension that stops the penalty from crushing useful uncertainty. Hafner style setups keep most gradient on the posterior side until clusters form, then let the prior catch up.

The Quick Version

Learning latent dynamics compresses frames into a small code and learns transitions between codes, so planning never renders pixels.
RSSM splits state into deterministic memory $h_t$ from a recurrent cell and stochastic belief $s_t$ for uncertainty.
The posterior $q(s_t \mid h_t, x_t)$ uses the frame while the prior $p(s_t \mid h_t)$ imagines without it, and training pulls them together.
In the worked 1D check the KL is $0.1681$ , reward loss is $0.0200$ , and ELBO is $-0.1881$ .
Keep a reconstruction loss and free bits, or the code drops rare but vital cues and imagination drifts.