World Model

The observation encoder is a neural network (often a convolutional or transformer-based encoder) that maps raw, high-dimensional sensory inputs (e.g., pixels from a camera) into a compact, structured latent representation or belief state. This compression discards irrelevant details while preserving information necessary for prediction and decision-making.

• Function: z_t = encoder(o_t)
• Purpose: Reduces dimensionality, extracts features, and provides a consistent input space for the dynamics model.
• Example: In DeepMind's Dreamer algorithm, this is a convolutional variational autoencoder that learns a stochastic latent representation of image observations.

The core predictive engine. The dynamics model (or transition model) learns the environment's rules. It takes the current latent state and a proposed action as input and predicts the next latent state and any intermediate reward.

• Function: (z_{t+1}, r_t) = dynamics(z_t, a_t)
• Purpose: Simulates the consequences of actions without interacting with the real world.
• Key Challenge: Model error—the discrepancy between predicted and real transitions—can lead to compounding error over long rollouts.
• Architectures: Often implemented as Recurrent State-Space Models (RSSM) to handle partial observability, or as ensembles of MLPs for uncertainty estimation.

While sometimes integrated into the dynamics model, the reward predictor is a distinct learned function that estimates the scalar reward associated with a given state or state-action pair. It allows the agent to evaluate the desirability of imagined future states.

• Function: r̂_t = reward(z_t) or r̂_t = reward(z_t, a_t)
• Purpose: Provides the objective signal for planning. In algorithms like MuZero, the reward predictor is a critical output of the model's forward pass.
• Training: Learned via supervised regression on collected (state, action, reward) tuples.

The latent state (often denoted z or s) is the compressed, internal representation of the agent's belief about the current world situation. It is the "memory" of the world model, summarizing the history of observations and actions.

• Properties: Should be Markovian (the future is conditionally independent of the past given the present state) and sufficient for accurate prediction.
• Types: Can be deterministic (a fixed vector) or stochastic (a distribution, e.g., z ~ N(μ, σ)), with the latter better capturing uncertainty.
• Role: Serves as the input and output for the dynamics model and the basis for the policy and value functions.

Though not part of the model per se, the policy (π) and value function (V) are trained exclusively using data generated by the world model through imagined rollouts. This is the key to sample efficiency.

• Process: The model is unrolled from a starting belief state, generating imagined trajectories of latent states, actions, and rewards. The policy is improved, and the value function is learned, via backpropagation through time (BPTT) on these synthetic trajectories.
• Algorithm Example: Dreamer uses an actor-critic framework where both actor (policy) and critic (value) are trained on latent imagination data.
• Advantage: Decouples costly environment interaction from policy optimization.

A robust world model must know what it doesn't know. Uncertainty estimation mechanisms quantify the model's confidence in its predictions, which is critical for safe planning and directed exploration.

• Epistemic Uncertainty: Uncertainty in the model itself due to limited data. Estimated via probabilistic ensembles (multiple models) or Bayesian Neural Networks (BNNs).
• Aleatoric Uncertainty: Inherent randomness in the environment.
• Use in Planning: Algorithms can perform pessimistic exploration by avoiding high-uncertainty states or using Model Predictive Control (MPC) with uncertainty-aware trajectory optimization.

Model-Based Reinforcement Learning (MBRL) is a paradigm where an agent learns an internal, predictive model of its environment's dynamics and reward function. This model is then used for planning and policy optimization, enabling more sample-efficient learning compared to model-free methods that learn purely from trial-and-error experience.

• Core Idea: Replace millions of real-world interactions with internal simulation.
• Key Trade-off: Balances the computational cost of model learning and planning against the reduced need for expensive environmental samples.
• Primary Use Case: Applications where real-world interaction is costly, risky, or slow, such as robotics, autonomous driving, and industrial control.

Dreamer is a seminal model-based RL algorithm that learns a latent dynamics model called a Recurrent State-Space Model (RSSM). Its key innovation is training the agent's policy and value function entirely through latent imagination—backpropagating gradients through sequences of imagined states and actions—without ever deploying the policy in the real environment during training.

• Architecture: Combines a variational autoencoder for perception, an RSSM for dynamics, and separate networks for policy and value.
• Training Loop: 1) Learn world model from real data, 2) Train actor and critic on imagined rollouts, 3) Deploy policy to collect new real data.
• Impact: Demonstrated that agents could learn complex behaviors from pixels using only a compact latent model for planning.

MuZero is a model-based RL algorithm that learns a value-equivalent model. Instead of predicting the true environment state, it learns a model that predicts future rewards, values, and policies—exactly the quantities needed for effective planning with a Monte Carlo Tree Search (MCTS). This abstraction allows it to master games like Go, Chess, and Shogi, and Atari, without any prior knowledge of the game rules.

• Model Outputs: Reward, policy (action probabilities), and value for a given state and action sequence.
• Planning: Uses MCTS guided by the learned model to choose actions.
• Significance: Shows that a perfect dynamics model is not necessary; a model accurate for planning is sufficient.

A Recurrent State-Space Model (RSSM) is a latent dynamics model architecture central to algorithms like Dreamer. It represents the world state as a combination of a deterministic recurrent state (for memory) and a stochastic latent variable (for uncertainty). This hybrid structure effectively models partially observable environments and temporal dependencies.

• Components:
  • Deterministic Path: An RNN (like a GRU) that accumulates information over time.
  • Stochastic Path: A latent variable sampled at each step to represent unpredictable aspects.
• Function: Predicts the next latent state and reward given the previous state and action.
• Benefit: Creates a compact, abstract state representation ideal for long-horizon imagination.

Model-policy co-adaptation is a critical failure mode in model-based RL where an agent's policy overfits to the biases and inaccuracies of its own learned world model. The policy learns to exploit the model's errors, producing behaviors that appear optimal in simulation but fail catastrophically when deployed in the real environment.

• Cause: The policy is trained exclusively on rollouts from an imperfect model, creating a feedback loop.
• Mitigation Strategies:
  • Uncertainty-Aware Planning: Using ensembles or Bayesian models to avoid uncertain states.
  • Regularization: Limiting the imagination horizon or blending real and imagined data.
  • Pessimistic Exploration: Penalizing actions that lead to states where the model is uncertain.

Model Predictive Control (MPC) is an online planning algorithm frequently used with learned world models. At each time step, MPC uses the model to simulate multiple potential action sequences over a finite planning horizon, selects the best sequence according to a reward/cost function, executes only the first action, and then replans from the new state. This closed-loop approach is robust to model inaccuracies.

• Online vs. Offline: MPC is an online planner; it does not learn a general policy but re-optimizes at every step.
• Advantage: Naturally compensates for model error by frequently re-grounding plans in real observations.
• Common Use: Robotics, process control, and any domain where conditions change rapidly.