Recurrent State-Space Models › Learning the Model: Posterior Inference and the ELBO

You now know what the parts are. This section explains how they are trained, and why the training objective has the peculiar two-part structure it does.

Consider a single training sequence of observations $o_{1:T}$ and actions $a_{1:T}$. The generative story is:

1. Start from an initial belief ($h_1$ can be a learned vector or fixed zeros, with the initial latent prior conditioned on it).
2. For each $t$: draw $z_t$ from the prior $p_\theta(z_t \mid h_t)$, emit $o_t$ from $p_\theta(o_t \mid h_t, z_t)$, then update $h_{t+1} = f_\theta(h_t, z_t, a_t)$.

The action-conditioned likelihood $p_\theta(o_{1:T}\mid a_{1:T-1})$ integrates over latent trajectories. For a general nonlinear neural model, that sum or integral is intractable; linear-Gaussian models are an important tractable exception. We use the filtering variational factorization $q(z_{1:T}\mid o_{1:T},a_{1:T-1})=\prod_t q_\phi(z_t\mid h_t,o_t)$, where each $h_t$ depends on earlier sampled latents. Jensen's inequality gives an evidence lower bound on the log likelihood. Written with all expectations visible, it is

$$\mathcal{L} = \sum_{t=1}^{T}\mathbb{E}_{q(z_{<t})}\!\left[ \underbrace{\mathbb{E}_{q_\phi(z_t\mid h_t,o_t)}[\log p_\theta(o_t\mid h_t,z_t)]}_{\text{reconstruction}} -\underbrace{\mathrm{KL}\!\left(q_\phi(z_t\mid h_t,o_t)\,\|\,p_\theta(z_t\mid h_t)\right)}_{\text{prior-posterior regularization}} \right]$$

The outer expectation averages over the earlier posterior draws that determine $h_t$; at $t=1$ there are no earlier latents. The inner expectation averages the current reconstruction likelihood. We estimate these expectations with a sampled trajectory, drawing $z_t=\mu_\phi+\sigma_\phi\odot\epsilon_t$, where $\epsilon_t\sim\mathcal{N}(0,I)$ is parameter-independent noise and $\sigma_\phi>0$. This reparameterization exposes a differentiable map from the mean and scale to the sample, so ordinary backpropagation supplies a pathwise Monte Carlo gradient under the usual differentiability and integrability conditions. The expected objective is a lower bound; an individual Monte Carlo estimate is not guaranteed to lie below the log likelihood.

Filtering uses only past and current observations. If a smoothing variational family contains this filtering family, its optimal ELBO cannot be lower, because it can use future observations too. That is a statement about the families' optima, not a promise about two fitted networks. One posterior trajectory is also not a particle filter with importance weighting and resampling. Finally, the decoder log density at a given sample may be analytic while its expectation through a nonlinear decoder is not. The conditional diagonal-Gaussian KL does have the closed form below; its outer expectation over earlier latent history is still estimated by sampling.

$$\mathrm{KL}(q \,\|\, p) = \sum_{i} \left[ \log \frac{\sigma_{p,i}}{\sigma_{q,i}} + \frac{\sigma_{q,i}^2 + (\mu_{q,i} - \mu_{p,i})^2}{2\sigma_{p,i}^2} - \frac{1}{2} \right]$$

where:

• $q$: shorthand for the posterior Gaussian $q_\phi(z_t \mid h_t, o_t)$, with mean $\mu_{q,i}$ and standard deviation $\sigma_{q,i}$ along latent dimension $i$
• $p$: shorthand for the prior Gaussian $p_\theta(z_t \mid h_t)$, with mean $\mu_{p,i}$ and standard deviation $\sigma_{p,i}$
• $i$: an index over the latent dimensions, so the sum runs over the coordinates of $z_t$

The three summands divide the work: the log-scale ratio and normalized variance compare their widths, the normalized squared mean gap penalizes mean disagreement, and $-1/2$ subtracts the posterior's expected standardized quadratic contribution. The terms cancel jointly: when the two Gaussians coincide the log ratio is zero and the variance term plus the constant $-1/2$ sum to zero, so the divergence vanishes exactly when the two Gaussians coincide. This closed form is what the code below implements directly. Notice also the asymmetry of the divergence: because it is the KL with $q$ first, the prior's standard deviation $\sigma_{p,i}$ appears in the denominator, so a prior that is too confident (small $\sigma_p$) pays a steep price when the posterior disagrees with it. For distinct prior-head parameters under this posterior-sampled objective, the $1/\sigma_p^2$ denominator punishes prior overconfidence steeply, and the KL is the only term whose gradient reaches the prior distribution's parameters (the mean and scale head of $p_\theta(z_t \mid h_t)$; the reconstruction gradient still reaches the deterministic transition $f_\theta$ through the sampled $z_t$). Its gradient pulls the posterior toward the prior too. That is the collapse route the stabilizers below are designed to suppress.

The asymmetry is worth seeing as a curve rather than a formula.

Read the vertical gap between the curves: at a mean gap of one, the narrowest prior already pays about 13.7 nats more than the widest one (14.11 versus 0.44) for the identical posterior. The KL direction therefore punishes prior overconfidence, which is exactly the behavior a dynamics-learning signal should have.

The two terms have distinct jobs. Their traces can suggest where to investigate a broken training run, but do not uniquely identify a failing component.

