Model-Based RL with Policies

From planning to policies: why open loop isn’t enough

The previous lecture ended on what Levine dubs “model-based RL version 1.5”: collect data, train a dynamics model, plan through the model, execute only the first planned action, append the observed transition to the buffer, and replan — retraining the model every so many steps. Replanning every step (model predictive control) compensates for model errors, and the planners from the optimal control lecture slot in as subroutines. But the simplest of those planners — random shooting, the cross-entropy method — are open loop: they optimize expected reward over a committed sequence of actions $\mathbf{a}_1, \ldots, \mathbf{a}_T$.

The closed-loop fix is to commit not to actions but to a policy $\pi(\mathbf{a}_t \mid \mathbf{s}_t)$, optimized under the trajectory distribution $p(\mathbf{s}_1) \prod_t \pi(\mathbf{a}_t \mid \mathbf{s}_t)\, p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$ — exactly the ordinary RL objective, except that we model $p(\mathbf{s}_{t+1} \mid \mathbf{s}_t, \mathbf{a}_t)$ explicitly. There is a choice of form for $\pi$: neural networks give global policies that produce decent actions everywhere they were trained, whereas iterative LQR yields a linear feedback controller that is technically closed loop but only locally good, near the plan it was derived around. This lecture is about training global policies with learned models.

The obvious idea — backpropagate through the model — and why it fails

If you knew deep learning but had “missed the entire first part of this course,” you would set up a computation graph: policy nodes map states to actions, dynamics nodes map state–action pairs to next states, reward nodes emit scalars. With a learned (deterministic) model $f$ and a known, differentiable reward, the negative sum of rewards is a perfectly valid PyTorch loss — call .backward() and do gradient ascent on $\theta$. (Gaussian stochasticity can be handled by the reparameterization trick, covered later in the course.) This is “model-based RL version 2.0”: collect data, train $f$, backpropagate through the graph into the policy, run the policy, retrain.

In general it does not work — and not because it is incorrect. As with shooting methods, early actions have compounding effects on the whole trajectory while late actions barely matter, so the policy parameters get hit with enormous gradients through the early steps and tiny ones through the late steps: a classically ill-conditioned problem. With trajectory optimization the cure was a second-order method, and LQR gave the Hessian-vector products almost for free by dynamic programming backward in time. That option is gone here: the policy parameters couple all the time steps together, so no backward recursion applies, and generic second-order optimizers for neural nets “tend to be kind of flaky.”

From the deep learning side, the failure is exactly naive backpropagation through time: the derivative of a late reward with respect to early parameters is a product of many Jacobians, and unless their eigenvalues sit near one, gradients explode or vanish.

Use the model as a simulator: model-free RL with a model

The perhaps surprising resolution: use model-free RL algorithms, with the model serving only to generate synthetic samples. If you learn a model, don’t exploit its known derivatives — treat it as a simulator of the environment and use it to manufacture data. To see why this is principled, compare the two gradient estimators. The policy gradient (the likelihood-ratio, or REINFORCE, estimator) is

$\nabla_\theta J(\theta) \approx \frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(\mathbf{a}_{i,t} \mid \mathbf{s}_{i,t})\, \hat{Q}^\pi_{i,t}$

and it is a perfectly good gradient estimator for any stochastic computation graph — nothing about it is intrinsically “reinforcement learning.” The transition probabilities never appear in the expression; they enter only through the sampling, and samples can come from a learned model just as well as from the real MDP. The backpropagation (pathwise) gradient, by contrast, is what the chain rule gives you:

$\nabla_\theta J(\theta) = \sum_{t=1}^{T} \frac{d\mathbf{a}_t}{d\theta} \frac{d\mathbf{s}_{t+1}}{d\mathbf{a}_t} \left( \sum_{t'=t+1}^{T} \frac{dr_{t'}}{d\mathbf{s}_{t'}} \left( \prod_{t''=t+2}^{t'} \frac{d\mathbf{s}_{t''}}{d\mathbf{a}_{t''-1}} \frac{d\mathbf{a}_{t''-1}}{d\mathbf{s}_{t''-1}} + \frac{d\mathbf{s}_{t''}}{d\mathbf{s}_{t''-1}} \right) \right)$

