World Models, SimPLe, and PlaNet › Variational latent dynamics

A latent dynamics model introduces an unobserved or inferred state and models its action-conditioned evolution. World Models and PlaNet aim to make their learned state more compact than pixels or easier to use for prediction and control, but neither property is guaranteed merely by introducing a latent variable. Formally, we imagine a latent state $z_t$, an observation $x_t$, an action $a_t$, and a reward $r_t$, and we model:

$$z_{t} \sim p_\theta(z_t \mid z_{t-1}, a_{t-1}), \qquad x_t \sim p_\theta(x_t \mid z_t), \qquad r_t \sim p_\theta(r_t \mid z_t).$$

where:

• $z_t$: the latent state at time $t$, evolved by the transition model and, in planning-based agents, rolled forward to evaluate candidate actions
• $a_{t-1}$: the action taken at the previous step, which conditions how the latent changes
• $p_\theta(z_t \mid z_{t-1}, a_{t-1})$: the transition (prior) model, parameterized by $\theta$, giving the distribution over the next latent given the current latent and action
• $x_t$: the observation at time $t$ (for example, an Atari frame)
• $p_\theta(x_t \mid z_t)$: the observation (decoder) model, rendering an observation from the latent
• $r_t$: the reward at time $t$
• $p_\theta(r_t \mid z_t)$: the reward model, predicting reward from the latent

The transition lets latent states evolve, the observation model renders observations, and the reward model predicts per-step rewards that accumulate into the return used for control. If we could reliably infer $z_t$ from the observation history, predict its transitions over the planning horizon, and estimate rewards accurately enough to compare candidate returns, we could plan entirely in latent space: roll the latent forward under candidate actions, evaluate predicted rewards, and pick the best action sequence. The obstacle is that we never observe $z_t$. We need an inference procedure. In the end-to-end variational formulation developed here, the observation model and inference network are optimized jointly; staged or pretrained alternatives are also possible.

It is worth pausing on why this factorization can be a reasonable modeling assumption. In image-based control, many pixel configurations can share a control-relevant situation: appearance and background variation may add visual complexity without changing the underlying dynamics. A latent can compress some of that redundancy and expose dynamics that are easier to model. In an Atari game, visible sprite positions and velocities are useful factors, but a sufficient emulator state may also contain timers, mode flags, lives, off-screen objects, and pseudorandom state. World Models and PlaNet aim to retain variables needed for prediction and control in a representation more compact than the image; this does not imply that every image-based method here, including SimPLe, feeds a compact latent to its decision mechanism or that the hidden state is always a handful of coordinates.

This is the setting of a variational autoencoder extended over time. We introduced latent-variable representations and temporal state in Part III. Here we reconnect those ideas to dynamics and control. For nonlinear transition and decoder networks, the exact sequence posterior is generally intractable. Let $T$ be the sequence length, let $h_t$ be the deterministic state computed from the sampled latent and action history before observing $x_t$, and let $q_\phi$ be the filtering variational posterior with inference-network parameters $\phi$. The approximate sequence posterior then factorizes as

$$q_\phi(z_{1:T} \mid x_{1:T}, a_{1:T-1}) = \prod_{t=1}^{T} q_\phi(z_t \mid h_t, x_t).$$

Assume the generative joint factorizes as $\prod_{t=1}^{T}p_\theta(z_t\mid h_t)p_\theta(x_t\mid h_t,z_t)p_\theta(r_t\mid h_t,z_t)$, so observation and reward are conditionally independent given $(h_t,z_t)$, with the deterministic recursion $h_t=f_\theta(h_{t-1},z_{t-1},a_{t-1})$ and fixed initial conditions. Under these assumptions, $\mathcal{L}(\theta,\phi)$ below is a sequential ELBO: it lower-bounds the conditional log evidence $\log p_\theta(x_{1:T},r_{1:T}\mid a_{1:T-1})$.

$$\mathcal{L}(\theta, \phi) = \mathbb{E}_{q_\phi(z_{1:T} \mid x_{1:T}, a_{1:T-1})} \left[ \sum_{t=1}^{T} \left( \log p_\theta(x_t \mid h_t,z_t) + \log p_\theta(r_t \mid h_t,z_t) + \log p_\theta(z_t \mid h_t) - \log q_\phi(z_t \mid h_t,x_t) \right) \right].$$

The outer expectation is over the whole sequential latent trajectory, so it also averages over the random histories that determine later priors. The two likelihood terms encourage the inferred state to retain information useful for predicting observations and rewards. They do not guarantee interpretable factors or prevent collapse. The log-prior minus log-posterior terms are the negative KL contributions. They encourage posterior mass, on average, to remain in regions assigned substantial density by the transition prior. Without that agreement, prior outputs can drift away from the states seen by the decoder and reward head, making multi-step prior rollouts progressively unreliable.

