Diffusion Model Predictive Control

TL;DR

D-MPC (Google DeepMind) factorizes offline model-based RL into two independently-trained multi-step diffusion models — a joint trajectory-level dynamics model $p(s_{t+1:t+H} \mid s_t, a_{t:t+H})$ and a multi-step action proposal $p(a_{t:t+H} \mid s_t)$ — and plans at run-time with a simple “sample, score, and rank” (SSR) procedure. This matches model-free offline RL (CQL/IQL) on D4RL locomotion / Adroit / Kitchen while retaining MPC’s three big advantages: optimize novel rewards at test time, fine-tune the dynamics model alone when hardware changes, and avoid compounding errors by predicting full trajectories rather than autoregressing one step at a time.

Key claims

Multi-step diffusion models for both dynamics and action proposal, combined via a simple SSR planner, significantly outperform MBOP (the single-step deterministic-MLP baseline) on D4RL: locomotion average 65.98 vs 33.13, Adroit cloned 39.28 vs 22.47 [Table 2].
D-MPC matches or exceeds model-free offline RL methods (CQL, IQL, BC, MOReL, DT) on D4RL despite carrying MPC’s extra flexibility for run-time reward/dynamics changes [Table 2].
Novel rewards can be optimized at test time by swapping the learned value model for any objective $r'$ in the SSR scoring step, without retraining the dynamics or action proposal — demonstrated by inducing lunging / balancing / jumping behaviors in a Walker2d trained only on medium-replay data [§4.2, Fig. 1].
Factorized dynamics-vs-proposal training enables hardware-defect adaptation: fine-tuning only the dynamics model on 100 episodes of “play” data lifts a defective-actuator Walker2d from 22.74 to 30.65, while Diffuser (which models joint state+action) drops from 25.85 to 6.8 under the same fine-tuning budget [§4.3].
Ablation: each component (multi-step diffusion dynamics, multi-step diffusion action proposal, SSR planner) helps independently, but together they reach 65.98 vs 33.13 for MBOP — with the multi-step diffusion dynamics specifically reducing compounding error in long-horizon rollouts compared to single-step diffusion and autoregressive transformer baselines [Table 3, §4.4, Fig. 2].
D-MPC can be distilled into a fast reactive MLP policy: a 6-domain locomotion student achieves 65.08 vs the planner’s 65.98 (and 41.92 for vanilla BC), recovering nearly all the benefit at MLP inference speed [§4.5].

Method

D-MPC has three learned components, all trained on a fixed offline dataset of (state, action, reward) triples: (1) a multi-step diffusion dynamics model $p_\theta(s_{t+1:t+H} \mid s_t, a_{t:t+H}, h_t)$ that jointly denoises an entire $H$ -step future state trajectory conditioned on past history and a proposed action sequence; (2) a multi-step diffusion action proposal $p_\phi(a_{t:t+H} \mid s_t, h_t)$ that proposes joint action chunks via behavior cloning on the offline data; and (3) a transformer-based value head $V(s_{t+H})$ that estimates reward-to-go at the leaf of the planning horizon. Training uses standard denoising score matching, no ensembles, no separate reward model. At control time, the SSR planner samples $M$ candidate action sequences from the proposal, rolls each through the dynamics model to predicted state sequences, scores them with $J = \sum_{\tau=t}^{t+H-1} \gamma^\tau r(s_\tau, a_\tau) + \gamma^H V(s_{t+H})$ , picks the best, executes the first action, and replans. Novel rewards $J'$ are dropped in by replacing $J$ at test time; novel dynamics are handled by SFT-ing only $p_\theta$ on play data from the new environment. The factorization $p(s, a) = p(s \mid a) \cdot p(a)$ — distinct from Diffuser’s joint $p(s, a)$ — is what enables independent dynamics fine-tuning.

Results

On D4RL locomotion (HalfCheetah / Hopper / Walker2d, medium and medium-replay), D-MPC averages 65.98 vs 33.13 for MBOP, 63.77 for Diffuser, 41.92 for BC, 65.45 for CQL, and 67.47 for IQL [Table 2]. On Adroit (pen / door / hammer, cloned), D-MPC reaches 39.28 vs 22.47 MBOP, 28.36 DT, 37.15 BC, 44.90 IQL [Table 2]. On Franka Kitchen (mixed / partial), D-MPC averages 70.42 vs 44.75 BC, 51.25 CQL, 48.65 IQL [Table 2]. Long-horizon dynamics-prediction accuracy: multi-step diffusion has the lowest median RMSD on non-velocity coordinates across horizons 5–50 on Walker2d medium-replay and expert, beating single-step diffusion and an autoregressive transformer baseline [Fig. 2]. Novel-reward generalization: from a single Walker2d-medium-replay-trained model, swapping in $r' = -10(h - h^*)^2$ with $h^* \in \{0.4, 1.1, 1.7\}$ produces lunge / balance / jump behaviors [Fig. 1, §4.2]. Distillation: 65.08 vs D-MPC’s 65.98 on the same 6 locomotion settings, ~3 orders of magnitude faster at inference [§4.5]. The main limitation, noted by the authors, is run-time cost — diffusion sampling at every replan step is much slower than a reactive MLP, motivating the distillation experiment.

Why it’s interesting

D-MPC is the cleanest filed instance of the factorized world model = dynamics × policy proposal pattern that the wiki’s World Foundation Models cluster keeps circling, but executed in the offline-RL low-dim-state regime rather than the pixel-video regime that dominates the cluster. It contrasts directly with Helios: Real Real-Time Long Video Generation Model and PAN: A World Model for General, Interactable, and Long-Horizon World Simulation (pixel-space joint world models) and with Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models (energy-based planners): D-MPC’s factorization is what enables independent dynamics fine-tuning under hardware drift, a claim the joint-modelling pixel-WFMs cannot make. It also sits adjacent to Generative World Renderer and MirageLSD: The First Live-Stream Diffusion AI Video Model on the question of whether the right “world model” output is a predicted future state or a predicted future pixel observation — D-MPC pushes hard on the state-prediction side, with distillation as the inference-cost answer. Worth tracking as a precedent now that pixel-space WFMs (Genie 3, Cosmos Predict) are starting to be used as planners themselves.