The problematic part is the giant product of $n \times n$ Jacobians — eigenvalues above one explode it, below one vanish it. (Technically the likelihood-ratio form wants stochastic policies and transitions and the pathwise form deterministic ones, but the reparameterization trick and near-deterministic limits bridge the gap; the fundamental difference is the Jacobian product.) There is no free lunch: the policy gradient pays for dropping the Jacobians by requiring samples — it has deep connections to finite differencing. But with a learned model, samples cost compute rather than physical interaction: “generating more samples is just a matter of sticking more GPUs in the data center.” A 2018 analysis by Parmas et al. works out the numerical stability comparison; the short version is that the model-free gradient can actually be better.

Hence “model-based RL version 2.5”: collect data, learn the model, sample many trajectories from the model with the current policy, improve the policy by policy gradient (with all the actor-critic tricks), repeat the inner loop several times, then collect fresh real data and retrain the model.

The curse of long rollouts, and the branched-rollout trick

Version 2.5 still has a flaw: long model-based rollouts. Recall the imitation learning lecture — a policy trained with supervised learning makes a small mistake, lands in an unfamiliar state, makes a bigger mistake, and errors compound; that was distributional shift. The same applies verbatim to learned models: run $\pi_\theta$ under the learned dynamics rather than the true dynamics and the model drifts into states it wasn’t trained on, where its errors grow. Worse, in version 2.5 you are also changing the policy away from the one that collected the data, exacerbating the shift. Even in the best case, the error accumulates as $O(\epsilon t^2)$ — the bound and its tightness follow by the same logic as the behavioral cloning analysis (Levine leaves the proof “as an exercise,” and it appears below). Long model-based rollouts are costly in accumulated error: the longer the rollout, the less its reward resembles the real one.

Simply truncating the horizon — rollouts of 50 steps in a 1,000-step MDP — changes the problem: a robot cooking a 30-minute meal never gets past putting the pot on the stove in a 5-minute rollout. The trick is to take a small number of full-length real rollouts, sample states uniformly along them, and branch short model-based rollouts from those states. Error stays low because rollouts are short, and every part of the horizon is visited because branch points come from everywhere in the real trajectories.

The price is a strange state distribution: you rolled into the branch point with the data-collection policy and roll out of it with the new policy, so the model-based states match neither — they are a mix. Small policy changes keep this benign (the advanced policy gradients analysis applies), but the whole point of model-based RL is to improve the policy a lot between collection rounds. So in practice this recipe works best with off-policy algorithms — Q-learning or Q-function actor-critic — rather than on-policy policy gradients.

Dyna, and the Dyna-style recipe behind MBPO

That gives “model-based RL version 3.0,” close to what people actually use: collect data, train the model, branch short rollouts (as short as one step; on the order of ten in practical algorithms) from buffer states, and improve the policy with off-policy RL on both real and model data. The classic ancestor is Dyna (Sutton, 1990s): online Q-learning where, after each real transition, you (1) update the model and reward function on that transition, (2) do a normal Q-learning update, then (3) repeat $K$ times: sample an old $(s, a)$ from the buffer and re-simulate $s' \sim \hat{p}(s' \mid s, a)$, $r = \hat{r}(s, a)$ for a model-based Q-update. Using the buffer’s own action and a single model step is peculiar — it squeezes little from the model — but it is statistically very safe: one step from a previously visited state–action pair incurs essentially no distributional shift.

Published variants of this recipe differ mainly in how model data feeds the Q-function: the procedure above is closest to model-based policy optimization (MBPO); model-based value expansion (MVE) uses model rollouts only to improve target values, not to train the Q-function directly; model-based acceleration (MBA) is a third variation. The trade-off is characteristic: these methods learn significantly faster — the model amplifies the dataset — but can plateau at lower final performance, because model bias “puts a ceiling on how tightly they can fit to the real-world dynamics.” Ensembles average out model errors and reduce exploitation, short rollouts limit compounding, and you cannot go too long without real data, or the branch-point states drift too far from what the current policy would actually visit.

Successor representations: the bare-minimum model

The last part steps back from models that predict $\mathbf{s}_{t+1}$ — this is current research, closer to ideas than finished algorithms, but suggestive of where model-based RL could go. In the three-box view of RL, the green box exists to let you evaluate the policy; evaluation is what enables improvement. So: what is the minimal model that evaluates a policy? Working with state-dependent rewards $r(s)$ for simplicity, expand the value function and swap the order of summation:

$V^\pi(s_t) = \sum_{t' = t}^{\infty} \gamma^{t' - t}\, \mathbb{E}\!\left[ r(s_{t'}) \mid s_t \right] = \sum_{s} \left( \sum_{t' = t}^{\infty} \gamma^{t' - t}\, p(s_{t'} = s \mid s_t) \right) r(s)$

