Recurrent State-Space Model (RSSM)

The deterministic recurrent state (often denoted h_t) is a hidden vector updated by a recurrent neural network (RNN), typically a Gated Recurrent Unit (GRU) or LSTM. It functions as a compressed memory of all past observations and actions, providing a stable, predictable pathway for temporal information flow.

• Purpose: Maintains a deterministic, evolving context for the stochastic latent variables.
• Mechanism: Updated as h_t = f_RNN(h_{t-1}, s_{t-1}, a_{t-1}), where s is the stochastic latent and a is the action.
• Role in Planning: Provides the consistent "backbone" for multi-step imagination rollouts, allowing gradients to flow effectively through time via backpropagation.

The stochastic latent state (often denoted s_t or z_t) is a random variable, typically modeled as a Gaussian, that represents the perceived, uncertain aspects of the environment at time t. It is sampled from a distribution conditioned on the current observation and the deterministic state.

• Purpose: Captures the partial observability and inherent randomness of the environment that cannot be perfectly memorized.
• Mechanism: s_t ~ p(s_t | h_t, o_t), where o_t is the observation. It is later used to predict the next observation and reward.
• Role in Imagination: During planning ("dreaming"), future s_t values are sampled from a prior p(s_t | h_t) without new observations, enabling the agent to simulate possible futures.

These components bridge the high-dimensional observation space (e.g., images) and the low-dimensional latent space.

• Encoder (Recognition Model): A neural network (e.g., CNN) that compresses a raw observation o_t into parameters (mean and variance) for the posterior distribution of the stochastic latent: q(s_t | h_t, o_t).
• Decoder (Observation Model): A neural network that reconstructs or predicts the observation from the latent state: p(o_t | h_t, s_t). Its reconstruction loss is a key part of the RSSM's training objective, forcing the latent space to retain all necessary information about the observation.
• Function: Enables the model to work with pixels directly, learning a semantically meaningful latent state without hand-engineered features.

These are simple MLP heads attached to the latent state that allow the model to simulate the outcomes of imagined actions.

• Reward Model: Predicts the expected immediate reward for a given latent state: r_t ~ p(r_t | h_t, s_t). This turns the RSSM into a reward-predictive world model.
• Terminal (Continue) Model: Predicts the probability of an episode terminating (or continuing) in a given latent state: c_t ~ p(c_t | h_t, s_t). This is crucial for accurate long-horizon value estimation in episodic environments.
• Critical for Planning: During imagination, these predictors allow the agent to evaluate the desirability of simulated trajectories without interacting with the real environment.

The RSSM is trained end-to-end to maximize a variational lower bound (ELBO) on the log-likelihood of observed sequences. The loss function has several key terms:

• Reconstruction Loss: log p(o_t | h_t, s_t) - Ensures the latent state contains information to decode the observation.
• KL Divergence Loss: D_KL( q(s_t | h_t, o_t) || p(s_t | h_t) ) - Regularizes the posterior (from encoder) to stay close to the prior (from dynamics).
• KL Balancing: A practical technique where the KL term is scaled (e.g., 0.1 to 0.5) or made free for parts of the training to prevent the model from ignoring the latent variables (posterior collapse) or the observations (over-regularization).

This is the primary use case of the trained RSSM. Instead of planning in pixel space, the agent's policy and value functions are trained entirely on imagined trajectories within the latent space.

• Process: Start from a real encoded state (h_t, s_t). For K steps, sample an action from the policy, predict the next deterministic state h_{t+1}, sample a prior latent s_{t+1} ~ p(s_{t+1} | h_{t+1}), and predict reward/continue. This creates a latent trajectory.
• Efficiency: Backpropagation through time (BPTT) can flow through these compact latent rollouts, making policy optimization via gradient descent highly sample-efficient.
• Algorithm Example: The Dreamer series of algorithms uses this exact paradigm: learn an RSSM (world model), then train an actor-critic agent using gradients derived from imagined latent rollouts.

A world model is an agent's internal, learned representation that predicts future environment states and rewards. It enables planning and imagination without direct, costly interaction with the real world. The RSSM is a specific, highly effective architecture for implementing a world model, particularly for high-dimensional observations like images.

• Core Function: Serves as a compressed, predictive simulator.
• Key Benefit: Drastically improves sample efficiency by allowing policy training via imagined rollouts.
• Example: In the Dreamer algorithm, the world model (an RSSM) is used to train the actor and critic entirely in latent space.

A latent dynamics model learns to predict future states in a compressed, abstract latent space, not in the raw, high-dimensional observation space (e.g., pixels). This is the primary function of the RSSM.

• Architecture: Typically uses an encoder to map observations to latents, a transition model to predict next latents, and a decoder to reconstruct observations.
• Advantage: Improves generalization, reduces computational cost, and filters out irrelevant sensory details.
• Contrast: Unlike a model that predicts raw pixels, a latent dynamics model learns the underlying state of the environment.

Dreamer is a seminal model-based reinforcement learning algorithm that popularized the RSSM architecture. It learns a world model (the RSSM) and then trains a policy and value function entirely through latent imagination—backpropagation through time on rollouts generated by the model.

• Three-Phase Loop: 1) Learn the RSSM world model from experience. 2) Train an actor and critic using trajectories imagined by the model. 3) Use the policy to collect new experience.
• Key Innovation: Demonstrates that effective policies can be learned purely from simulated experience within a learned latent space.
• Impact: Established a strong baseline for sample-efficient, visual model-based RL.

The RSSM explicitly separates its latent state into stochastic and deterministic components. This design is critical for modeling partial observability and complex temporal dependencies.

• Deterministic Path: A recurrent neural network (e.g., GRU) that maintains a deterministic memory of the past. It represents predictable, sequential information.
• Stochastic Path: A latent variable sampled from a distribution at each step. It represents the unpredictable aspects of the environment and future.
• Interaction: The deterministic state conditions the prior for the stochastic variable, which is then updated by the current observation. This creates a rich, hybrid state representation for robust long-horizon prediction.

Model-policy co-adaptation is a critical failure mode in model-based RL that architectures like the RSSM must guard against. It occurs when a policy overfits to the specific biases and inaccuracies of its own learned dynamics model.

• The Problem: The policy learns to exploit flaws in the model, achieving high reward in simulation but failing catastrophically in the real environment.
• RSSM's Mitigation: By using a stochastic latent variable, the RSSM inherently represents uncertainty. Training the policy over many imagined trajectories with different stochastic samples encourages robustness.
• Related Solution: Pessimistic exploration or planning with uncertainty quantification (e.g., using an ensemble) are other methods to prevent this issue.

Imagined rollouts (or simulated experience) are sequences of states, actions, and rewards generated by unrolling a learned dynamics model from a starting state. They are the primary data source for training the policy in algorithms like Dreamer that use an RSSM.

• Process: Starting from a real encoded state, the model (RSSM) and a candidate policy are used to simulate multiple steps into the future in latent space.
• Purpose: Provides a cheap, abundant source of training data for the actor and critic networks via backpropagation through time (BPTT).
• Trade-off: The quality of these rollouts is limited by model error and compounding error, where small inaccuracies grow over long simulation horizons.