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.
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.
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.
- 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)\}$.
- 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$$
- 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:
- Data coverage matters enormously. The model is only accurate in regions of state-action space covered by the training data $\mathcal{D}$. If the base policy $\pi_0$ only visits a small part of the state space, the model will make poor predictions elsewhere—and the RL algorithm may exploit these inaccuracies. This is analogous to the distribution mismatch problem in imitation learning (DAgger).
- Modeling difficulty is domain-dependent. Some environments have simple, low-dimensional dynamics that are easy to learn (e.g., a pendulum). Others involve complex contact physics, deformable objects, or high-dimensional observations (e.g., images), making accurate modeling far more challenging. A maze with discrete transitions is trivial to model; a humanoid robot performing dexterous manipulation is not.
- Model inaccuracies compound. Even with a good model, running RL inside the learned simulator is nontrivial. Small per-step prediction errors can compound over long rollouts, causing the simulated trajectories to diverge from reality. An RL agent trained purely in a slightly inaccurate simulator may learn to exploit model errors rather than solve the real task.
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:
- End-to-end learning in observation space. Directly learn a mapping $f_\phi: \mathcal{S} \times \mathcal{A} \to \mathcal{S}$ from the current state and action to the next state. When states are low-dimensional vectors (e.g., joint angles and velocities), a standard neural network with mean squared error loss works well. When observations are high-dimensional (e.g., images), one might use convolutional architectures or even video prediction models.
- Latent-space models. Learn a compact, low-dimensional state representation $\mathbf{z} = g(\mathbf{s})$, then learn the dynamics in this latent space: $\hat{\mathbf{z}}_{t+1} = h(\mathbf{z}_t, \mathbf{a}_t)$. This approach can be more data-efficient and avoids the need to predict every pixel of a high-dimensional observation. World models (discussed briefly later) follow this paradigm.
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)$$- Run some policy (e.g., random) to collect data $\mathcal{D} = \{(\mathbf{s}, \mathbf{a}, \mathbf{s}')_i\}$.
- Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
- 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}$.
- Sample many candidate sequences $\mathbf{A}_1, \ldots, \mathbf{A}_N$ from some distribution (e.g., uniform over the action space).
- 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'})$.
- 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.
Initialize a sampling distribution $p(\mathbf{A})$ (e.g., Gaussian with broad variance). Repeat:
- Sample many action sequences $\mathbf{A}_1, \ldots, \mathbf{A}_N$ from $p(\mathbf{A})$.
- Evaluate each: $J(\mathbf{A}_i) = \sum_{t'=t}^{t+H} r(\mathbf{s}_{t'}, \mathbf{a}_{t'})$.
- 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$.
- 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) | |
|---|---|---|
| Strengths | Scales to high dimensions; works well in overparameterized regimes | Highly parallelizable; requires no gradient information; simple to implement |
| Weaknesses | Requires differentiable model and smooth optimization landscape | Scales 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.
- Run some policy (e.g., random) to collect initial data $\mathcal{D} = \{(\mathbf{s}, \mathbf{a}, \mathbf{s}')_i\}$.
- Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
- Use the model to plan an action sequence (via backpropagation or CEM).
- Execute the planned actions in the real environment, appending the resulting tuples $(\mathbf{s}, \mathbf{a}, \mathbf{s}')$ to $\mathcal{D}$.
- 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.
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.
- 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\}$.
- Learn model $f_\phi(\mathbf{s}, \mathbf{a})$ to minimize $\sum_i \|f_\phi(\mathbf{s}_i, \mathbf{a}_i) - \mathbf{s}'_i\|^2$.
- Use model $f_\phi(\mathbf{s}, \mathbf{a})$ to plan an optimal action sequence over a horizon $H$.
- Execute only the first planned action $\mathbf{a}_t$, observe the resulting state $\mathbf{s}'$.
- Append $(\mathbf{s}, \mathbf{a}, \mathbf{s}')$ to dataset $\mathcal{D}$.
- 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.
Summary of Planning with Learned Models
The planning-based approach to model-based RL has several appealing properties and limitations:
| Advantages | Disadvantages |
|---|---|
| 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.
- Advantage: No longer compute-intensive at test time—the policy network produces actions in a single forward pass.
- Limitation: Still limited to short-horizon problems, since the planner it imitates can only plan over short horizons.
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:
- 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.
- 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.
- 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.
Repeat:
- Collect real data: Run current policy $\pi_\phi$ in the real environment, add transitions to $\mathcal{D}_{\text{env}}$.
- Update model: Train dynamics model $p_\theta(\mathbf{s}' \mid \mathbf{s}, \mathbf{a})$ using $\mathcal{D}_{\text{env}}$.
- 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}}$.
- 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}}$).
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
- Sample efficiency. Model-based methods can be dramatically more data-efficient than model-free methods, particularly when the dynamics are relatively easy to learn. In robotics, this can mean learning a task in hours rather than days.
- Self-supervised learning. The dynamics model can be trained without reward labels—it only needs $(s, a, s')$ tuples. This means data collection for the model does not require a reward function, and in principle, the same model can be reused across different tasks.
- Task-agnostic models. Because the dynamics model captures how the world works rather than what the task is, a well-trained model can sometimes be transferred across different reward functions or goals. This is particularly appealing in robotics, where the physics are shared across tasks.
Disadvantages
- Model error and the objective mismatch problem. The dynamics model is typically trained to minimize prediction error (e.g., MSE on next-state predictions), not to maximize task reward. A model might be very accurate in irrelevant parts of the state space and inaccurate in the parts that matter most for the task. This is known as the objective mismatch problem.
- Sometimes harder to learn than a policy. Predicting the full next state is in some sense a harder problem than learning a good policy. The model must accurately predict all state dimensions, while a policy only needs to know which action is best. For high-dimensional observation spaces (e.g., images), learning an accurate forward model can be prohibitively difficult.
- Added complexity. Model-based methods require training and tuning an additional model, introducing more hyperparameters (model architecture, ensemble size, rollout length, model update frequency, etc.) and more compute.
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:
- Inverse model: $p(\mathbf{a}_t \mid \mathbf{s}_t, \mathbf{s}_{t+1})$ — predicts what action was taken to cause a particular transition. Useful for learning from observation-only demonstrations.
- Multi-step inverse model: $p(\mathbf{a}_t \mid \mathbf{s}_t, \mathbf{s}_{t+n})$ or $p(\mathbf{a}_{t:t+n} \mid \mathbf{s}_t, \mathbf{s}_{t+n})$ — predicts actions needed to reach a distant state.
- Future prediction without actions: $p(\mathbf{s}_{t+1:t+n} \mid \mathbf{s}_t)$ — predicts future states without conditioning on actions. Useful for learning representations of environment dynamics.
- Video interpolation: $p(\mathbf{s}_{t+1:t+n} \mid \mathbf{s}_t, \mathbf{s}_{t+n+1})$ — fills in intermediate states given start and end states.
- Joint transition distribution: $p(\mathbf{s}_t, \mathbf{a}_t, \mathbf{s}_{t+1})$ — models the full joint distribution of transitions rather than just the conditional.
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:
- State space: hand joint positions and velocities, plus object positions and orientations.
- Action space: 24-dimensional continuous control of the five-fingered hand.
- Reward: tracking error to a target object trajectory, plus a penalty for dropping the object.
Method: PDDM
The proposed method, called PDDM (Planning with Deep Dynamics Models), uses a model-based planning approach with several carefully chosen design decisions:
- Model: An ensemble of 3 neural networks, each with 2 hidden layers of size 500. The ensemble helps estimate model uncertainty—when the ensemble members disagree, the model is likely inaccurate in that region of state space.
- Planner: A modified version of the cross-entropy method (CEM), incorporating softer reward-weighted means (rather than hard elite selection) and temporal smoothing on the action sequences. These modifications produce smoother, more physically plausible action sequences.
- Training loop: The system alternates between collecting approximately 30 trajectories using the planner and then updating the dynamics model on the accumulated dataset.
Results and Comparisons
The authors compared PDDM against both model-free and model-based baselines in simulation:
- Model-free methods: SAC (actor-critic) and NPG (natural policy gradient)—these required far more data to achieve comparable performance.
- Model-based baselines: MBPO (RL with model-generated data), PETS (a CEM-based planner without the PDDM modifications), and Nagabandi et al. (random shooting without ensembles)—PDDM outperformed all of these.
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:
- Model capacity matters. A sufficiently large model (2 hidden layers of size 500) was needed; smaller models (e.g., 2 x 64) performed significantly worse.
- Ensembles are important. At least 3 ensemble members were needed for reliable performance. A single model was noticeably worse, likely because it cannot capture model uncertainty.
- Planning horizon is a trade-off. Too short a horizon (2 steps) limited the planner's ability to reason about the future, while too long a horizon (15 steps) introduced compounding model errors. An intermediate horizon (5–7 steps) worked best.
- Modified CEM is crucial. The reward-weighted mean and temporal smoothing in PDDM's modified CEM significantly outperformed both standard CEM and random shooting. This underscores that the planner's optimization quality matters as much as the model's accuracy.
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.
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:
- Learn a dynamics model $f_\phi(\mathbf{s}, \mathbf{a})$ from environment interaction data, using supervised learning (MSE regression).
- Use the model for planning—optimizing action sequences via gradient-based (backpropagation) or sampling-based (CEM) methods.
- Address distribution mismatch by iteratively collecting data under the current plan and retraining the model (analogous to DAgger).
- Use model-predictive control (MPC) to replan at each time step, gaining robustness to model errors at the cost of increased test-time computation.
- 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.