The bracketed quantity is (up to a constant) a probability of landing in $s$ “in the future.” Normalizing with $1 - \gamma$ so it sums to one defines $p_\pi(s_{\text{future}} = s \mid s_t)$: sample a future offset from a geometric distribution with parameter $\gamma$ and ask where the policy is then — equivalently, terminate with probability $1 - \gamma$ each step and ask where you stop.

Successor representations obey a Bellman equation with a pseudo-reward $(1 - \gamma)\, \delta(s_t = i)$:

$\mu^\pi_i(s_t) = (1 - \gamma)\, \delta(s_t = i) + \gamma\, \mathbb{E}_{\mathbf{a}_t \sim \pi,\, s_{t+1}}\!\left[ \mu^\pi_i(s_{t+1}) \right]$

so vectorized Bellman backups (value iteration and friends) can train them. Two worries remain: it is not obvious this is easier than model-free RL when the reward is already known, and one entry per state cannot scale — for continuous states the delta is always zero and a density version is needed. Scaling is handled by projecting onto a feature basis: given features $\phi_j(s)$, the successor features are $\psi^\pi_j(s_t) = \sum_s \mu^\pi_s(s_t)\, \phi_j(s)$ — each $\phi_j$ acting as a pseudo-reward with $\psi_j$ its pseudo-value — and 100 features can replace backups over millions of states.

The intended use: train $\psi^\pi$ once with Bellman backups; when someone hands you reward samples, solve a least-squares regression $\phi(s_i)^\top \mathbf{w} \approx r_i$ and immediately read off $Q^\pi$ — one representation reused across many rewards. The caveat: acting by $\arg\max_{\mathbf{a}}$ on this $Q^\pi$ is one step of policy iteration — better than $\pi$, not optimal. A stronger variant trains successor features for many policies $\pi_k$ and acts by $\arg\max_{\mathbf{a}} \max_k Q^{\pi_k}(s, \mathbf{a})$, which provably does as well as or better than the best single $\pi_k$ — better, because it can pick different policies in different states (see Successor Features for Transfer in Reinforcement Learning). For continuous states, frame it as classification: train a classifier to distinguish “futures” $s_{\text{future}} \sim p_\pi(\cdot \mid s_t, \mathbf{a}_t)$ (positives) from states drawn from the policy’s marginal (negatives). The Bayes-optimal classifier is $p(F = 1 \mid s_t, \mathbf{a}_t, s_{\text{future}}) = \frac{p_\pi(s_{\text{future}} \mid s_t, \mathbf{a}_t)}{p_\pi(s_{\text{future}} \mid s_t, \mathbf{a}_t) + p_\pi(s_{\text{future}})}$, so the ratio of the two classifier outputs recovers the future-state density up to a constant independent of $s_t$ and $\mathbf{a}_t$ — harmless when maximizing over actions. Training is a cross-entropy loss, positives sampled at a geometric future offset along the trajectory; the off-policy extension is in the C-learning paper.

Check yourself

Exercises in the lecture’s spirit — derivations and small calculations, on paper.

Problem 12.1 Jacobian product blowup

A deterministic learned model has state Jacobians $d\mathbf{s}_{t+1}/d\mathbf{s}_t$ whose eigenvalues all equal $\lambda$. In the pathwise gradient, the term coupling the action at time $t$ to the reward at time $t' = t + 100$ contains a product of roughly 100 such Jacobians. Estimate the magnitude of that product for $\lambda = 1.05$ and $\lambda = 0.95$, and explain why the LSTM remedy from RNN training cannot be applied.

Show solution

The product’s eigenvalues are $\lambda^{100}$: for $\lambda = 1.05$, $1.05^{100} \approx 131$ — gradients from distant rewards are amplified two orders of magnitude and explode further with horizon; for $\lambda = 0.95$, $0.95^{100} \approx 0.006$ — distant rewards effectively vanish from the gradient. Either way the early-time-step parameters receive wildly different gradient magnitudes than late ones, producing severe ill-conditioning. An LSTM fixes this in RNNs by choosing the recurrent function so its Jacobian is dominated by the identity (eigenvalues near one). In model-based RL the recurrent function is the dynamics model $f$, which must be trained to match the real dynamics — whose Jacobians are dictated by the environment, not by us. Changing $f$ to be easier to differentiate would make it a wrong model, and the resulting policy would fail on the real system.

Problem 12.2 rollout error bound

