Recurrent State-Space Models

Watch a soccer ball roll behind a sofa. You can often predict where it will reappear, provided nothing hidden changes its motion. That familiar inference suggests a concrete engineering problem: maintain an internal state that keeps evolving when observations stop. A robot cannot pause its predictions whenever a camera is occluded. It must project forward from the information it has, while retaining uncertainty about what happens out of sight. In this chapter we build a small model that filters noisy observations and then predicts without seeing the future. We will also check where that model falls short.

In Part II: Inference, Decisions, and Control you met the classical answer, the Bayesian filter. A Kalman filter represents a Gaussian belief by its mean and covariance. With known linear transition and observation models, a Gaussian initial belief, and independent Gaussian process and measurement noise, its updates give the exact state posterior and minimum-mean-squared-error state estimate. Outside those conditions, ordinary Kalman updates generally cease to be exact Bayesian inference, although optimal linear estimation and other tractable special cases remain possible. In Part III: Representing Agents and Worlds you saw learned latent representations. The recurrent state-space model (RSSM), used by PlaNet and the Dreamer family, combines learned recurrent memory with stochastic latent dynamics.

An RSSM keeps two components of internal state. A deterministic recurrent state $h_t$, carried by a GRU or LSTM, provides memory. Deterministic describes its update given its inputs; it does not mean that the model is certain about the world. A stochastic latent $z_t$ is sampled from a conditional distribution. During filtering, a posterior network uses the current observation to infer that distribution. A prior network predicts it without the current observation, using past information through $h_t$. A KL divergence encourages the prior to match the observation-conditioned posterior and regularizes how much that posterior can differ. During acting, the agent can still filter its current observation; only its imagined future lacks future observations and therefore uses the prior alone.

The filtering/imagination distinction is the central operating boundary in the chapter, not a train/test split. Both regimes can be evaluated on held-out data, and filtering also occurs during deployment. The inference model asks, "what latent states explain what I have just observed?" The dynamics model asks, "what latent states should I predict before looking?" Training encourages consistency between them, but it cannot make unavailable future information available. A good predictive prior can therefore differ from the posterior even when neither is defective.

This section builds the two ingredients separately before combining them. PlaNet's experiments motivate the combination, but neither ingredient is a universal requirement for a capable world model. We need to distinguish how a state is computed from what distribution its outputs can represent.

Recall from Bayesian Filtering and Belief States that an observation-action history contains the information available for inferring hidden state. Storing that raw history costs more as an episode grows. Some observations are redundant, while repeated noisy measurements can still help reduce uncertainty. A belief $b_t = p(s_t \mid o_{\le t}, a_{<t})$ is the distribution of hidden environment state $s_t$ given observations through $t$ and earlier actions. It is a sufficient summary for the usual known-model POMDP prediction and decision problem, but a general belief distribution need not have an exact finite-dimensional representation. Classical filters exploit tractable special families. A learned world model instead fits a compact representation intended to retain the information its prediction tasks need.

There are two computational choices to separate:

• A deterministic hidden vector is a fixed function of its inputs and parameters. It can store point estimates, probabilities, variances, or other learned features.
• A stochastic latent is a random variable. Its distribution represents modeled alternatives; one sample is only one realization, not the distribution itself.

A Kalman filter is a useful counterexample to the idea that deterministic memory means certainty: its mean and covariance are updated deterministically, yet jointly describe an uncertain state. Likewise, a neural vector can encode the weights, means, and scales of a mixture distribution. The RSSM chooses a different implementation: recurrent memory plus sampled latent states. That choice can help retain information and generate alternative trajectories, but it does not by itself guarantee good inference or calibrated uncertainty.

A simple learned filter is a recurrent neural network that ingests observations and actions and maintains a fixed-size vector:

where:

• $h_t$: the learned hidden state at time $t$, a fixed-size vector that summarizes the observation and action history
• $f_\theta$: the learned recurrence, for instance a GRU or LSTM cell, with parameters $\theta$
• $a_{t-1}$: the action taken at the previous step
• $o_t$: the observation arriving at the current step

A related recurrent-memory pattern appears in Ha and Schmidhuber's World Models (2018). Its LSTM consumes VAE codes and actions and outputs a mixture distribution, not a point prediction. In our simpler recurrence, $h_t$ is a learned predictive summary, not a guaranteed sufficient statistic. A fixed-width representation puts pressure on the network to retain task-relevant information; it may also forget information that later turns out to matter. The prediction objective determines what is rewarded, and held-out tests tell us whether that training signal was enough.

Deterministic state has real strengths:

