Lecture 11

Model-Based RL

Learning dynamics models for planning: world models, model-predictive control, Dyna-style methods, and model-based policy optimization.

World Models MPC Dyna Learned Dynamics
Original PDF slides

Why Model-Based RL?

Throughout this course, we have studied a rich landscape of reinforcement learning algorithms—from imitation learning and offline RL to on-policy and off-policy methods like SAC, DQN, and PPO. All of the RL methods we have seen so far are model-free: they learn a policy or value function directly from interaction data without ever attempting to understand how the environment works. Model-based reinforcement learning takes a fundamentally different approach by asking: can we learn a predictive model of the environment—a "learned simulator"—and then use it to make better decisions?

The central motivation is sample efficiency. Model-free methods, particularly on-policy algorithms like REINFORCE, require enormous amounts of environment interaction. In domains like robotics, where every real-world sample involves physical wear on hardware, or in medicine, where each sample corresponds to a patient outcome, this data hunger is a critical bottleneck. If we can learn an accurate dynamics model from a small amount of data, we can simulate many trajectories internally and dramatically reduce the number of real-world interactions needed.

Definition
Model-Based Reinforcement Learning. A family of RL methods that explicitly learn a dynamics model $\hat{f}(\mathbf{s}, \mathbf{a}) \approx \mathbf{s}'$ (or more generally, $p_\theta(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$) of the environment and use it to improve decision-making, either through planning, synthetic data generation, or both.

The idea of learning a simulator is broadly applicable. In robotics and physical systems, the model captures the physics of the system—predicting how joint angles, object positions, and contact forces evolve over time. In finance, a model might predict market dynamics conditioned on trading actions. In games, the model encodes the rules of the game (though in some games like chess, the rules are already known and need not be learned). Modern video generation models like Sora and Veo 2 can even be viewed as learned simulators that predict future visual observations conditioned on actions.

