12.Model-Based RL & World Models

12.1 What a model buys

A model predicts what the environment will do: $\hat p(s' \mid s,a)$, and usually a reward model $\hat r(s,a)$ as well. Given one, three things become possible that were not before.

Generate synthetic experience. Chapter 7's Dyna already did this with a tabular model. Each real transition can be mined many times over by training on imagined ones.

Plan at decision time. Optimize a short action sequence against the model, execute the first action, re-plan. This is model predictive control, and it needs no policy network at all.

Transfer across tasks. The dynamics of a robot arm do not change when the goal moves. A model learned for reaching is immediately useful for pushing — whereas a value function must be relearned from scratch. For a robot expected to do many things, this is the deepest argument for model-based methods.

The empirical payoff is large. Model-based methods routinely reach good policies in 10–100× fewer environment steps than model-free ones. PILCO famously learned cart-pole swing-up in under 20 seconds of real robot interaction, at a time when model-free methods needed hours.

12.2 Learning dynamics, and knowing what you do not know

The naive approach — train a network to minimize $\|\hat f(s,a) - s'\|^2$ — produces a model that is confidently wrong off the training distribution, which is exactly where a planner will take it. The planner is an adversary that seeks out wherever the model is optimistic.

The fix is to model uncertainty explicitly, and to distinguish two kinds.

Aleatoric uncertainty is noise inherent in the system: sensor noise, contact stochasticity, unmodelled disturbances. More data does not reduce it. Capture it by predicting a distribution — a Gaussian with learned mean and variance — rather than a point.

Epistemic uncertainty is ignorance: the model has not seen this region of state space. More data does reduce it, and this is the kind that matters for planning, because it tells you where not to trust yourself. Capture it with an ensemble of $K$ models trained on bootstrapped data with different initializations; where they agree, the model is confident, and where they disagree, it is guessing.

The law of total variance separates them cleanly:

$\underbrace{\operatorname{Var}[s']}_{\text{total}} = \underbrace{\mathbb{E}_k\!\left[\sigma_k^2\right]}_{\text{aleatoric}} + \underbrace{\operatorname{Var}_k\!\left[\mu_k\right]}_{\text{epistemic}},$

where $k$ indexes ensemble members. The first term is the average predicted noise; the second is the disagreement between members.

12.3 Why imagination must be short

Here is the central difficulty, and it is not subtle once stated.

Suppose the model has per-step error at most $\varepsilon$, and the true dynamics are Lipschitz with constant $L$. Roll the model forward $H$ steps, feeding each prediction back in as the next input. The error after $H$ steps is bounded by

$\varepsilon_H \le \varepsilon \cdot \frac{L^H - 1}{L - 1}.$

For $L > 1$ — which holds for any system with unstable or chaotic modes, meaning essentially every interesting robot — this is geometric in the horizon. Doubling the rollout length does far more than double the error.

Every ensemble member in that widget has small per-step error. The fan is not caused by bad models; it is caused by feedback. Each prediction becomes the next input, so the rollout steadily walks off the data distribution into a region where all bets are off.

12.4 Planning with a learned model: CEM and MPC

Given a model, one can skip the policy entirely and optimize actions directly at each step.

Random shooting samples $N$ action sequences, evaluates each through the model, and executes the first action of the best. Simple, embarrassingly parallel, and weak in high dimensions.

The cross-entropy method iterates that idea. Sample sequences from a Gaussian over action trajectories, keep the top-performing "elite" fraction, refit the Gaussian to the elites, repeat. It is importance-sampled optimization: each iteration reweights the sampling distribution toward the high-return region, and the elite-refitting step is the maximum-likelihood update of that reweighting.

Model predictive control wraps either in a receding horizon: plan $H$ steps, execute one, discard the rest, re-plan from the state you actually reached. Re-planning is what makes MPC robust — the model only needs to be accurate for the few steps before the next correction, which is exactly the regime §12.3 says it can manage.

PETS combines all of this: an ensemble of probabilistic models, CEM planning, and trajectory sampling that propagates particles through ensemble members so that the planner's return estimates account for both uncertainty types.

use burn::prelude::*;
 
pub struct CemPlanner {
    horizon: usize,
    population: usize,
    elites: usize,
    iterations: usize,
    action_dim: usize,
}
 
impl CemPlanner {
    /// Returns the first action of the best plan — MPC discards the rest.
    pub fn plan<B: Backend>(
        &self,
        ensemble: &DynamicsEnsemble<B>,
        state: &Tensor<B, 1>,
        reward: impl Fn(&Tensor<B, 2>, &Tensor<B, 2>) -> Tensor<B, 1>,
        rng: &mut impl rand::Rng,
    ) -> Vec<f32> {
        // Sampling distribution over action SEQUENCES, refit each iteration.
        let mut mean = vec![0.0f32; self.horizon * self.action_dim];
        let mut std = vec![0.5f32; self.horizon * self.action_dim];
 
        for _ in 0..self.iterations {
            let mut candidates: Vec<(f32, Vec<f32>)> = Vec::with_capacity(self.population);
 
            for _ in 0..self.population {
                let seq: Vec<f32> = mean
                    .iter()
                    .zip(&std)
                    .map(|(&m, &s)| m + s * gaussian(rng))
                    .collect();
 
                // Roll the plan forward, switching ensemble member each step so
                // the return averages over epistemic uncertainty (TS-∞).
                let mut s = state.clone();
                let mut total = 0.0f32;
                for t in 0..self.horizon {
                    let a = slice_action(&seq, t, self.action_dim);
                    let member = rng.gen_range(0..ensemble.len());
                    let next = ensemble.predict(member, &s, &a);
                    total += reward(&s.clone().unsqueeze(), &a.clone().unsqueeze())
                        .into_scalar()
                        .elem::<f32>();
                    s = next;
                }
                candidates.push((total, seq));
            }
 
            // Keep the elites and refit — the cross-entropy step.
            candidates.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
            let elite = &candidates[..self.elites];
            for i in 0..mean.len() {
                let vals: Vec<f32> = elite.iter().map(|(_, s)| s[i]).collect();
                let m = vals.iter().sum::<f32>() / vals.len() as f32;
                let v = vals.iter().map(|x| (x - m).powi(2)).sum::<f32>() / vals.len() as f32;
                mean[i] = m;
                std[i] = v.sqrt().max(0.05); // floor keeps the search from collapsing
            }
        }
        mean[..self.action_dim].to_vec()
    }
}

12.5 Latent world models

Everything above predicts in state space. For a robot with cameras, the observation is a 64×64×3 image, and predicting pixels is both expensive and wasteful — most pixel-level detail is irrelevant to control.

Latent world models learn a compact latent state $z_t$ in which dynamics are simple, and predict there. The Dreamer family learns four components jointly: an encoder $q(z_t \mid z_{t-1}, a_{t-1}, o_t)$, a transition model $p(z_t \mid z_{t-1}, a_{t-1})$, a reward predictor, and a decoder used only as a training signal.

Training maximizes an evidence lower bound with the familiar two-term structure:

$\mathcal{L} = \underbrace{\mathbb{E}_q\left[\log p(o_t \mid z_t) + \log p(r_t \mid z_t)\right]}_{\text{reconstruction}} - \underbrace{\beta\, D_{\text{KL}}\!\left(q(z_t \mid \cdot) \,\|\, p(z_t \mid z_{t-1}, a_{t-1})\right)}_{\text{consistency}}.$

The reconstruction term forces the latent to retain what matters; the KL term forces the transition model to predict it. The payoff is that the policy is then trained entirely inside the latent model — thousands of imagined trajectories per real episode, at negligible cost, with no rendering.

12.6 Model-based versus model-free, decided

The trade is clean enough to state as a rule.

Model-based methods win when samples are expensive and dynamics are learnable — a real robot, a slow simulator, smooth dynamics without hard contact. Model-free methods win when samples are cheap or dynamics are hard to model — massively parallel simulation, contact-rich manipulation where the model would have to capture exactly the discontinuities it is worst at.

Hybrids increasingly win outright: learn a model, use it to generate short imagined rollouts, and train a model-free algorithm on the mixture. MBPO is that recipe, and the boundary between the two families is genuinely blurring.

There is also a robotics-specific option this chapter has been circling: rather than learning the dynamics from scratch, fit the parameters of a physics simulator you already trust — masses, friction coefficients, motor constants. That is system identification, it needs far less data than learning a network, and Chapter 15 builds it.

12.7 Chapter bridge

Part II is complete. We can approximate value functions (Chapter 8), train deep value-based agents on discrete actions (Chapter 9), optimize policies directly with stable on-policy updates (Chapter 10), learn off-policy in continuous action spaces with maximum entropy (Chapter 11), and learn dynamics models to rehearse in imagination (this chapter).

Every one of those chapters treated the robot as an abstract environment — a thing that returns $s'$ and $r$. Part III opens the box. Chapter 13 asks what is actually inside when the environment is a physical machine: kinematics, Jacobians, the manipulator equation, and the classical controllers — PID, LQR, operational-space control — that RL must be measured against rather than assumed to beat. Reacher stops being a task and becomes a mechanism, with dynamics we derive from first principles.

References