Latent Dynamics

The Encoder Is Not Enough: We Need to Predict the Future

The ball-on-table example has run through both previous pages: a single frame under-determines velocity, and an encoder compressing that frame is not guaranteed to keep the velocity it needs. A VAE can compress ${\mathbf{o}}_t$ into ${\mathbf{z}}_t$, but a world model must also use history to recover hidden variables and predict the future:

In latent space, given the current state ${\mathbf{z}}_t$ and action ${\mathbf{a}}_t$, predict the next state ${\mathbf{z}}_{t+1}$.

This prediction capability lets the agent "simulate" the future internally, enabling planning without actually executing actions in the environment. This is the key reason world models reduce sample complexity.

The Simplest Dynamics Model: GRU

The Gated Recurrent Unit (GRU) is a foundational tool for sequence modeling. As a dynamics model, the GRU takes $({\mathbf{z}}_t,{\mathbf{a}}_t)$ and predicts the next latent state:

$${\mathbf{z}}_{t+1}=\text{GRU}({\mathbf{z}}_t,{\mathbf{a}}_t.\theta )$$

📖 GRU internals (brief): A GRU controls information flow through two gates: the reset gate decides "how much of the past to forget", and the update gate decides "how much of the old state to retain vs. how much new information to write in". Gate values lie between 0 and 1, determined jointly by the current input and the previous hidden state. This allows the GRU to selectively retain long-term dependencies while discarding irrelevant information, making it better at handling longer sequences than a plain RNN. Compared to an LSTM, the GRU has one fewer gate (no separate memory cell), fewer parameters, and trains faster.

The GRU's strengths are simplicity and stable training. Its limitation is that it produces deterministic predictions and cannot express uncertainty. In real environments, the same action can lead to multiple different outcomes: pushing the ball from the earlier pages might send it rolling cleanly, or it might strike something just outside the frame and stop short, and a single deterministic prediction cannot represent both.

MDN-RNN: Modeling Uncertainty

MDN-RNN (Mixture Density Network + RNN), proposed in Ha & Schmidhuber (2018), models uncertainty over the next state using a mixture of Gaussians:

$$p({\mathbf{z}}_{t+1}|{\mathbf{z}}_t,{\mathbf{a}}_t)=\sum _{k=1}^K{\pi }_k\cdot \mathcal{N}({\mathbf{z}}_{t+1}.{\mu }_k,{\sigma }_k^2)$$

• $K$ Gaussian components, each with its own mean ${\mu }_k$ (center of the distribution) and variance ${\sigma }_k^2$ (width of the distribution)
• Mixture weights ${\pi }_k$: the probability mass of the $k$-th Gaussian component, satisfying $\sum _{k=1}^K{\pi }_k=1$, ${\pi }_k\ge 0$. These can be read as "the probability that the $k$-th future occurs". The network outputs ${\pi }_k$ values and normalizes them through softmax to ensure the weights sum to 1.

MDN-RNN can capture multimodal distributions: the environment may transition to several distinct next states, and the model can represent all of them.

MDN-RNN architecture from Ha & Schmidhuber (2018): the RNN hidden state is passed through a fully connected layer to produce K parameter groups (π_k, μ_k, σ_k), representing mixture weights, means, and variances, which together define the Gaussian mixture distribution over the next latent state.

RSSM: Separating Deterministic and Stochastic Components

The RSSM (Recurrent State Space Model) was introduced in PlaNet and later became the core state model of the Dreamer series. It splits the state into two parts:

• Deterministic hidden state ${\mathbf{h}}_t$: maintained by an RNN, aggregating information from the historical trajectory, with no stochasticity
• Stochastic latent state ${\mathbf{z}}_t$: sampled from a distribution conditioned on ${\mathbf{h}}_t$, expressing current uncertainty

Core equations of the RSSM:

📖 Subscript $\varphi$ (phi): the subscript $\varphi$ in $f_{\varphi }$, $p_{\varphi }$, $q_{\varphi }$ denotes "this function has parameters $\varphi$", i.e., the learnable weights of the neural network. $f_{\varphi }(\cdot )$ is read as "function $f$ parameterized by $\varphi$". During training, gradient descent updates $\varphi$ so that the predictions of these functions become increasingly accurate. Similarly, $\theta$ (theta) that appears later is another commonly used symbol for a distinct set of learnable parameters.

$${\mathbf{h}}_t=f_{\varphi }({\mathbf{h}}_{t-1},{\mathbf{z}}_{t-1},{\mathbf{a}}_{t-1})\text{(deterministic update, GRU/RNN)}$$$${\mathbf{z}}_t\sim p_{\varphi }({\mathbf{z}}_t\mid {\mathbf{h}}_t)\text{(prior: no access to real observations, infers current state from history }{h}_t\text{ alone. Used for pure imagination/prediction)}$$$${\mathbf{z}}_t\sim q_{\varphi }({\mathbf{z}}_t\mid {\mathbf{h}}_t,{\mathbf{o}}_t)\text{(posterior: corrects the prior using real observation }{o}_t\text{. Used during training)}$$

