Lecture 11 · Part 3

Model-Based RL

Learn the dynamics, then plan through them with last lecture's tools. Levine builds the algorithm in versions — 0.5 fails from distributional shift, 1.0 fixes it DAgger-style, 1.5 adds replanning (MPC) — then shows why the planner exploits an overconfident model, why epistemic (not aleatoric) uncertainty is the cure, and how bootstrap ensembles deliver it. The lecture closes with model-based RL from images via latent state-space models.

▶ Lecture videos (5 parts) Slides (current revision, PDF) ↗

Why learn the model, and the naive recipe

The premise of model-based RL: if we had an estimate of the dynamics — a function $f(s_t, a_t) = s_{t+1}$ , or stochastically $p(s_{t+1} \mid s_t, a_t)$ — we could use all the planning tools from the previous lecture to control the system instead of doing model-free RL at all. Everything in this lecture learns only a model and plans through it; using models to learn policies comes next lecture.

The most obvious algorithm — Levine calls it model-based RL version 0.5, “not quite the thing you want to use” — is three steps. Run a basic exploration policy $\pi_0$ (uniformly random is fine) to collect transitions $\mathcal{D} = \{(s_i, a_i, s'_i)\}$ . Fit the dynamics by supervised learning:

$\min_f \sum_i \left\| f(s_i, a_i) - s'_i \right\|^2$

(squared error being negative log-likelihood under a Gaussian; cross-entropy for discrete states). Then plan through $f$ to choose actions.

Does this work? Sometimes, very well: it is essentially how system identification works in classical robotics. There you already know the equations of motion and are fitting a handful of unknowns — masses, friction or drag coefficients — so if the true model is in your class, modest data pins down those few numbers (given a base policy that excites the system’s different modes). The recipe is effective precisely when you can hand-engineer the dynamics representation. With high-capacity models like deep nets, it breaks — for a reason we have seen before.

Distributional shift: DAgger for models

Picture trying to reach the top of a mountain. A random walk $\pi_0$ explores one slope, and from that data the model learns a very reasonable inference: the more right you go, the higher you get. Plan through that model and you walk right — off the far edge.

The data came from the state distribution $p_{\pi_0}(s_t)$ induced by the base policy. But planning through $f$ is itself a policy — call it $\pi_f$ , not a network but a planning algorithm run on top of the model — with its own distribution $p_{\pi_f}(s_t)$ . The problem is exactly that $p_{\pi_f}(s_t) \neq p_{\pi_0}(s_t)$ : the model is valid where $\pi_0$ went, but the planner seeks out the actions the model scores highest, those lead to states it was never trained on, it makes an erroneous prediction there, feeds it back into itself, and predicts something even more wrong at the next step.

The fix is borrowed from DAgger: collect more data where the learned object will actually operate. Here it is far easier than DAgger, because no human expert is needed — to label a state–action pair with its next state, you just take the action and watch nature. Version 1.0: learn the model, plan through it, execute those plans, append the resulting transitions to $\mathcal{D}$ , retrain, repeat. “It’s essentially just like DAgger, only for models” — with the anachronistic caveat that this procedure long predates DAgger. In principle, version 1.0 works.

Replanning: model predictive control

But what if you make a mistake? Some mistakes are cliffs — sudden, unrecoverable — but many are not. Say your model predicts you’ll drive straight when the wheel is two degrees left. Execute the open-loop plan and you drift left, a little more each step, and stay wrong until retraining eventually fixes the model. We can learn faster by fixing mistakes as they happen: replan from the state you actually reached instead of the one the plan predicted.

Code walkthrough — The three versions, as one loop

# v0.5: steps 1-3 once.  v1.0: loop 1-4.  v1.5: MPC inner loop.
D = run_base_policy()                    # 1. random pi_0 collects (s, a, s')
while True:
    f = fit_dynamics(D)                  # 2. supervised learning on all data
    for step in range(N):                # outer retraining period
        plan = optimize_actions(f, s)    # 3. CEM / random shooting, horizon H
        a = plan[0]                      # 4. execute ONLY the first action
        s_next = env.step(a)             #    observe what really happened
        D.append((s, a, s_next))
        s = s_next                       # ...and replan from the real state

Version 1.5 is much more computationally expensive — you plan at every time step — but it does much better with a much worse model: once the car has actually drifted left, the error is big enough that even an imperfect model knows to steer right.

The intuition for how to replan: “the more you re-plan, the less perfect each individual plan needs to be.” In practice people use much shorter horizons than for one open-loop plan, and even random shooting often works — last lecture’s closing demonstration was exactly MPC with a short horizon, leaning almost entirely on replanning. Homework 4 involves implementing essentially this procedure.

The planner exploits your model

So version 1.5 pretty much always works — but how well? In a Berkeley experiment on half-cheetah, version 1.5 with a deep-net model learned quickly but plateaued around reward 500 — moving forward, just slowly — while a model-free learner bootstrapped from it eventually reached about 5,000. Why the tenfold gap?

Iterative data collection poses a unique challenge: the model must not overfit early, when data is scarce, yet needs the capacity to do well later, when data is plentiful. Deep nets struggle in the low-data regime, so they do poorly early, which means poor exploration, which means getting stuck. And it is worse than ordinary overfitting:

Uncertainty estimation is the fix. Suppose you want to walk to the edge of a cliff for the view. If the model predicts a distribution over next states and you plan for expected reward, then near the edge — where the model is uncertain — substantial probability mass falls off the cliff and incurs a huge negative reward, so the planner automatically stays back. No explicit pessimism is engineered in; taking the expected value under an uncertain model already makes the planner hedge its bets across the possible worlds consistent with the data, and as data accumulates near the edge, uncertainty shrinks and it walks closer. Caveats Levine flags: you still need to explore (excess caution can keep you from high-reward regions forever), and the expected value is neither the pessimistic value (a lower confidence bound, reasonable for safety) nor the optimistic value (which favors exploration) — but it “is a pretty good place to start.”

Two kinds of uncertainty, and how to estimate the right one

The obvious idea — have the network output a distribution over $s_{t+1}$ (a softmax, or a Gaussian mean and variance) and use its entropy — is, Levine says flatly, “a bad idea; this does not work.” An overfitted model matches its training points almost exactly, so its optimal output variance is nearly zero: it will be extremely confident and completely wrong out of distribution. The entropy of the output is measuring the wrong thing.

Being uncertain about the model means being uncertain about the parameters $\theta$ . Instead of the usual $\arg\max_\theta \log p(\theta \mid \mathcal{D}) = \arg\max_\theta \log p(\mathcal{D} \mid \theta)$ , estimate the full posterior $p(\theta \mid \mathcal{D})$ — its entropy is the model uncertainty — and predict by integrating the parameters out:

$\int p(s_{t+1} \mid s_t, a_t, \theta)\, p(\theta \mid \mathcal{D})\, d\theta$

Done exactly, this is intractable for neural nets, so we approximate. Bayesian neural networks put a distribution over every weight; the common tractable approximation factorizes the posterior into independent marginals, $p(\theta \mid \mathcal{D}) = \prod_i p(\theta_i \mid \mathcal{D})$ with $p(\theta_i \mid \mathcal{D}) = \mathcal{N}(\mu_i, \sigma_i)$ — each weight gets an expected value and an uncertainty. It’s crude (weights interact tightly), and training it needs variational inference, covered next week; see Blundell et al., “Weight Uncertainty in Neural Networks” and Gal et al., “Concrete Dropout”.

Today’s simpler method — which in Levine’s experience “actually works a little bit better” in model-based RL — is bootstrap ensembles: train many networks that agree on the training data but make different mistakes outside it, and let their disagreement measure uncertainty.

Bootstrap ensembles approximate the parameter posterior

An ensemble of $N$ models represents the posterior as a mixture of Dirac deltas, $p(\theta \mid \mathcal{D}) \approx \frac{1}{N} \sum_i \delta(\theta_i)$ , so the predictive integral becomes an average over models:

$\int p(s_{t+1} \mid s_t, a_t, \theta)\, p(\theta \mid \mathcal{D})\, d\theta \approx \frac{1}{N} \sum_i p(s_{t+1} \mid s_t, a_t, \theta_i)$

Crucially, you average the probabilities, not the means: $N$ Gaussians with different means form a mixture of Gaussians, not one Gaussian at the average mean — a warning that matters for homework 4. The classical bootstrap trains each $\theta_i$ on $\mathcal{D}_i$ , resampled from $\mathcal{D}$ with replacement, to make the datasets statistically independent. In deep learning the resampling is “usually unnecessary”: SGD and random initialization already make the models sufficiently independent, even trained on identical data.

The approximation is crude — training ten nets costs ten times the compute, so ensembles usually have fewer than 10 members — but it is enough.

Planning under an uncertain model

Before, the planner optimized $\sum_{t} r(s_t, a_t)$ with $s_{t+1} = f(s_t, a_t)$ . With $N$ models, choose the action sequence that maximizes the average reward over models:

$J(a_1, \ldots, a_H) = \frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{H} r(s_{t,i}, a_t), \qquad s_{t+1,i} = f_i(s_{t,i}, a_t)$

The general recipe, for any candidate action sequence: sample $\theta \sim p(\theta \mid \mathcal{D})$ (for an ensemble, pick one of the $N$ models at random); step it forward, sampling $s_{t+1} \sim p(s_{t+1} \mid s_t, a_t, \theta)$ at each time step; sum the rewards; repeat and average. With a small ensemble you can sum over all models instead of sampling; fancier alternatives like moment matching over the state distribution exist. Then optimize over action sequences with your favorite method — random shooting or CEM — and run it all inside the version 1.5 MPC loop.

The approach transfers to the real world: a robotic hand with an ensemble-based model, learning directly from physical interaction, performs a full in-palm rotation of two objects after about three hours and does it reliably after four. For the foundations, Levine recommends PILCO (Deisenroth, ~2011) — Gaussian processes rather than neural nets, but one of the papers that established why epistemic uncertainty matters in model-based RL.

Model-based RL with images

What about image observations — Atari frames, a robot’s camera? Three obstacles: images are high-dimensional, which makes prediction hard; they are redundant; and they are usually partially observed — one frame of Breakout doesn’t tell you the ball’s velocity. So we adopt the POMDP model: dynamics $p(s_{t+1} \mid s_t, a_t)$ over an unobserved state, and an observation model $p(o_t \mid s_t)$ . The appealing separation of concerns: the observation model handles the high-dimensional stuff but no temporal dynamics; the latent dynamics handle time but stay low-dimensional. A reward model $p(r_t \mid s_t, a_t)$ joins them, since reward depends on a state we never see.

With states observed we would maximize $\sum_i \sum_t \log p_\phi(s_{t+1,i} \mid s_{t,i}, a_{t,i})$ . With states latent, we need the expected log-likelihood under the posterior over states given the observations — as in hidden Markov models. So we learn an approximate posterior, the encoder $q_\psi$ , with a spectrum of choices: from a full smoothing posterior $q_\psi(s_t, s_{t+1} \mid o_{1:T}, a_{1:T})$ — most accurate, hardest to train — down to a single-step encoder $q_\psi(s_t \mid o_t)$ , simplest and crudest, appropriate when the problem is nearly fully observed. Stochastic encoders need variational inference (next week); today’s limiting case is a deterministic single-step encoder, $s_t = g_\psi(o_t)$ , a delta function that collapses the expectation entirely. Substitute it in:

$\max_{\phi, \psi} \sum_{i} \sum_{t} \log p_\phi\!\left(g_\psi(o_{t+1,i}) \mid g_\psi(o_{t,i}), a_{t,i}\right) + \log p_\phi\!\left(o_{t,i} \mid g_\psi(o_{t,i})\right)$

The second term is an autoencoder — encode $o_t$ , reconstruct it back; the first forces encoded states to obey the learned dynamics; a reward term $\log p_\phi(r_t \mid g_\psi(o_t), a_t)$ is added the same way. Everything is differentiable, so train it all jointly with backprop, and plug the result straight into version 1.5: collect $(o, a, o')$ tuples, train encoder + dynamics + reward + observation models, plan in latent space, execute the first action, replan.

Published instances: Embed to Control (Watter et al.) uses a stochastic encoder and LQR in the latent space — its learned embeddings visibly recover the 2D structure of a navigation task and the cyclic structure of a pendulum. A successor regularizes the latent space to be locally linear, making it ideal for iterative LQR, and runs on real robots: stacking Lego blocks and pushing objects from images, with a human supplying reward labels by keypress, after roughly 20 minutes of training — far more efficient than the model-free alternatives.

And one final twist: you can skip the latent space altogether and model dynamics directly in observation space, $p(o_{t+1} \mid o_t, a_t)$ (with recurrence over past observations to handle partial observability). When the scene has dozens of objects and no obvious compact state, predicting pixels directly can work very well — a robot arm using such a model plans to move a designated image point to a target location, reaching out and pushing a stapler where it was told.

Back of the chalk tray — what to remember

Model-based RL v0.5 (random data, supervised model, plan) is system identification — it works with a few physical parameters, and fails with deep nets because $p_{\pi_f}(s_t) \neq p_{\pi_0}(s_t)$ : distributional shift, the imitation-learning problem transplanted onto the model.
v1.0 fixes it DAgger-style, and more cheaply than DAgger: nature labels next states for free — plan, execute, append, retrain.
v1.5 (MPC) executes only the first planned action and replans from the real state; it costs compute but tolerates much worse models, and shorter horizons plus random shooting often suffice.
The planner adversarially exploits positive errors in the model, so what you need is epistemic uncertainty — uncertainty about which model is right — not the aleatoric noise a maximum-likelihood output distribution captures.
Bootstrap ensembles approximate $p(\theta \mid \mathcal{D})$ as a mixture of deltas; average the predicted probabilities (not the means) and plan for expected reward across models — on half-cheetah this lifts version 1.5 from ~500 to over 6,000 reward.
With images, learn a latent state-space model — observation model, latent dynamics, reward model, encoder, trained jointly (deterministic encoder in the simplest case) — or predict directly in observation space when no compact state exists.

Check yourself

Conceptual and derivation exercises in the lecture’s spirit — work them before opening the solutions.

Problem 11.1 distributional shift diagnosis

A colleague fits a deep-net dynamics model to 10,000 transitions collected with a random policy on a driving simulator, verifies a tiny held-out prediction error on transitions from the same policy, then plans through the model and reports the car does something absurd. (a) Explain precisely why the small held-out error did not predict planning performance, naming the two distributions involved. (b) Why would the same procedure likely have worked if the model were a known vehicle model with three unknown coefficients? (c) What is the minimal change to the algorithm that fixes the failure asymptotically?

Show solution

(a) The held-out set is drawn from $p_{\pi_0}(s_t)$ , the state distribution of the random data-collection policy, so it certifies accuracy only there. Planning through the model defines a new policy $\pi_f$ with its own distribution $p_{\pi_f}(s_t)$ , and the planner specifically selects actions whose predicted outcomes look best — steering toward states that are out of distribution, where errors are large and compound: an erroneous predicted state is fed back into the model, producing a still more erroneous one. Since $p_{\pi_f}(s_t) \neq p_{\pi_0}(s_t)$ , held-out error under $p_{\pi_0}$ says nothing about the states the planner actually visits. (b) With three physical coefficients, the hypothesis class is so constrained that if the true model lies in it, almost any exciting data pins the parameters down; the model then extrapolates correctly even off-distribution. This is classical system identification. Expressive models fit tightly to the training distribution and extrapolate arbitrarily badly. (c) Close the loop on data: plan through the model, execute those plans, append the resulting transitions to the dataset, retrain, and repeat (version 1.0 — DAgger for models). The training distribution then converges toward the distribution the planner induces.

Problem 11.2 replanning trade-off

Your model has a small systematic bias: executing its “drive straight” action actually turns the car 2 degrees left per step. Compare open-loop execution of a 50-step plan (version 1.0) against MPC (version 1.5) on this system. (a) Roughly how far off-heading is the open-loop car at the end? (b) Why does MPC largely eliminate the problem even without retraining the model? (c) Why can MPC get away with much shorter planning horizons and cruder optimizers like random shooting?

Show solution

(a) The heading error accumulates roughly linearly: about $2° \times 50 = 100°$ by the end — the car finishes nearly perpendicular to its intended direction, since the open-loop plan never consults the real state. (b) MPC executes one action, then observes the actual resulting state and replans from it. After one step the car is only 2 degrees off; the planner, starting from the true state, immediately compensates. The model is still wrong about the fine effect of small steering inputs, but once a deviation is large enough to matter, correcting it no longer requires a model accurate to within 2 degrees — the required correction is coarse. Errors are repeatedly zeroed out instead of integrated. (c) Because “the more you re-plan, the less perfect each individual plan needs to be”: each plan only has to be roughly right for one executed step before it is thrown away and recomputed. Mistakes in the tail of the horizon are never executed, so a short horizon and a noisy optimizer such as random shooting suffice — the closed loop supplies the precision the individual plans lack, at the price of planning at every time step.

Problem 11.3 classifying uncertainty

For each scenario, state whether the relevant uncertainty is aleatoric or epistemic, whether collecting more data reduces it, and what a maximum-likelihood model that outputs a Gaussian $\mathcal{N}(\mu(s,a), \sigma(s,a))$ would report there: (a) predicting the sum of two dice about to be rolled, after observing a million rolls; (b) predicting the next state at the edge of a cliff the agent has never approached; (c) an overfitted dynamics model evaluated exactly on its own training points.

Show solution

(a) Aleatoric: the world itself is random. More data does not reduce it — the sum of two dice stays a nondegenerate distribution forever. The ML model correctly reports a wide output distribution here; this is the one case output entropy gets right. (b) Epistemic: nothing is necessarily noisy about the cliff; the agent simply does not know which dynamics function is correct there because it has no data. More data (approaching gradually) shrinks it to zero. The ML model reports whatever its extrapolation happens to be, typically with confidently small $\sigma$ — an erroneous mean and an erroneous variance, for the same reason: nothing trained the uncertainty head out of distribution either. (c) Neither, really — but the ML model reports near-zero variance, because on training points its means match the targets almost exactly, so zero variance maximizes likelihood. This is exactly why output entropy is the wrong tool: the model is maximally confident precisely where overfitting makes it least trustworthy off the training points. What is needed instead is the posterior $p(\theta \mid \mathcal{D})$ — “the model is certain about the data, but we are not certain about the model.”

Problem 11.4 ensemble mixture arithmetic

A two-model ensemble predicts the next (scalar) state as $\mathcal{N}(+1, 0.1)$ and $\mathcal{N}(-1, 0.1)$ . (a) Write the correct ensemble predictive distribution and contrast it with the “average the means” shortcut. (b) Compute the mean and variance of the correct predictive distribution. (c) A planner gets reward $1$ if the next state lands in $[-0.5, +0.5]$ and $0$ otherwise. Compare the expected reward under the correct mixture with the reward the averaged-mean prediction implies, and explain what this says about disagreement as an uncertainty signal.

Show solution

(a) The posterior is approximated as $p(\theta \mid \mathcal{D}) \approx \tfrac{1}{2}(\delta(\theta_1) + \delta(\theta_2))$ , so the predictive distribution is the mixture $\tfrac{1}{2}\mathcal{N}(+1, 0.1) + \tfrac{1}{2}\mathcal{N}(-1, 0.1)$ — two sharp bumps at $\pm 1$ . The shortcut averages the means into a single $\mathcal{N}(0, 0.1)$ : one confident bump at a state neither model predicts. (b) Mixture mean: $\tfrac{1}{2}(+1) + \tfrac{1}{2}(-1) = 0$ . Mixture variance, by the law of total variance: within-component variance $0.1$ plus between-component variance $\tfrac{1}{2}(1-0)^2 + \tfrac{1}{2}(-1-0)^2 = 1$ , giving $1.1$ — eleven times the shortcut’s variance. The disagreement between the models is the epistemic term. (c) Under the correct mixture, essentially all mass sits near $\pm 1$ , outside $[-0.5, 0.5]$ , so expected reward is close to $0$ . The averaged-mean Gaussian puts nearly all its mass inside the interval, promising reward close to $1$ . The shortcut manufactures confidence out of disagreement and would lure the planner into exactly the exploitation failure uncertainty estimation exists to prevent; mixing the probabilities keeps the hedge, so the planner correctly treats this action as unpromising in expectation.