Recurrent State-Space Model (RSSM)

The RSSM explicitly separates its latent state into two components: a deterministic recurrent state and a stochastic latent variable. The deterministic RNN (often an LSTM or GRU) maintains a memory of past information, while the stochastic component models the inherent uncertainty and partial observability of the environment. This hybrid structure allows the model to capture predictable dynamics in the recurrent state while representing aleatoric uncertainty in the stochastic part.

At its core, the RSSM is a sequential latent variable model trained with variational inference. It defines a generative process (the prior) that predicts the next latent state given the previous state and action, and an inference model (the posterior) that updates the belief about the current latent state based on the new observation. This framework, using the Evidence Lower Bound (ELBO), allows the model to learn a compact state representation from pixels or other sensory data in an unsupervised manner.

A key application of the RSSM is latent imagination. Once trained, the model can generate realistic rollouts entirely within its learned latent space, without decoding to pixels at each step. Planning algorithms, like the Cross-Entropy Method (CEM) used in Dreamer, perform these rollouts to evaluate action sequences. This makes planning computationally efficient, as it avoids expensive rendering of high-dimensional observations during search.

The RSSM is designed for Partially Observable Markov Decision Processes (POMDPs), where the agent receives incomplete sensory data. The recurrent network aggregates a history of observations and actions into the deterministic state, forming a belief state. This allows the agent to maintain an internal estimate of the true world state, which is crucial for tasks where memory is required, such as navigating mazes or tracking objects that become occluded.

The RSSM serves as the dynamics model in Model-Based Reinforcement Learning (MBRL). It is trained to predict future rewards and episode continuation (discount) signals alongside reconstructing observations. These predictions enable value functions and policies to be trained on imagined trajectories, leading to the high sample efficiency characteristic of agents like Dreamer. The model learns a self-consistent world where rewards and states evolve realistically.

The RSSM is a deep learning generalization of classical state-space models like the Kalman filter. Both frameworks maintain an estimate of a hidden state that evolves over time. Key differences include:

• The RSSM uses neural networks for non-linear transition and observation models.
• It employs amortized variational inference instead of analytical updates.
• The latent state is learned rather than engineered. This connection grounds the RSSM in decades of state estimation theory while leveraging the representational power of deep networks.

A world model is a learned or engineered internal representation of an environment that enables an agent to predict future states and outcomes without direct interaction. It serves as a foundational component for planning and model-based reinforcement learning.

• Function: Acts as a "simulator in the mind," allowing the agent to perform mental rollouts.
• Contrast with RSSM: An RSSM is a specific, popular neural network architecture for implementing a world model, combining deterministic recurrence with stochastic latents.
• Example: In the Dreamer agent series, the world model (implemented as an RSSM) learns to predict future image observations and rewards from past images and actions.

A Partially Observable Markov Decision Process (POMDP) is the mathematical framework that formally defines the problem an RSSM is designed to solve. It extends the standard MDP by acknowledging the agent does not have direct access to the true environment state.

• Core Challenge: The agent only receives partial, noisy observations (e.g., pixels from a camera).
• RSSM's Role: The RSSM learns to infer a compact belief state from this sequence of observations and actions, effectively performing state estimation within the POMDP framework.
• Key Components: Includes a state transition function, an observation function, and a reward function, all of which are learned by the RSSM.

Model-Based Reinforcement Learning (MBRL) is the overarching paradigm where an agent learns an explicit model of the environment's dynamics and uses it for planning or policy optimization. The RSSM is a premier example of a learned dynamics model in deep MBRL.

• Advantage: Dramatically improves sample efficiency compared to model-free methods, as the model can be queried extensively without real-world interaction.
• RSSM's Contribution: Provides a tractable latent dynamics model suitable for high-dimensional observations like images.
• Planning Loop: In agents like Dreamer, the learned RSSM is used with a planner (e.g., a learned policy or cross-entropy method) to select actions that maximize predicted reward.

A belief state is a probability distribution over possible true states of the environment, representing the agent's internal estimate given its history of actions and observations. In an RSSM, the stochastic latent variable is a compressed, distributed representation of this belief.

• From POMDPs: The optimal solution to a POMDP is a policy that maps belief states to actions.
• RSSM Implementation: The RSSM's recurrent network maintains a deterministic hidden state that summarizes history, which then conditions the prior and posterior distributions of the stochastic latent state.
• Function: This latent belief state is the input to the transition predictor and decoder, enabling consistent long-horizon predictions.

A disentangled representation is a latent state encoding where distinct, semantically meaningful factors of variation in the data are separated into independent or semi-independent dimensions. While not explicitly enforced, a well-trained RSSM often learns a form of disentanglement.

• Benefit for Planning: Disentanglement can improve generalization and interpretability. If factors like object position, color, and type are separated, the model can more robustly recombine them for novel situations.
• RSSM Mechanism: The separation of a deterministic RNN state (for modeling temporal dependencies) and a stochastic latent state (for modeling uncertainty and perceptual details) encourages a degree of factorization.
• Related Technique: Methods like β-VAE explicitly add a loss to encourage disentanglement in the latent space.

Model-Predictive Control (MPC) is an online planning and control method that repeatedly solves a finite-horizon optimization problem using a dynamics model. The RSSM can serve as the dynamics model within an MPC loop for robotic control.

• Process: At each timestep, MPC uses the model (RSSM) to simulate multiple action sequences, evaluates them with a reward/cost function, executes the first action of the best sequence, and then replans.
• Advantage with RSSM: Provides real-time, adaptive control by leveraging the learned world model to predict outcomes of candidate actions.
• Contrast with Learned Policy: Dreamer agents typically learn an explicit policy network from imagined RSSM rollouts, while MPC uses planning at inference time, which can be more flexible but computationally heavier.