Model-Predictive Control and Planning

Model-predictive control is planning that does not trust its own plan: optimise a short action sequence in the model, execute only the first action, and optimise again from wherever you actually landed.

A world model gives you a function to call instead of the world. Dyna spends those calls offline, distilling them into a value function. MPC spends them online, at every decision, and throws most of the result away. This post builds three sampling-based planners on a 2-D point mass, then measures the three knobs that decide whether MPC works: the horizon, how often to replan, and how to know when the model itself is guessing. All numbers come from running the code shown with NumPy.

The system and the objective#

A point mass in the plane with state $x = (p_x, p_y, v_x, v_y)$, a force $u \in [-1, 1]^2$, linear drag, and a time step of 0.1. The task is to bring it from $(2, -1.5)$ to the origin and stop there.

import numpy as np
DT, DRAG = 0.1, 0.5
 
def true_step(x, u):
    """Point mass in 2-D. x = [px, py, vx, vy], u = force in [-1, 1]^2. Batched over leading dims."""
    p, v = x[..., :2], x[..., 2:]
    v2 = v + DT * (u - DRAG * v)
    return np.concatenate([p + DT * v2, v2], -1)
 
def cost(x, u):
    return (x[..., :2] ** 2).sum(-1) + 0.1 * (x[..., 2:] ** 2).sum(-1) + 0.01 * (u ** 2).sum(-1)
 
def rollout_cost(model, x0, U):
    """U: (K, H, 2) candidate action sequences. Returns (K,) total imagined cost."""
    x = np.repeat(x0[None], len(U), 0); J = np.zeros(len(U))
    for t in range(U.shape[1]):
        x = model(x, U[:, t]); J += cost(x, U[:, t])
    return J

rollout_cost is the only place the model is used. It pushes $K$ candidate action sequences through $H$ model steps in parallel and returns each candidate's total cost. Every planner below is a different answer to one question: which $K$ sequences should be evaluated?

Random shooting#

The simplest answer: sample $K$ sequences uniformly and keep the best.

def random_shooting(model, x0, H, K, rng, **_):
    U = rng.uniform(-1, 1, (K, H, 2))
    return U[np.argmin(rollout_cost(model, x0, U))]

The receding-horizon loop wraps any planner. It plans, executes replan_every actions in the real system, and plans again. When it replans, it warm-starts from the remainder of the previous plan, shifted forward in time.

def run_mpc(planner, model, H=15, K=300, T=60, replan_every=1, seed=0, warm=True):
    rng = np.random.default_rng(seed)
    x = np.array([2.0, -1.5, 0.0, 0.0]); total = 0.0; plan, k = None, 0
    for t in range(T):
        if plan is None or k >= replan_every:
            mu = None
            if warm and plan is not None and planner is not random_shooting:
                mu = np.concatenate([plan[k:], np.zeros((k, 2))])   # shift the old plan
            plan = planner(model, x, H, K, rng, mu=mu); k = 0
        u = plan[k]; k += 1
        x = true_step(x, u); total += cost(x, u)
    return total, np.linalg.norm(x[:2])

The bottleneck of random shooting is dimensionality. A plan has $2H$ numbers. With $H = 15$ that is a 30-dimensional box, and 300 uniform samples cover it about as well as 300 darts cover a 30-dimensional room. The best sample is merely the least bad.

The cross-entropy method#

CEM turns the sample-and-pick step into an iterative search. Keep a Gaussian over action sequences, sample from it, refit the Gaussian to the elite fraction, repeat. To make the comparison fair, the same total budget of $K$ rollouts is split across the iterations.

