Model-Based Reinforcement Learning (MBRL)

The dynamics model is the core predictive component of an MBRL system. It is a function, typically parameterized by a neural network, that approximates the environment's transition function: s_{t+1} = f_θ(s_t, a_t). Its purpose is to predict the next state (and often the reward) given the current state and action.

• Types: Can be deterministic (single prediction) or stochastic (predicts a distribution).
• Training: Learned from collected experience (s_t, a_t, s_{t+1}, r_t) tuples via supervised learning, minimizing prediction error.
• Challenge: Model bias—inaccuracies can compound during long rollouts, leading to catastrophic planning failures.

The planning algorithm uses the learned dynamics model to simulate potential future trajectories and select optimal actions. It is the mechanism that converts predictions into decisions.

• Online Planning: Performed at each timestep. Examples include Model-Predictive Control (MPC), which solves a short-horizon optimization problem on-the-fly.
• Offline Planning: Uses the model to generate synthetic data for training a separate policy network, as seen in algorithms like Dreamer.
• Core Methods: Includes shooting methods, cross-entropy method, and Monte Carlo Tree Search (MCTS) for discrete action spaces.

The policy (or controller) is the function that ultimately selects actions. In MBRL, it can take two primary forms:

• Model-Free Policy: A separate neural network trained on data generated by the model (e.g., via backpropagation through time). This amortizes the cost of planning.
• Model-Based Controller: The planning algorithm itself acts as the controller, as in MPC, where no separate policy network is stored.

The choice impacts the trade-off between computational latency (planning online) and sample efficiency (amortizing via a learned policy).

State representation refers to the format of the input s_t for the dynamics model. A good representation is compact, informative, and conducive to accurate long-horizon prediction.

• Raw Observations: High-dimensional pixels or sensor readings. Often requires learning a latent state, as in the Recurrent State-Space Model (RSSM).
• Abstract Latent States: Learned low-dimensional vectors that capture essential information, discarding irrelevant details. This is crucial for sample efficiency.
• Object-Centric or Disentangled Representations: Aim to factor the state into independent, semantically meaningful components (e.g., object position, velocity), improving generalization and interpretability.

Managing model uncertainty is critical for robust MBRL. Inaccurate models lead to poor plans. Systems estimate two types of uncertainty:

• Aleatoric Uncertainty: Inherent randomness in the environment. Modeled by predicting a distribution.
• Epistemic Uncertainty: Uncertainty due to limited data. A primary method is Ensemble Dynamics, training multiple models. Disagreement among the ensemble signals areas where the model is less certain.

This uncertainty is used for uncertainty-aware planning (e.g., avoiding uncertain states) or directed exploration (seeking out uncertain states to improve the model).

