Recurrent State Space Models
Recurrent State Space Models (RSSM) are a hybrid latent dynamics architecture that simultaneously maintains a deterministic recurrent state for temporal consistency and a stochastic latent variable for uncertainty representation — combining the memory of RNNs with the probabilistic expressiveness of VAEs to model both the reliable patterns and the inherent randomness of real-world environments — introduced as the core of the Dreamer agent and now the dominant architecture for learning dynamics models in model-based reinforcement learning from high-dimensional observations.
What Is the RSSM?
- Two-Path Design: The RSSM maintains two parallel state components at each timestep: a deterministic recurrent hidden state (from a GRU cell) and a stochastic latent variable (drawn from a learned Gaussian distribution).
- Deterministic Path: The GRU hidden state h_t captures a summary of all past observations and actions — providing temporal consistency, long-range memory, and a stable context for dynamics prediction.
- Stochastic Path: The latent variable z_t is sampled from a distribution conditioned on h_t — capturing environmental stochasticity, multimodal futures, and inherent uncertainty not resolved by past context.
- Prior vs. Posterior: During imagination (no observations), z_t is sampled from the prior p(z_t | h_t). During training with observations, z_t is sampled from the posterior p(z_t | h_t, o_t) — a richer estimate given the observation.
- Together: The full latent state (h_t, z_t) captures both what has happened (deterministic) and what is happening right now with uncertainty (stochastic).
RSSM Equations
The RSSM update at each step t given action a_{t-1} and observation o_t:
- Deterministic recurrence: h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1})
- Prior (for imagination): z_t ~ p(z_t | h_t) — predicted stochastic state without observation
- Posterior (for training): z_t ~ q(z_t | h_t, e_t) where e_t = Encoder(o_t) — refined with current observation
- Observation model: o_t ~ p(o_t | h_t, z_t) — reconstruction for training signal (DreamerV1/V2)
- Reward model: r_t ~ p(r_t | h_t, z_t) — used for policy learning
Training uses ELBO: reconstruction + reward prediction + KL(posterior prior).
Why The Two-Path Design?
| Property | Deterministic Path | Stochastic Path |
|---|---|---|
| Purpose | Long-range memory, temporal context | Uncertainty, multimodal futures |
| Update | Always updated from previous state + action | Sampled from distribution |
| During Imagination | Used directly | Sampled from prior |
| Information Flow | Carries all past context forward | Captures current randomness |
A purely deterministic model can't represent stochastic environments. A purely stochastic model (VAE at each step) loses temporal context. RSSM combines both strengths.
Evolution Across Dreamer Versions
- DreamerV1: Continuous Gaussian stochastic state, GRU deterministic — image reconstruction training.
- DreamerV2: Replaced continuous Gaussian with discrete categorical latent (32 groups × 32 classes) — better for representing sharp multimodal futures, enabling human-level Atari.
- DreamerV3: Added symlog predictions, free bits KL balancing, and robust normalization — enabling the same RSSM to work across 7+ domains without tuning.
RSSM is the workhorse of world-model-based RL — the architectural insight that bridging deterministic memory and stochastic uncertainty produces a dynamics model expressive enough to learn the structure of diverse real and simulated environments from raw sensory observations.