def cem(model, x0, H, K, rng, iters=5, elite=0.1, mu=None, **_):
    mu = np.zeros((H, 2)) if mu is None else mu; sd = np.ones((H, 2)) * 0.6
    n_el = max(2, int(K / iters * elite))
    for _ in range(iters):                                    # same total budget K
        U = np.clip(mu + sd * rng.standard_normal((K // iters, H, 2)), -1, 1)
        el = U[np.argsort(rollout_cost(model, x0, U))[:n_el]]
        mu, sd = el.mean(0), el.std(0) + 1e-3
    return mu

Each iteration concentrates samples where the previous one found low cost. The returned plan is the elite mean, not the single best sample, which averages out some sampling noise. CEM is the planner inside PETS (Chua et al., 2018, arXiv:1805.12114) and PlaNet (Hafner et al., 2018, arXiv:1811.04551).

MPPI#

CEM throws away every non-elite sample. Model-predictive path integral control (Williams et al., 2017) uses all of them, weighting each perturbation by the exponentiated negative cost:

The temperature $\lambda$ interpolates between "average everything" (large $\lambda$) and "take the best" (small $\lambda$).

def mppi(model, x0, H, K, rng, lam=0.5, sigma=0.5, mu=None, **_):
    mu = np.zeros((H, 2)) if mu is None else mu
    eps = sigma * rng.standard_normal((K, H, 2))
    U = np.clip(mu + eps, -1, 1)
    J = rollout_cost(model, x0, U)
    w = np.exp(-(J - J.min()) / lam); w /= w.sum()         # subtract min for numerical safety
    return np.clip(mu + (w[:, None, None] * eps).sum(0), -1, 1)

MPPI does a single weighted update per call and relies on the warm start to accumulate improvement across time steps. That makes it cheap per decision and smooth across decisions, which is why it is popular on real vehicles.

With the true model, $H = 15$, $K = 300$, 60 steps, mean over 5 seeds:

All three reach the goal, because replanning forgives a great deal. The cost column shows the difference: random shooting wanders; CEM and MPPI take near-direct paths.

The horizon trade-off#

A quadratic cost toward the origin rewards greed, so to make the horizon matter, add a forbidden strip that the straight-line path crosses:

def cost_wall(x, u):
    """Quadratic cost plus a forbidden strip at 0.8 < px < 1.2 for py < 0.5."""
    px, py = x[..., 0], x[..., 1]
    inside = (px > 0.8) & (px < 1.2) & (py < 0.5)
    return (x[..., :2] ** 2).sum(-1) + 0.01 * (u ** 2).sum(-1) + 50.0 * inside
 
def biased_model(x, u, drag=0.0, gain=1.6):
    """A 'learned' model that got drag and actuator gain wrong."""
    p, v = x[..., :2], x[..., 2:]
    v2 = v + DT * (gain * u - drag * v)
    return np.concatenate([p + DT * v2, v2], -1)

First, how wrong is the biased model? Feed it and the real system the same random actions from the same 1000 random states, and measure the mean position gap:

That is compounding error in a continuous system: velocity errors integrate into position errors, so the gap grows faster than linearly. Now sweep the horizon, with cost replaced by cost_wall, 80 steps, 10 seeds:

Read the true-model columns first. At $H = 3$ the planner cannot see past the wall and grinds along it: myopia. At $H = 10$ both planners find the detour. Beyond that, cost rises again even with a perfect model, because the search space grows with $H$ while the budget does not. Random shooting suffers most, as the dimensionality argument predicts.

The biased columns are noisier and not monotone, and that is itself the lesson. Replanning every step feeds back the true state, which absorbs much of a smooth, systematic bias. The biased CEM run even scored better at $H = 40$: a wrong model can accidentally discourage aggressive plans. A plan's cost under the model and its cost in reality are different numbers, and only the second one counts.

When to replan#

Feedback is what makes the bias tolerable, so remove it gradually. Plan with a milder bias, biased_model(x, u, drag=0.1, gain=1.3), CEM, $H = 20$, the original cost without the wall, and execute $k$ actions of each plan before replanning:

Executing the whole plan open-loop costs 45% more and ends far from the target. The rule is simple: replan as often as compute allows, and if you cannot, replan at least as often as the model's error grows past what the task tolerates, which the drift table above tells you how to estimate.

Knowing when the model is guessing#

A single model returns a confident prediction even far from its data. PETS answers this with an ensemble: train several models on bootstrap resamples and treat their disagreement as epistemic uncertainty. A linear ensemble on the point mass shows the mechanism:

def fit_ensemble(n_models, n_data, rng):
    """PETS-flavoured: bootstrap linear models on a small random dataset."""
    X = rng.uniform(-2, 2, (n_data, 4)); U = rng.uniform(-1, 1, (n_data, 2))
    Y = true_step(X, U) + 0.01 * rng.standard_normal((n_data, 4))
    Z = np.concatenate([X, U, np.ones((n_data, 1))], 1); Ws = []
    for _ in range(n_models):
        i = rng.integers(n_data, size=n_data)
        Ws.append(np.linalg.lstsq(Z[i], Y[i], rcond=None)[0])
    return Ws

Five members, evaluated one step ahead on 500 held-out states:

Disagreement tracks error across two orders of magnitude, without access to the true dynamics. That is what makes it usable inside a planner: PETS propagates particles through randomly chosen members ("trajectory sampling") and averages the plan's cost over them, so a plan that passes through regions of disagreement is scored against several divergent futures instead of one confident one.

What to try next#

1. Add disagreement as a penalty, J + beta * std, and rerun the horizon sweep with an ensemble trained on 12 transitions. Long horizons should stop paying off once disagreement accumulates.
2. Replace the linear members with small MLPs and the point mass with a pendulum. Nonlinearity is where ensembles start to disagree in useful places.
3. Replan only when the observed state leaves the predicted state by more than a threshold. Measure compute saved against cost lost.
4. Move the planner into a learned latent space, as PlaNet does, or distil it into a policy, as Dreamer does. The trade is per-step planning compute against amortised policy quality.
5. Score the model by the cost of the plans it produces, not by one-step error. That is the planning-centric view of evaluating a world model, and the reason large models like Genie are judged by what can be done inside them.