Diffusion Model Predictive Control
Authors: Guangyao Zhou
Sivaramakrishnan Swaminathan
Rajkumar Vasudeva Raju
J Swaroop Guntupalli
Wolfgang Lehrach
Joseph Ortiz
Antoine Dedieu
Miguel Lazaro-Gredilla
Kevin Patrick Murphy
Venue: Accepted by TMLR ()

Abstract: We propose Diffusion Model Predictive Control (D-MPC), a novel MPC approach that learns a multi-step action proposal and a multi-step dynamics model, both using diffusion models, and combines them for use in online MPC. On the popular D4RL benchmark, we show performance that is significantly better than existing model-based offline planning methods using MPC (e.g. MBOP) and competitive with state-of-the-art (SOTA) model-based and model-free reinforcement learning methods. We additionally illustrate D-MPC’s ability to optimize novel reward functions at run time and adapt to novel dynamics, and highlight its advantages compared to existing diffusion-based planning baselines.

Published in Transactions on Machine Learning Research (05/2025)
Diffusion Model Predictive Control
Guangyao Zhou' Google Mind
stannis Mlanniss@ooogl..com com com
Sivaramakrishman Swaminathan Google Deep
sivark@google.co
Rajkumar Vasudeva Raju Google Mind
vujeraju@@google.com
J. Swaroop Guntupalli Google Mid
swaroopgj@google.com
Wolfgang Lehrach Google Mind
upl@google.com com
Joseph Ortiz Google DeepMind
jocortiz@google.cm com com
Antoine Dedicu Google Mind
udedien@@goolee.com
Miguel Gredilla Google DeepMind
Atzaroreddilaeggogee.comm
Kevin Murphy Google DeepMind
kpmourphhy@ggoogll.com
Reviewed on OpenReview: https: // openreview. net/ forum? id= putgffHHJm
Abstract
We propose Diffusion Model Predictive Control (D-MPC), a novel MPC approach that learns multi-step action proposal and multi-step dynamics model, both using diffusion models, and combines them for use in online MPC. On the popular D4RL benchmark, we show performance that is significantly better than existing modellbased offline planning methods using MPC (e.g. MBOP(Argensse (Argenson & Dulac-Arnold] 2021) and competitive with stato-of-the-art (SOTA) model-based and model-frce reinforcement learning methods. We additionally illustrate D-MPO's ability to optimize novel reward functions al run lime and adapt to novel dynamics, and highlight its advantages compared existing diffusion-bbar planning basclines.
1 Introduction
Model predictive control (MPC), also called receding horizon control, uses a dynamics model and an action selection mechanism (planner) to construct "agents" that can solve wide variety of tasks by means of maximizing known objective function (see e.g., Schwenzer et al. (2021) for review of MPC). More precisely, we want to design an agent that maximizes an objective function J [ F-1, St41++F) over planning horizon
Corresponding author: stamiss@gnogle.com

---

Published in Transactions on Machine Learning Research (05/2025)
F from the current timestep
Outt F-1 arg max E
(1)
F
where h2 = 1111-1,,aaa)) is the history. MPC thus factorizes the problem that the agent needs to solve into two pieces: a modeling problem (representing the dynamics model Pa($t-1o=P(s1n 01++F-1), which in this paper we learn from offline trajectory data) and a planning problem (finding the best sequence of actions for a given reward function). Once we have chosen the action sequence, we can execute the first action a (or chunk of actions) and then replan, after observing the resulting next state, thus creating a closed loop policy.
The advantage of this MPC approach compared to standard policy learning methods is that can easily adapt to novel reward functions at test time, simply by searching for state-action trajectories with high reward. This makes the approach more flexible than policy learning methods, which are designed to optimize fixed reward function!" [ In addition, learning a dynamics model is often more sample-efficint than learning a policy directly (Zhu et al... This is because dynamics model training is essentially a supervised regression problem, predicting the next state given the current state and action a mapping that is typically well-behaved and near-deterministic. Policy learning, however, involves predicting actions, where the optimal behavior may be multimodal (multiple good actions exist) and require accurate long-horizon credit assignment, making it a more complex learning task given the same data budget. Finally, dynamics models can often be adapted more easily than policies to novel environments, as we show in our experiments.
However, to make MPC effective in practice, we have to tackle two main problems. First, the dynamics model needs to be accurate to avoid the problem of compounding errors, where errors in next state prediction accumulate over time as the trajectory rolled out (Venkatraman et al., 2015 Asadi et al. 2019 Xiao et al. 2019; Lambort et al. 2022p. p. To avoid compounding errors, multi-step models are profcrable. However, these require model class capable of capturing the complex, multimodal distribution of entire trajectories. This motivates use of diffusion models. Second, the planning algorithm needs to be powerful enough to select good sequence of actions, avoiding the need to exhaustively scarch through a large space of possible actions.
We tackle both problems by using diffusion models to learn joint trajectory-level representations of (1) the world dynamics, Pu($++1+++ ac1+F-1, f2,), which we learn using offline "play" data (Cui et al., 2023); and (2), an action sequence proposal distribution, p(aar+F-1ll (h),, (h,,), which we can also learn offline using behavior cloning on some demonstration data. Although such a proposal distribution might suggest actions that are not optimal for solving new rewards that were not seen during training, we show how to compensate for this using a simple sampling- bamel planner, variant of the random shooting method that uses a multi-step diffusion model trained on offline datasets as the action proposal and a simple alternative to more complex methods such as trajectory optimization or cross-entropy method. We call our overall approach Diffusion Model Predictive Control (D-MPC).
We show experimentally (on a variety of state-based continuous control tasks from the D4RL benchmarks (Fu al. 2020)) that our proposed D-MPC framework significantly outperforms existing model-based offline planning methods, such as MBOP (Argenson & Dulac-Arnold 2021), which Iearns single-step dynamics model and a single-step action proposal, and hence sullers from the compounding error problem. By contrast, D-MPC learns more accurate trajectory-Ievel models, which avoid this issue, and allow our model-based approach to match (and sometimes exceed) the performance of model-free offlino-RL. mothods. We also show that our D-MPC method can optimize novel rewards at test time, and that we are able to fine-tune the dynamics model to a new environment (after a simulated motor defect in the robot) using a small amount of data. Finally we perform an ablation analysis of our method, and show that the different pieces namely the use of stochastic multi-step dynamics, multi-step action proposals, and the sampling- -aaee planner with learned action proposals
each individually valuable, but produce even more benefits when combined.
In summary, our key contributions are:
I 'Coal-conditioned reinforcement learning (RL) increase the flexibility of the policy, but the requested goal (or something very similar to it) needs have been seen in the training set, so cannot optimize completely novel reward functions that are specified at run time.
2

---

Published in Transactions on Machine Learning Research (05/2025)
1. We introduce Diffusion Model Predictive Control (D-MPC), combining multi-step action proposals and dynamics models using diffusion models for online MPC.
We show D-MPC outperforms existing model-based offline planning methods on D4RL benchmarks, and is competitive with SOTA reinforcement learning approaches.
3. We demonstrate D-MPC can optimize novel reward functions at runtime with simple sampling-based planner, and adapt to novel dynamics through fine-tuning.
4. Through ablations, we validate the benefits of our method's key components individually and in combination.
2 Related work
Related work can be structured hierarchically as in Table [ Model-based methods postulate a particular dynamics model, whereas model-free methods, whether more traditional behaaioral cloning (BC), con- servative Q learning (CQL) (Kunar et al. 2020), implicit Q-learning (IQL) (Kostrikov et al.] 2021 etc
or diffusion based- - diffusion policy (DP) (Chi et al. 2023), diffuusion BC (DBC) (Pearce et al.) 2023
simply learn policy. This can be done either by regressing directly from expert data or with some variant of Q-leaning. Model-basod approaches can be further divided according to how they use the model: Dyna-style [Sutton] 1991) approaches use it to learn a policy, either online or offline, which they can use al runtime, whereas MPC-style approaches use the full model at runtime for planning, possibly with the guidance of a proposal distribution
It is possible to model the dynamics of the model and the proposal jointly using Py (, a) or factorized distribution pa())))p(a). The latter allows for extra flexibility, since both pieces can be fine-tuned or even re-learned independently. Finally, we can categorize these models as either single-step (SS) or multi-step (MS). SS methods model the dynamics as pa(s h4, a1+1) (where h, = (St-Ha: ac-H:t) is the history of length H), and the proposal as p(a, so we predict (a distribution over) the next time step, conditioned on past obsorvations (and the next action, for the dynamics). We extend this to the whole planning horizon of length F by composing it in an autoregressive form, as a product of one-step conditionals, i.., (($ 0144F-1) II. F- pa((s+1lss a,, h.). (Note that this might result in compounding errors, even if h4 contains the entire past history, i.c., is non-Markovian.) By contrast, multi-step (MS) methods model the joint distributions at trajectory level. Thus the MS dynamics model represents pa($++ia+ r/s h4, 0+++F-1) and the MS proposal represents p(
Several recent papers follow the SS Dyna framework. Some using traditional dynamics modeling (c.g., MOREL (Kidambi et al. 2020). MOPO
Yu et al. 2020), COMBO (Yu et al., 2021). ) RAMBO-RL (Rigter 2020)). ; and others using diffusion. The latter includes "Diffusion 2024h (previously called "Diffusion World Models" (Alonso et al. 2023)). ) "UniSim" " paper (Yang et al. 2024 and the "SynthER" paper (Lu et al. 2024p. These are then used to generate samples from the model at training time in order to train a policy with greater data efficiency than standard model-free reinforcement learning (RL) mothods. Some other recent papers - such as "Diffusion World Model" (Ding et al. 2024) "PolyGRAD" (Rigter et al. 2022)) and "Policy-Guided Diffusion" (Jackson) et al. 2024) have proposed to use diffusion for creating joint multi-step (trajectory-level) dynamics models. However, being part of the Dyna framework, they are not able to plan run-time, like D-MPC does.
et al., 2022) and Dreamer (Hafner et al. for World Modeling' paper (Alonso et al.
There are many papers with a model-based approach. Probably the closest to our work is "Diffuser" (anner et al. 2022,, which uses diffusion to fit joint (state, action) multi-step model P)( a1:7) using offline trajectory data. They then use classifeer-guidance to steer the sampling process to generate joint sequences that optimize a novel reward at test time. The main difference to our method is that we represent the joint as a product of two models, the dynamics (((1:T(a117) and the policy / action proposal, p(a).r). As we show in Section 4.3. this factorization allows us easily adapt to changes in the world (e.g., due to hardware defects) from small amount of new data, whereas Diffuser struggles in this context. In addition, we propose
Note that proposal distribution (which we denote by p(a)) is different than a policy (which we denote by x(a)), since rather than determining the next. best action directly, it helps accelerate the search for one.
3

---

Published in Transactions on Machine Learning Research (05/2025)
Factored: pa(ele) pla)
Joint: py(*,a)
Model-free: x(e)
Dyna
MPC
MPC
(single-step)
(multi-step)
(single-step)
(multi-step)
(multi-step)
MOReL etc., Dreamer, DWMS,
DWM, PolyGRAD PGD
BC, CQL, IQL, DD, DP, IH, DBC
Examples
MBOP
D-MPC
Difluser, DT, TT
UniSim,
SynthER
Ran-time planning
Ran-time novel rewards
/
Novel dynamics
/
Non-expert data Speed as runtime
/ Med.
Fast
Fast
Slow
Slow
Fast
Table 1: A tale of three distributions; comparing properties across offline RL methods. The methods we
mention are defined as follows: MOREL etc. Kidambi et al. (2020); Yu et al.
(2020)
2021 b: Rigter et al.
2022)
Dreamer Hafner et al. (2020), DWMS (Diffusion World Models) Alonso et al.
[023)_ , UniSim Yang et al.
2024
SynthER Lu et al. [2024), DWM (Diffusion World Model) Ding et al. (2024). PolyGRAD Rigter et al. (2024)
PGD (Policy-Guided Diffusion) Jackson et al. (2024). Diffuser Janner et al.
(2022). DT (decision transformer)
Chen et al.
2021], TT (Irajectory transformer Janner et al. (2021), BC (behavior cloning). CQL (conservative
Q learning) Kumar et al. (2020), IQL (implicit Q learning) Kostrikov al. (2021) DD (Decision Diffuser)
Ajay et al.
2023), DP (Diffuson Policy) Chi et al. (022)). IH (Imitiating Humans) Pearce et al. (2023
DBC
(Diffusion BC)
Wang et al. (2023)
simple sampling- based planner that does not rely classifier guidance. Other works using MS with joint proposal are decision transformer (DT) (Chen et al. 2021), and trajectory transformer (Janner et al. 2021
Similarly, the "Decision Diffuser" paper (Ajay et al. 2023 learns a trajectory distribution over states, and uses classifier-free guidance to generate trajectories that have high predicted reward; the state sequence is then converted into an action sequence using a separately trained inverse dynamics model (IDM). However, this approach does not allow for run-time specification of new reward functions.
MPC has also been applied in model-based RL, with TD-MPC (Hansen al. 2022 and TD-MPC2 (Hansen et al. 2023) being the representative methods. D-MPC differs from the TD-MPC line of work in that D-MPC uses multi-step diffusion models for both action proposal and dynamics model, while the TD-MPC line of work uses single-step MLPs. In addition, the TD-MPC line of work focuses on online learning with environment interactions while in D-MPC we focus on learning from offline data and then use the learned models for doing MPC the environment.
The model-based offline planning or MBOP paper (Argenson & Dulac-Arnold 2021) was the original inspiration for our method. In contrast with the previous MPC methods, it factorizes the dynamics models and the action proposal model, which are learned separately and used at planning to optimize for novel rewards. The main difference with our work is that they use ensembles of one-step deterministic MLPs for their dynamics models and action models, whereas we use a single stochastic trajectory level diffusion model. In addition they use a somewhat complex trajectory optimization method for the action selection, whereas we use a simple sampling-bast based planner. Finally, we also study adapting the model to novel dynamics.
Several recent works are closely related to D-MPC. "Model-based Diffusion' (MBD) (Pan et al., 2024) also utilizes diffusion for trajectory-level search, but it assumees known environment dynamics, unlike our setting. "DyDiff' Zhao et al. (2024) employs a diffusion-basee dynamics model with structure similar to D-MPC's, conditioning on action sequences to generate state sequences. However, DyDiff's primary focus is generating synthetic on-policy data by modeling the interaction between a given single-step policy and the multi-step diffusion dynamics model, differing from our goal of planning with learned dynamics and action proposals.
D-MPC is a novel combination of MPC, factorized dynamics/action proposals, and MS diffusion modeling. This allows us to be able to adapt to novel rewards and dynamics and avoid compounding errors.
4

---

Published in Transactions on Machine Learning Research (05/2025)
3 Method
We will now describe our new D-MPC method. Our approach can be seen as multi-step diffusion extension of the model-based offline planning (MBOP) paper (Argenson & Dulac- 2021 with few other modifications and simplifications.
3.1 Model predictive control
D-MPC first learns the dynamics model Pslas action proposal and heuristic value function J (see below), in offline phase, discuss in Section 3.2 and then proceeds to alternate between taking action in the environment with planning the next sequence of actions using planner, we discuss in Section 3.3 The overall MPC pseudocode is provided in Algorithm
3.2 Model learning
We assuine access to an offline dataset of trajectories, consisting (state, action, reward) triples: D rMM
{sl.n, ai fin, , T. model Pa($t+1t++ h ax(+F-1) and another diffusion- these models, minimize the denoising score matching loss in the usual way. We include detailed review of diffuusio model training in Appendix A, and refer the readers to e.g. Karras et al. (2022) for additional discussions.
M a pm T........ 8 TTT, a1 ri:T>) )- We then use this to difflusion-based dynamics action proposal p(Qm=4F1lss, h.). To fit
We also define a function that approximates the reward-to-ge given any proposed sequence of states and actions:
+F--1 R ax) V((++))
J ((ar-))
aca4r-1) E
Σ
(2)
Here 7 the discount factor, and Vss) represents the value function from state 8 (i.e., estimate of future reward the leaves of this search process). We also use transformer to learn / (although we can also just compute J directly, if the reward function R is known, and V) by regressing from (St.+F (x(-F-1) to the discounted future reward in Eq. (2 We use L2 loss for the regression. We use J as the objective function for optimization in MPC, and way to specify novel tasks. Refer to Appendix for additional details on model architectures and hyperparameters.
admissible lower bound (such as 0)
Algorithm 1: Main MPC loop.
Input: D = offline datasct, N = mum. samples, F forecast horizon, H history length
(pa. p, J)- = train(D)
8o = env env.init()
ho (so)
for t 0o do
a: planner.plan(s, hr, J,N, F,H)
(8t 1, 1(+1 env.step(s;;@c ar)
h = append(a: T(41)
h suffix(h,, H)
Note that unlike MBOP (Argenson & Dulac- Arnold 2021) we do not need to (rain ensembles, since diffusion models are expressive enough to capture the richness of the respective distributions directly. Also, contrast Argenson Dulac-Arnold (2021), we do not need train a separate reward function: we estimate our value function the beginning the planning horizon, for a given sequence of states and actions along that horizon. In this way, our objective function J already includes the estimated reward along the horizon.
5

---

Published in Transactions on Machine Learning Research (05/2025)
3.3 Planning
D-MPC is compatible with any planning algorithm. When the action space is discrete, can solve this optimization problem using Monte Carlo Tree Search, used in the MuZero algorithm (Schrittwieser et al., 2020
Here we will only consider continuous action spaces.
We propose a simple sampling- based planner, depicted as Algorithm 2 In order to plan, given the current state 8, and history hi, we use our diffusion action proposal p to sample N action sequences, we predict the corresponding state sequences using Pola: we score these state/action sequences with the objective function J, we rank them to pick the best sequence, and finally return the first action in the best sequence, and repeat the whole process. We show empirically that this outperforms more complex methods such as the Trajectory Optimization method used in the MBOP paper, which we describe in detail in the Appendix (Algorithm 5 We believe this is because the diffusion model already reasons at the trajectory level, and can natively generate a diverse set of plausible candidates without the need for additional machinery.
3.4 Adaptation
As with all MPC approaches, our proposed D-MPC is more computationally expensive than methods that use a reactive policy without explicit planning. However, one of the main advantages of using plamning-based methods in the offline setting is that they can easily be adapted to novel reward functions, which can be different from those optimized by the behavior policy that generated the offline data. In D-MPC, we can easily incorporate novel rewards by replacing Vn in Alg. 2 by Vn = KJ(BFFF An,I:)) + (J An,I:F), where the novel objective J(s1.r, Am,1:F)= } Ei= fuovel(St, An,t): fnovel is novel reward function, and K, R are weights of the original and novel objectives, respectively. We demonstrate this approach in Section 4.2 Of course, if the new task is very different from anything the agent has seen before, then the action proposal may be low quality, and more search may be needed.
If the dynamics of the world changes, we can use supervised fine tuning of Pala on a small amount of exploratory "play" from the new distribution, and then MPC before. We demonstrate this in Section 4.3
4 Experiments
In this section, we conduct various experiments to evaluate the effectiveness of D-MPC. Specifically we seek answer the following questions with our experiments:
Does our proposed D-MPC improve performance over existing MPC approaches (which learn the model offline), and it perform competitively with standard model-based and model-free offline RL methods?
2. Can D-MPC optimize novel rewards and quickly adapt to new environment dynamics at run time?
3. How do the different components of D-MPC contribute to its improved performance?
Can we distill D-MPC into a fast reactive policy for high-frequency control?
Algorithm 2: Sampling based planner with learned multi-step diffusion action proposals
Def a Planner(So. ho, Pd, P, J, N, H):
2 for n = 1:N do
0n,1:F ~ p(-180: ho)
$1.F ~ Sp. heg In.1:F)
V. = J(SIFF, an.1:F)
arg max, Vn
Return 02,1

---

Published in Transactions on Machine Learning Research (05/2025)
4.1 For fixed rewards, D-MPC is comparable to other offline RL methods
We evaluate the performance of our proposed D-MPC on various DARL (uu el al. 2020 lasks. Planning-based approaches are especially beneficial in cases where we do not have access to expert data. As a result, we focus our comparisons on cases where we learn with sub-optimal data. We experiment with locomotion tasks for Halfcheetah, Hopper and Walker2D for levels medium and medium-replay, Adroit tasks for pen, door and hammer with cloned data, and Franka Kitchen tasks with mixed and partial data.
Our results are summarized in Table 2 We see that our method significantly beats MBOP, and a behavior cloning (BC) baseline. It also marginally beats Diffuser, strong model-based offinne RL approach that uses classifier guidance during sampling for planning instead of MPC. In contrast to Diffuser's single diffusion model for joint state-action sequence generation, our approach decouples the dynamics model and action proposal. We note that the total compute of our two transformers is designed to be roughly equivalent to Diffuser's 's single model; we achieve this by using a smaller transformer for the dynamics model, reflecting its relative simplicity. We additionally compare to other popular model-free offline RI. methods, like conservative Q-learning (CQL) (Kumar al., 2020) and implicit Q-learning (IQL) (Kostrikov et al.) 2021), well model-based RL methods like MOReL (Kidambi et al., 2020), and mequence-models like Decision Transformer (DT) (Chen et al., 2021). These methods camnot adapt to novel rewards at test time (unlike D-MPC, MBOP and Diffuser), but we include them to give a sense of the SOTA performance on this benchmark. We see that our method, despite its extra flexibility, can still match the performance of these existing, but more restrictive, methods.
Domain
Level
MOReL
MBOP
D-MPC (ours)
Diffuser
DT
BC
CQL
IQL
halfcheetah medium
42.10
44.60
46.00
(=0.17)
44.20
42.60
42.60
44.00
47.40
hopper
medium
95.40
48.80
61.24
(+2.30)
58.50
67.60
52.90
58.50
66.30
walker2d
medium
77.80
41.00
76.21
(#2.67)
79.70
74.00
75.30
72.50
78.30
halfcheetah
medium-replay
40.20
42.30
41.12
(+0.31)
42.20
36.60
36.60
45.50
44.20
hopper
medium-replay
93.60
12.40
92.49
(12.23)
96.80
82.70
18.10
95.00
94.70
walker2d
medium-replay
49.80
9.70
78.81
(+4.11)
61.20
66.60
26.00
77.20
73.90
Locomotion
Average
66.48
33.13
65.98
63.77
61.68
41.92
65.45
67.47
Domain
Level
MBOP
D-MPC (ours)
DT
BC
CQL
IQL
adroit-pen
cloned
63.20
89.22
(=12.57)
71.17 (+2.70)
99.14 (=12.27)
14.74 (#2.31)
114.05 (#4.78)
adroit-door
cloned
0.00
16.36
(#:220)
11.18 (20.96)
3.40 (..95)
-0.08 (+0.13)
9.02 (+1.47)
aaroitthhammr
cloned
4.20
12.27
(1:358)
2.74 (1:0.22)
8.90
(14.04)
0.32 (10003)
11.63 (1:1.7)
Adroit Average
22.47
39.28
28.36
37.15
4.99
44.90
Domain
Level
D-MPC (ours)
BC
CQL
IQL
kitchen
mixed
67.50
(+2.099)
51.50
52.40
51.00
kitchen
partial
73.33
(+1.64)
38.00
50.10
46.30
Kitchen Average
70.42
44.75
51.25
48.65
Table 2: Performance comparison of D-MPO with various model-based and model-free offine RL methods across different domains. Baseline results are obtained from existing papers (Ajay et al. 2023 Tarasov et al., 2024). Performance is measured using normalized scores [Fu al. 2020). For D-MPC, we report the and standard error of normalized scores over 30 episodes with different random initial environment conditions. Following (Kostrikov al. 2021). highlight in bold scores within 5% of the maximum per task. Baseline numbers from (Ajay et al. 2023) do not have associated standard errors. We include the standard errors for baseline numbers from
Tarasov al. 2024) when they are present.
4.2 Generalization novel rewards
In Fig. D we demonstrate how novel rewards can be used to generate interesting agent behaviors. We first trained the dynamics, action proposal, and value models for D-MPC on the Walker2d medium-replay dataset.

---

Published in Transactions on Machine Learning Research (05/2025)
We then replaced the trained value model with novel objective Vn for scoring and ranking trajectories in our planning loop, using (8t, A,) eexp(-(h, Maagget) / where h4 is the dimension of the state 9r that corresponds to the height of the agent, Marget is the target height, x² = x 10-4, K = 0 and = (so we only use the new J and ignore the original J). By using this simple method, we were able to make the agent lunge lorward and keep its head down (Marget = 0.9), balance itself' (Marget = 1.2), and repeatedly jump (Marget = 1.4). Note that for the experiments presented in Fig. [ we utilize a pure novel reward setting 0 and & 1. This means the value function used for planning, Va, is solely determined by the novel reward function, J. The original reward function, J. derived from the pre-training data, has no influence on the agent's is behavior in these cases. The resulting lunge, balance, and jump behaviors are therefore direct consequences of optimizing and are not present in the original dataset.
by setting
We also compared D-MPC with Diffuser on these novel reward tasks. Diffuser can also optimize these rewards, and we did not observe a significant performance difference. However, D-MPC's key advantage lies in its ability to adapt novel dynamics, as discussed Scction
trained value function
target height 0.9 m
target height 1.2 m
target height 1.4 m
height the walker (m)
1.4
NmM w
1.4
1.0-
1.0-
200
400
600
800
1000
50
150
200
200
800
1000
100
200 timestep
timestep
timestep
timestep
100
50
175
reward
2
Novel reward
Novel reward
Novel roward
Enwv' reward
Erv reward
Env reward
Env reward
0
200
400
600
800
1000
0
50
100
150
200
200
400
600
800
1000
100
200
300
timestep
timestep
timestep
timestep
Figure 1: Novel reward functions can generate interesting agent behaviors. The leftmost column shows an example episode generated by D-MPC trained on the Walker2d medium-replay dataset, using the trained value function in the planner. The remaining three columns present individual examples of behaviors generated using a height-based novel objective in the planner, with each column corresponding to a different target. height. The top row of each column displays the agent's ' height al each timestep within the episode. The middle row shows two snapshots of the agent per episode, while the bottom row graphs the novel reward (targeted by the reward received by the agent at each timestep. This figure serves
planner) and the actual environment-proviae as a qualitative demonstration how novel rewards can employed to produce interesting behaviors.
4.3 Adaptation to novel dynamics
In this section, we demonstrate our model's ability to adapt to novel dynamics at test time with limited experience. The need for such adaptions to novel dynamics is common in real world robotics applications where wear and tear or even imperfect calibrations can cause hardware defects and changed dynamics at test time. Because of our factorized formulation, where we separate dynamics Psla from policy na we can leverage a small amount of "play" data collected with the hardware defects, and use fine-tune our multi-step diffusion dynamics model while keeping our action proposal and trained value functions the same.
8

---

Published in Transactions on Machine Learning Research (05/2025)
We demonstrate this Walker2D. We train the original models the medium dataset and simulate hardware defect by restricting the torque executed by the actions on foot, joint (action dimension 2). On the original hardware, without the defect, trained D-MPC achieves a normalized score of 76.21 (#2.67). When executing this model on defective hardware, performance drops to only 22.74 (#1.41). Performance of our implementation of Diffuser in the same setup when deployed on defective hardware drops from 72.91 (+ 3.47) to 25.85 (+1.0
To compensate for the changed system dynamics, collect 100 episodes of "play" data on the defective hardware by deploying the original D-MPC trained the medium-replay dataset. We use this small dataset to fine-tune our multi-step diffusion dynamics model, while keeping the policy proposal and value model fixed. Post-finetuning, performance improves to 30.65 (#1.89). Since diffuser jointly models state and action sequences, there no way to independently finetune just the dynamics model. We instead fine-tune the full model with the collected "play" data. After fine-tuning, diffuser performance actually drops to 6.8 (=0.6)). See Table 3a for summary.
4.4 Ablation studies
In this section, we conduct a series of detailed ablation studies to illustrate how different components in D-MPC contribute to its good performance. In particular, we investigate the effect of using diffusion for action proposals, and the impact of using single-step vs. multi-step models both for the action proposals and for the dynamics models. We evaluate all variants on D4RL locomotion tasks. Soe Table 3b for a high-level summary of the results, and Table 4 in the appendix for detailed performances of different D-MPC variants individual D4RL domains and levels.
4.4.1 Diffusion action proposals improve performance and simplify the planning algorithm
Existing model-based offline planning methods, such as MBOP, typically use a single-step deterministic MLP policy for action proposals, an ensemble of single-stop deterministic MLP models to emulate a stochastic dynamics, and rely on trajectory optimization methods for planning. This MBOP method achieves an average score of 33.13 on the locomotion tasks.
We can significantly improve this bascline MBOP score, and simplify the algorithm, by (1) replacing their single-step MLP action proposal with single-step diffusion proposal, and (2) replacing their TrajOpt. planner with our simpler sampling- based planner. This improves performance to 52.93. Replacing their MLP dynamics with a single-step diffusion dynamics model further provides a minor improvement, to 53.32.
4.4.2 Multi-step diffusion action proposals contribute to improved performance
D-MPC uses multi-step diffusion action proposals. In this section, we illustrate how this further improves performance when compared with single-step diffuusio action proposals.
We start with the same single-step MLP dynamics model as in Section 4.4.1 We then replace the single-step diffusion action proposal with multi-step diffusion action proposal that jointly samples a chunk of actions. This improves average performance from 52.93 to 57.14. We then repeated this experiment on top of the single-step diffusion dynamics, and improved performance from 53.32 to 57.81.
4.4.3 Multi-step diffusion dynamics models contribute to improved performance
D-MPC uses a multi-step diffusion dynamies model. In this section we illustrate how this reduces compounding error in
dynamics prediction and contributes to improved performance.
We first measure the accuracy of long- horg-horizon dynamics predictions of single-step diffusion, multi-step diffusson and auto-regressive transformer (ART) dynamics models (described in Appendix E.3), independent of the planning loop. We train the dynamiis models on medium datasets from the respective domains, and measure the accuracy of long-horizon dynamics prediction, using statc/action soquences sampled from the modium (training data), medium-replay (lower quality data) and expert (higher quality data) datasets. Concretely, we follow Schubert et al. (2023) and calculate the median root mean square deviation (RMSD) the non-velocity
9

---

Published in Transactions on Machine Learning Research (05/2025)
Action
Dynamics Model
Diffuser
D-MPC
Proposal
SS Diff
SS MLP
ART
MS Diff
Original
79.60
76.21 (+2.67)
Pre-FT w/ defect
25.85 (+1.08)
22.74(+141)
SS Diff
53.32
52.93
N/A
Post-FT w/ defect
6.8(+0.86)
30.65(+1.89)
MS Diff
37.81
57.14
59.83
65.98
(b)
Table 3: (a) Performance Walker2D before and after simulated hardware defect, followed by fine-tuning (FT) on play data. (b) Average performances D-MPC variants on D4RL locomotion tasks. The full D-MPC method is bottom right. MS: multi-step, SS: single step, Diff: diffusion, ART: auto-regressive transformer. Our baseline MBOP uses ensembling and MPPI trajectory optimization with SS MLP dynamics models and SS MLP action proposals, and achieves
of 33.13.
Walker2d
Hopper
HalfCheetah
10
107
107
medium
10.1
10:
10
10
10-2
10aa
16
64
256
16
64
256
16
64
256
medium-repla replay
10a
10a
10a
10:1
10
10
10
16
64
256
16
64
256
16
64
256
expert
RMSD pred.
10-1
10-1
10.4
256
64
256
64
256
Rollout steps
Diffusion 32-step
Diffusion 1-step
MLP ART
Figure 2: Accuracy of long-horizon dynamics prediction. We train the dynamics models the medium dataset and evaluate on medium (training data), medium-replay (lower-quality data), and expert (higher-quality data) datasets. Prediction errors measured by the median root square deviation (RMSD) velocity coordinates based on 1024 sampled state action sequences of length 256. Plots show median + 10 percentile bands. The multi-step diffusion dynamics model incurs significantly lower prediction data while maintaining superior generalization abilities, outperforming other single-step and auto-regressive alternatives. The ruto-vegressive transformer (ART) dynamics model outperforms the single step diffusion dynamics model. The single-step MLP dynamics model exhibits compounding errors that grow rapidly for long-horizes dynamics predictions.
training
coordinates with increasing trajectory length. Fig. summarizes the results. From the figure the multi-step diffusion dynamics model reduces compounding errors in long-horizon dynamics predictions compared to single-step and auto-regressive alternatives, while maintaining superior generalization abilities, especially for action distributions that are not too far from training distributions.
how
We then evaluate the quality of these dynamics models when used inside the D-MPC planning loop with a multi-step diffusion action proposal. When using the ART dynamics model, we get a score of 59.83, but when using the multi-step diffusion dynamics model, we get 65.98. We believe this difference is due to the fact that
10

---

Published in Transactions on Machine Learning Research (05/2025)
the transformer dynamics model represents the sequence level distribution as product of one-step (albeit Non-Markovian) conditionals, ie., Pa($++1++r] $1:t 01:++--1) = I Pa($:+1|$1:: d1:--1, az). By contrast, the diffusion dynamics model is an 'a-causal" joint distribution that goes from noise to clean trajectories, rather than working left to right. We conjecture that this enables diffusion to capture global properties of a signal (e.g., predicting if the final state corresponds to the robot falling over) in a more faithful way than a causal-in-time model.
While some D4RL tasks may not present explicit, immediate penalties for short-term planning, achieving high scores even on benchmarks like Walker2D and Hopper often requires anticipating delayed consequences of actions (e.g., an excessively fast gait leading to later instability). This section demonstrates D-MPC's ability to mitigate compounding error through accurate long-horizon predictions, crucial capability for both maximizing performance within D4RL and for tackling more complex, real-world scenarios where long-term planning is essential.
4.4.4 Necessity of Expressive Models for Multi-Ster Step Prediction
To demonstrate the importance of model expressiveness for effective multi-step prediction, we replaced the multi-step dilfusion models within D-MPC (for both action proposal and dynamics) with multi-step MLPss These MLPs were two-layer networks with 4096 hidden units, using flattened multi-step states and actions as inputs. The D-MPC configuration using diffusion models achieves a normalized score of 65.98 (Table 3b). In contrast, the version using multi-step MLPs for both components only achieved a score of 50.01. This substantial performance drop underscores the necessity of employing expressive models, such as diffusion models, to accurately capture the complex distributions inherent in multi-step planning.
4.5 D-MPC can be distilled into fast reactive policy for high-frequency control
Due to the use of diffusion models, D-MPC has slower runtime. In Appendix J we include a detailed runtime comparison between D-MPC and other methods. However, if high-froquency control is important, we distill the D-MPC planner into a fast lask-specific MLP policy, similar to what is done in MoREL (Kidambi et al. 2020 or MOPO YYu et al.) 2020). Concretely, we train an MLP policy on offline datasets, using the planned actions from pretrained D-MPC as supervision. We do this for the 6 D4RL locomotion domain and level combinations we use in our ablation studics, and compare performance with both the vanilla MLP policy and D-MPC. We train all models for 1M steps, and evaluate the last checkpoint for the distilled MLP policy.
We observe that the distilled MLP policy achieves an average normalized score of 65.08 across the 6 D4RL locomotion domain and level combinations, which is only slightly worse than D-MPC's average normalized score of 65.98, and greatly outperforms the vanilla. MLP policy's average normalized score of 41.92. In addition, after distillation we just have an MLP policy, and it runs at the same speed as the vanilla MLP policy.
5 Conclusions
We proposed Diffusson Model Predictive Control (D-MPC), which leverages diffusson models to improve MPC by learning multi-step action proposals and multi-step dynamics from offline datasets. D-MPC reduces compounding errors with its multi-step formulation, achieves competitive performance on the D4RL benchmark, and can optimize novel rewards at run time and adapt to new dynamics. Detailed ablation studies illustrate the benefits of different D-MPC components.
One disadvantage of our method (shared by all MPC methods) is the need to replan at each step, which is much slower than using a reactive policy. This is particularly problematic when using diffusion models, which are especially slow to sample from. In the future, we would like to investigate the use of recently developed speedup methods from the diffusion literature, such as distillation (see e.g., Chang et al. (2023) ) Furthermore, while our ablation studics demonstrate the surprising effectiveness of our simple sampling- planner, incorporating guided sampling techniques (as suggested Janner et al. (2022) could offer a path towards greater efficiency by combining the strengths of model-based and model-free approaches.
11

---

Published in Transactions on Machine Learning Research (05/2025)
Another limitation of the current D-MPC is we only explored setups where directly have access to the low-dimensional states, such as proprioceptive sensors on robot. In the future, we plan to extend this work to handle pixel observations, using representation learning methods that extract abstract latent representations, which can form the imput to our dynamics models, similar to existing latent-space world modeling approaches such as the Dreamer line of work, but an MPC context, rather than a Dyna context.
Like all offline RL methods, D-MPC's performance is influenced by the distribution of behaviors in the training dataset. When offline datasets lack behaviors relevant to the target task, the generalization capabilities of any method are inherently constrained without additional data collection. While this does present limitation for D-MPC, it is not unique to our approach but rather a fundamental challenge in offline RL. Within the scope of available data, D-MPC excels at optimizing and adapting to novel rewards and dynamics, which represents the realistic scenario for offline RL applications. Our approach's ability to effectively leverage the existing behavioral distribution is a significant strength. Future work could explore techniques to encourage broader exploration within the constraints of offline data, potentially expanding the applicability of D-MPC and similar methods to even wider range of scenarios.
References
Amurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In ICLR, 2023. URL https //openreview. http://wwww...eeefoooood/dssssss
Eloi Alonso, Adam Jelley, Anssi Kanervisto, and Tim Pearce. Diffusion world models. October 2023. URL
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari, 2024.
Arthur Argenson and Gabriel Dulac-Arnold. Model-Based offline planning. In ICLR, 2021. URL https:
05556
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel Littman. Combating the Compounding- Error Error
problem with multi-step model. arXiv [cs. May 2019. URL httpp:://arrd...rg//bb////////// 13320
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Ziyi Chang, George A Koulieris, and Hubert P H Shum. On the design fundamentals of diffusion models: A survey. June 2023. URL http::/ wwwpeiiii.orraaaee a30/ 04542
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084 15097, 2021.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In RSS, March 2023. URL http:
//arxiv. org. hargg/abs//// abs//303..0. 041377
Zichen Jeff Cui, Yibin Wang, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. From play to policy: Conditional behavior generation from uncurated robot data. In ICLR, 2023. URL httt:://axxiv.oot hrrg/ abs/2210.10047
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv [cs.LG], February 2024. URL https:://arx/ htt///rxxi.oorg/bbb../ hbs///bs/2200 attgg//2s/4402....
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXio preprint arXiv:2004.07721, 2020.
Danijar Hafner, Timothy Lillicaa,, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent. imagination. In ICLR, 2020. URL http:///ariiv.orr//////1111.
12

---

Published in Transactions on Machine Learning Research (05/2025)
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXix:2003.044 04955, 2022.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpo2: Scalable, robust world models for continuous control. Xrviv preprint ar1o:23100.68825 2023.
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. Coft, abs/1606.00 1606.08445,, 2016. URL
wwtp://wrrbbbbbb...../o// 08415
Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Foerster. Policy-guided diffusion. ar preprint urXiv:2/04.06356 2024.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural informalion processing systems, $4:1273-1288 2021.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine Planning with diffusion for flexible behavior synthesis. In ICML, May 2022. URL http:://arrv..orr//bb//222......
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space Diffusion-Based generative models. In NIPS, June 2022. URL http://arxiv. http://aaaiiv.oorr/bbb//2200...
Rahul Kidambi, Aravind Rajeswaran, Pranceth Netrapalli, and Thorsten Joachims. Morel: Model-based oflline reinforcement learning. Advances in neural information processing systemss, 33:21810 21823, 2020.
Πya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. International Conference on Learning Representations, 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforce- ment learning. Advances in Newral Information Processing Systems, 33:1179--199) 2020.
Nathan Lambert, Kristofer Pister, and Roberto Calandra. Investigating compounding prediction errors learned dynamics models. March 2022. URL http://arxiv. wwtp:/orr..bbbbb..oooomm
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprinl Xiv: 03003, 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. +pm-solver+++ Fast solver for guided sampling of diffusion probabilistic models. ar preprint arXi::2211001095, 2022.
Cong Lu, Philip Ball, Yee Whye Teh, and Jack Parker-Holder Synthetic experience replay, Advances in Neural Information Processing Systems, 36, 2024.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171 PMLR, 2021.
Chaoyi Pan, Zoji Yi, Guanya Shi, and Guannan Qu. Model-based diffusion for trajectory optimization. Advances in Neural Information Processing Systems, 37:57914-57943, 2024.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In ICLR, January 2023. URL http:://axxw..rrg///////
10677
Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. Advonces in neural information processing systems, 35:16082 16097, 2022.
Marc Rigter, Jun Yamada, and Ingmar Posner. World models via. Policy-Guided trajectory diffusion. TMLR, 2024. URL http:://axiv..o/// http://aaxiv.oorrab///////////// hww.g//ab//....../ 08533
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint ar 00150 2022.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 2020. URL http:
//arxiv. wwexiv.org/aae abs www.b//9...0............111
Ingmar Schubert, Jingwei Zhang, Jake Bruce, Sarah Bechtle, Emilio Parisoltoo, Martin Riedmiller, Jost Tobias Springenberg, Arunkumar Byravan, Leonard Hasenclever, and Nicolas Heess, A generalist dynamics model for control. May 2023. URL http:///axxvv...gabbbssss/
10912
Max Schwenzer, Muzalfer Ay, Thomas Bergs, and Dirk Abel. Review on model predictive control: an engineering perspective. Int. J. Adv. Manuf. Technol. 117(5))1322-1344) 1349, November 2021. URL https: 000ol.orr//0011007/0111000012010
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. Advances in Newral Information Processing Systems, 36, 2024.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference Learning Representations, 2020.
Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. ar preprint arXiv:2310. 14189, 2023.
Yang Song, Prafulla Dhariwal, Mark Chen, and Tlya Sutskever. Consistency models. arXiv preprint arXiv:2303.01468 2023.
Richard S Sutton. Dyna, an integrated architecturre for learning, planning, and reacting. SIGART Bull., 2(4):
160-163, July 1991. URL
Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov. Corl: Research-oriented doep offline reinforcement learning library. Advances in Newral Information Processing Systems, 36, 2024.
Arun Venkatraman, Martial Hebert, and Bagnell. Improving Multi-Step prediction of learned time series models. AAAI, 29(1), February 2015. URL httpss:/ojs aaaoo.saaaa...................aa oro//irgg... wttp:pp////aticlee/////ooooaaaee 9590
Hsiang-Chie Chun Wang, Shang-Fu Fu Chen, Ming-Hos Hao Hsu, Chun-Mao Lai, and Shao-Hua Sun. Diffusion model- augmented behavioral cloning. Xixie preprint orXio:2302.11 13335, 2023.
Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning to combat Compounding Error in Model-Based reinforcement learning. December 2019. URL http::/arxiv.orr http//awwvvv.r.rreee./// orggabs//912/ 19tp/1992..11111 11206
Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. Em distillation for one-step diffusion models. arXiv preprint XrXiv:2405.16833 2024.
Ruibin Xiong, Yunchang Yang, Di Hc, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu. On layer normalization in the transformer architecture. CoRR,
2020. URL https: http:////aarr..oorr///////00.0. wwtp:o/g/abbb/...////2/ 04745
Sherry Yang, Yilu Du, Scyed Kamyar Soyed Ghasemipour, Jonathan Tompson, Leslic Pack Kaclbling, Dale Schurmans, and Pieter Abbeel. Learning interactive Real- World simulators. In ICLR, 2024. URL
httts://openreveee
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances Neural Information Processing Systems, 15:14129-141-1 2020.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 34:28954 28967, 2021.
Hanye Zhao, Xiaoshen Han, Zhengbang Zhu, Minghuan Liu, Yong Yu, and Weinan Zhang. Long-horizon rollout via dynamics diffusion for offiine reinforcement learning. arXiv preprint or arXiv:2405. 191 2024.
Hanlin Zhu, Baihe Huang, and Stuart Russell. On representation complexity of based and model-fiee free reinforcement learning. In The Twelfth International Conference on Learning Representations.
15

---

Published in Transactions on Machine Learning Research (05/2025)
A
Algorithms for model learning
In Algorithm
3 and Algorithm we present the multi-step and single-step training used in our experiments.
Diffusion models are generative models that define a forward diffusion process and reverse denoising process. The forward process gradually adds noise to the data, transforming it into a simple Gaussian distribution. The reverse process, which is learned, denoises the data step by step to recover the original data distribution.
In D-MPC, we model the action proposals and dynamics models as conditional diffusion models. Formally, let 2o, be the original data, where zo represents the data we want to model and y represents the conditioning variable. Let vk be the data al step
the diffusion process. The forward process is defined as (r=-)) (TRi vag (1-ax)I) where ar determines the variance schedule. The reverse process aims to approximate po(#x-1|2z) = (= 1: (oo (I k. y), Σn) where N(u, Σ) denotes Gaussian distribution with mcan μ and covariance matrix E. For suitably chosen 0g and large enough K, IK ~ NN(.,) I). Starting with Ganssian noise, we iteratively denoise to generate the samples.
We train diffusion models using the standard denoising score matching loss. Concretely, we start by randomly sampling ummodified original Zo. y from the dataset. For each sample, we randomly select a step in the diffusion process, and sample the random noise 6 with the appropriate variance for the diffusion time step k. We train a noise prediction network Eg by minimize the mean-squared-error loss MSE(<,, ex + Ek. k. v)). po(0 [at; y) can be calculated function of ca(*,, k; y), which allows us to draw samples from the trained conditional diffusion model.
In the D-MPC case, for learning the action proposal p, the future actions 0ct+F-1 are the clean data To. and the current state and history h, are the conditioning variable y. For learning the multi-step dynamics model Pa: the future states $t41:44/ are the clean data Ho, and the current state st, the history fy and the future actions Fct+F-1 are the conditioning variable y. In both cascs, the noise prediction network co is a transformer. Details of the transformer are given in Appendix E.T
Algorithm 3: Model learning (one step models).
Def .M- = train(D, F, H):
2 Create dataset of tuples: D' {(s,, hr, ar 8t41, Tt, Gt)). he = (0-- G. = Ej
Fit pi (11411188, h4, at) using MLE on D', set Pa = II ++8- =
a,,
Fit Pi (r h2) using MLE on D', set p -= [1;= 1
P pi(a,ls,,h;)
Fit R = E(r, 88, his at) using regression on D'
Fit = ([Ge/s,, he. a.) using regression on D
Algorithm 4: Model learning (multi-step models).
Def M train(D, F, H)): Create dataset of tuples:
D' = {(st, he, Qcc=F-1:
G.)}. h; 1S--Hi-L.a-HH-- G,
Fit pa($441244 124 8t, he, 0-14F-1) using diffusson on D'
Fit p( p(am.4F-ilst, he) using diffusion on D'
Fit R- = E(r, St, he, at) using regression on D
Fit V = B(G, he, a:) using regression on D
B The MBOP-TrajOpt Algorithm
In Algorithm 5 we include the complete MBOp-TrajOpt algorithm from Argenson & Arnold 20021 adapted to notations reference.
16

---

Published in Transactions on Machine Learning Research (05/2025)
Algorithm 5: MBOP-TrajOpt
Def a = MBOP-TrajOPt(ss, arrev ,,, N. F,K,
R =0
AN =0
for n N do
[= n mod K
81 = su. dg d grevy
R. =
for / - : F do
~~ o2)
a Sprop (c: (1,a))1) +e
10
An, = (1 B)a:
Drev prev Fin(t,F-1)
81+1 = fstates (5t, (
11
R = R+ + K Ei=1 freward (8t, Ant)
1
12
Rn R + Li=l fratue (Salle ($F41, An.F)
13
Σ
uB Ano-t Vt [0
14 02
Bn
15 Return a
C Detailed ablation study results
In Table we present detailed performances of the D-MPC variants studied in Section high-level summary.
See Table
3b for
Doumin Name
Level
MBOP
D-MPC
MS Duff Action Proposal
SS DIE Action Propoeni
Dall Dyuamics
MLP Dyaacnics
ART Dynamic
Dill Dyuamies
S5 MLP
Dymoniee
ballchedaah
mediam
11.61
26.00
(20.17)
21.30
[1 0.18)
44.78
(1 0.13)
45.17 (# 0.15)
16.51
[1: 0.17)
$1.88
(1: 0.18)
happer
medium
2.38)
50.66
1.77)
walker2d
mediam
11.00
76.21
2.67)
72.23
[1 2:29)
77.09
(1 1.62)
73.16
(4: 2.97)
75.59
[1: 2.85]
76.41
[1 1.79)
balfcheetali
mediun-replan replay
0.31)
42.85
0.19)
hopper
medinn--relay
12:40
92.49
(422223)
74.45
[1: 1.44)
76.38
(1: 3.83)
79.84
(4: 3.93)
70.17
[1: 5.56)
08.31
(1: 5.60)
wallerrdd
medium- mediun-replay
78.81
(4:1.19)
4.79)
52.33
4.85)
4.43)
5.62)
5.51)
Averago
33.13
65.98
57.81
57.14
50.83
53.32
52.03
Table d: Detailed performances of the D-MPC variants studied in the ablation studies on different DARL domains and levels. MS multi-step, SS single step, Diff diffusion, ART auto-regressive transformer.
D Normalizing state coordinates
Following Ajay et al. (2023), we normalize the states that are input to our models by using the the empirical cumulative distribution function (CDF) to remap cach coordinate to lie uniformly in the range [-1,
Given an offline dataset of trajectories, consisting of (state, action, reward) triples
M D {sim) ain rim si. aim, 1 Tim.:... ST ai TilM
M M
let S (m=1i
(sx) m be the accumulated corpus for the k-th coordinate of each state. We can define the empirical CDF for the k-th coordinate of the state by
1 N
(t)
Σ 1 si
for st = N St.
N
where ly is the indicator function for event Y.
17

---

Published in Transactions on Machine Learning Research (05/2025)
For the state vector 5 consisting of coordinates 8k, the relation with the normalized state coordinate is then given by 8): 2 (8k) - 1. The states output from our dynamics model are unnormalized by the inverse relation Sk- = Fx 182
E
Model architectures and training details
E.1
Diffusion models
In this paper, we train 4 kinds of diffusion models: single-step diffusion action proposals, single-step diffusion dynamics models, multi-step diffusion action proposals, and multi-step diffusion dynamics models.
We implement all 4 models as conditional diffusion models. For single-step diffusion action proposals, we use diffusion to model (s); for single-step diffusion dynamics models, for multi-step diffusion action proposals, dynamics models use diffusion to model p(5t+1:1+ Se, Oca+F-1).
diffusion to model p(s+ 1 1| ae); diffusion to model p((estt-1/st): and for multi-step diffusion
Our diffusion implementation uses DDIM Song et al. 2020) with cosine schedule Nichol & Dhariwal 2021 We use transformers as our denoising network: for each conditional diffusion model, we embed the diffusion time using sinusoidal embeddings, project the time embeddings and each state and action (both for states states/actions /actions that are used as conditioning and for states/actions that are being modelled) to a shared token space with tokens of dimension 256, add Fourier positional embeddings with 16 Fourier bases to all tokens, and pass all the tokens through number of transformer layers. We take the output tokens that correspond to the states/actions are predicting, and project them back to the original ilate/action spaces.
For all our transformer layers, we use multi-headed attention with 8 heads, and 1024 total dimensions for query, key and value, and 2048 hidden dimensions for the MLP.
For all of our single-step diffusion action proposals, we use diffusson timesteps and 2 transformer layers for the denoiser.
For all of our single-step diffusion dynamics models, use 3 diffusion timesteps and 2 transformer layers for the denoiser.
For all of our multi-step diffusion action proposals, we use 32 diffusion timesteps and 5 transformer layers for the denoiser.
For all of our multi-step diffusion dynamics models, we use 10 diffusion timesteps and transformer layers for the denoiser.
E.2 One-step MLP dynamics models
We follow Argenson & Dulac-Arnoll (2021 for the multi-layer perceptron (MLP) architecture of our one-step dynamics model, and train to approximate 8t41 f(s,, We use only a single MLP. Hyperparameters for the model and training summarized in Table
Hyperparameter
Value
Number FC layers
2
Size of FC layers
512
Non-linearity
ReLU
Batch size
256
Loss function
Mean square error
Table 5: Hyperparameters for the MLP one-step dynamics model and training.
E.3 Auto-Re Sett-Reegeessve Transformer dynamics model (ART)
We follow Chen et al. 2021 in our transformer dynamics model, leaving out the rewards tokens and the time-step embedding. Each state and action mapped into a single token with sepa- rate linear layer, namely embed, and embed,, respectively. This results in the following tokens: T
18

---

Published in Transactions on Machine Learning Research (05/2025)
(embed,(s1) (anbeda(a)) embed,(s2), embed,(az)....) These then normalized using LayerNorm Ba et al. (2016) then mapped with causal transformer to series of output tokens 02,02, The loss function is then:
L Σ (sE(pred((O2-1))) a.) ) (eE(pred,(Oz))) 84+1)
(3)
where pred, and pred, are again linear prediction layers mapping from the output token the state action and MSE is the mean squared error.
The hyperparameters used are summarized in Table 6
Hyperparameler
Value
Encode dimension
512
Number layors
3
Number of heads
MLP size per head
512
Attention window
Position embbedding Dropout
Fourier embedding with 16 basis functions.
used.
LayerNorm location
Before attention block Xiong al 2020
Non-lincarity
GeLU Hendrycks & Gimpel 201h
Table 6: Hyperparameters used in the ART model.
E.4 Model architectures for the objective function
Our objective function takes as input future action proposals F- and future states Fta++F, F, and regresses the discounted future reward as defined in Eq. (2) We again implement our objective function using a transformer, using the same transformer layer as in Appendix F.T] We project all states and actions to a shared token space with tokens of dimension 256, and specify an additional learnable token for the discounted future reward. We add Fourier positional embedding with 16 Fourier bases all tokens, pass all tokens through transformer of 10 layers, take the token that corresponds to the discounted future reward, and read out the discounted future reward prediction. We train the objective function J using an L2 loss.
In our experiments for all domains except Hopper, we used a discount factor of 0.99. For Hopper we used a discount factor 0.997. For Walker2D and Hopper, episodes can terminate early due to the agent falling down. For episodes that terminate carly, we include an additional 100 termination penalty for the last stop as reward, and calculate the discounted future reward taking into account the termination penalty.
E.5 Training setups and hyper-pane pappr-parameters
For all of our model training, we use the Adam optimizer for which the learning rate warms up from 10-4 over 500 steps and then follows cosine decay schedule from 10-4 to 10 We train all models for 2 108 steps. We use gradient clipping at norm and uses EMA with decay factor of 0.99. All of our evaluations are done using the EMA parameters for the models.
E.6 Compute Resources
We train and evaluate all models A100 GPUs. We use single A100 GPU for each training run, and separate worker with a single A100 GPU for evaluation. Training for 2 x 106 steps for each variant takes about 2 days, for variants we considered.
F
for the sampling-based planner
For all of our experiments, we use forecast horizon F 32, number of samples N 64, and a history length H 1. A forecast horizon F 32 already works well since our trained objective function J predicts discounted future rewards.
19

---

Published in Transactions on Machine Learning Research (05/2025)
G Long-horizon dynamics prediction
Following Schubert et al. 2023 wC mcasure prediction errors by the median Root Mcan Square Deviation (RMSD) on the non-velocity coordinates, as depicted in Figure 2 While this metric allows us to directly analyze the effectiveness of the dynamics model, it is a somewhat crude metric of the correctness or usefulness of the prediction. For example, the following predictions would all produce errors 1.0 compared to the correct ones: treating each walker as a bundle of limbs with the same center-of-mant mass, predicting states that would trigger termination criteria (unreasonable joint angles), or predicting an upside down position for the HalfChectah (a state from which it cannot recover). This metric is more effective probe of dynamics model quality in the regime where the errors are smaller. We see in Figure 2 that the multi-step diffusion dynamics model in particular has low prediction errors even for long rollouts, indicating that it performs well in these situations.
H Generalization to novel rewards
For the examples in Fig. we first trained the dynamics, action proposal and value components of D-MPC on the Walker2d medium--eppay dataset. The leftmost column shows the agent's 's height and rewards allained during an example episode generated by D-MPC.
To incorporate novel rewards, we replaced the the trained value model in the planning loop with a novel objective based solely on the height of the agent. For this objective, we used (8t, At) exp((((h, Margee)) )2/2o") where h2 is the dimension of the state 8; that corresponds to the height of the agent, harget is the desired target height, o² = 5 x 10-4, was chosen to roughly match the maximum reward attainable in the environment.
0 and ă = 1. The scale factor of 5 in the reward function fnovel
In each episode, the agent starts at a height of 1.25 (with uniform noise in the range of [-5 x 10-A, 5 x 10-4] added for stochasticity). Figure [ demonstrates the agent's behavior for different target heights. For htrget 1.2, which is close to the initial position, the agent maintains the desired height for an extended duration. For Marget 0.9, the agent lowers its torso to achieve the target height, but eventually leans too far forward resulting in carly episode termination. For Marget = 1.4, the agent has to jump achieve the desired height, which can only be momentarily allained due to the environment's physics. In the example shown, the agent jumps three times before falling over, leading to early episode termination.
Adaptation to novel dynamics
We used D-MPC and our implementation of the diffuser models trained on Walker2D medium datasct as our pre-trained models. To simulate defective hardware, we modified the action being executed in the environment. Specifically, we clip the action component corresponding to the torque applied on the right foot rotor to [-0.5, 0.5] vs the original [1.
To collect the play data on this defective hardware, we run our D-MPC model trained on medium-replay dataset, and collect data for 100 episodes. It has total of 30170 transitions (steps) and an average normalized episode reward of 23.14 (#: 2.31). Note that the actions in this dataset correspond to the actions output by the model and not the clipped actions.
For fine-tuning, we load the pre-trained diffuser and D-MPC models and train them on this dataset with same training parameters as the original training. For D-MPC, we only train the dynamics model and frecze other components. For the diffuser, we fine-tune the full model, since it is a joint model. During online evaluation, we sampled 256 trajectories in both the models and picked the one with the best value execute the action at each step. We report the maximum scores for both approaches.
J Timing measurements
To illustrate the differences in rum times between different methods, we list in Table the measured wall clock planning time (along with standard errors) per execution step, on single A100 GPU, for each algorithm, in
20

---

Published in Transactions on Machine Learning Research (05/2025)
three D4RL locomotion environments. The measurements illustrate that simple MLP policy is the fastest, followed by an MBOP-like setup, followed by D-MPC, followed by MPC with an autoregressive transformer dynamics model. If tasks require faster control loops, D-MPC could be sped up in a few different ways such as amortizing the planning over a larger chunk of actions executed in the environment (since the planner
naturally generates long-horizon plans), using accelerated diffusion sampling strategies Lu el al. (2022) ; Shih
ct al. (2024), and distilling the diffusion models Song al. (2023): Song & Dhariwal (2023) Salimans Ho
(2022
Liu al. ) Xie al. (2024). We leave this exploration for future work.
DARL domain
MLP policy
MBOP
D-MPC
ART-MPC
Walker2d
3.51(40.06) x 10
4.77(#0.033) x 10
9.37(+0.03) x 10
3.37(#0.00) x 102
HalfChoetah
3.73(10.06) 10
3.93(10.10) 10
9.24(10.07) 10*
2.62(10.04) 10
Hopper
3.49(10.06) x 10
5.07(110.8) 10
9.48(10.04) x 10
3.51(11.007 x 10
Table 7: Wall-clock planning time milliseconds) per environment step for different algorithms, measured on single A100 GPU.
21


Reviews:
[{"review_id": "B4oJaatJMX", "reviewer": "Reviewer_AvyK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths)**\n- Straightforward to understand\n- Due to the nature of the MPC algorithm, it can be easily adapted to novel rewards and novel dynamics\n- Interestingly, the D-MPC algorithm can be distilled into a single-step MLP policy, which can achieve similar results compared to the original one. It allows the algorithm to be applied to real-world settings where high-frequency action selection is required.\n\n**Weakness**)\n- Analysis of the effect of hyperparameters: While the authors fixed several crucial hyperparameters such as $F, N,$. It would be nice to include the effect of each hyperparameter. For example, if we use small $F$, does the performance drop? Is there any relationship between the minimum requirements for $N$ and the dimension of actions space? Analysis of those parts makes readers understand this work more thoroughly.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose D-MPC, a novel approach that learns the dynamics model using diffusion models and replaces the trajectory planner with the sampling from diffusion models. D-MPC alleviates the error accumulation issue raised by single-step dynamics and mitigates the complexity of the planner in selecting an action. Experiment results validate that D-MPC can be easily adapted to novel rewards or novel dynamics due to the nature of the MPC algorithm.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "As stated above, it would be better to include an analysis of several decision choices.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xUt4poQLCg", "reviewer": "Reviewer_tGYL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The baseline choices look solid and are a fair comparison.\n- The experiments cover a broad spectrum of offline learning problems. \n\nWeaknesses:\n- The paper is overall difficult to read and incoherent.\n- The motivation is not clear. Why are diffusion models better than existing methods? The paper claims that they are more expressive and other methods can face compounding errors, however they are unrelated and neither are highlighted in the experiments section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper explores using diffusion models as the policy and dynamics model in offline learning and used online with MPC. The method, D-MPC is evaluated on d4rl locomotion manipulation and kitchen tasks and compared to other offline learning methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "- The overall flow and phrasing can be greatly improved.\n- The claims and results do not align well. The abstract and introduction claims significantly better performance however the results are marginally better than Diffuser.\n- Results in Section 4.2 and Figure 1 do not corroborate with the earlier claim in Table 1 of Run-time novel rewards. The proposed Kappa weighting suggests that the behaviors can be slightly altered, rather than novel.\n- Section 4.4.3 highlights long horizon planning, which is not needed in the most of the evaluated tasks.\n- More motivation and context for why diffusion models are considered should be added. Section J shows why they should not be used for control, taking 300x the time to perform inference. For example, environments where a slow inference and the expressiveness of diffusion models can be leveraged would make the method look more convincing.", "broader_impact_concerns": "- The paper proposes to use more computationally expensive method for marginal improvements.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vJgy3JtWSz", "reviewer": "Reviewer_CH73", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n1. The writing is very clear, and the introduction and categorization of previous works are appropriately structured.\n\n2. The proposed algorithm aligns well with intuition and can be easily implemented with existing open-source codebases.\n\n3. The authors effectively motivate the use of diffusion models as both the dynamics model and the action proposal model. And both designs are supported by corresponding ablation studies.\n\n4. D-MPC demonstrates strong performance in adapting to new rewards and environment dynamics, highlighting its potential in sim-to-real applications.\n\n### Weaknesses\n\n1. Compared to directly generating state-action sequences with one diffusion model (like Diffuser), this decoupled modeling comes with additional computational cost. Also, the sampling and ranking approach is less efficient than existing methods that optimize decision sequences through guided sampling. In fact, the action proposal model is very similar to the action model in [1], where with appropriate conditioning or guidance, it can produce sufficiently good policies in most cases. It would be an interesting direction to explore whether the strengths of model-based and model-free methods can be effectively integrated into a unified framework.\n\n2. Section 4.2 lacks comparative experiments with Diffuser as those in Section 4.3. When combined with a value function, Diffuser can also perform a similar search process for novel rewards, though it samples in the joint space. I am particularly curious whether the decoupled approach of D-MPC offers advantages in terms of search efficiency and stability.\n\n[1] Chi, C., Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., ... & Song, S. (2023). Diffusion policy: Visuomotor policy learning via action diffusion. *The International Journal of Robotics Research*, 02783649241273668.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces D-MPC, an MPC algorithm based on diffusion models. The core idea is to utilize trajectory-level diffusion models to implement both the dynamics model and the action proposal model in MPC. Experimental results demonstrate that D-MPC achieves comparable performance to model-free methods. Moreover, due to the decoupling of policy and model, it exhibits better generalization capabilities than previous diffusion-based methods like Diffuser, enabling adaptation to new reward functions and efficient fine-tuning for novel environmental dynamics.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "There are some details in the paper that I found unclear:\n\n1. In the introduction section, the authors mention that \"learning a dynamics model usually requires less data than learning a policy.\" This seems somewhat counterintuitive. Although the supervisory signal is high-dimensional, the model's output space is also larger and harder to fit. Therefore, it feels like more data would be required instead. Is there any reference that can support this claim?\n\n2. In Section 3.2, the authors state: \"To learn the value function, we use a transformer to regress from $(s_{t:t+F} , a_{t:t+F −1})$ to the discounted future reward in Eq. (2).\" They also mention, \"We also use a transformer to learn $J$.\" Does this mean that two separate transformers are used to fit $V$ and $J$ independently? Given that only $J$ is required in MPC optimization, why is it necessary to fit $V$ as well?\n\nI believe some recent studies are closely related to D-MPC, and the authors may consider discussing and comparing them:\n\n1. Model-based Diffusion (MBD) [2] also leverages diffusion for efficient trajectory-level search, but it operates under the assumption that the environment dynamics are known.\n\n2. DyDiff [3] employs a diffusion-based dynamics model with the structure quite similar to the one in D-MPC, where both condition on action sequences to generate state sequences. The key difference is that DyDiff focuses on generating synthetic on-policy data by modeling the interaction between a given single-step policy and the multistep diffusion model.\n\n[2] Pan, C., Yi, Z., Shi, G., & Qu, G. (2025). Model-based diffusion for trajectory optimization. *Advances in Neural Information Processing Systems*, *37*, 57914-57943.\n\n[3] Zhao, H., Han, X., Zhu, Z., Liu, M., Yu, Y., & Zhang, W. DyDiff: Long-Horizon Rollout via Dynamics Diffusion for Offline Reinforcement Learning.", "broader_impact_concerns": "The authors did not discuss broader impacts, nor do I think it is necessary.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces Diffusion Model Predictive Control (D-MPC), which combines diffusion models for multi-step action proposal and dynamics modeling with model predictive control. The reviewers had mixed opinions: Reviewer CH73 recommended "Accept," Reviewer AvyK "Leaning Accept," and Reviewer tGYL "Leaning Reject."

After weighing these recommendations, the AE find the positive aspects outweigh the concerns. Reviewer CH73 highlighted the clear writing, intuitive algorithm design, and strong performance in adapting to new rewards and dynamics. Reviewer AvyK appreciated the extensive experiments and the potential for distillation into a single-step policy. While Reviewer tGYL raised valid concerns about computational cost and marginal performance improvements over some baselines, the authors provided convincing responses, including:

- Demonstrating how D-MPC can be distilled into a fast MLP policy for real-time applications
- Providing additional hyperparameter analysis showing the importance of proper tuning
- Clarifying that their significant performance improvements were over comparable MPC-based methods (e.g., MBOP), while being competitive with other state-of-the-art approaches
- Explaining the importance of long-horizon planning even for seemingly simple locomotion tasks
Decision: Accept (poster)