Diffusion Model Predictive Control — Planning by Painting the Future.

Suppose you're a robot arm and you want to do something useful — say, pick up a cup. You could learn a reflex: state goes in, action comes out, no thinking. That's model-free RL. Fast, but brittle, and if the goal changes you retrain. Or you could imagine: picture several ways the next few seconds might go, score each imagined future, and act on the best plan. That's model predictive control (MPC) — and it's flexible because the "what do I want" part is decided at run time, not baked into the weights.

The classic problem with MPC is that both halves are hard to learn well from offline data: the dynamics model ("if I take this action, what happens next?") and the action proposal ("what actions are even worth considering?"). Errors compound over a multi-step rollout, and the proposals are often too narrow.

"Diffusion Model Predictive Control" (D-MPC) by Guangyao Zhou and colleagues at Google DeepMind (arXiv:2410.05364, published at TMLR) replaces both halves with diffusion models — the same generative machinery behind image generators — and gets a planner that matches state-of-the-art offline RL while staying adaptable at run time.

The two things MPC needs

MPC runs a loop: at the current state, propose candidate action sequences, simulate each one forward with a dynamics model, score the resulting trajectories with a reward function, take the first action of the best plan, then re-plan from the new state. So you need:

1. An action proposal — a distribution over plausible multi-step action sequences to evaluate. (Random actions waste your whole compute budget on nonsense.)
2. A dynamics model — predicts the future states that a proposed action sequence would produce.

D-MPC learns both as multi-step diffusion models. That word "multi-step" is the load-bearing part.

Why diffusion, and why multi-step

A diffusion model learns to turn noise into a sample from some data distribution by gradually denoising. For images it paints a picture; here it "paints" a chunk of a trajectory — a whole sequence of actions, or a whole sequence of future states — in one coherent shot.

The standard alternative is a single-step dynamics model: predict state $s_{t+1}$ from $(s_t, a_t)$, then feed that prediction back in to get $s_{t+2}$, and so on. The problem is compounding error: each small prediction mistake becomes the input to the next prediction, and errors snowball over a long horizon until the imagined future is garbage.

A multi-step diffusion dynamics model instead generates the whole future window jointly, conditioned on the action sequence. Because it models the joint distribution over the trajectory rather than chaining one-step guesses, it sidesteps a lot of that snowball.

The core bet: model trajectories as a joint object you sample all at once, not as a chain of one-step predictions you stitch together. Diffusion is good at exactly that — generating high-dimensional, coherent samples in a single conditioned pass.

optional read — single-step vs multi-step factorization

A single-step model factorizes a length-$H$ rollout autoregressively:

$$p(s_{1:H} \mid a_{1:H}, s_0) = \prod_{t=1}^{H} p(s_t \mid s_{t-1}, a_{t-1})$$

Each factor's error feeds the next factor's input, so the variance of $\hat{s}_t$ grows with $t$. D-MPC's dynamics diffusion model instead samples the block jointly, $s_{1:H} \sim p_\theta(\cdot \mid a_{1:H}, s_0)$, learning intra-trajectory correlations directly. The action proposal is likewise a diffusion model over $a_{1:H}$ rather than a per-step policy, so candidate plans are coherent sequences instead of independent jittered actions.

The D-MPC loop

Because the reward function is plugged in during planning rather than learned into a policy, you can change the goal at run time and the planner adapts immediately — no retraining. That's the structural advantage of MPC, and D-MPC keeps it while fixing the quality problem that usually makes learned MPC lose to model-free RL.

Results on D4RL

D4RL is the standard offline-RL benchmark suite — locomotion and manipulation tasks where you learn purely from a fixed dataset of logged experience, with no further environment interaction during training. The headline findings:

• Significantly beats prior model-based offline MPC methods like MBOP. The multi-step diffusion dynamics model is the difference — less compounding error means the plans it scores are actually trustworthy.
• Competitive with SOTA model-based and model-free RL. A planner that imagines matches methods that bake the policy into weights — while keeping run-time flexibility they don't have.
• Beats prior diffusion-based planning baselines (e.g. Diffuser-style approaches) — so the win isn't "diffusion is magic," it's the specific propose-and-simulate-with-two-diffusion-models design.

The part that makes it more than a benchmark number

The paper specifically demonstrates two run-time superpowers that model-free policies don't have:

1. Optimize novel reward functions at run time. Give it a reward it never trained on, and it plans toward that new objective on the fly. The dynamics model doesn't care what you're optimizing — it just predicts futures; the reward does the steering.
2. Adapt to novel dynamics. When the world changes, the planner can adjust because planning is explicit search over imagined futures, not a frozen reflex.

This is the real argument for MPC over end-to-end policies: generality. A model-free policy is fast but committed — it does the one thing it was trained to do. A planner is slower but rethinks from scratch every step, so it bends to new goals and new conditions. The historical knock was that learned planners weren't accurate enough to compete. D-MPC's claim is that multi-step diffusion finally makes them accurate enough.