• Repeatable updates. With fixed inputs, initialization, parameters, and deterministic numerical execution, the same recurrence returns the same state. This removes latent-sampling variance from that update, not representation or floating-point error.
• Direct differentiation. Ordinary backpropagation follows the recurrence without a fresh sampling operation on each edge. Gradients can still vanish or explode.
• Simple inference. The point-output baseline below needs one recurrent pass and no variational posterior.

The limitation appears when we attach a single point output and train it with mean squared error. If the next observation has two plausible modes, that output cannot report both. Conditional squared risk decomposes into the irreducible conditional variance plus the squared distance from the conditional mean, so the MSE-optimal point forecast is the mean. For a balanced symmetric bimodal distribution, that mean lies between the modes. A forecast of a ball on the left or right of a pillar can therefore average into a low-probability location in the middle. The cause is the point-output objective, not deterministic recurrent memory: a distributional head can report both alternatives.

The failure is easiest to see with numbers rather than words, so we make it concrete before moving on.

The plot illustrates mode averaging: the conditional mean lies in the low-density valley between the modes, not at the global minimum of the density. Enlarging a model while retaining an MSE point-output objective does not change the optimum into two alternatives. Changing the output distribution or the prediction objective addresses a different problem than adding hidden-state capacity.

A point output also leaves uncertainty unreported. Consider two grasps with the same predicted mean outcome but different failure risks. A planner receiving only those mean outcomes cannot distinguish them by uncertainty. Risk-sensitive decisions or information gathering may need a distribution, quantiles, or other uncertainty-sensitive information. These can come from deterministic memory with a probabilistic head, from ensembles, or from stochastic latent rollouts. A point forecast can still be useful for decisions; it simply does not answer every decision question.

One way to generate alternatives is to make part of the latent state stochastic. Instead of reporting only a point, the model predicts a distribution and samples from it. One trajectory may go left and another right if its learned distribution represents both outcomes. Sampling represents what the model believes, not a certificate that the belief is correct. Calibration must be checked against held-out outcomes.

Stochastic latent sequence models precede RSSMs: variational RNNs (Chung et al., 2015), stochastic RNNs (Fraccaro et al., 2016), deep variational Bayes filters (Karl et al., 2017), and structured inference networks (Krishnan et al., 2017) develop related latent-state and inference ideas. In these models, distinguish the approximate latent distribution from a sample drawn from it. For an RSSM, the joint distribution of $(h_t,z_t)$ matters: $h_t$ varies across previously sampled latent histories, even though each recurrent update is deterministic conditional on its inputs.

Stochastic recurrence can make retaining information and optimizing long sequences difficult. PlaNet reported that its combined recurrent and stochastic state performed substantially better than the particular single-path variants it tested. That result does not establish a universal law that each additional sample increases gradient variance. For example, in $z_t=z_{t-1}+\theta+\epsilon_t$ with a terminal linear loss $z_T$, the pathwise derivative is exactly $T$ for every noise draw. Actual gradient variance depends on the transition, objective, and estimator. The RSSM's deterministic path offers a route for recurrent information without requiring a fresh sample on every edge.

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.

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

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.

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:

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.

Once trained, the model is used by imagining: rolling the latent state forward under actions, without observations. This section explains how rollouts work, when their error can grow, and what keeps them useful.

A rollout has two phases.

Burn-in. Run posterior updates over an observed prefix to initialize an approximate filtered belief. A known initial environment state can instead inform a suitable initializer; an unconditional constant is not a substitute for arbitrary known states. How much history is needed depends on observability and the model. To represent uncertainty over the installed belief, run several independently sampled posterior histories through the same observations and actions. Repeating one sampled endpoint would condition away the differences between those histories.

Imagination. From the installed belief, iterate the dynamics with the prior only, conditioned on a candidate action sequence:

where:

• $k$: the number of steps beyond the last observed time $t$; the pair of updates is iterated for $k = 1, 2, \ldots, H$ over the imagination horizon $H$
• $a_{t+k-1}$: the candidate action supplied by the planner at step $t+k-1$
• $p_\theta(z_{t+k} \mid h_{t+k})$: the prior, now the only source of the stochastic state because no observation arrives to correct it

The two updates play the same roles as during training: the GRU refreshes the memory from the previous belief and action, and the prior samples a fresh uncertain state. The agent then decodes whatever it needs from $(h_{t+k}, z_{t+k})$ — typically only the reward inside the planning loop, with observations decoded for inspection or evaluation rather than per candidate plan.

Two properties make latent imagination useful in PlaNet and Dreamer. Their action evaluation or actor training can operate through latent reward and value predictions without decoding images at every step. This avoids the image decoder's computation in those implementations. PlaNet evaluates candidate action sequences with a planner; Dreamer instead learns an actor through imagined trajectories and value bootstrapping. A single sampled rollout is one trajectory. Repeated independent posterior and prior noise draws approximate the model's conditional distribution of futures, which still needs validation. We will build the decision-making machinery in Part VII: Planning and Agency.

