Recurrent State-Space Model (RSSM)

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.

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.