Model-Based Deep RL: Sample Efficiency and PETS

The tabular planning chapter made one argument worth carrying into deep RL: planning and learning are the same backup machine, differing only in whether the experience they consume is real or simulated by a model. Dyna kept a table of observed transitions and replayed them; MCTS grew a search tree from a known simulator. Both assumed the model was either exact (the game rules) or tiny (a lookup table). Neither assumption holds for pixels, continuous control, or dynamics the agent is never handed.

Model-based deep RL asks the harder question: what if you must learn the model, in high dimensions, from the same limited experience you are trying to act on? Get it right and the payoff is large: a learned model lets you generate unlimited imagined experience and plan against it, so each real interaction goes much further. Get it wrong and the model's errors compound over the planning horizon into confident nonsense. This lesson is about that trade and the most direct way to win it: PETS — learn a probabilistic model, then plan against it while distrusting it where it is unsure. A companion lesson takes the two more radical lines: World Models / Dreamer (learn a compact latent world and imagine inside it) and MuZero (learn a latent model that predicts only what planning needs).

Why model-based: the sample-efficiency argument

Model-free methods learn a value function or policy directly from reward, discarding each transition after a gradient step (or, with replay, after a few). They are simple and unbiased by any model's mistakes, but they are sample-hungry: DQN needed tens of millions of Atari frames, and model-free continuous control routinely spends millions of environment steps. When a step is a real robot moving, or a real dollar spent, that price is the whole problem.

A model changes this arithmetic. A single real transition $(s,a,r,s^{\prime })$, once absorbed into a learned dynamics model $\hat{p}(s^{\prime }\mid s,a)$, can seed thousands of imagined transitions that never touch the environment. Planning against those imagined transitions is planning against a compressed summary of everything the agent has seen. That is the sample-efficiency case in one line: a good model turns experience into imagined planning, and imagined planning is nearly free.

The catch is in the same figure. A model-free method fits its value directly to returns, so it is asymptotically unbiased; a model-based method inherits whatever bias its model carries, and if the model is imperfect the planned policy can plateau below what a patient model-free learner would reach. The engineering of the whole field is about getting the early steepness without inheriting too much of the late gap.

The compounding-error problem

The reason a learned model is not simply a free win is that planning uses it recursively. To imagine a trajectory of length $H$, you feed the model's own output back as its next input: ${\hat{s}}_1=\hat{p}(\cdot \mid s_0,a_0)$, then ${\hat{s}}_2=\hat{p}(\cdot \mid {\hat{s}}_1,a_1)$, and so on. Each step adds a little error, and each subsequent step is now conditioned on an already-wrong state, so the errors do not merely add; they compound.

For example, suppose the model's one-step prediction is unbiased but noisy, adding an independent error of standard deviation $\sigma =0.1$ (in state units) at each step, and — the pessimistic-but-realistic case — that the dynamics are mildly expansive, so an error at one step is amplified by a factor $L=1.2$ before the next error is added. After $H$ steps the accumulated error has standard deviation

At $H=1$ this is just $0.1$. At $H=5$ it is $0.1\sqrt{(1.2^{10}-1)/(1.44-1)}=0.1\sqrt{(6.19-1)/0.44}=0.1\sqrt{11.8}=0.34$. At $H=15$ it is $0.1\sqrt{(1.2^{30}-1)/0.44}=0.1\sqrt{(237.4-1)/0.44}=0.1\sqrt{537}=2.3$ — a $23\times$ blow-up of the per-step error over a horizon of fifteen. Even with a contractive system ($L<1$) the sum converges but the error still grows before it saturates; with $L\ge 1$ it grows without bound. This is the arithmetic behind the figure and the reason every method below either keeps $H$ short (so ${\sigma }_H$ stays near $\sigma$), quantifies the growing uncertainty so the planner discounts far-horizon rollouts, or sidesteps state prediction entirely.

Three responses run through the rest of the lesson. Quantify the uncertainty so planning distrusts the model where it is unsure — the ensemble idea behind PETS. Shorten the effective horizon by replanning every step from the true state instead of committing to a long imagined plan — model-predictive control, also in PETS. And change what the model predicts: instead of reconstructing full future states (pixels), predict only the quantities planning consumes — reward and value — so per-step errors in irrelevant detail never enter the loop, which is the move MuZero makes.