The trade is compute for flexibility. D-MPC pays planning cost every step (sampling and scoring trajectories) in exchange for being able to change its mind about the goal at run time. For many real control problems — where the objective genuinely shifts — that's the trade you want.

A vocabulary refresher (so the rest lands)

Three families of methods, one quick map:

• Model-free RL — learn a policy $\pi(a \mid s)$ directly: state in, action out. Fast at run time, but the goal is baked into the weights. Change the reward and you retrain. Examples: SAC, TD3, CQL.
• Model-based RL with a learned policy — learn a dynamics model, use it to train a policy. Better sample efficiency, still a frozen policy at the end.
• Planning / MPC — learn a dynamics model, but don't distill a policy; instead search over action sequences at run time using the model. The goal lives in a reward function you pass in during planning, so it's swappable. D-MPC lives here.

The reason planning lost popularity is precisely the reason D-MPC is interesting: learned dynamics models were too inaccurate over multi-step rollouts to plan against, so model-free reflexes won on raw performance. Fix the accuracy and planning's flexibility becomes free upside.

Why diffusion beats the usual proposal tricks

The action proposal is doing more work than it looks. A planner is only as good as the candidates it bothers to evaluate — score a thousand nonsense action sequences and you learn nothing. Prior planners proposed candidates in limited ways:

• Gaussian noise around a mean (CEM / MPPI style) — jitter each action independently. Produces incoherent sequences: step 3 doesn't "know" what step 2 did, so the candidate trajectories look like twitchy noise rather than purposeful motion.
• Autoregressive policy rollout — sample actions one at a time from a learned policy. Coherent, but narrow — it only proposes what the policy already prefers, so the planner can't discover anything the base policy wouldn't have done.
• Diffusion proposal (D-MPC) — generate a whole multi-step action sequence as one joint sample. Coherent and diverse: the model has learned the manifold of plausible action sequences from the offline data, so its candidates are realistic multi-step behaviours, not jitter and not policy-clones.

Same story on the dynamics side. The multi-step diffusion dynamics model generates the future state window jointly, so the candidate futures it hands the reward function are internally consistent — no snowballing one-step error to mislead the scoring.

The trick for run-time rewards: guidance

How does a model that never trained on your new reward suddenly optimize it? The mechanism is the same one image diffusion models use to follow a text prompt: guidance. During the denoising that generates a trajectory, you nudge each step toward regions of higher reward. The unconditional diffusion model supplies "what trajectories are physically plausible"; the reward gradient supplies "which of those I actually want." Multiply them together and you sample plausible and high-reward trajectories.

That factorization is the whole reason the reward can be swapped at run time. Plausibility lives in the (fixed) learned diffusion model; preference lives in the (swappable) reward. You never retrain the expensive part to chase a new objective.

optional read — reward guidance during denoising

A diffusion model samples a trajectory $\tau$ by reversing a noising process, following the score $\nabla_\tau \log p(\tau)$ at each denoising step. To bias sampling toward high reward $r(\tau)$, add the reward gradient:

$$\nabla_\tau \log\big[p(\tau)\, e^{\,\alpha\, r(\tau)}\big] \;=\; \underbrace{\nabla_\tau \log p(\tau)}_{\text{plausible (fixed model)}} \;+\; \underbrace{\alpha\, \nabla_\tau r(\tau)}_{\text{preference (swappable)}}$$

The temperature $\alpha$ trades realism against greedy reward-chasing. Because the first term is the frozen learned model and the second is just the reward you pass in, a brand-new $r(\cdot)$ at run time changes only the second term — no retraining. This is also how D-MPC adapts to novel dynamics: re-condition the dynamics diffusion on observed transitions.

What D4RL actually measures

D4RL (Datasets for Deep Data-Driven RL) is the field's standard offline-RL gym. The relevant slices here are continuous-control locomotion (HalfCheetah, Hopper, Walker2d) under dataset qualities like medium, medium-replay, and medium-expert — i.e. you learn from logged trajectories of varying skill, never touching the live environment during training. It's a brutal test of a dynamics model: the offline data has gaps, and a planner that hallucinates good outcomes in those gaps gets punished.

D-MPC's results there — beating prior MPC methods like MBOP and matching SOTA model-based and model-free RL — are the evidence that the multi-step diffusion dynamics model is accurate enough to plan against. And beating earlier diffusion planners (Diffuser) shows the gain is the specific propose-and-simulate design, not just "add diffusion."

Honest limitations

• Planning is expensive. Sampling from diffusion models every control step costs more than a single forward pass of a reflex policy. Real-time, high-frequency control is where this bites, and is an active engineering frontier (faster samplers help).
• Offline data still bounds it. The dynamics and proposal models only know what the offline dataset showed them. Plan into a region the data never covered and the imagined futures get unreliable — the usual offline-RL distribution-shift caveat.

arXiv:2410.05364 — Diffusion Model Predictive Control, Guangyao Zhou et al. (TMLR). Presented at the YC Paper Club.