Key Insight
Not all domains require learning a model. In games like chess or Go, the transition dynamics are perfectly known (the rules of the game). In such cases, model-based methods can use the known model directly for planning (as in AlphaZero's Monte Carlo Tree Search) without needing to learn anything about the dynamics. The need to learn a model arises when the dynamics are unknown or too complex to specify analytically.

The Learned Simulator Algorithm

The most basic model-based RL approach follows a conceptually simple three-step recipe. We collect transition data from the environment, fit a neural network to predict next states, and then use this learned model as a stand-in for the real environment.

Notional Learned Simulator Algorithm
  1. Collect data. Run some base policy $\pi_0(\mathbf{a} \mid \mathbf{s})$ (e.g., a random policy) to collect a dataset $\mathcal{D} = \{(\mathbf{s}_i, \mathbf{a}_i, \mathbf{s}'_i)\}$.
  2. Learn a "simulator." Train a model $f_\phi(\mathbf{s}, \mathbf{a})$ to predict the next state by minimizing the mean squared error: $$\min_\phi \sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$$
  3. Optimize inside the learned simulator. Run your favorite RL algorithm or planning method using $f_\phi$ as the environment dynamics.

This algorithm is appealingly simple, but several critical challenges arise in practice. Understanding these challenges is essential for making model-based RL work reliably.

Key Challenges

Three fundamental issues can cause this naive approach to fail:

Example
The cliff analogy. Imagine a drone navigating along a cliff edge. A base policy (e.g., random exploration) might only generate trajectories in a safe area away from the cliff. The learned model has no data about what happens near the edge, so it might predict that going right (toward the cliff) leads to gentle, safe terrain. An RL agent optimizing in this model might confidently plan a trajectory right off the cliff. The data distribution of $\pi_0$ does not match the distribution induced by the final policy $\pi_f$: $p_{\pi_0}(\mathbf{s}) \neq p_{\pi_f}(\mathbf{s})$.

Learning Dynamics Models

Before discussing how to use a learned model, we briefly survey the spectrum of approaches for obtaining one. The right approach depends heavily on how much prior knowledge is available about the domain.

Known Dynamics

In some domains, the dynamics are fully specified. Board games like chess and Go have exact, deterministic rules. Certain card games and combinatorial optimization problems also fall into this category. When the model is known, we can skip the learning step entirely and proceed directly to planning.

Partially Known Dynamics

In many physical systems, we know the general form of the dynamics (e.g., Newton's laws, rigid-body mechanics) but not all the parameters (e.g., friction coefficients, mass distributions, damping constants). In these cases, we can use the known structure and fit the unknown parameters from data. This approach, sometimes called system identification, is common in classical control and robotics. It typically requires far less data than learning dynamics from scratch because the model class is constrained to physically plausible dynamics.

Unknown Dynamics: End-to-End Learning

In the most general—and most common—setting, we have no prior knowledge of the dynamics and must learn the model entirely from data. Two main strategies exist:

Key Insight
In addition to learning the dynamics, model-based RL often requires learning a reward model as well. If the reward function is not known analytically (which is common when rewards depend on complex sensory inputs), the agent must learn to predict rewards from $(s, a)$ pairs alongside learning the transition dynamics.

Planning with Learned Models

Once we have a learned dynamics model $f_\phi(\mathbf{s}, \mathbf{a})$, the most direct way to use it is for planning—searching over action sequences to find ones that maximize cumulative reward over a finite horizon $H$. The planning problem can be stated as:

$$\max_{\mathbf{a}_{t:t+H}} \sum_{t'=t}^{t+H} r(\mathbf{s}_{t'}, \mathbf{a}_{t'})$$

where each next state is obtained from the model: $\mathbf{s}_{t'+1} = f_\phi(\mathbf{s}_{t'}, \mathbf{a}_{t'})$. This optimization can be performed using either gradient-based or sampling-based methods.

Approach 1a: Gradient-Based Planning (Backpropagation)

If the learned model $f_\phi$ is a differentiable neural network and the reward function $r$ is also differentiable, we can backpropagate through the model to optimize the action sequence. Concretely, we treat the actions $\mathbf{a}_t, \mathbf{a}_{t+1}, \ldots, \mathbf{a}_{t+H-1}$ as free parameters, compute the total reward by unrolling the model forward, and then take gradient ascent steps on the actions:

$$\nabla_{\mathbf{a}_{t:t+H}} \sum_{t'=t}^{t+H} r\bigl(f_\phi(\mathbf{s}_{t'}, \mathbf{a}_{t'}),\; \mathbf{a}_{t'}\bigr)$$
Planning via Backpropagation
  1. Run some policy (e.g., random) to collect data $\mathcal{D} = \{(\mathbf{s}, \mathbf{a}, \mathbf{s}')_i\}$.
  2. Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
  3. Backpropagate through $f_\phi(\mathbf{s}, \mathbf{a})$ to optimize the action sequence via gradient ascent on cumulative reward.

Gradient-based planning is scalable to high-dimensional action spaces and works especially well in overparameterized regimes where the optimization landscape is smooth. However, it requires a differentiable model and reward, and it can struggle with rugged loss landscapes containing many local optima.

Approach 1b: Sampling-Based Planning

An alternative is to use sampling-based (zeroth-order) optimization, which requires no gradient information. These methods evaluate the objective for many candidate action sequences and iteratively refine the search distribution.

Random Shooting

The simplest sampling-based method is random shooting. Denote an action sequence as $\mathbf{A} := \mathbf{a}_t, \ldots, \mathbf{a}_{t+H}$.

  1. Sample many candidate sequences $\mathbf{A}_1, \ldots, \mathbf{A}_N$ from some distribution (e.g., uniform over the action space).
  2. For each candidate, roll out the model to compute $J(\mathbf{A}_i) = \sum_{t'=t}^{t+H} r(\mathbf{s}_{t'}, \mathbf{a}_{t'})$.
  3. Return $\mathbf{A}_i$ with the highest $J(\mathbf{A}_i)$.

Random shooting is extremely simple, but the quality of the solution depends entirely on the coverage of the random samples. In high-dimensional action spaces or long horizons, the probability of sampling a good sequence by chance becomes vanishingly small. Can we do better by iteratively improving the sampling distribution?

The Cross-Entropy Method (CEM)

The cross-entropy method (not to be confused with the cross-entropy loss in classification) is an iterative refinement of random shooting that progressively concentrates the sampling distribution around high-performing action sequences.

Cross-Entropy Method (CEM) for Action Optimization

Initialize a sampling distribution $p(\mathbf{A})$ (e.g., Gaussian with broad variance). Repeat:

  1. Sample many action sequences $\mathbf{A}_1, \ldots, \mathbf{A}_N$ from $p(\mathbf{A})$.
  2. Evaluate each: $J(\mathbf{A}_i) = \sum_{t'=t}^{t+H} r(\mathbf{s}_{t'}, \mathbf{a}_{t'})$.
  3. Select elites: pick the top-$M$ sequences $\mathbf{A}_{i_1}, \ldots, \mathbf{A}_{i_M}$ with the largest $J(\mathbf{A}_i)$, where $M < N$.
  4. Refit the distribution $p(\mathbf{A})$ to the elites (e.g., fit a Gaussian to $\mathbf{A}_{i_1}, \ldots, \mathbf{A}_{i_M}$).

After several iterations, return the best action sequence found.

Gradient-Based vs. Sampling-Based: Trade-offs

These two families of optimization methods have complementary strengths:

Gradient-Based (1st order)Sampling-Based (0th order)
StrengthsScales to high dimensions; works well in overparameterized regimesHighly parallelizable; requires no gradient information; simple to implement
WeaknessesRequires differentiable model and smooth optimization landscapeScales poorly to high dimensions (both horizon $H$ and action dimension $|\mathbf{a}|$)

Distribution Mismatch and Online Model Learning

The basic planning algorithm described above can fail catastrophically due to distribution mismatch. The model is trained on data from the base policy $\pi_0$, but the planner finds an action sequence that visits states far outside this training distribution. In these out-of-distribution regions, the model's predictions are unreliable, and the planner may exploit these errors.

Formally, the state distribution under the base policy $p_{\pi_0}(\mathbf{s})$ differs from the state distribution under the final planned policy $p_{\pi_f}(\mathbf{s})$:

$$p_{\pi_0}(\mathbf{s}) \neq p_{\pi_f}(\mathbf{s})$$

This is exactly the same distribution shift problem that arises in behavioral cloning (addressed by DAgger in Lecture 2). The solution is analogous: iteratively collect new data under the current planner, retrain the model, and replan.

Planning with Online Model Learning
  1. Run some policy (e.g., random) to collect initial data $\mathcal{D} = \{(\mathbf{s}, \mathbf{a}, \mathbf{s}')_i\}$.
  2. Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
  3. Use the model to plan an action sequence (via backpropagation or CEM).
  4. Execute the planned actions in the real environment, appending the resulting tuples $(\mathbf{s}, \mathbf{a}, \mathbf{s}')$ to $\mathcal{D}$.
  5. Go to step 2.

By alternating between planning and data collection, the model gradually becomes accurate in the regions of state space that the planner actually visits. This is conceptually identical to DAgger's iterative data aggregation strategy, but applied to dynamics learning rather than policy cloning.

Example
Revisiting the cliff. With online model learning, the agent initially plans conservatively (the model is uncertain about the cliff region). After executing its plan and collecting new data near the cliff edge, it learns that going too far right leads to falling. Over subsequent iterations, the model becomes accurate in the relevant regions, and the final policy learns to go to the top of the cliff and stop—achieving the optimal behavior.

Model-Predictive Control (MPC)

The planning approaches described so far are open-loop: the planner computes an entire action sequence $\mathbf{a}_{1:H}$ upfront and then executes it without feedback. This is fragile. If the model is even slightly inaccurate, errors accumulate over the horizon, and the executed trajectory diverges from the plan. In stochastic environments, this problem is even worse—the plan cannot adapt to unexpected state transitions.

The solution is closed-loop planning, where the agent observes the actual state after each action and replans. This is precisely the idea behind model-predictive control (MPC), a cornerstone technique from classical control theory that adapts naturally to the learned-model setting.

Model-Predictive Control (MPC)
  1. Run a base policy $\pi_0(\mathbf{a}_t \mid \mathbf{s}_t)$ (e.g., random) to collect initial dataset $\mathcal{D} = \{(\mathbf{s}, \mathbf{a}, \mathbf{s}')_i\}$.
  2. Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
  3. Use model $f_\phi(\mathbf{s}, \mathbf{a})$ to plan an optimal action sequence over a horizon $H$.
  4. Execute only the first planned action $\mathbf{a}_t$, observe the resulting state $\mathbf{s}'$.
  5. Append $(\mathbf{s}, \mathbf{a}, \mathbf{s}')$ to dataset $\mathcal{D}$.
  6. Go to step 2 (periodically retrain model) or step 3 (replan from current state).

The key insight of MPC is that by replanning at every time step (or every few steps), the agent can correct for model errors in real time. Even if the model's predictions for $\mathbf{s}_{t+5}$ are inaccurate, the agent only commits to $\mathbf{a}_t$ and then replans from the actual state $\mathbf{s}_{t+1}$. This makes MPC remarkably robust to model inaccuracies.

Key Insight
MPC trades compute for robustness. By replanning at every step, it avoids committing to a long sequence of actions based on potentially inaccurate long-horizon predictions. The downside is that planning is computationally expensive, and it must be performed at every time step during execution—not just during training.

Summary of Planning with Learned Models

The planning-based approach to model-based RL has several appealing properties and limitations:

AdvantagesDisadvantages
Simple and modular—any model can be paired with any planner Compute-intensive at test time (must plan at every step)
Easy to plug in different goals or rewards, even at test time Only practical for short-horizon problems or very shaped reward functions
MPC naturally corrects for model errors Model accuracy degrades over long horizons

The short-horizon limitation arises for two reasons: (a) planning over long horizons is computationally expensive (the search space grows exponentially with $H$), and (b) model prediction errors compound over many steps, making long-horizon model predictions unreliable. This naturally raises the question: can we train a policy using a learned model, so that we do not need to plan at test time?

Model-Based Policy Optimization

Rather than using the model purely for planning at test time, an alternative is to use it to generate training data for a model-free RL algorithm. This combines the sample efficiency of model-based methods with the flexibility and long-horizon capability of model-free policy learning.

Option 1: Distilling the Planner into a Policy

The simplest approach is to run the MPC planner to generate actions across many states, then train a policy network $\pi_\phi(\mathbf{a} \mid \mathbf{s})$ to imitate the planner's actions via supervised learning. This is essentially behavioral cloning where the "expert" is the planner itself.

To solve longer-horizon problems, we need a different strategy. Two key ideas are: (1) plan with a terminal value function that estimates long-term reward beyond the planning horizon, and (2) augment model-free RL methods with data generated from the model. The second approach leads to the family of algorithms known as model-based policy optimization.

Option 2: Data Augmentation via Model Rollouts

The central idea of model-based policy optimization (MBPO) is to use the learned model to generate synthetic transition data that augments the real environment data. The model-free RL algorithm (e.g., SAC) then trains on both real and synthetic data, gaining the sample efficiency of model-based methods without being limited to short planning horizons.

A critical design question is how to generate the synthetic data. Three natural strategies reveal increasing sophistication:

  1. Full trajectories from initial states. Reset to an initial state and roll out the model for many steps. The problem: the model may not be accurate over long horizons, and the resulting trajectories may be unrealistic.
  2. Short trajectories from initial states. Generate partial rollouts from initial states to stay within the model's accuracy window. The problem: later states in the trajectory are underrepresented, leading to poor coverage.
  3. Short trajectories ("branched rollouts") from all states in the real data. For each state visited in real environment interactions, generate a short model rollout of $k$ steps using the current policy. This provides good coverage across the entire trajectory while keeping rollouts short enough for the model to remain accurate.

The third strategy is the key insight behind MBPO (Janner et al., 2019). By branching short rollouts from every state in the real dataset, the method achieves broad coverage without requiring the model to be accurate over long horizons.

Model-Based Policy Optimization (MBPO)

Repeat:

  1. Collect real data: Run current policy $\pi_\phi$ in the real environment, add transitions to $\mathcal{D}_{\text{env}}$.
  2. Update model: Train dynamics model $p_\theta(\mathbf{s}' \mid \mathbf{s}, \mathbf{a})$ using $\mathcal{D}_{\text{env}}$.
  3. Generate synthetic data: For states $\mathbf{s}$ sampled from $\mathcal{D}_{\text{env}}$, roll out $\pi_\phi$ in the model $p_\theta$ for $k$ steps. Add the resulting transitions to $\mathcal{D}_{\text{model}}$.
  4. Update policy: Perform model-free RL updates (e.g., SAC) on policy $\pi_\phi$ and critic $Q$ using data from $\mathcal{D}_{\text{model}}$ (and optionally $\mathcal{D}_{\text{env}}$).
Key Insight
MBPO is compatible with any model-free RL algorithm in step 4. The model serves purely as a data-generation engine—a way to amplify a small amount of real experience into a much larger synthetic dataset. The rollout horizon $k$ is a tunable hyperparameter that controls the trade-off between data quantity (longer rollouts produce more data) and data quality (shorter rollouts are more accurate).
Why branched rollouts help with coverage

Consider a real trajectory $\mathbf{s}_1 \to \mathbf{s}_2 \to \cdots \to \mathbf{s}_T$. If we only generate synthetic rollouts from the initial state $\mathbf{s}_1$, the model must be accurate for up to $T$ steps to produce useful training data for the later parts of the trajectory. With branched rollouts, we generate short $k$-step rollouts from every state $\mathbf{s}_t$ in the trajectory. This means the policy receives training signal for states near $\mathbf{s}_T$ without the model ever needing to predict more than $k$ steps ahead. The model error in each synthetic transition is bounded by the single-step prediction error, and the coverage of the synthetic data mirrors the coverage of the real data, modulo $k$ steps of model drift.

When to Use Model-Based RL

Model-based RL offers significant advantages but also introduces new challenges. The decision of whether to use a model depends critically on the problem domain.

Advantages

Disadvantages

Key Insight
Whether to use a model depends on how hard it is to learn. In domains with low-dimensional state spaces and smooth dynamics (e.g., robotic control with joint-angle states), model-based methods often shine. In domains with high-dimensional observations and complex, discontinuous dynamics (e.g., contact-rich manipulation from pixels), model-free methods may be more practical. The best choice is always domain-dependent.

Other Kinds of Models

The standard forward dynamics model $p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$ is not the only useful type of model. Several alternative formulations serve different purposes:

Each of these model types captures different aspects of the environment structure and can be useful for specific downstream tasks, from goal-conditioned planning to representation learning.

Case Study: Dexterous Manipulation with PDDM

To make the ideas of model-based RL concrete, we examine a landmark case study: Deep Dynamics Models for Learning Dexterous Manipulation (Nagabandi, Konoglie, Levine, and Kumar; Google Brain, 2019). This work remains one of the most impressive demonstrations of model-based RL on five-fingered dexterous manipulation and provides valuable ablation studies on what design choices matter most.

Problem Setup

The task is to control a five-fingered robotic hand (the ShadowHand, with 24 degrees of freedom) to manipulate objects—including rotating Baoding balls and tracking a pencil along a target trajectory. The problem setup is:

Method: PDDM

The proposed method, called PDDM (Planning with Deep Dynamics Models), uses a model-based planning approach with several carefully chosen design decisions:

Results and Comparisons

The authors compared PDDM against both model-free and model-based baselines in simulation:

The key findings: model-based methods were more sample-efficient than model-free methods, and PDDM was more performant than other model-based methods on these dexterous tasks.

Ablation Studies

The ablation studies reveal several important practical lessons for model-based RL:

Real-World Results

Perhaps the most compelling result was the transfer to a real ShadowHand robot performing Baoding ball rotation. Sample efficiency was critical here: the hardware is fragile and expensive, making data collection costly. PDDM learned the Baoding ball rotation task in approximately 4 hours of real-world interaction—a feat that would have been impractical with model-free methods. The ball was reset using a separate robot arm to enable autonomous training.

Example
Why sample efficiency matters in the real world. The ShadowHand costs over $100,000 and has delicate tendons that wear with use. Model-free methods like SAC might require millions of time steps (days of continuous operation) to learn dexterous tasks, during which the hardware could degrade or break. By learning from just a few hours of data, model-based methods make such experiments feasible on real hardware.

Summary and Looking Ahead

Model-based reinforcement learning offers a compelling alternative to model-free methods when sample efficiency is important and the dynamics are learnable. The key ideas covered in this lecture form a coherent progression:

  1. Learn a dynamics model $f_\phi(\mathbf{s}, \mathbf{a})$ from environment interaction data, using supervised learning (MSE regression).
  2. Use the model for planning—optimizing action sequences via gradient-based (backpropagation) or sampling-based (CEM) methods.
  3. Address distribution mismatch by iteratively collecting data under the current plan and retraining the model (analogous to DAgger).
  4. Use model-predictive control (MPC) to replan at each time step, gaining robustness to model errors at the cost of increased test-time computation.
  5. Use the model for data generation (MBPO) to augment model-free RL with synthetic rollouts, combining the sample efficiency of model-based methods with the long-horizon capability of model-free policy learning.

The fundamental trade-off in model-based RL is between the benefits of increased sample efficiency and task-agnostic learning, and the costs of model error, objective mismatch, and added complexity. The PDDM case study demonstrates that with careful engineering—ensembles, modified CEM, appropriate horizon lengths—model-based methods can achieve impressive results on challenging real-world robotics tasks.

In the next lecture, we turn to multi-task and goal-conditioned RL, examining how to train generalist policies that can perform many tasks rather than specializing in a single one. Model-based methods will make a natural appearance there as well, since a learned dynamics model is inherently task-agnostic and can be shared across goals.