Planning with Learned Models

Learn a world model, roll out candidate action plans inside it, run the best first step, then replan from the real next state.

CEM keeps two elites at 4.0 and 3.2, refits the mean to 0.52 and 0.56, and MPC runs only the first action 0.52

Why Does This Exist?

A fixed policy memorizes what worked during training. Move the goal, add an obstacle, or change the reward, and the memorized mapping goes stale until thousands of fresh gradients rewrite it. That lag hurts in robotics and control where targets shift at test time.

A learned dynamics model offers a faster route. Learn how states evolve and what rewards look like, then search for good action sequences inside that model whenever a new goal appears. No weight rewrite needed for each new target.

Long open loop plans break on their own error. Each predicted step adds a little model mistake, and after 20 or 30 steps the imagined state sits far from the real one. The tail of the plan then commands actions for a world you are not in.

This page covers the standard fix: Model Predictive Control with receding horizon plus the Cross Entropy Method for search. Optimize over horizon H, run only the first action, observe the true next state, and plan again. For learning the model itself, see model based deep RL. For compact latent models that make this cheap, see learning latent dynamics. For a full latent planner, see Dreamer.

You need Gaussian sampling before this page. A Gaussian has a mean mu for the center and a variance sigma squared for the spread. CEM samples action plans from one, keeps the winners, and moves the center toward them.

Think of It Like This

Helm checks every wave

You pilot a small boat through chop. Every few seconds you picture ten ways to work the rudder over the next ten seconds, keep the three smoothest, and average them into a plan. Then you run only the first second of it.

At the next crest you look at the real water, not your memory of it, and picture ten fresh plans. The old nine seconds go in the bin.

Where the analogy stops: you use feel and sight. CEM uses numbers, with N sampled plans, K elites, and a smoothing factor beta that stops the distribution from snapping shut too fast. And your boat answers slowly, while a robot may need this whole loop inside 20 milliseconds.

How It Actually Works

Write s t for the true state, H for horizon, A for a plan of H actions, p hat for the learned dynamics, r hat for the learned reward, gamma for discount, N for candidates, K for elites, beta for smoothing. Plans roll forward as s hat next equals f hat of s hat and a, and the plan score is the discounted sum of predicted rewards plus an optional terminal value.

MPC runs five steps: solve for the best plan A star from s t, run only a t star, throw away the rest, observe the true s t plus 1, and solve again from there. Grounding every step in truth is what stops compounding error.

The CEM search

CEM is derivative free, which helps when learned dynamics make the return surface jagged. Each planning call runs M rounds:

Sample N plans from a Gaussian with mean mu and variance sigma squared, clipped to actuator limits.
Roll each plan through p hat and score its predicted return J.
Sort, keep the top K elites, and compute their mean mu elite and variance sigma elite squared.
Smooth with Polyak averaging: mu next equals beta times mu elite plus (1 minus beta) times mu, and likewise for variance.
After M rounds, return the first entry of mu as a t star and shift the rest forward to warm start the next call.

Worked example

Horizon 2, one action dim, N 4, K 2, beta 0.8, prior mean [0.0, 0.0], prior variance [1.0, 1.0]. Four candidates score through the model:

A1 [0.5, 0.8] scores 3.2
A2 [minus 0.5, 0.2] scores 1.0
A3 [0.8, 0.6] scores 4.0
A4 [minus 1.0, minus 0.5] scores minus 0.5

Elites are A3 and A1. Their mean is [(0.8 plus 0.5) divided by 2, (0.6 plus 0.8) divided by 2], which is [0.65, 0.70]. Their variance is [0.0225, 0.0100] because the squared gaps are 0.0225 each on step one and 0.0100 each on step two, halved across two elites.

Smooth with beta 0.8: mu1 equals 0.8 times [0.65, 0.70] plus 0.2 times [0.0, 0.0], which is [0.52, 0.56]. Variance becomes 0.8 times [0.0225, 0.0100] plus 0.2 times [1.0, 1.0], which is [0.2180, 0.2080]. MPC runs a0 star 0.52, bins 0.56 except as a warm start of [0.56, 0.00] next step. The diagram shows the four dots, the two elites circled, and the arrow landing on 0.52.

Code

This snippet runs the exact refit from the example.

1import numpy as np2
3
4def cem_refit(5    candidates: np.ndarray, returns: np.ndarray, prior_mean: np.ndarray, prior_var: np.ndarray, elites: int = 2, beta: float = 0.86) -> tuple[np.ndarray, np.ndarray, np.ndarray]:7    """One deterministic CEM update matching the worked example."""8    idx = np.argsort(returns)[-elites:]9    picked = candidates[idx]10    elite_mean = np.mean(picked, axis=0)11    elite_var = np.mean((picked - elite_mean) ** 2, axis=0)12    new_mean = beta * elite_mean + (1.0 - beta) * prior_mean13    new_var = beta * elite_var + (1.0 - beta) * prior_var14    return new_mean, new_var, new_mean[0]15
16
17cands = np.array([[[0.5], [0.8]], [[-0.5], [0.2]], [[0.8], [0.6]], [[-1.0], [-0.5]]], dtype=float)18rets = np.array([3.2, 1.0, 4.0, -0.5], dtype=float)19mu0 = np.zeros((2, 1))20var0 = np.ones((2, 1))21mu1, var1, a0 = cem_refit(cands, rets, mu0, var0)22print([f"{v:.2f}" for v in mu1.flatten()])23# -> ['0.52', '0.56']24print([f"{v:.4f}" for v in var1.flatten()])25# -> ['0.2180', '0.2080']26print(f"{float(a0):.2f}")27# -> 0.52

Watch Out For

Control latency that eats the plan

The symptom is a robot that jitters or lags even though simulation looks clean. With N 500, H 20, and M 5, each step needs 50,000 model calls. At 20 to 50 milliseconds of GPU time, the command lands on a stale state and the loop goes unstable.

Fix it by warm starting mu from the shifted prior plan so M drops to 1 or 2, by planning in a compact latent state instead of pixels, and by distilling the planner into a fast actor for routine steps while keeping CEM for tricky ones. Measure end to end latency before trusting any horizon.

Collapsing onto the first good plan

The symptom is a planner that finds one decent trajectory early, shrinks variance to near zero, and never improves. A small obstacle shift then breaks it because no samples land near the new route.

The cause is raw elite refits with beta at 1.0 or tiny elite sets. Fix it by keeping beta near 0.8, keeping elites near 10 percent of N, and flooring variance so sampling stays alive. If variance still dives, raise N before raising M.

The Quick Version

Planning with learned models searches action sequences at test time instead of memorizing one policy.
MPC plans over horizon H, runs only the first action, and replans from the true next state.
CEM samples N plans, keeps the K best, and smooths the Gaussian toward them without gradients.
The worked step keeps [0.8, 0.6] and [0.5, 0.8], refits the mean to [0.52, 0.56] with variance [0.2180, 0.2080], and runs 0.52.
Warm starts and latent models keep this inside real time limits, and variance floors stop early collapse.