In System Identification, you saw why evaluating a one-step model recursively is a different test from predicting with fresh observations. An imperfect prediction can become an input to the next transition and affect later errors. Stable dynamics may attenuate that error; unstable dynamics may amplify it. Neither error growth nor monotonic widening is guaranteed merely by rolling a model forward.

Two contributions are useful to distinguish when interpreting forecast error:

• True conditional variability. Even a perfect conditional-mean predictor makes errors against individual stochastic futures. Process noise and uncertainty left after filtering affect both point forecasts and sampled forecasts. For a stable linear system, predictive covariance can approach a finite stationary level rather than grow forever.
• Model and inference error. Misspecified dynamics, imperfect optimization, a poor initial belief, or evaluation outside the training distribution can add error. Uncertainty about fitted parameters is one part of this problem, often called epistemic uncertainty; it is not a synonym for every approximation error. Recursion may amplify or attenuate these defects.

Accurately represented environmental variability should not be removed just to make intervals narrower. An inaccurate model may instead need better data coverage, inference, or dynamics fitting; shorter rollouts reduce reliance on extrapolation without necessarily fixing it. A point-output model has no interval, but that absence is not a declaration of confidence. Conversely, drawing a broad stochastic ensemble does not establish that the model knows its errors.

Use matched reference predictions when the environment is known. Below we implement the exact linear-Gaussian filter and its open-loop conditional mean and variance under the same observed prefix and actions. This gives a reference for irreducible squared error rather than comparing a random forecast sample with a conditional mean. We also measure interval coverage and width. Empirical 95% coverage at one horizon is only a marginal diagnostic: an overly wide interval can buy coverage, biases can be hidden across subgroups, and finite validation sets give noisy proportions. Neither coverage nor an RMSE curve alone decomposes model error from process noise.

Useful tools for managing rollout degradation include:

• Short horizons. PlaNet's reported setup plans 12 steps; the original Dreamer imagines 15 steps and bootstraps values beyond that window. PlaNet replans as observations arrive, while Dreamer acts through a learned policy conditioned on an updated latent state.
• Free nats and KL balancing. These adjust information regularization and gradient allocation. Free nats explicitly stop KL pressure below a threshold; neither technique guarantees a smaller prediction gap.
• Latent overshooting. PlaNet compares multistep prior predictions with posteriors inferred from real observations, adding direct multistep consistency pressure beyond one-step training.
• Stochastic latents. Sampling can retain modeled alternative futures instead of feeding back an intermodal mean. A sample lies in the learned prior's support, not necessarily in a well-trained region or in the true environment's support.
• Limited planner pressure. Optimizing strongly against an imperfect model can exploit its errors. Shorter horizons and uncertainty-sensitive objectives are possible defenses, not guarantees; Part XII: Reliable World Models develops this issue.

Evaluate the operating regimes separately. Posterior-pass error and prior-only rollout error depend on several components, but they ask different questions about access to observations. Reporting only posterior reconstruction can conceal open-loop weaknesses because each fresh observation offers a chance to correct the model. A point-error curve alone cannot assess distribution quality.

We now build an RSSM for a small controlled linear-Gaussian environment. The experiment exposes true hidden states and the data-generating equations, so we can check filtering, open-loop forecasts, and uncertainty against a known-model reference. It does not illustrate every difficulty of world modeling: in particular, this environment's conditional futures are Gaussian, not multimodal. Training is CPU-sized; runtime depends on hardware and numerical libraries.

The environment is a discrete spring-damper system with position and velocity as its two state variables. Position is observed with sensor noise; velocity is never an input or target during world-model training. Recorded force actions are also supplied. Position, velocity, actions, and process noise jointly determine future motion. With the known model parameters and Gaussian initial distribution, a Kalman filter gives the exact posterior. The position observation map is $C=[1,0]$, and the two-step observability matrix has rows $C$ and $CA=[1,0.1]$, so its determinant is $0.1$: position history contains information about velocity, even though noise prevents exact recovery.

The experiment asks four questions, each mapping to a section of the chapter:

1. Can the model filter? Score posterior decoder means against clean positions, alongside raw-observation echo, a memoryless shrinkage estimate, and a known-model Kalman filter. Beating echo shows lower estimation error, not by itself use of history.
2. Can the model imagine? Observe 10 positions, then predict 40 steps under the recorded future actions. Compare the RSSM ensemble mean with the deterministic baseline and the Kalman conditional mean using the same point-RMSE calculation.
3. How informative are its intervals? Score central sample-quantile intervals against clean future positions, report width and sampling uncertainty, and compare with exact Kalman intervals. These are diagnostics, not a proof of full calibration.
4. What information is recoverable? Fit a linear probe on training representations and test it on held-out sequences. A high velocity $R^2$ supports linear decodability; it does not prove an identifiable physical coordinate, causal use, or Markov sufficiency.