📖 Prior vs. posterior: these are fundamental concepts in Bayesian statistics. The prior is "belief before seeing data", the RSSM's guess about the current state $z_t$ based on historical memory $h_t$. The posterior is "belief updated after seeing data", refining the prior with real observation $o_t$ to obtain a more accurate estimate. During training, the posterior generates $z_t$ and the KL loss is computed (measuring the gap between prior and posterior). During inference and imagination, only the prior is available (there is no real $o_t$), so the RSSM rolls forward using the prior alone.

Why separate them?

State	Role	Property
${\mathbf{h}}_t$	Memory	Deterministic, aggregates history
${\mathbf{z}}_t$	Perception	Stochastic, expresses uncertainty

After separation, the model can roll forward using only the prior $p({\mathbf{z}}_t|{\mathbf{h}}_t)$ without real observations, enabling planning purely in imagination. This is the fundamental reason for Dreamer's sample efficiency.

The PlaNet paper (Hafner et al., ICML 2019) verified this design through ablation studies (systematically removing one component of the model and observing the change in performance, thereby confirming the component's necessity): a purely stochastic path (no deterministic $h_t$) struggles to reliably retain information across multiple steps, and training optimization may fail to find solutions where some dimensions collapse to near-zero variance to store long-term information. A purely deterministic path (no stochastic $z_t$) cannot express the inherent stochasticity of the environment, and the distribution gap between imagined and real trajectories grows larger. Both paths are indispensable. The observation model is therefore conditioned on both $h_t$ and $z_t$: $o_t\sim p(o_t|h_t,z_t)$, with deterministic memory and stochastic perception jointly determining the reconstructed image.

Worked example (simplified single-step transition): the real $f_{\varphi }$ is a gated GRU, which is not convenient to compute by hand, so this uses a simplified version with the gates removed, keeping only a linear combination, to demonstrate the same mechanism. Suppose ${\mathbf{h}}_{t-1}=[0,0]$ (start of the sequence, hidden state initialized to zero), $z_{t-1}=0.50$, $a_{t-1}=1.0$, and the deterministic update uses weights $w_z=[0.4,0.2]$, $w_a=[0.3,0.1]$, bias $b=[0.1,0]$:

$${\mathbf{h}}_t=w_z\cdot z_{t-1}+w_a\cdot a_{t-1}+b=[0.4,0.2]\times 0.50+[0.3,0.1]\times 1.0+[0.1,0]=[0.60,0.20]$$

Using the prior network (also simplified to a linear layer) with $w_{\mu }=[0.5,0.5]$ and a fixed ${\sigma }_{\text{pr}}=0.20$, compute the prior's mean:

$${\mu }_{\text{pr}}=w_{\mu }\cdot {\mathbf{h}}_t=0.5\times 0.60+0.5\times 0.20=0.40$$

Finally, sample $z_t$ using the reparameterization trick (sampling noise $\epsilon =0.50$):

$$z_t={\mu }_{\text{pr}}+{\sigma }_{\text{pr}}\cdot \epsilon =0.40+0.20\times 0.50=0.50$$

This chain shows the complete data flow of one RSSM transition step: $({\mathbf{h}}_{t-1},z_{t-1},a_{t-1})\to {\mathbf{h}}_t\to {\mu }_{\text{pr}}\to z_t$. In the real code, $f_{\varphi }$ is PyTorch's GRUCell (with reset and update gates, see the GRU callout above), and the prior network is more than a single linear layer, so the actual numbers will differ, but every step the data flows through is exactly the same as here. P02 runs the same numbers through this simplified chain to verify the formula and the code agree.

Comparison of Three Dynamics Models

Model	Uncertainty Modeling	Memory Mechanism	Primary Use
GRU	None (deterministic output)	Fixed-dimension hidden state $h_t$	Simple sequence prediction, rapid prototyping
MDN-RNN	Mixture of Gaussians (multimodal)	Fixed-dimension hidden state $h_t$	Multimodal uncertainty, Ha & Schmidhuber M-module
RSSM	Separated prior/posterior (Gaussian)	Dual-track: deterministic $h_t$ + stochastic $z_t$	Core of Dreamer, supports pure-imagination planning

The three form a progression: GRU establishes the foundation for sequence modeling, MDN-RNN introduces uncertainty, and RSSM further decouples "memory" from "perceptual uncertainty", enabling the model to roll forward and plan without real observations.

Next, read Training Distributions and Free Rollouts to examine what happens when these models repeatedly consume their own predictions. Doing so before P02 makes the distinction between one-step error and horizon drift concrete.