Recurrent State-Space Models › The Recurrent State-Space Model

This section assembles the architecture, name by name, and explains the wiring that makes the posterior and prior cooperate. Read it slowly: every design choice here has a reason, and the reasons compound.

An RSSM for a partially observable environment with observations $o_t$ and actions $a_t$ consists of:

• Deterministic recurrence (the memory path). A GRU (or LSTM) updates $h_t$ from the previous deterministic state, the previous sampled latent, and the previous action:

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

• Prior (transition model). Before seeing anything at time $t$, the model guesses the stochastic state:

  $$z_t \sim p_\theta(z_t \mid h_t) = \mathcal{N}\big(\mu_\theta(h_t),\ \mathrm{diag}(\sigma^2_\theta(h_t))\big)$$

• Posterior (inference model). During training and filtering, an encoder sees the current observation and re-places $z_t$ where the evidence says it should be:

  $$z_t \sim q_\phi(z_t \mid h_t, o_t) = \mathcal{N}\big(\mu_\phi(h_t, o_t),\ \mathrm{diag}(\sigma^2_\phi(h_t, o_t))\big)$$

• Observation model (decoder). Reconstructs or predicts the observation from the joint state:

  $$o_t \sim p_\theta(o_t \mid h_t, z_t)$$

• Optional heads. In the Dreamer line, a reward head $p_\theta(r_t \mid h_t, z_t)$ and a discount or continuation head for episode boundaries (the original Dreamer already modeled early termination). We will use them in Part VII: Planning and Agency; here we focus on the observation model.

where:

• $f_\theta$: the GRU or LSTM update, a neural network with parameters $\theta$ applied identically at every step
• $\mu_\theta(h_t)$, $\sigma^2_\theta(h_t)$: the prior's mean vector and per-dimension variances, produced from $h_t$ by a small network head
• $\mu_\phi(h_t, o_t)$, $\sigma^2_\phi(h_t, o_t)$: the posterior's mean vector and per-dimension variances, produced from the pair $(h_t, o_t)$ by the encoder with parameters $\phi$
• $\mathrm{diag}(\cdot)$: a covariance matrix carrying the given variance vector on its diagonal and zeros elsewhere, so the Gaussian treats the coordinates as independent conditional on its inputs
• $r_t$: the reward at time $t$

Notice the timing convention, because it trips everyone up at least once. The GRU input at time $t$ is $(z_{t-1}, a_{t-1})$, not $z_t$. So $h_t$ is the model's prediction of the situation at time $t$ before the observation $o_t$ arrives. The observation never touches $h_t$ directly; it enters the model only through the posterior, which shapes $z_t$, which enters the recurrence at the next step. All per-step information about the current observation must therefore flow through the stochastic latent. This is deliberate, and it is what makes the KL term meaningful, as you will see below. If the observation could leak into $h_t$ by some side channel, the prior itself could use the observation it is supposed to predict, and the KL would no longer measure a real prediction gap. By forcing all observation information through the sampled bottleneck $z_t$, the architecture guarantees that the only way the model can use the present is by encoding it in exactly the place the prior is trying to predict.

The figure below shows the wiring unrolled over three steps. Read the middle column: the prior flows down from $h_t$ into $z_t$; the posterior flows up from $o_t$ into $z_t$; the decoder emits a predicted observation from $(h_t, z_t)$; and the GRU carries $(h_{t-1}, z_{t-1}, a_{t-1})$ into $h_t$. The third column is drawn faded deliberately: it is imagination, where the posterior arrows do not exist and the prior must carry the sequence on its own.

The division of labor between $h$ and $z$ is the design insight that makes the architecture work, so it deserves a careful statement.

The deterministic path provides a direct recurrent route for information and gradients. With differentiable gates, backpropagation through a GRU is an ordinary chain-rule calculation conditional on its inputs. This can help preserve information that would be difficult to carry through repeated stochastic transitions. It does not guarantee small gradient variance or stable optimization: the inputs still include sampled latents, and the recurrent Jacobians can attenuate or amplify gradients. PlaNet's ablation is evidence about its tested architecture and tasks, rather than a proof that stochastic-only models cannot learn.

The stochastic path supplies sampled latent variation. In the graph defined here, $o_t$ enters the joint state through the posterior distribution for $z_t$, then influences later memory through the recurrence. The KL compares distributions, not two sampled values: even identical distributions can generate different samples, and an observation can change posterior variance without changing its mean. Uncertainty can also reside in distribution parameters carried by memory, in the observation model, and in uncertainty about fitted parameters. No single latent draw contains all of those uncertainties.

The Kalman analogy concerns the order of computation: predict a state distribution, then update it using an observation. A Kalman filter represents uncertainty with both a mean and a covariance, and its innovation is an observation residual, not a KL divergence. Under the linear-Gaussian assumptions its update is exact; an RSSM uses learned nonlinear dynamics and approximate inference. Comparing filtering with imagination can help locate a failure, but does not isolate one guilty component: the encoder, transition, decoder, and available information all influence the observed errors.

DreamerV2 (Hafner et al., 2021) uses 32 categorical variables with 32 classes each, sampled with straight-through gradients; DreamerV3 retains categorical stochastic latents. Categoricals provide discrete alternatives in latent coordinates. A diagonal Gaussian is unimodal conditional on its context in those coordinates, but a nonlinear decoder can still turn Gaussian latents into multimodal or correlated observations. The high-level prior/posterior/recurrence graph remains recognizable across the variants, while the distribution parameterization, KL calculation, and gradient estimator change. Straight-through differentiation is a biased surrogate; its practical usefulness does not imply that score-function estimators inevitably have exploding variance.