There is a well-known tradeoff hiding in the KL term. If we weight reconstruction heavily, the model may allocate capacity according to pixel likelihood rather than task relevance, for example to background animation or texture. If we weight the prior too heavily, the latent may carry too little information and both observation and reward prediction can degrade. The β-VAE objective also changes the KL coefficient, and Burgess et al. analyze its capacity–reconstruction tradeoff from a rate–distortion perspective; the same coefficient-level reasoning applies to this sequential objective. PlaNet proposed latent overshooting as an additional multi-step regularizer: from an earlier posterior state, the transition is rolled forward for several steps and its multi-step predictive prior is matched by KL to later posterior beliefs. The released implementation treats those posterior targets as stopped-gradient references by default. Those extra terms encourage consistency at several prediction distances. The paper's reported final RSSM used three free nats and did not require overshooting; Appendix D reports a slight RSSM performance decrease when overshooting was added. The released configuration likewise defaults to three free nats and an overshooting loss scale of zero. These two settings agree between the final paper agent and released defaults, even though the paper's proposed overshooting objective was evaluated separately.

To make the tradeoff concrete, think of the KL term as a communication-channel budget. The posterior can carry information about the current observation beyond what the prior predicts, though it may also ignore that observation. Averaged over data, its KL to the transition prior upper-bounds this conditional information and also penalizes mismatch between the aggregate posterior and the prior. With natural logarithms, the KL is an expected log-density excess in nats; under a common discretization or an ideal coding construction, this corresponds to expected excess code length. Divide by $\ln 2$ to express nats in bits. Too little regularization can let the posterior encode detail that prior-only rollouts cannot reproduce. Too much can leave the state too vague for observation or reward prediction. The useful operating point depends on the task and on the relative loss weights.

A second design decision is how $h_t$, the history summary, is formed. A feedforward latent derived from one frame is Markov only if that frame is predictively sufficient for action-conditioned future observations and rewards. This can fail when histories with the same current frame imply different futures—for example, if persistent velocity or hidden object state can be inferred from earlier frames. Chronologically, World Models came first, PlaNet followed, and SimPLe came later. Architecturally, they offered different answers: World Models encodes each frame and adds separate recurrent memory; PlaNet integrates deterministic recurrence with stochastic state in an RSSM; SimPLe predicts in observation space from four-frame inputs, while its selected video model also maintains an internal recurrent spatial state and a separate latent-bit LSTM (released configuration; video-model implementation). The SimPLe policy consumes frame stacks, not those internal model states. These papers belong to overlapping observation-space and latent-state research lines rather than a single progression. RSSM-style recurrent latent models became one influential modern lineage, alongside transformer-based observation-token models and other alternatives.

The single-frame Markov issue deserves emphasis because it can silently merge histories that share a frame but imply different future distributions. Consider a game in which two visually identical situations have different futures because one has an enemy moving toward the agent and the other has an enemy moving away. A model that sees only the current frame cannot distinguish them. If both histories have positive conditional probability under that frame and the same action, the true next-step predictive distribution mixes their possible outcomes; a misspecified model need not cover both. Recurrence can mitigate the ambiguity when the history contains enough identifying information: the recurrent state can accumulate evidence about which situation the agent is in, and the transition can condition on that summary. It does not solve partial observability by itself. Finite capacity, aliased histories, and imperfect optimization can still leave different situations indistinguishable. This resembles belief-state methods in partially observable control; PlaNet's RSSM supplies a learned approximate state summary rather than an exact belief.

PlaNet's RSSM is a variational latent dynamics architecture that Dreamer later reused, so we describe it carefully. The RSSM splits the latent state into two parts. First, the deterministic recurrent state $h_t$ summarizes the past. It follows the recurrence

$$h_t = f_\theta(h_{t-1}, z_{t-1}, a_{t-1}),$$

where:

• $h_t$: the deterministic recurrent state at time $t$, carrying memory of the past
• $f_\theta$: the recurrent network (for example, a GRU) parameterized by $\theta$
• $h_{t-1}$: the previous recurrent state
• $z_{t-1}$: the previous stochastic latent state
• $a_{t-1}$: the action taken at the previous step

This component is deterministic and carries memory. Second, the stochastic state $z_t$ is sampled from a diagonal Gaussian. It can represent per-step conditional variation, but the architecture does not require it to contain only information absent from $h_t$.

Why split the latent at all, rather than using a single stochastic latent? The two paths provide different inductive biases. The finite recurrent state is a fixed-size learned summary of history that can retain useful temporal information without resampling every component at every step. It may discard task-relevant information when capacity or training is inadequate; for some tasks, a finite state can instead be sufficient. The stochastic state represents conditional variation through a distribution. A diagonal Gaussian supplies a mean and variance, but its presence does not by itself imply calibrated uncertainty, and a single conditional Gaussian latent density is unimodal. If calibrated predictive uncertainty is claimed, the induced predictive distribution should be evaluated against held-out outcomes under a stated calibration target and metric. Richer one-step multimodality in the latent transition requires a richer latent distribution; a nonlinear decoder can nevertheless induce a multimodal observation distribution, and multi-step propagation can produce complicated trajectory distributions.

The transition prior, conditioned on the history, is:

$$p_\theta(z_t \mid h_t) = \mathcal{N}\!\big(\mu_\theta^p(h_t),\, \operatorname{diag}(\sigma_\theta^p(h_t)^2)\big),$$

where:

• $h_t$: the deterministic recurrent state summarizing the preceding inferred history and action $a_{t-1}$, before the current observation $x_t$ is incorporated
• $\mu_\theta^p(h_t)$: the prior mean, a function of the history produced by the transition network
• $\sigma_\theta^p(h_t)^2$: the vector of coordinatewise prior variances, also a function of the history
• $\mathcal{N}(\cdot, \cdot)$: a multivariate Gaussian specified by its mean and covariance matrix, diagonal here with coordinatewise variances

The approximate posterior, used during training and whenever we want to incorporate a new observation, is:

$$q_\phi(z_t \mid h_t, x_t) = \mathcal{N}\!\big(\mu_\phi^q(h_t, x_t),\, \operatorname{diag}(\sigma_\phi^q(h_t, x_t)^2)\big),$$

where:

• $x_t$: the current observation, the extra input the posterior receives
• $\mu_\phi^q(h_t, x_t)$: the posterior mean, a function of both the history and the current observation
• $\sigma_\phi^q(h_t, x_t)^2$: the vector of coordinatewise posterior variances

The posterior sees the current observation, the prior does not. During imagination, we sample $z_t$ from the prior; during inference from real observations we may sample from the posterior to stay grounded. The KL term now compares two Gaussians over the stochastic state, and because both are diagonal, it has a closed form:

$$\begin{aligned} D_{\mathrm{KL}}\!\big( q \,\|\, p \big) &= \frac{1}{2} \sum_{k=1}^{K} \Big( \frac{\sigma_{q,k}^2 + (\mu_{q,k} - \mu_{p,k})^2}{\sigma_{p,k}^2} - 1 + \log \sigma_{p,k}^2 - \log \sigma_{q,k}^2 \Big) \end{aligned}$$

where:

• $q$: the approximate posterior Gaussian over the stochastic state (mean $\mu_q$, variance $\sigma_q^2$)
• $p$: the transition prior Gaussian over the stochastic state (mean $\mu_p$, variance $\sigma_p^2$)
• $K$: the dimensionality of the stochastic latent $z_t$
• $\mu_{q,k}, \sigma_{q,k}^2$: the posterior mean and variance of the $k$-th latent dimension
• $\mu_{p,k}, \sigma_{p,k}^2$: the prior mean and variance of the $k$-th latent dimension

This expression measures, dimension by dimension, how much the posterior differs from the prior. Holding the positive prior variance $\sigma_{p,k}^2$ fixed, the mean-mismatch contribution grows quadratically with $|\mu_{q,k}-\mu_{p,k}|$. The variance ratio and logarithm must be read together: their contribution is minimized when the posterior and prior variances match and increases when the posterior variance is either larger or smaller. Because both distributions are diagonal, the total separates into per-dimension contributions.

The closed form matters practically as well as conceptually. The analytic Gaussian KL is differentiable and cheap to compute, but that does not imply a gradient everywhere in a thresholded training objective. The released PlaNet objective subtracts three free nats and clips the divergence penalty at zero. Where the standard one-step penalty has a nonzero gradient, it trains the prior toward the posterior and regularizes the posterior toward the prior; below the free-nats threshold, it supplies no such gradient. Optional multi-step overshooting instead uses stopped-gradient posterior targets and is disabled by default in the released configuration. Active KL pressure can encourage prior-only rollouts (which planning uses) to remain closer to posterior-inferred latents seen during training, but it does not prevent multi-step divergence.

Two structural advantages follow. First, because $h_t$ is deterministic and recurrent, the model need not be Markov in the current frame alone: it can retain evidence about velocity or partially observed variables when the history and learned capacity are sufficient. Second, because $z_t$ is stochastic, the model can represent conditional variation rather than being restricted to one deterministic next state. That is not the same as representing arbitrary multimodality or calibrated epistemic uncertainty. The combination of a deterministic memory path and a stochastic path recurs in later RSSM-based world models, including the Dreamer family.

In PlaNet, the observation and reward heads consume the combined state $[h_t,z_t]$. The paper evaluated six pixel-based continuous-control tasks from the DeepMind Control Suite, using $64\times64$ RGB observations and scalar rewards (Hafner et al., 2019). The observation decoder is not needed inside candidate rollouts at control time: CEM advances the prior and queries predicted rewards without rendering future frames. The encoder and posterior incorporate the retained real observation at each decision step, after the task's configured action-repeat interval.

It is worth stating explicitly what these components do. The recurrent state $h_t$ summarizes the preceding inferred history through $z_{t-1}$ and $a_{t-1}$, and parameterizes the time-$t$ prior $p_\theta(z_t\mid h_t)$ before $x_t$ is incorporated. The stochastic state $z_t$ supplies the sampled per-step component. PlaNet's observation and reward heads use their combination; CEM consumes the returns produced by rolling that combined state forward rather than passing $z_t$ through a learned controller. The decoded observation supplies a training signal but is not rendered during planning.