The reconstruction term trains the posterior, decoder, and recurrent transition to make observations predictable. It encourages the posterior to encode current information that helps the decoder, but does not force use of every latent dimension. The current reconstruction has no dependency on $h_{t+1}$: it is later losses that reward retaining useful information there. In this observation-only example, reconstruction is the direct incentive to encode observations; an agent with reward or other auxiliary heads would have additional incentives.

The KL term trains the prior and regularizes the posterior. Either distribution can move to reduce it. Its initial value and subsequent trend depend on both networks, so neither a large starting KL nor a monotonic decline is guaranteed. Residual KL is not sensor-noise variance. In a scalar exact Gaussian update with prior variance $P$ and measurement variance $R$, expected information gain is $\tfrac12\log(1+P/R)$, while innovation variance is $P+R$. Increasing sensor noise can therefore decrease expected information gain while increasing innovation variance. Nats and squared observation units describe different quantities.

Collapse happens when the KL cost of encoding information exceeds the reconstruction benefit of using it. The ELBO literature developed several remedies, and the RSSM line uses two that you should know by name.

Free nats replace the per-step KL by $\max(\mathrm{KL},\epsilon_{\mathrm{free}})$. PlaNet and the original Dreamer use 3 nats; the cited DreamerV3 paper uses 1 nat. Below the threshold, this term is constant and supplies no gradient through either KL argument. It removes marginal regularization pressure in that region; it does not measure irreducible surprise or guarantee that the posterior uses the available capacity. A collapsed posterior equal to its prior also has a clipped value at the floor. We use a 1-nat threshold below and log both raw and clipped KL to keep these cases distinguishable. Reconstruction and shared upstream parameters can still change the distributions while the clipped KL's own gradient is zero.

KL balancing (DreamerV2, kept in DreamerV3): split the KL gradient by stopping gradients through one argument at a time:

$$\mathcal{L}_{\mathrm{dyn}} = \lambda_{\mathrm{dyn}}\, \mathrm{KL}\big(\mathrm{sg}[q] \,\big\|\, p\big), \qquad \mathcal{L}_{\mathrm{rep}} = \lambda_{\mathrm{rep}}\, \mathrm{KL}\big(q \,\big\|\, \mathrm{sg}[p]\big)$$

where:

• $q$ and $p$: shorthand for the per-step posterior $q_\phi(z_t \mid h_t, o_t)$ and prior $p_\theta(z_t \mid h_t)$
• $\mathrm{sg}[\cdot]$: stop-gradient, meaning the bracketed distribution is held fixed, so that term's gradient can only reach the other argument
• $\lambda_{\mathrm{dyn}}$ and $\lambda_{\mathrm{rep}}$: the loss weights. DreamerV2 introduced the balance with the factor $\alpha = 0.8$, weighting the dynamics term four times the representation term; the DreamerV3 paper v2 applies loss scales of $\lambda_{\mathrm{dyn}} = 1.0$ and $\lambda_{\mathrm{rep}} = 0.1$ to the two terms

Both terms are minimized alongside the negative reconstruction objective. The first differentiates through the prior argument with the posterior argument stopped; the second reverses those roles. Shared upstream parameters, including the recurrent context, may influence both distributions after an update. Larger dynamics coefficients emphasize direct prior-fitting gradients, but coefficient ratios are not gradient-norm ratios and do not guarantee an optimizer's outcome. The aim is to reduce the incentive to discard useful observation information merely to make matching easier. DreamerV3 reports results across its evaluated domains with a shared configuration and several coordinated design choices; that empirical result should not be attributed to KL balancing alone.

Loss units matter for both images and low-dimensional observations. A sum over many pixels differs from a mean over pixels, and the observation likelihood's scale determines reconstruction strength. For a fixed scalar Gaussian decoder variance $\sigma_o^2$, negative log likelihood equals squared error divided by $2\sigma_o^2$, plus a parameter-independent constant. Thus squared error plus $\beta\,\mathrm{KL}$ is, up to positive scaling and an additive constant, a Gaussian negative-ELBO objective with $\sigma_o^2=\beta/2$. Here $\beta=0.1$ corresponds to variance $0.05$ before clipping. Free-nat clipping modifies that objective; a learned variance would also introduce scale-dependent terms. Scalar loss magnitudes alone do not establish gradient dominance, especially when a large clipped term is locally constant.

During filtering, the posterior uses the currently available observation; during imagination, future states come from the prior. At fixed recurrent context $h$, averaging $\mathrm{KL}(q(z\mid h,o)\|p(z\mid h))$ over a specified observation distribution gives two terms: the mutual information between $z$ and $o$ under that joint distribution, plus the KL between the observation-averaged posterior and the prior. Consequently, expected KL upper-bounds that observation information but is not equal to it in general. It mixes useful or irrelevant encoded information with prior mismatch; it is not a direct measure of forecast error.

Keep observation-conditioned filtering and open-loop prediction separate when evaluating a model. The experiment below measures both; it does not train a policy or measure closed-loop decision usefulness. A good posterior reconstruction can coexist with poor imagination because observations supply information that an open-loop forecast lacks. Approximation and training failures can widen that difference too. A high held-out KL is a reason to investigate, not proof of a failed prior: check predictions, interval quality, and which observations change the inferred state.