Levine leaves as an exercise: show that if the learned model’s per-step error is $\epsilon$ (probability of a “mistake,” in the sense of the behavioral cloning analysis, when queried in-distribution), the expected accumulated error of a model-based rollout of length $T$ is $O(\epsilon T^2)$, and argue the bound is tight.

Show solution

Follow the behavioral cloning argument with the model in place of the policy. At each step the model errs with probability at most $\epsilon$ while it remains in-distribution; once it has erred, it may be out of distribution and we assume the worst (error every subsequent step). The probability of having made the first mistake by step $t$ is at most $\epsilon t$ (union bound), and each of the remaining $T - t$ steps then contributes error. Summing the per-step error probabilities: $\sum_{t=1}^{T} \epsilon t = \epsilon\, T(T+1)/2 = O(\epsilon T^2)$. Tightness follows from the same construction as the imitation learning tightrope example: an MDP where a single wrong transition leaves the data distribution permanently, after which the model gets no credit for recovery — the first mistake at time $t$ genuinely costs $O(T - t)$, and the quadratic sum is realized. The takeaway: doubling the rollout length quadruples the accumulated error, which is why practical algorithms keep model rollouts to a handful of steps.

Problem 12.3 branched-rollout distribution

In a branched rollout scheme, states are sampled from real trajectories collected under an old policy $\pi_{\text{old}}$, and short model rollouts are generated from those states with the current policy $\pi_{\text{new}}$. (a) Whose state distribution do the resulting model-based transitions follow? (b) Why does this favor off-policy algorithms like Q-learning over policy gradients? (c) Why must you still collect real data periodically, even with a perfect off-policy learner?

Show solution

(a) Neither policy’s. Branch points are distributed as $\pi_{\text{old}}$‘s visitation distribution; the few steps after the branch are governed by $\pi_{\text{new}}$ under the learned dynamics. The mixture matches $\pi_{\text{old}}$ early in each branch and drifts toward $\pi_{\text{new}}$‘s local behavior, so overall it is a mix of the two. (b) On-policy methods like the policy gradient require samples from the current policy’s trajectory distribution; a mixed distribution biases the gradient unless policy changes stay small — but the whole point of model-based acceleration is large policy improvement between collection rounds. Off-policy methods are, in principle, indifferent to which distribution generated the transitions. (c) “In principle” hides a state- coverage requirement: the Q-function must be accurate on states the new policy will actually visit at deployment. If the buffer states (and short branches from them) no longer resemble those states, training concentrates accuracy in the wrong region. Refreshing the buffer keeps the branch-point distribution close to reality.

Problem 12.4 successor representation calculation

A deterministic two-state MDP under a fixed policy $\pi$ cycles $A \to B$ and $B \to B$ (B is absorbing). With discount $\gamma$, compute the successor representation vectors $\mu^\pi(A)$ and $\mu^\pi(B)$ from the definition, verify the Bellman equation at $A$, and recover $V^\pi(A)$ for the reward $r(A) = 0$, $r(B) = 1$.

Show solution

From $B$: the state at every future offset is $B$, so $\mu^\pi_B(B) = (1-\gamma)\sum_{k \ge 0} \gamma^k = 1$ and $\mu^\pi_A(B) = 0$. From $A$: offset 0 is $A$, all later offsets are $B$, so $\mu^\pi_A(A) = (1-\gamma)\gamma^0 = 1 - \gamma$ and $\mu^\pi_B(A) = (1-\gamma)\sum_{k \ge 1} \gamma^k = \gamma$. Bellman check at $A$ for entry $B$: $(1-\gamma)\,\delta(A = B) + \gamma\, \mu^\pi_B(B) = 0 + \gamma \cdot 1 = \gamma$ — matches. Value recovery: $V^\pi(A) = \frac{1}{1-\gamma} \mu^\pi(A)^\top \mathbf{r} = \frac{1}{1-\gamma}\left[(1-\gamma)\cdot 0 + \gamma \cdot 1\right] = \frac{\gamma}{1-\gamma}$, which is exactly the discounted return $\sum_{k \ge 1} \gamma^k$ of reaching $B$ one step from $A$ and staying forever. Note $\mu^\pi$ never referenced the reward — with a different $\mathbf{r}$ (say $r(A) = 1$, $r(B) = 0$) the same $\mu^\pi$ instantly gives $V^\pi(A) = \frac{(1-\gamma) \cdot 1}{1-\gamma} = 1$, illustrating reuse across rewards.