World Model

The encoder compresses high-dimensional sensory inputs (e.g., pixels from images) into a low-dimensional latent state representation (z_t). This serves as the agent's internal belief about the current world state.

• Function: Maps observations → latent states.
• Purpose: Dimensionality reduction and feature extraction, creating a compact, information-rich code for planning.
• Example: A convolutional neural network (CNN) that processes an 84x84 game frame into a 64-dimensional vector.

The dynamics model predicts the next latent state (z_{t+1}) and immediate reward (r_t) given the current latent state and a proposed action (a_t). It learns the environment's rules.

• Function: Models p(z_{t+1}, r_t | z_t, a_t).
• Purpose: Enables rollout or imagination of possible futures without interacting with the real environment.
• Critical Role: Forms the core of the agent's predictive understanding, allowing it to simulate consequences of actions.

The decoder reconstructs the original observation from the latent state. It is often used during training to ensure the latent space retains meaningful information but may be omitted at inference time for efficiency.

• Function: Maps latent states → reconstructed observations.
• Purpose: Provides a reconstruction loss that forces the encoder to learn a useful, information-preserving representation.
• Note: In some architectures (e.g., Dreamer), the decoder is only used for training the encoder.

The latent space is the compressed, abstract representation where the world model operates. It is the cornerstone of the model's efficiency and generalization.

• Characteristics: Lower-dimensional, continuous, and structured.
• Benefits:
  • Enables fast planning by operating on small vectors.
  • Can learn to ignore irrelevant sensory details (e.g., background textures).
  • Often exhibits smooth transitions, where similar latent states lead to similar futures.

While sometimes integrated into the dynamics model, the reward predictor is a distinct component that forecasts the scalar reward signal. Accurate reward prediction is essential for evaluating the quality of imagined trajectories.

• Function: Estimates r_t = f(z_t, a_t).
• Purpose: Provides the reinforcement signal for planning within the model.
• Challenge: Often harder to predict accurately than state transitions, as rewards can be sparse and discontinuous.

The controller or policy (π) is the agent's decision-making module. In world model architectures, it is typically trained entirely within the latent space using trajectories generated by the dynamics model.

• Training Paradigm: The policy proposes actions, the dynamics model simulates outcomes, and a latent critic (value function) evaluates them. Gradients are backpropagated through the learned model to the policy.
• Key Advantage: This latent imagination allows for massive, cost-effective trial-and-error without real-world interaction.

Model-Based Reinforcement Learning is a paradigm where an agent learns an explicit model of the environment's dynamics (the transition function) and reward function. This learned model is then used for planning and policy improvement, either by generating synthetic rollouts for training a policy or by searching for optimal action sequences. MBRL contrasts with model-free RL, which learns a policy or value function directly from experience without an internal model. Key advantages include improved sample efficiency and the potential for safe exploration in a simulated environment.

Dreamer is a seminal model-based reinforcement learning agent that learns a world model from high-dimensional image inputs. Its architecture consists of three components:

• A representation model that encodes images into a compact latent state.
• A dynamics model that predicts future latent states.
• A reward model that predicts expected rewards. The agent's actor-critic policy is trained entirely within this learned latent imagination (latent dreaming), where trajectories are generated by the world model. Experiences from the real environment are stored in a replay buffer to train the world model, creating a tight integration with experience replay mechanisms.

MuZero is a model-based reinforcement learning algorithm that learns a model implicitly, solely for the purpose of planning. Unlike traditional world models, it does not explicitly predict future observations or environment states. Instead, it learns:

• A representation function to map observations to hidden states.
• A dynamics function to predict the next hidden state and immediate reward.
• A prediction function to estimate the policy and value from a hidden state. It uses a Monte Carlo Tree Search (MCTS) planner with this learned model to choose actions. All components are trained via gradient descent using data stored in a large replay buffer, making it a masterclass in combining learned models with planning and experience replay.

A latent state space is a compressed, abstract representation of the environment learned by a model (like a world model's encoder). This space is designed to contain all information necessary for predicting future observations and rewards while discarding irrelevant details. Working in a latent space provides significant computational and learning benefits:

• Dimensionality Reduction: Planning and prediction occur in a smaller, denser space than raw pixels.
• Temporal Abstraction: The latent state aims to be Markovian, where the current state contains all relevant history.
• Improved Generalization: The model learns invariances and key features. World models like those in Dreamer perform all future prediction and policy training within this learned latent space.

In the context of world models, planning refers to the process of using the internal model to simulate sequences of future states and actions before committing to a real action. The agent searches for action sequences that maximize predicted cumulative reward. Common planning methods include:

• Random Shooting: Evaluating many randomly generated action sequences.
• Cross-Entropy Method (CEM): Iteratively refining a distribution over promising action sequences.
• Monte Carlo Tree Search (MCTS): A heuristic search algorithm that builds a look-ahead tree, used by MuZero. Planning with a world model allows for deliberative behavior and can solve tasks with long-term dependencies more effectively than reactive, model-free policies.

Model Predictive Control is a planning and control strategy frequently used with learned world models. At each timestep, MPC:

1. Uses the current world model to plan an optimal sequence of actions over a finite planning horizon.
2. Executes only the first action from the planned sequence.
3. Observes the new state, re-plans from scratch, and repeats. This receding horizon control makes MPC robust to model inaccuracies, as the plan is constantly re-evaluated. It is a primary method for converting a passive world model (which only predicts) into an active controller. MPC is widely used in robotics and process control, where safety and real-time adjustment are critical.