Learn a model, then plan: PETS

The most direct model-based recipe is the one Dyna sketched: fit a dynamics model to observed transitions, then use it to choose actions. PETS, Probabilistic Ensembles with Trajectory Sampling, Chua et al. (2018), is the version that made this competitive with model-free methods on continuous control, and it does so by taking the compounding-error problem seriously.¹

Two design choices carry it. First, the model is probabilistic: each network outputs not a point prediction but a Gaussian over the next state, ${\hat{p}}_{\varphi }(s^{\prime }\mid s,a)=\mathcal{N}({\mu }_{\varphi }(s,a),{\Sigma }_{\varphi }(s,a))$, capturing the environment's own stochasticity (aleatoric uncertainty). Second, the model is an ensemble of $B$ such networks, trained on different bootstraps of the data; the disagreement between ensemble members captures uncertainty about the model itself (epistemic uncertainty), which is large exactly where data is scarce and the model should not be trusted.

With a model in hand, PETS does not train a policy at all. It plans online by model-predictive control (MPC): at the current real state $s_t$, it searches over short action sequences, evaluates each by rolling it forward through the ensemble and summing predicted rewards, executes only the first action of the best sequence, then throws the plan away and replans from the next real state. Replanning every step is what limits compounding error — the imagined rollout only has to be trusted for $H$ steps, after which the true state resets the plan.

The evaluation of a candidate action sequence must fold in the ensemble's uncertainty rather than trusting one network. PETS uses trajectory sampling: to score a sequence, it propagates several particles forward, each particle at each step drawing a successor from a (bootstrap-chosen) ensemble member's Gaussian. The particles fan out where the ensemble disagrees, so a sequence that scores well only under one over-confident member earns a poor average return. The action search itself is a sampling optimizer — the cross-entropy method (CEM), which repeatedly samples action sequences, keeps the top-scoring fraction, and refits the sampling distribution to them.

One cross-entropy-method iteration on a scalar toy problem shows the loop; the same loop runs (vectorized over the action sequence) inside PETS. Say the planner is choosing a single action $a\in \mathbb{R}$ and the true reward peaks near $a=2$. CEM starts with a sampling distribution $\mathcal{N}({\mu }_0,{\sigma }_0^2)=\mathcal{N}(0,2^2)$ and draws, say, six candidates: $\{-2.1,-0.4,0.9,1.8,2.3,3.5\}$. Rolling each through the model and scoring by predicted reward, the two best turn out to be the ones nearest the peak, $\{1.8,2.3\}$ (the elite set). CEM refits the sampling distribution to the elites: the new mean is their average, ${\mu }_1=(1.8+2.3)/2=2.05$, and the new standard deviation is their spread, ${\sigma }_1=\sqrt{\frac{1}{2}[(1.8-2.05{)}^2+(2.3-2.05{)}^2]}=0.25$. The next round samples from $\mathcal{N}(2.05,0.25^2)$ — already tightly centered on the optimum — and after a few such rounds the distribution collapses onto the best action. The whole method is: sample, keep the elite fraction, refit to them, repeat. It needs no gradients, which is what lets it optimize through a black-box learned model.

Empirically, on continuous-control benchmarks like the half-cheetah, PETS matched the asymptotic performance of strong model-free methods (SAC, PPO) while using roughly an order of magnitude fewer environment steps to get there.¹ The probabilistic ensemble was the essential part: ablating it to a single deterministic network collapsed the gains, because without a calibrated estimate of its own uncertainty the planner exploits the model's errors.

Where this leaves us

PETS is the most literal reading of the model-based promise: fit a dynamics model to observed transitions, then plan against it. Two ideas make it work despite compounding error. The probabilistic ensemble gives the planner a calibrated estimate of model uncertainty — members agree where data is plentiful and spread out where it is scarce, so trajectory sampling naturally distrusts rollouts into unexplored regions. And model-predictive control replans from the true state every step, so any imagined rollout only has to be trusted for a short horizon before reality resets the plan.

PETS still predicts full next states in the environment's native coordinates. The two methods in the companion lesson push on that. World Models and Dreamer learn a compact latent state and do almost all their learning by imagining inside it; MuZero goes further and predicts neither states nor pixels, only the reward, value, and policy that planning actually reads. Both continue in Model-Based Deep RL: World Models, Dreamer, and MuZero.