The experience replay buffer stores past interactions (s, a, s', r). Its management is pivotal for learning both the model and the policy.

• Model Learning: The buffer provides the supervised learning dataset for the dynamics model.
• Policy Learning: For algorithms that train a policy, the buffer may also store synthetic data generated by model rollouts.
• Strategic Sampling: Techniques like prioritized experience replay can focus on rare or high-error transitions to improve model accuracy in critical regions. Balancing real and model-generated data is key to preventing model exploitation.

The Dreamer agent is a seminal MBRL algorithm that learns a Recurrent State-Space Model (RSSM) from image inputs. Its architecture combines:

• A deterministic recurrent path to track temporal dependencies.
• A stochastic latent variable model to capture uncertainty and partial observability.
• A latent imagination core where policy and value networks are trained entirely on compact state sequences predicted by the world model, achieving high sample efficiency. It demonstrates that effective policies can be learned purely from simulated experience within a learned latent space.

MuZero is a model-based algorithm that learns an implicit model of the environment's dynamics without being given the rules. It jointly learns:

• A representation function to encode observations into a hidden state.
• A dynamics function that predicts the next hidden state and immediate reward given a state and action.
• Policy and value functions to guide a Monte Carlo Tree Search (MCTS).

Unlike explicit forward models, MuZero's learned dynamics are tailored for accurate value and policy prediction, enabling superhuman performance in board games, Atari, and Go. It exemplifies value-equivalent model learning.

Model-Predictive Control (MPC) is a foundational online planning method widely used in robotics and process control. At each timestep, MPC:

1. Uses a dynamics model (learned or analytical) to simulate multiple potential action sequences over a finite horizon.
2. Evaluates these trajectories using a cost function.
3. Executes only the first action of the optimal sequence.
4. Repeats the process from the new observed state.

This receding horizon control approach is robust to model inaccuracies and disturbances. In MBRL, the dynamics model is often a learned neural network, and planning is performed via random shooting or gradient-based optimization.

The Imagination-Augmented Agent (I2A) architecture enhances a standard model-free agent with 'imagination' from a learned world model. Its key components are:

• An imagination core that runs multiple rollouts from the current state using the learned model.
• An imagination encoder that aggregates these rollout trajectories into a context vector.
• A model-free backbone (e.g., an A3C agent) whose policy and value networks receive the current observation and the imagination context.

This hybrid approach allows the agent to supplement direct experience with simulated foresight, improving sample efficiency and planning capabilities on complex tasks like Sokoban.

Probabilistic Ensembles with Trajectory Sampling (PETS) is an algorithm that addresses model uncertainty in MBRL. Its core principles are:

• Ensemble Dynamics: Learning an ensemble of neural network dynamics models to capture epistemic uncertainty (uncertainty due to lack of data).
• Probabilistic Prediction: Each model outputs a probability distribution over next states and rewards.
• Uncertainty-Aware Planning: Trajectories are sampled from the ensemble, and planning (e.g., via Cross-Entropy Method) can incorporate uncertainty, often leading to pessimistic or risk-averse behavior in uncertain regions of the state space. This makes it particularly effective for safe exploration and offline RL.

This architectural pattern separates the world model learning from the controller optimization, a concept popularized by Ha & Schmidhuber. The system has two main components:

• A Vision Model (V) and Memory Model (M) that compress high-dimensional inputs (like images) into a compact latent state representation and learn to predict future latent states.
• A simple Controller (C), often a linear or small neural network, that is trained via evolutionary strategies or reinforcement learning to maximize expected reward using rollouts performed entirely inside the learned world model.

This decoupling allows for fast, low-cost controller training in simulation and demonstrates the power of a accurate, compressed world model as a substrate for planning.

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 the core predictive engine in MBRL, allowing for planning and rollout simulations. Unlike a simple dynamics model, a world model often aims to capture a more abstract, compact, and actionable representation of the agent's experience.

• Function: Enables mental simulation of possible futures.
• Architectures: Include Recurrent State-Space Models (RSSM) and variational autoencoders.
• Use Case: The Dreamer agent uses a learned world model from pixels to plan via latent imagination.

Model-Predictive Control (MPC) is an online planning and control method that repeatedly solves a finite-horizon optimization problem using a dynamics model. At each step, it plans a sequence of actions, executes the first one, and then re-plans, making it a cornerstone technique for applying learned models in MBRL to real-time systems.

• Core Loop: Plan → Execute first action → Re-plan with new state.
• Model Usage: Can use either analytical/physics-based models or learned neural network models.
• Application: Widely used in robotics, process control, and autonomous driving for its robustness to model inaccuracies.

Monte Carlo Tree Search (MCTS) is a heuristic search algorithm that builds a look-ahead tree by recursively simulating trajectories (rollouts) to evaluate the potential of different action sequences. It is a powerful planning algorithm often paired with a model in MBRL to select optimal actions.

• Four Steps: Selection, Expansion, Simulation, Backpropagation.
• Model Role: The dynamics model is used to generate the simulated rollouts during the expansion and simulation steps.
• Famous Example: AlphaGo and MuZero use MCTS combined with learned models to achieve superhuman performance in games.

Ensemble dynamics refers to training multiple (e.g., 5-10) neural network models to predict environment dynamics. The disagreement, or epistemic uncertainty, among the ensemble members is a critical signal in MBRL for managing the reliability of the learned model.

• Primary Use: Uncertainty quantification. High variance in predictions indicates areas where the model has not been trained sufficiently.
• Algorithmic Applications: Used in PETS (Probabilistic Ensembles with Trajectory Sampling) and other MBRL methods to guide exploration (favor uncertain states) or enable risk-averse planning (avoid uncertain states).

Model-Based Value Expansion (MBVE) is a hybrid technique that blends model-free and model-based RL. It uses short-horizon rollouts from a learned dynamics model to generate more accurate target values (Q-targets) for training a model-free Q-function or value function.

• Mechanism: Instead of a 1-step TD target, it computes an n-step target using model-generated transitions.
• Benefit: Reduces the bias in Q-learning targets, leading to more stable and sample-efficient learning.
• Distinction: Unlike pure planning methods, MBVE uses the model only to improve the targets for a model-free critic, not for direct action selection.

The sim2real gap is the performance discrepancy between a system trained in simulation and its performance when deployed in the real world. This is a major challenge for MBRL, as inaccuracies in the learned or analytical model are magnified during planning, leading to failures.

• Causes: Unmodeled physics, sensor noise, actuator delays, and visual domain differences.
• MBRL Impact: A poor model leads to invalid planning and compounding errors over long rollouts.
• Mitigation Strategies: Domain randomization, system identification, online model adaptation, and the use of differentiable physics simulators for gradient-based policy learning.