We retain a deterministic GRU with a point decoder as an illustrative baseline. The RSSM has 20,625 parameters and 1,500 updates; the baseline has 10,689 parameters and 800 updates, and their objectives and observation timing differ. This is not a controlled ablation isolating stochastic state. The shared evaluation cutoff, actions, and point-error metric make the resulting forecasts comparable, but cannot attribute their differences to one architectural choice.

The simulator below uses an explicit-Euler spring-damper transition with step size $0.1$, plus discrete process noise. It is the discrete model being studied, not the exact sampled transition of a continuous-time stochastic oscillator. Observation $o_t$ measures the position in state $s_t$; action $a_t$ affects $s_{t+1}$. We use consistent arbitrary position and time units, so velocity has position-per-time units. The learned models receive positions and actions, not the true state, simulator coefficients, or individual noise realizations.

The model receives the noisy position observations shown here and their recorded actions. The true state curves are for our inspection. Try estimating the direction of motion from a short noisy history: a single position generally does not determine velocity, while earlier positions and applied forces provide additional information. Useful predictive memory need not contain a coordinate explicitly named velocity; equivalent learned features can carry motion information.

The RSSM below mirrors the equations of the architecture section exactly: a GRU cell for the deterministic path, two small heads for the prior and posterior Gaussians, and a decoder. Note that the posterior head is the only component that receives the observation as direct input; the decoder touches the observation only indirectly, through z_t sampled from the posterior.

The gauss_head splits each head into a mean and log standard deviation. Exponentiation makes the scale positive; clamping the log scale to $[-5,2]$ bounds it to approximately $[0.0067,7.39]$. The mean is not bounded, and this clamp does not prevent posterior collapse or every numerical failure. The GRU uses $(h_{t-1},z_{t-1},a_{t-1})$ before the current observation is encoded, matching the timing convention above.

The KL helper implements the analytic conditional diagonal-Gaussian expression. The loop updates $h_t$, computes both heads, samples the posterior, decodes, and accumulates squared reconstruction error plus weighted clipped KL. This is a single-trajectory estimate of a modified variational objective, not the literal unmodified ELBO displayed earlier.

The free-nat floor is 1.0, applied after summing latent dimensions but separately for each sequence and timestep. The KL coefficient is 0.1. With our scalar squared-error reduction, its fixed-Gaussian interpretation uses decoder variance $0.05$ before free-nat clipping. This is a chosen observation-model scale, not an estimate of the simulator's sensor variance $0.0625$. We log raw KL as well as the clipped loss term. Neither their scalar ratio nor convergence near the threshold establishes that the prior has learned all predictable information.

Read raw and clipped KL separately. A clipped value exactly at 1 would only imply that every contributing raw value was at or below 1; it would not prove optimal prediction or identify a noise floor. Posterior reconstruction can fall below the sensor variance because its encoder sees the same noisy observation used as its target. That is allowed by the objective, not an accidental input leak. By contrast, a history-only prediction of a fresh observation has expected squared error at least the independent sensor variance. Clean-position evaluation below asks whether the posterior output denoises, and the Kalman reference shows how far it is from the known-model optimum.

In the seeded reproduction, reconstruction fell to 0.0175 per step. The logged clipped KL began near 1, rose, and finished at 1.101 nats; raw KL finished at 1.017 nats. The final values are last-minibatch diagnostics, not held-out averages. The initial increase is another reason not to interpret every training trace as a prior steadily catching up to a fixed posterior.

The baseline GRU receives $(o_{t-1},a_{t-1})$ and predicts $o_t$ with a squared-error point decoder. Its hidden memory could encode motion information or distribution parameters, but this particular output head returns only a point. It differs from the mixture-density RNN in World Models. We leave its small training budget explicit rather than describing the comparison as an isolated test of stochasticity.

Teacher forcing feeds the real previous observation during baseline training; recursive evaluation feeds the prediction instead. This change can expose compounding errors, but does not guarantee that every horizon worsens. The RSSM also changes operating regime when it replaces future posterior updates with prior draws. It supplies an ensemble while this baseline supplies no interval. That difference in output format does not establish that the ensemble is well calibrated.

Now the two evaluation regimes. The filter runs the posterior at every step: the model always gets to look. The imagination burns in the belief on the first BURN_IN real observations, then rolls out with the prior only, under the true future actions, sampling N_SAMPLES independent futures. The deterministic baseline rolls out by feeding its own predictions back in place of observations, the classic deterministic world-model rollout.