Model-Based Policy Optimization (MBPO)

The dynamics model (or transition model) is a neural network, typically an ensemble of probabilistic networks, trained to predict the next state and reward given the current state and action: s', r = f_θ(s, a). This model serves as a learned simulator of the environment. MBPO uses short, fixed-length imagined rollouts from this model, starting from states sampled from a real experience buffer, to generate synthetic training data. The use of an ensemble helps with uncertainty quantification; high disagreement among ensemble members signals areas where the model is unreliable, which informs rollout horizon limits.

To mitigate compounding error—where inaccuracies in the dynamics model lead to increasingly unrealistic states over long rollouts—MBPO strictly limits the length of imagined trajectories. The algorithm uses a fixed rollout horizon (e.g., 1-5 steps). This is a critical hyperparameter: too short, and the model provides limited benefit; too long, and the policy trains on erroneous synthetic data, leading to model-policy co-adaptation. The policy is trained on a mixture of these short synthetic rollouts and real data, which provides a robust training signal while preventing exploitation of model flaws.

Unlike planning algorithms like Model Predictive Control (MPC) that re-plan online, MBPO uses its model solely for data augmentation. The synthetic rollouts are added to a replay buffer. A standard model-free RL algorithm, such as Soft Actor-Critic (SAC) or Proximal Policy Optimization (PPO), is then used to train the policy from this combined buffer. This decoupling allows MBPO to leverage the sample efficiency of model-based data generation while retaining the strong asymptotic performance and stability guarantees of modern model-free algorithms.

A core challenge is preventing the policy from exploiting model bias—systematic errors in the learned dynamics. MBPO employs several strategies:

• Ensemble-Based Uncertainty: Using an ensemble of models and measuring their disagreement.
• Conservative Rollout Horizon: The short horizon acts as a regularizer against compounding error.
• Data Mixing: Training on a blend of real and synthetic data prevents the policy from drifting into regions of state space where the model is catastrophically wrong. This approach is less conservative than pessimistic exploration used in offline RL but is designed for the online setting where the agent can continually gather new real data to correct the model.

MBPO operates via a parallel, asynchronous loop between three processes:

1. Real Data Collection: The current policy interacts with the real environment, storing transitions (s, a, r, s') in a real replay buffer.
2. Model Training: The dynamics model ensemble is periodically retrained on all real data collected so far.
3. Policy Training via Imagination: In parallel, the policy and value networks are updated using batches of data sampled from a combined replay buffer that contains both real data and short synthetic rollouts generated on-demand from the latest model. This loop maximizes hardware utilization (e.g., using GPUs for policy/model training while CPUs collect environment samples).

MBPO vs. Pure Planning (e.g., MPC): MBPO learns a general policy, while MPC solves for optimal actions online at every step. MBPO is more computationally efficient at deployment.

MBPO vs. Latent Imagination (e.g., Dreamer): Algorithms like Dreamer train the policy via backpropagation through time (BPTT) on latent rollouts. MBPO uses simpler, model-free policy gradients on decoded state rollouts, which can be more stable and easier to implement.

MBPO vs. Value-Equivalent Models (e.g., MuZero): MuZero learns a model that predicts future rewards, values, and policies, not accurate state transitions. MBPO's model aims for accurate state prediction, making its synthetic data usable by any model-free algorithm.

Model-Based Reinforcement Learning (MBRL) is a paradigm where an agent learns an internal model of its environment's dynamics and reward function. This model is then used for planning and policy optimization, with the primary goal of improving sample efficiency compared to model-free methods. The core challenge is managing model error to prevent the policy from exploiting model inaccuracies.

A world model is an agent's internal, learned representation that predicts future states and rewards based on current states and actions. It enables planning and imagination without direct, costly interaction with the real environment. In MBPO, the learned dynamics model acts as a world model for generating short, synthetic rollouts to train the policy.

• Function: Compresses experience into a predictive representation.
• Architecture: Often implemented as a Recurrent State-Space Model (RSSM) or a probabilistic ensemble of neural networks.

Model Predictive Control (MPC) is an online planning algorithm that repeatedly solves a finite-horizon optimal control problem using a learned model, executing only the first planned action before re-planning from the new state. It contrasts with MBPO, which uses the model to generate data for training a separate, amortized policy network.

• MPC: Plans from scratch at each timestep (high online compute).
• MBPO: Trains a policy network offline using model data (amortizes compute into training).

Dreamer is a leading model-based RL algorithm that learns a latent dynamics model (an RSSM) and uses it to train policies and value functions entirely via latent imagination—backpropagation through time on imagined rollouts. Unlike MBPO, which uses model rollouts to create a dataset for a model-free algorithm like SAC, Dreamer optimizes the policy directly through gradients propagated through the model.

Model-based offline reinforcement learning involves learning a dynamics model from a static, pre-collected dataset without any online interaction. The model is then used to train a policy via planning or synthetic data generation. MBPO can be adapted for offline settings, but requires techniques like pessimistic exploration to avoid exploiting model errors in unseen state-action regions, which is a critical challenge distinct from the online MBPO setting.

Uncertainty quantification is the process of estimating the epistemic (model) and aleatoric (environmental stochasticity) uncertainty in a learned dynamics model's predictions. It is critical for robust planning and safe exploration in MBRL.

• Techniques: Bayesian Neural Networks (BNNs), probabilistic ensembles, and bootstrap methods.
• Role in MBPO: MBPO uses short rollouts to limit the impact of compounding error, which arises when model inaccuracies accumulate over long simulations. Advanced variants explicitly quantify uncertainty to dynamically adjust rollout horizons.