Model-based Diffusion for Trajectory Optimization

Authors: chaoyi  pan
zeji  yi
guanya  shi
guannan  qu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:28:00

Abstract: Recent advances in diffusion models have demonstrated their strong capabilities in generating high-fidelity samples from complex distributions through an iterative refinement process. Despite the empirical success of diffusion models in motion planning and control, the model-free nature of these methods does not leverage readily available model information and limits their generalization to new scenarios beyond the training data (e.g., new robots with different dynamics). In this work, we introduce Model-Based Diffusion (MBD), an optimization approach using the diffusion process to solve trajectory optimization (TO) problems without data. The key idea is to explicitly compute the score function by leveraging the model information in TO problems, which is why we refer to our approach as model-based diffusion. Moreover, although MBD does not require external data, it can be naturally integrated with data of diverse qualities to steer the diffusion process. We also reveal that MBD has interesting connections to sampling-based optimization. Empirical evaluations show that MBD outperforms state-of-the-art reinforcement learning and sampling-based TO methods in challenging contact-rich tasks. Additionally, MBD’s ability to integrate with data enhances its versatility and practical applicability, even with imperfect and infeasible data (e.g., partial-state demonstrations for high-dimensional humanoids), beyond the scope of standard diffusion models. Videos and codes are available in the supplementary materials.

Main paper: 
 problem statement and background: 
 Notations: We use lower (upper) scripts to specify the time (diffusion) step, e.g., x t , u t , y t represent the state, control and state-control pair at time t, and Y (i) represents the diffusion state at step i.This paper focuses on a class of trajectory optimization (TO) problems whose objective is to find the sequences {x t } and {u t } that minimize the cost function J(x 1∶T ; u 1∶T ) subject to the dynamics and constraints. The optimization problem 1 can be formulated as follows:EQUATIONs.t. x 0 = x init (1b) x t+1 = f t (x t , u t ), ∀t = 0, 1, . . . , T -1, (1c) g t (x t , u t ) ≤ 0, ∀t = 0, 1, . . . , T -1.(1d) where x t ∈ R nx and u t ∈ R nu are the state and control at time t, f t ∶ R nx × R nu → R nx represents the dynamics, g t ∶ R nx × R nu → R ng are the constraint functions and l t ∶ R nx × R nu → R are the stage costs. We use Y = [x 1∶T ; u 1∶T ] to denote all decision variables. Traditionally, TO is solved using nonlinear programming, which faces challenges such as non-convex problem structures, nonlinear or discontinuous dynamics, and high-dimensional state and control action spaces. Recently, there has been a growing interest in bypassing these challenges by directly generating samples from the optimal trajectory distribution using diffusion models trained on optimal demonstration data [12, 40, 46, 61] .To use diffusion for TO, (1) is first transformed into a sampling problem. The target distributionEQUATIONObtaining the solution Y * from the TO problem in Eq. ( 1) is equivalent to sampling from Eq. ( 2) given a low temperature λ → 0. In fact, in Appendix A.2, we prove that the distribution of J(Y ) with Y ∼ p 0 (⋅) converges in probability to the optimal value J * as λ → 0, under mild assumptions. However, it is generally difficult to directly sample from the high-dimensional and sparse target distribution p 0 (⋅). To address this issue, the diffusion process iteratively refines the samples following a backward process, which reverses a predefined forward process as shown in Fig. 1 . The forward process corrupts the original distribution p 0 (⋅) to an isotropic Gaussian p N (⋅) by incrementally adding small noise to it and scaling it down by √ α i to maintain an invariant noise covariance (see Fig. 2 (b) for an example). Mathematically, this means we iteratively obtainY (i) ∼ p i (⋅) with p i|i-1 (⋅|Y (i-1) ) ∼ N ( √ α i Y (i-1) , (1 -α i )I).Because the noise at each time step is independent, the conditional distribution of Y (i) |Y (i-1) also leads to that of Y (i) |Y (0) :EQUATIONThe backward process p i-1|i (Y (i-1) |Y (i) ) is the reverse of the forward process p i|i-1 (Y (i) |Y (i-1) ), which removes the noise from the corrupted distribution p N (⋅) to obtain the target distribution p 0 (⋅).The target distribution p 0 (⋅) in the diffusion process is given by:EQUATIONEQUATIONStandard diffusion models [33, 40, 61] , which we refer to as Model-Free Diffusion (MFD), solve the backward process by learning score function merely from data. In contrast, we propose leveraging the dynamics model to estimate the score to improve the generalizability of the model and allow a natural integration with diverse quality data.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: This paper includes theoretical results, and all the assumptions and proofs are provided in the main paper. Guidelines:• The answer NA means that the paper does not include theoretical results. 
model-based diffusion: 
 In this section, we formally introduce our MBD algorithm that leverages model information to perform backward process. To streamline the discussion, in Section 4.1, we first present MBD with Monte Carlo score ascent to solve simplified and generic unconstrained optimization problems. In Section 4.2, we extend MBD to the constrained optimization setting to solve the TO problem given complex dynamics and constraints. Lastly, in Section 4.3, we augment the MBD algorithm with demonstrations to improve sample quality and steer the diffusion process. We consider the reverse process for an unconstrained optimization problem min Y J(Y ), where thetarget distribution is p 0 (Y (0) ) ∝ exp (-J(Y (0) ) λ). In our MBD framework, "model" implies that we can evaluate J(Y (0) ) for arbitrary Y (0) , enabling us to compute the target distribution up to the normalizing constant.MBD uses Monte Carlo score ascent instead of the commonly adopted reverse SDE approach in MFD. Specifically, when denoising from i to i -1, MBD performs one step of gradient ascent on log p i (Y (i) ) and then scales the sample by the factor 1 √ αi as defined in the forward process:EQUATIONCritically, with the model-informed p 0 (Y (0) ), we can estimate the score function ∇ Y (i) log p i (Y (i) ) by connecting p i (Y (i) ) to p 0 (Y (0) ) via Bayes' rule:EQUATION= -EQUATIONBetween Eq. (7a) and Eq. (7b), we use the forward Gaussian density in Eq. (3):p i|0 (Y (i) | Y (0) ) ∝ exp(-1 2 (Y (i) - √ ᾱiY (0) ) ⊺ (Y (i) - √ ᾱiY (0) ) 1-ᾱi ). Its log-likelihood gradient is ∇ Y (i) p i|0 (Y (i) | Y (0) ) = -1 1-ᾱi (Y (i) - √ ᾱi Y (0) )p i|0 (Y (i) | Y (0)). Given Y (0) as the integration variable in Eq. (7c),p i|0 (Y (i) | Y (0)) is evaluated as a function of Y (0) parameterized by Y (i) . Based on that, we define the function ϕ i (Y (0) ) as:EQUATION(Y (0) -Y (i) √ ᾱi ) ⊺ (Y (0) -Y (i) √ ᾱi ) 1-ᾱi ᾱi ) ∝ N ( Y (i) √ ᾱi , I ᾱi -I) (8)This finding enables the Monte-Carlo estimation for computing the score function. We collect a batch of samples from ϕ i (⋅) which we denote as Y (i) and approximate the score as: Comparison between MFD and MBD. Table 1 highlights the key differences between MBD and MFD, which originate from two assumptions made in MBD: (a) a known target distribution p 0 (Y (0) ) given the model; (b) the objective of sampling Y (0) from the high-likelihood region of p 0 (Y (0) ) to minimize the cost function. For (a), MBD leverages p 0 to estimate the score following Eq. (9a), whereas MFD learns that from the data. For (b), MBD runs Monte Carlo score ascent in Eq. ( 6) to quickly move the samples to the high-density region as depicted in Fig. 2 (c), while MFD runs reverse SDEEQUATIONY (i-1) = 1 √ αi (Y (i) + 1-αi 2 ∇ Y (i) log p i (Y (i) )) + √ 1 -α i z i, where z i is Gaussian noise, to maintain the sample diversity. Given low temperature λ, it can be shown that ∇ log p i (Y (i) ) ≈ -1(1-ᾱi) (Y (i) -arg max p i (⋅))2 , i.e., the function log p i (Y (i) ) is 1 (1-ᾱi) -smooth. Therefore, choosing the step size (1 -ᾱi ) in Eq. ( 6) is considered optimal, as for L-smooth functions, O( 1 L ) is the step size that achieves the fastest convergence [64] . How diffusion helps? The diffusion process plays an important role in helping Monte Carlo score ascent overcome the local minimum issue in highly non-convex optimization problems, as shown in Fig. 2(a) . Compared with optimizing a highly non-convex objective, Monte Carlo score ascent is applied to the intermediate distribution p i (⋅) = ∫ p 0 (Y (0) )p i|0 (⋅)dY (0) , which is made concave by convoluting p 0 (⋅) with a Gaussian distribution p i|0 (⋅), as shown in Fig. 2(b) . Starting from the strongly concave Gaussian distribution p N ∼ N (0, I) with scale ᾱN → 0, the density is easy to sample. The covariance of the sampling density Σ ϕi = ( 1 ᾱi -1)I is large when i = N , implying that we are searching a wide space for global minima. In the less-noised stage, the intermediate distribution p i (⋅) is more peaked and closer to the target distribution p 0 (⋅), and ᾱi → 1 produces a smaller sampling covariance Σ ϕi to perform a local search. By iteratively running gradient ascent on the smoothed distribution, MBD can effectively optimize a highly non-convex objective function as presented in Fig. 2 . The MBD algorithm is formally depicted in Algorithm 1.Connection with Sampling-based Optimization. When diffusion step is set to N = 1, MBD effectively reduces to the Cross-Entropy Method (CEM) [11] for optimization. To see this, we can plug the estimated score Eq. (9b) into the Monte Carlo score ascent Eq. ( 6) and set N =1: Y (0) = ᾱ1 α1 Ȳ (0) (Y (1) ) = Ȳ (0) (Y (1) ) = ∑ Y (0) ∈Y (1) Y (0) w(Y (0) ) ∑ Y (0) ∈Y (1) w(Y (0) )where w(Y (0) ) = p 0 (Y (0) ) ∝ exp(-J(Y (0) ) λ ) and Y (1) ∼ N ( Y (1)α0 , ( 1 α0 -1)I). This precisely mirrors the update mechanism in CEM, which aims to optimize the objective function f CEM (Y (0) ) = J(Y (0) ) and determine the sampling covariance Σ CEM = ( 1 α0 -1)I, thus linking the sampling strategy of CEM with the α schedule in MBD. The advances that distinguish MBD from CEM-like methods are (1) the careful Algorithm 1 Model-based Diffusion for Generic Optimization1: Input: Y (N ) ∼ N (0, I) 2: for i = N to 1 do 3: Sample Y (i) ∼ N ( Y (i) √ ᾱi-1 , ( 1 ᾱi-1 -1)I) 4: Calculate Eq. (9b) Ȳ (0) (Y (i) ) = ∑ Y (0) ∈Y (i) Y (0) p 0 (Y (0) ) ∑ Y (0) ∈Y (i) p 0 (Y (0) )Estimate the score Eq. (9a):∇ Y (i) log pi(Y (i) ) ≈ -Y (i) 1-ᾱi + √ ᾱi 1-ᾱi Ȳ (0) (Y (i) ) 6:Monte Carlo score ascent Eq. ( 6):Y (i-1) ← 1 √ α i (Y (i) + (1 -ᾱi)∇ Y (i) log pi(Y (i))) 7: end for scheduling of α and (2) the intermediate refinements on p i , both following the forward diffusion process. This allows MBD to optimize for smoothed functions in the early stage and gradually refine the solution for the original objective. On the contrary, CEM's solution could either be biased given a large Σ CEM which overly smoothes the distribution as in p 20 , p 100 of Fig. 2 (b), or stuck in local minima with a small Σ CEM as in p 1 of Fig. 2(b ) where the distribution is highly non-concave.For TO, we have to accommodate the constraints in Eq. (1) which change the target distribution top 0 (Y (0) ) ∝ p d (Y (0) )p J (Y (0) )p g (Y (0) ). Given that p d (Y (0)) is a Dirac delta function that assigns non-zero probability only to dynamically feasible trajectories, sampling from ϕ i (Y (0) ) could result in low densities. To enhance sampling efficiency, we collect a batch of dynamically feasible samplesY (i) d from the distribution ϕ i (Y (0) )p d (Y (0)) with model information. Proceeding from Eq. (9a), andEQUATION≈ -EQUATION= -EQUATIONwhereEQUATIONThe model plays a crucial role in score estimation by transforming infeasible samples Y (i) from Line 3 in Algorithm 2 into feasible onesY (i) d . The conversion is achieved by putting the control part U = u 1∶T of Y (0) = [x 1∶T ; u 1∶T ]into the dynamics Eq. (1c) recursively to get the dynamically feasible samples Y (0) d (Line 4), which shares the same idea with the shooting method [29] in TO. MBD then evaluates the weight of each sample with p g (Y (0) )p J (Y (0) ) in Line 5. One common limitation of shooting methods is that they could be inefficient for long-horizon tasks due to the combinatorial explosion of the constrained space p g (Y ) ∝ ∏ T t=1 1(g t (x t , u t ) ≤ 0), leading to low constraint satisfaction rates. To address this issue, we will introduce demonstration-augmented MBD in Section 4.3 to guide the sampling process from the state space to improve sample quality.With the ability to leverage model information, MBD can also be seamlessly integrated with various types of data, including imperfect or partial-state demonstrations by modeling them as noisy observations of the desired trajectory p(Y demo | Y (0) ) ∼ N (Y (0) , σ 2 I). Given suboptimal demonstrations, sampling from the posterior p( ∼ N (0, I)Y (0) | Y demo ) ∝ p 0 (Y (0) )p(Y demo | Y (0) ) could lead to poor solutions as the demonstration likelihood p(Y demo | Y (0) ) could dominate the model-based distribution p 0 (Y (0) ) ∝ p d (Y (0) )p J (Y (0) )p g (Y (0)2: for i = N to 1 do 3: Sample Y (i) ∼ N ( Y (i) √ ᾱi-1 , ( 1 ᾱi-1 -1)I) 4:Get dynamically feasible samples:Y (i) d ← rollout(Y (i) ) 5:Calculate Ȳ (0) with Eq. (10d) (model only) or Eq. ( 13) (model + demonstration) 6:Estimate the score Eq. (10c):∇ Y (i) log pi(Y (i) ) ≈ -Y (i) 1-ᾱi + √ ᾱi 1-ᾱi Ȳ (0)Monte Carlo score ascent Eq. ( 6):Y (i-1) ← 1 √ α i (Y (i) + (1 -ᾱi)∇ Y (i) log pi(Y (i))) 8: end for cess. Rather, we assess Y (0) using p(Y demo | Y (0) ), employing a similar technique to interchange the distribution's parameter with the random variable, as demonstrated in Eq. ( 8), to establishp demo (Y (0) ) ∝ p(Y demo | Y (0) ) ∝ N (Y (0) | Y demo , σ 2 I).To accommodate demonstrations of varying qualities, instead of fixing target to p 0 (Y (0) )p(Y demo | Y (0) ), we propose seperating the p 0 (Y (0) ) from p demo (Y (0) ) to form a new target distribution3 :EQUATION)where η is a constant to balance the model and the demonstration. Here, we have introduced two extra constant terms p J (Y demo ) and p g (Y demo ) to ensure that the demonstration likelihood is properly scaled to match the model likelihood p 0 (Y (0) ). With these preparations, we propose to adaptively determine the significance of the demonstration by choosing η as follows:EQUATIONWhen samples have a high model-likelihood p 0 , we ignore the demonstration and sample from the model. Otherwise, we trust the demonstration. With the demonstration-augmented target distribution, we modify the way to calculate Ȳ (0) in Eq. (10d) as follows to obtain the score estimate:EQUATION) 
experimental results: 
 The experimental section will focus on demonstrating the capabilities of MBD in: (1) its effectiveness as a zeroth-order solver for high-dimensional, non-convex, and non-smooth trajectory optimization problems, and (2) its flexibility in utilizing dynamically infeasible data to enhance performance and regularize solutions. Our benchmark shows that MBD outperforms PPO by 59% in various control tasks with 10% computational time.Beyond control problems, in Appendix A.3, we also show that MBD significantly improves sampling efficiency by an average of 23% over leading baselines in high-dimensional (up to 800d) blackbox optimization testbeds [23, 18, 56, 42, 41, 44] . We also apply MBD to optimize an MLP network with 28K parameters in a gradient-free manner, achieving 86% accuracy within 2s for the MNIST classification task [2] , which is comparable to the gradient-based optimizer (SGD with momentum, 93% accuracy). To further extend MBD to closed-loop control, we employ receding horizon strategy to MBD in Appendix A.6 to update control sequence at each timestep, further improving the performance of MBD by 9.6% in terms of reward.To test the effectiveness of MBD as a trajectory optimizer for systems involving non-smooth dynamics, we run MBD on both locomotion and manipulation tasks detailed in Appendix A. tasks includes hopper, half-cheetah, ant, walker2d, humanoid-standup, and humanoid-running. The selected manipulation task, pushT [15] , presents its own challenges due to the complexity introduced by contact dynamics and the long-horizon nature of the task. These tasks are widely considered difficult due to their hybrid nature and high dimensionality.MBD is compared with the state-of-the-art zeroth-order optimization methods, including CMA-ES [7] , CEM [11] , and MPPI [59] , as well as reinforcement learning (RL) algorithms (e.g., PPO [47] and SAC [22] ) on these tasks. Please note RL is only used for performance reference not as there is no existing TO method that can solve such high-dimensional discontinuous tasks as we have shown in the experiments. Model-free RL, especially PPO/SAC, is widely used in such tasks and is considered the SOTA method. The difference between RL and MBD is further discussed in Appendix A.8 and A.6. The RL implementation follows the high-performance parallelized framework from Google Brax [20] elaborated in Appendix A.5.3. For the zeroth-order optimizer, we match the iteration and sample number with the MBD. All the experiments were conducted on a single NVIDIA RTX 4070 Ti GPU. Quantitative metrics including the average step reward and the computational time tested over 50 steps repeated for 8 seeds are reported in Tables 2 and 3 . MBD substantially outperforms zeroth-order optimization methods and even outperforms RL in most tasks. Specifically, for the pushT task, MBD achieves a significantly higher reward than the RL algorithm thanks to its iterative refinement process, which effectively explores the full control space while keeping fine-grained control to precisely push the object. Compared with the computationally heavy RL algorithms, MBD only requires one-tenth of time, which is similar to other zeroth-order optimization methods. The optimization process of MBD is visualized in Fig. 3 , where the iterative refinement process with the model plays a crucial role in optimizing high-dimensional tasks. We also evaluate the performance of MBD with data augmentation on the Car UMaze Navigation and Humanoid Jogging tasks to see how partial and dynamically infeasible data can help the exploration of MBD and regularize the solution by steering the diffusion process.For Car UMaze Navigation, the map blocked by U-shaped obstacles is challenging to explore given a nonlinear dynamics model. Therefore, random shooting has a low chance of reaching the goal region.To sample with loosened dynamical constraints, we augment MBD with data from the RRT [38] algorithm through goal-directed exploration with simplified dynamics. Fig. 4 (b) shows the difference between data-augmented MBD and data-free one: the former can refine the infeasible trajectory and further improve it to reach the goal in less time, while the latter struggles to find a feasible solution.The reason is that the infeasible trajectory from RRT serves as a good initialization for MBD, which can be further refined to minimize the cost function with MBD.For Humanoid Jogging, we aim to regularize the solution for the task with multiple solutions to the desired one with partial state data. Due to the infinite possibilities for humanoid jogging motion, the human motion data provide a good reference to regularize MBD to converge to a more human-like and robust solution instead of an aggressive or unstable one [26, 45] . We use data from the CMU Mocap dataset [1], from which we extract torso, thigh, and shin positions and use them as a partial state reference. Fig. 4 (a) demonstrates a more stable motion generated by data-augmented MBD.The Humanoidrun task is a 3D locomotion task where the agent is required to run forward. The reward is composed of the forward velocity of the agent and standing upright. Here we also modify the simulation substeps from 10 to 20 for longer planning horizon. The control dimension is 17.6. Humanoidstandup: The Humanoidstandup task is a 3D locomotion task where the agent is required to stand up. The reward is the upright torso position of the agent. The control dimension is 17.7. PushT: The PushT task is a 2D manipulation task where you can apply force to a sphere to push the T-shaped object to the target location. The reward is composed of the distance between the target and the object and orientation difference between the target and the object.To make the task more challenging, we randomize the target location 20cm away from the initial position and make sure the rotational angle is greater than 135 degrees, which makes it hard to solve the task with single continous contact policy. The control dimension is 2.8. Car2D: We implement a 2D car task with standard bicycle dynamics model, where state is x = [x, y, θ, v, δ], and action is u = [a, δ]. The dynamics is defined asẋ = f (x, u) = [v cos(θ), v sin(θ), v L tan(δ), a, δ].The constraints are defined as the U-shape area in the middle of the map, where the car cannot enter. The reward is composed of the distance between the target and the car and the control cost. The control dimension is 2. 
conclusion and future work: 
 This paper introduces Model-Based Diffusion (MBD), a novel diffusion-based trajectory optimization framework that employs a dynamics model to approximate the score function. MBD not only outperforms existing methods in terms of sample efficiency and generalization, but also provides a new perspective on trajectory optimization by leveraging diffusion models as powerful samplers. Future directions involve theoretically understanding its convergence, optimizing the standard Gaussian forward process using the model information, adapting it to online tasks with receding horizon strategies, and exploring advanced sampling and scheduling techniques to further improve performance.A Appendix / Supplemental Material We first give the definition of the volume of the sub-level set for cost J.Definition 1. Let F ∶ R d → R be a measurable function. Define the volume of the sub-level set for a given level t as:V F (t) = ∫ R d χ {Y ∈R d ∶F (Y )≤t} (Y ) dY,where χ denotes the indicator function.The volume function V J (t) plays a crucial role in linking geometric properties with probabilistic outcomes in optimization and learning algorithms. This function provides a quantitative measure that helps us to understand how changes in parameters like λ influence the distribution and concentration of probability mass.The interplay between geometry and probability, represented by V J (t), is crucial for evaluating the convergence and stability of algorithms. It provides a significant method for utilizing the PDF of the random variable Y to constrain the CDF, thereby facilitating convergence in distribution. < α 0 < α 1 < ⋯ < α M < ∞, It follows that: lim λ→0 J(Y ) p → J(Y * ) = 0.The cost value J(Y ) converges in probability to J(Y * ) as λ → 0.The condition on the polynomial bounds of V J (t) is generally not restrictive. For instance, considerJ = η c ∥Y -Y * ∥ m, where Y * is the optimal point and η c > 0 is any constant. In this case,V J (t) = Ct d m, where C is a constant, meets the constraint in a straightforward way. This condition can be extended beyond this simple scenario, as even if J has multiple modes, it can still adhere to this polynomial constraint.Proof. The convergence in distribution of Y towards Y * as λ → 0 is established by analyzing the behavior of the probability density function, defined up to a multiplicative constant. Consider the density Y 0 ∼ p λ (Y ) approximating Y * when λ approaches zero.EQUATIONEQUATIONEQUATIONwhere Eq. ( 14c) is valid since P (Y ) ∝ exp (-J(Y ) λ ) and J(Y ) represents the sufficient statistics of the distribution. We can obtain Eq. (14d) by computing the derivative of V J (x) based on the volume definition as shown in Definition 1.We denote J min = min Y J(Y ) = 0 and J max = max Y J(Y ) with J max satisfying 0 ≤ J max ≤ +∞. We proceed to analyze Eq. (14d) by performing integration by parts as shown in Eq. (15a).EQUATIONTo establish convergence in probability, we need to demonstrate that for any small ϵ > 0 and δ > 0, there exists sufficiently small λ > 0, such thatEQUATIONwhere the equality is due to Eq. (14d). Setting δ ′ = 1-δ δ , it suffices to show that:EQUATION)Assuming without loss of generality that J max = ∞, becuase dV J (t) ≥ 0, exp(-t λ ) > 0, we have:EQUATIONThis ratio as in Eq. ( 17) can be expanded using the integral bounds and the polynomial approximations for V J (t),then it suffices to show thatEQUATIONBy inserting Eq. (15b) into both the numerator and denominator on the LHS of Eq. ( 19), we obtain∫ ϵ 0 exp(-t λ )dV J (t) ∫ ∞ ϵ exp(-t λ )dV J (t) = exp(-ϵ λ )V (ϵ) + 1 λ ∫ ϵ 0 exp(-t λ )V J (t)dt -exp(-ϵ λ )V (ϵ) + 1 λ ∫ ∞ ϵ exp(-t λ )V J (t)dt (20a) ≥ ∫ ϵ 0 exp(-t λ )V J (t)dt ∫ ∞ ϵ exp(-t λ )V J (t)dt (20b) ≥ ∫ ϵ 0 exp(-t λ ) ∑ M k=0 c l k t α k dt ∫ ∞ ϵ exp(-t λ ) ∑ M k=0 c u k t α k dt (20c)To bound the expression in Eq. (20c), we first derive the following integrals by utilizing a change of variables x = t λ , which simplifies the expressions:EQUATIONEQUATIONFor these transformed integrals, we can observe that ∫ ∞ 0 exp(-x)x α k dx = Γ(α k + 1), the gamma function, which is well-defined for all non-negative α k . Given that δ ′ is a function of δ , by applying the intermediate value theorem and definition of the limit of the integral, we can choose ϵ k in such a way that:∫ ϵ k 0 exp(-x)x α k dx ≥ c k δ ′ 1 + c k δ ′ Γ(α k + 1), where c k = c l k c u kdenotes the ratio of coefficients in polynomial lower and upper bounds for V J (t). By selecting ϵ max = max ϵ 0 , ϵ 1 , ⋯, ϵ M to be the maximum of all such ϵ k , ensuring coverage for all polynomial terms up to M , we establish that:EQUATIONBy ensuring that λ ≤ ϵ ϵmax , we can conclude:EQUATIONThus, the condition specified in Eq. ( 16) is satisfied, validating that the distribution of Y converges in distribution to Y * as λ approaches zero.By adding another mild assumption regarding the landscape of J near the global optimum, we can demonstrate the convergence of the random variable Y itself, rather than the convergence of J(Y). Definition 3. We denote the minimum of the complementary set of neighborhood as:J * B (δ) = min ∥Y -Y * ∥>δ J(Y ) -J(Y * ).Proposition 4. Given the context and conditions specified in Definitions 1 and 3 and Proposition 2, and given that J has only one golbal minimizer Y * , i.e. there exist small δ * , that for δ ∈ (0, δ * ], J * B (δ) is strictly increasing, and J * B (δ * ) < ∞. It follows that:lim λ→0 Y p → Y * .The random variable Y converges in probability to Y * as λ → 0.Proof. In order to prove that lim λ→0 Y p → Y * . We need to prove that for any sufficient small γ > 0 and δ > 0, there exists small λ > 0, such thatEQUATIONFrom Definition 3 and due to the strict increase ofEQUATIONwhere0 < δ ≤ δ * . Because if ∥Y -Y * ∥ > δ, J(Y ) -J(Y * ) < J * B (δ) = min ∥Y -Y * ∥>δ J(Y ) -J(Y * ) contradicts Definition 3.Given that lim λ→0 J(Y ) p → J(Y * ). and any sufficient small ϵ, γ > 0.EQUATIONTherefore, ∃λ > 0, such thatEQUATIONAnd From Eq. ( 25), we have thatEQUATIONWe have that Y converges in probability to Y * ,i.e, lim λ→0 Y p → Y * .Proposition 5. Given the context and conditions specified in Propositions 2 and 4 and the way we define the forward process as in Eq. (3). The diffused Y i converge in density to a Gaussian distribution. lim λ→0Y (i) d → N ( √ ᾱi Y * , √ 1 -ᾱi I),where Y (i) ∼ p i (⋅) as in Eq. (3).Proposition 5 is derived by using Slutsky's theorem on Proposition 4 and offers insight into choosing the stepsize as discussed in Section 4.1.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Appendix A.5 provides all the necessary information to understand the experimental results.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Main quantitative results in Table 2 report the mean and standard deviation over 8 runs given different random seeds.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: The paper does not address societal impact. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not release data or models that have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We mentioned the package used in the paper on Appendix A.5.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
: 
 1 Introduction Model information x t+1 = f t (x t , u t ) x 0 x 1 u 1 y (0) 1 x T u T y (0) T … y (i-1) 1 y (i-1) T … y (i) 1 y (i) T … y (N) 1 y (N) T … Y (0) Y (i-1) Y (i) Y (N )Model-based diffusion in humanoid running tasks TrajectoryY (i) = αY (i-1) + 1 -αw Y (i-1) = 1 α i ( Y (i) + (1 -ᾱi ) ∇ Y (i) log p i (Y (i)) )Model-based score approximation Trajectory optimization (TO) aims to optimize the state and control sequence to minimize a cost function while subject to specified dynamics and constraints. Given nonlinear, non-smooth dynamics and non-convex objectives and constraints, traditional optimization methods like gradient-based methods and interior point methods are less effective in solving TO problems. In response, diffusion models have emerged as a powerful tool for trajectory generation in complex dynamical systems due to their expressiveness and scalability [12, 54, 34, 33, 40, 5] .problems, where gradient-based optimizers such as Augmented Lagrangian [32] , Interior Point [36] , and Sequential Quadratic Programming [3, 48] are employed. To leverage the parallelism of modern hardware and improve global convergence properties, sampling-based methods like Cross-Entropy Motion Planning (CEM) [37] and Model Predictive Path Integral (MPPI) [58, 62] have been proposed to solve TO by sampling from target distributions. To solve stochastic optimal control problems, trajectory optimization has also been framed as an inference problem in a probabilistic graphical model, where system dynamics defines the graph structure [35, 39] . This perspective extends methods such as iLQG by integrating approximate inference techniques to improve trajectory optimization [55] .The connection between diffusion and optimal control has been explored in [10] , which motivates us to use diffusion models as solvers for trajectory optimization.Diffusion for Planning. Diffusion-based planners have been used to perform human motion generation [12, 54] and multi-agent motion prediction [34] . Diffusion models are capable of generating complete trajectories by folding both dynamics and optimization processes into a single framework, thus mitigating compounding errors and allowing flexible conditioning [33, 40, 5] . In addition, they have been adeptly applied to policy generation, enhancing the capability to capture multimodal demonstration data in high-dimensional spaces for long-horizon tasks [46, 15] . These works assume no access to the underlying dynamics, limiting the generalization to new environments. To enforce dynamics constraints, SafeDiffuser [60] integrates control barrier functions into the learned diffusion process, while Diffusion-CCSP [61] composes the learned geometric and physical conditions to guarantee constraint compliance. Our approach uses diffusion models directly as solvers, rather than simply distilling solutions from demonstrations.Langevin-based Markov Chain Monte Carlo for Global Optimization. Gradient-based sampling algorithms have been widely used in global optimization, where the energy function J is optimized by sampling from the Boltzmann distribution p ∝ exp(-J λ ) [57, 43] . By annealing the temperature λ to zero, the sampling process converges to the global minimum of the energy function J [31] . The convergence of Langevin-based MCMC methods has been well studied in both continuous and discretized settings [16, 21] , showing that the distribution will converge in probability to the target distribution with certain decreasing schedule of the step size and temperature λ. In practice, the most common Langevin-based MCMC methods are unadjusted Langevin Monte Carlo (ULMC) [17] and Underdamped Langevin Monte Carlo (UdLMC) [14] , with convergence rates of O(1 ϵ 2 log( 1 ϵ )) and O( 1 ϵ ) given strongly convex and smooth energy functions, respectively. Recently, Langevin-based MCMC methods have been integrated into diffusion processes to improve global convergence and sampling efficiency [30, 6] , where the score function is estimated by Monte Carlo to accelerate and stabilize the diffusion process. Our work differs from these methods in that we aim to sample from the high-probability region of the target distribution without assuming access to the gradient of the energy function, and without assuming the energy function is smooth or convex. 
 
 Appendix: As a zeroth order optimizer, MBD is capable of addressing both trajectory optimization and broader, high-dimensional unconstrained optimization challenges. Such black-box optimization tasks are universally acknowledged as difficult [7, 63] . We first show superior performance of MBD within this black-box optimization context. In such settings, the Bayesian Optimization technique struggles due to the computational intensity required to develop surrogate models and identify new potential solutions [18] . Alternative black-box optimization strategies [23] are not limited by computational issues but tend to be less efficient because they do not estimate the black-box function as accurately. MBD's effectiveness is evaluated using two well-known highly non-convex black-box optimization benchmarks: Ackley [4] and Rastrigin [9] , each tested across three different dimensionalities. Comparisons were made with CMA-ES [23] , TuRBO [18] , LA-MCTS [56] , HesBO [42] , Shiwa [41] , and BAxUS [44] . -7000 -6000-5000-4000Rastrigin-400d 0 1000 2000 3000 4000 5000 6000 .-14000 -12000-10000-8000TuRBO-1 LA-MCTS HesBO CMA-ES Shiwa BAxUS MBD Fig. 5 shows the performance of MBD on the Ackley and Rastrigin benchmarks. MBD demonstrates superior performance over other algorithms for several reasons. Firstly, the implementation of a scheduled forward process that determines the total number of samples consequently boosts sample efficiency. Secondly, the application of Monte Carlo score ascent on various log p i (Y (i) ) facilitates its escape from local optima of varying scales. It is important to acknowledge that comparing computational efficiency may not be entirely fair, given that black-box optimization problems typically involve functions that are costly to evaluate. However, MBD markedly outperforms other Gaussian Process-based Bayesian Optimization approaches, achieving computational time savings of more than twentyfold, similar to the improvements observed with different evolutionary optimization strategies.Here are the implementation detail for the benchmarks. For the BO benchmarks, the experiments were conducted on an A100 GPU because of the high computational demands of the Gaussian Process Regression Model it incorporates.TuRBO: TuRBO is implemented based on tutorials from Botorch [9] .LA-MCTS: LA-MCTS, we refer to authors' reference implementations, and use TuRBO as its local BO solver [56] .HesBO: For HesBO, we refer to authors' reference implementations [42] . We transformed default GP component into Gpytorch version for faster inference speed on GPU. We set the embedding dimension to 20 for all tasks CMA-ES: We use pycma 4 to implement CMA-ES, and use default setting except setting population size eqauls to batch size.Shiwa: We use Nevergrad5 to implement Shiwa, and use default setting to run experiments.We refer to the authors' reference implementations [44] .To further demonstrate the effectiveness of MBD in high-dimensional systems, we apply MBD to optimize an MLP network for MNIST classification [2] without access to the gradient information. MBD achieve 85.5% accuracy with 256 samples within 2s, which is comparable to the performance of the SGD optimizer with momentum (92.7% accuracy). We use MLP with 2 hidden layers, each with 32 neurons, and ReLU activation function. The input is flattened to 784 dimensions, and the output is a 10-dimensional vector. We use cross-entropy loss as the objective function. The network has 27, 562 parameters in total, which makes sampling-based optimization challenging. MBD can effectively optimize the network with a small number of samples, demonstrating its effectiveness in high-dimensional black-box optimization tasks.Data-augmented MBD calculate the score function with demostration as follows:EQUATIONEQUATION)where We want MBD converge to the optimal point 8 with the help of demonstration data •. Although the demostration point is not optimal, MBD can still converge to the optimal point with the guidance of the demonstration data. Here data serves as a regularization term to guide the diffusion process to the negative optimal point while allowing to use model further to refine the solution.EQUATIONwhere demonstrate likelihood term w demo (Y (0) ) will draw samples towards data without considering the model. Given w = w demo ,Ȳ (0) = ∑ Y (0) ∈Y (i) d Y (0) p demo (Y (0) ) ∑ Y (0) ∈Y (i) d p demo (Y (0) ) = Y demo .The score function would be∇ Y (i) log p i (Y (i) ) = -Y (i) 1-ᾱi + √ ᾱi1-ᾱi Y demo , which means the score function is a linear combination of the current sample and the demonstration data.If we don't use Eq. ( 11) and employ the posterior distribution p(Y (0) |Y demo ) ∝ p 0 (Y (0) )p demo (Y (0) ), it will yields update weights w = w demo w model , which will draw samples to both model and demonstration data. If the demonstration data is not optimal, the final solution will be a compromise between the model and demonstration data. In Fig. 6 , the resulted solution will lie between optimal point 8 with the help of demonstration data •.Using the max function in w can aviod this issue. In the early stage while p J (Y (0) ) is low due to poor sample quality, w demo will dominate thanks to the high p J (Y demo ). This will draw samples towards the demonstration data as shown in the earlier stage of Fig. 6 . As the sample quality improves and p J (Y (0) ) > p J (Y demo ), w model will dominate and the sample will converge to the optimal point.We leverage the GPU-accelerated simulator Google Brax [20] to design the locomotion and manipulation tasks. All task is set to use positional backend in Brax except for the pushT task, which uses the generalizable backend for better contact dynamics simulation. Here we provide a brief description of each task implementations:1. Ant: The Ant task is a 3D locomotion task where the agent is required to move forward as fast as possible. The reward is composed of the forward velocity of the agent and control cost, same as the original Brax implementation. The control dimension is 8. 2. Hopper: The Hopper task is a 2D locomotion task where the agent is required to jumping forward as fast as possible. We use the same reward function as the original Brax implementation. We modify the simulation substeps from 10 to 20 for longer planning horizon given the same control node. The control dimension is 3. 3. Walker2d: The Walker2d task is a 2D locomotion task where the agent is required to walk forward. The reward is composed of keep the agent upright and moving forward. The control dimension is 6. 4. Halfcheetah: The Halfcheetah task is a 2D locomotion task where the agent is required to run forward. The reward is composed of the forward velocity of the agent and control cost. We follow the same reward function as the original Brax implementation. The control dimension is 6.In general, MBD is very little hyperparameters to tune compared with RL. We use the same hyperparameters for all the tasks, with small tweaks for harder tasks. 4 : MBD hyperparameters for various tasks For diffusion noise schedulling, we use simple linear scheduling β 0 = 1 × 10 -4 and β N = 1 × 10 -2 , and the diffusion step number is 100 across all tasks. Each step's α i is calculated as α i = 1 -β i .For reinforcement learning implementation, we strictly follow the hyperparameters and implementation details provided by the original Brax repository, which optimize for the best performance. For our self-implemented PushT task, the hyperparameters is ported from Pusher task in Brax for fair comparison. The hyperparameters for the RL tasks are shown in Table 5 and Table 6 .For the zeroth order optimization tasks, we the same hyperparameters as the MBD algorithm.For RRT algorithm in Car2D task, we set the max step size to 0.2, and the max iterations to 1000 given the maximum episode length is 50.For the demonstration collection in Humanoid Jogging task, we first download the mocap data which contains each joints' position in the world frame. Then we use the joint data to calculate the position of torso, thigh and shin position as partial state reference for our task. Even though MBD is designed as a trajectory optimization algorithm, it can be naturally extended to receding horizon control as shown in Algorithm 3. By conditioning the planning on each step's observation, MBD's MPC extension can further improve the performance of MBD by 9.6% in terms of reward, especially given control noise as shown in Fig. 7 .The online running frequence of MBD MPC is shown in Table 7 on RTX 4070Ti GPU. Please note that the frequency is calculated under the assumption of solving the whole 50 steps TO problem without reduced model at each iteration. Besides, the Brax code environment we used is not optimized for GPU, so the actual frequency could be higher with optimized environment. In our work we just use Brax as a simple and easy-to-use option. As the major computation time of MBD is spent on the forward dynamics simulation, it can be further improved by using more efficient physics engine.Frequency Given that sampling-based optimization is the core of MBD, we ablate the sample number in Fig. 8 . We can see that MBD converges to the optimal point with as few as 128 samples. The harder the task is, the larger performance gap between MBD and other TO solvers.1: Initialize: Optimize trajectory x 0∶T , u 0∶T -1 with MBD 2: for t = 0 to tfinal do 3:Observe the state xt 4:Optimize trajectory x t∶t+T , u t∶t+T -1 with single-step MBD 5:Apply the first control input ut to the system 6:Shift trajectory for next initialization: x t+1∶t+T , u t+1∶t+T -1 7: end for Given that RL and TO have different objective settings, especially the horizon difference, we conducted an ablation study by swapping the optimization objectives of MBD and RL. RL optimizes for a longer horizon discounted reward J = ∑ HRL t=0 γ t r t , H = 1000, γ < 1 while MBD optimizes for a shorter horizon undiscounted cumulative rewards J = ∑ HMBD t=0 r t , H = 50, γ = 1. Figure 9 shows the performance of MBD and RL under each other's optimization objectives.MBD outperforms RL by 44.5% under the RL objective and 805.5% under the TO objective. The results demonstrate that MBD's superior performance is attributed to its better diffusion-style iterative optimization process compared with RL's random exploration. Figure 9 : Ablation study on RL and TO objective function. γ = 1, H = 50 is TO objective, γ < 1, H = 1000 is RL objective, where γ is the discount factor and H is the horizon. For both objective function, MBD outperforms RL by 44.5% and 805.5% on average. We also find RL's overall performance is better given the discounted objective while MBD's performance is better given the total objective, which is consistent with our main results. should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.The paper provides open access to the data and code, with dedicated README files and instructions to reproduce the main experimental results comes with the supplemental material.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.We assume deterministic dynamics for simplicity to sample the dynamically feasible trajectory. The extension to stochastic dynamics is straightforward.See more elaborations in Appendix A.2A comparison between the demonstration-augmented MBD and the vanilla MBD is illustrated in Fig.6with detailed breakdowns in Appendix A.4.https://github.com/CMA-ES/pycmahttps://github.com/facebookresearch/nevergrad 


Decision: ### Summary
The paper introduces a method called Model-Based Diffusion (MBD) for solving trajectory optimization (TO) problems. Unlike traditional model-free diffusion approaches, MBD leverages explicit model information to compute the score function in the diffusion process. This allows it to efficiently solve TO problems without relying on pre-existing data. The method also has the flexibility to incorporate data of varying quality, enhancing its applicability in diverse tasks. Experimental results demonstrate that MBD surpasses state-of-the-art methods, including reinforcement learning and sampling-based optimization, particularly in complex tasks involving contact-rich environments. The paper also highlights MBD's ability to integrate incomplete or imperfect data, offering practical advantages over standard diffusion models.

### Decision
The paper is well-written, and the application of model-based diffusion for the trajectory optimization problems is novel and interesting. The reviewers initially had concerns about the paper on the applicability of the proposed method to real-world problems, experiments (in particular, comparison to simple RL baseline like PPO), and clarity of the theory section. However, I think during the rebuttal period, the authors have done a good job addressing most of these concerns. The paper was borderline. However, the approach is novel, and there are no critical problems, the concerns raised by the reviewers could be addressed with a minor modification of the paper. I recommend the authors to carefully read the points made by the reviewers and clarify those points for the camera-ready version of the paper.

Reviews: review: paper_summary: paper_summary: The paper presents Model-based Diffusion (MBD), a framework to train diffusion models to solve trajectory optimization (TO)  problems. In particular, the method assumes that the cost function associated with a TO problem can be evaluated for any trajectory of decision variables (states and control variables), which it then leverages to compute the target distribution of optimal solutions ($p_0(Y^{(0)})$) up to a normalizing constant. The denoising process in MBD performs a gradient ascent on intermediate distributions which are smoothed, easier to optimize versions of $p_0(Y^{(0)})$.

The method deals with unconstrained and constrained TO problems and is also able to incorporate demonstrations. MBD outperforms sota RL and also  sampling-based TO methods  in contact-rich tasks.
strength_weakness: strength_weakness: - While the comparison against RL is valuable and gives insights into the overall performance of MDB, the paper should better emphasize that the rewards reported in Table 2 are obtained by rolling out RL policies in an online closed-loop fashion, while the reward associated with MBD is computed according to an open loop execution.  Given that the paper lists the adaptation of MDB to online setup as future work, and that the rewards are computed in different ways, simply stating that MDB outperforms PPO by 59% can be misleading.  Qualifying such statements with an explanation of the different ways that rewards are computed would further improve the quality of the submission. 

[Minor] 
Typo L213 esitimation -> estimation
questions: questions: - What are the challenges of adapting your method to online tasks with receding horizon strategies?  Is it a matter of training MBD to solve TO problems starting from more initial states? 

-Once MDB has been trained, how long does it take to run the denoising process to get a new trajectory?
limitations: limitations: Limitations of the method are mentioned in the paper as potential lines of future work.
paper_summary: paper_summary: This paper presents Model Based Diffusion (MBD), a novel approach to trajectory optimization that leverages a diffusion process. Unlike traditional diffusion models where denoising networks are trained on data using a score matching loss, MBD directly computes the score function for a given target distribution. Specifically, for a trajectory optimization problem, the unnormalized target distribution is formulated as a product of three terms: an optimality term that depends on the cost function, a dynamical feasibility term, and a constraint satisfaction term. The authors demonstrate that in each iteration of the reverse diffusion process, the score function can be approximated using Monte-Carlo estimation – a weighted sum of samples drawn from a Gaussian distribution parameterized by the current trajectory, where the weight of each sample is the target density evaluated at that sample. Additionally, they show how their approach can be extended to handle scenarios with noisy observations of trajectories from the target distribution. Finally, they present experimental results on locomotion and manipulation tasks, demonstrating that their method outperforms other trajectory refinement techniques and RL algorithms in both average step reward and computational time.
strength_weakness: strength_weakness: The requirement for explicit specification of the cost function, dynamics, and constraints can be a limitation in real-world applicatioms. It would also be interesting to see how MBD would perform in the sparse reward or goal completion settings.
The performance of MBD presumably heavily relies on the accuracy of the provided dynamics model. Inaccurate models could lead to suboptimal trajectories and compounding errors. It would be useful to study how robust MBD is to varying degrees of model inaccuracy.
questions: questions: - The paper primarily evaluates MBD on tasks with relatively short horizons (e.g., 50 steps for most locomotion tasks). It would be valuable to understand how the performance of MBD scales with increasing horizon lengths. Do the computational benefits over RL persist for longer horizons? 
- How does the performance of MBD vary as a function of the number of samples used in the Monte Carlo approximation?
- For the RL agents, it looks like the full episode length (e.g. 1000 steps for locomotion tasks like hopper, halfcheetah) is used. The RL agents are presumably trained to maximize rewards over the entire episode, while is MBD optimized for a much shorter horizon. Is this then a fair comparison? The paper could benefit from a discussion of this comparison or additional experiments with RL agents trained on shorter horizons.
- How do you envision incorporating conditioning signals in MBD? This would be important for adapting MBD to a receding horizon strategy, where the trajectory is re-planned at each time step based on new observations.
limitations: limitations: The authors have not discussed the limitations of their work.
paper_summary: paper_summary: This work introduces model-based diffusion (MBD) to solve trajectory optimization (TO) problems without relying on data. The key idea is to explicitly compute the score function by leveraging the model information in TO problems. This paper also shows that MBD has interesting connections to sampling-based optimization. The empirical evaluations demonstrate that MBD outperforms state-of-the-art TO methods.
strength_weakness: strength_weakness: - The theoretical derivation and notation in the methodology are quite confusing and hard to follow.
- The comparison to some simple baseline of RL, such as PPO, cannot demonstrate the real advantages.  
- In line 137 and equation $(3)$, there may be an incorrect formula: 
$$
p_{i|i-1}(\cdot, Y^{(i-1)}) \sim \mathcal{N}(\sqrt{\alpha_i}Y^{(i-1)},1-\alpha_iI)
$$
and 
$$
p_{i|0}(\cdot, Y^{(0)}) \sim \mathcal{N}(\sqrt{\overline{\alpha}_i}Y^{(0)},1-\overline{\alpha}_iI)
$$

- In equation $(8)$, there may be an incorrect formula:
$$
\phi_i(Y^{(0)})\propto\cdots\propto\mathcal{N}(Y^{(0)}-\frac{Y^{(i)}}{\sqrt{\overline{\alpha}_i}},\frac{I}{\overline{\alpha}_i}-I)
$$

- In all algorithms, it is evident that the calculation step can be simplified: 

$Y^{(i-1)} \leftarrow \frac{1}{\sqrt{\alpha_i}}(Y^{(i)}-Y^{(i)}+\sqrt{\overline{\alpha}_i}\overline{Y}^{(0)}(\mathcal{Y}^{(i)})) = \frac{\sqrt{\overline{\alpha}_i}}{\sqrt{\alpha_i}}\overline{Y}^{(0)}(\mathcal{Y}^{(i)})=\sqrt{\overline{\alpha}_{i-1}}\overline{Y}^{(0)}(\mathcal{Y}^{(i)})$

Note that this coefficient $\sqrt{\overline{\alpha}_{i-1}}$ is canceled at the first step. Also, note that in equations $(9b)$, $(10d)$, and $(13)$, it is a weighted average of the samples from $\mathcal{Y}^{(i)}$. It may be clearer if the calculation steps are more simplified and easy to understand.

- In Table 3, the computational times of the CEM and MC-MBD algorithms are close. I'm wondering if this table includes the computation time for environment interaction. It might be fairer to compare just the computation times for the algorithms.
questions: questions: Please refer to the Weaknesses.
limitations: limitations: Limitations are not discussed in the manuscript.
paper_summary: paper_summary: The paper proposed MBO, a novel method to solve trajectory optimization problems with diffusion models. MBD formulates the optimization problem as a sampling problem and estimates the score function using the dynamic model.
strength_weakness: strength_weakness: - In the introduction, the authors claimed that existing diffusion-based approaches often require high-quality demonstration data. However, I cannot find any evidence of experiment results supporting this claim in the manuscript. Diffusion-based planners are not in the baselines of the experiments, and there are several model-free diffusion-based methods that can efficiently generate dynamically feasible and high-rewarding trajectories in RL fields. Could authors validate the reason why model-free diffusion planners are not included in the baselines?

- The authors also claimed that one of the advantages of model-based diffusion over model-free diffusion is its generalization ability (e.g., generalization to new tasks). However, I cannot find empirical evidence for this claim. Could the authors elaborate more on this claim?
questions: questions: There are some minor questions for better understanding.

- For model-free diffusion, we need to train the score function estimator through the data. In model-based diffusion, which part should be parametrized by the neural network and trained with the data?

- MBD is compared with the baselines, including zeroth-order optimization and RL algorithms. Is there any other recently proposed method for trajectory optimization? 

- For data-augmented MBD, there are also several works that utilize model-free diffusion models to augment data given sub-optimal offline datasets. They claim that through stitching ability, model-free diffusion can also generate novel trajectories [1, 2]. Could authors validate empirically that MBD is more powerful compared to model-free diffusion models given partial and dynamically infeasible datasets?

[1] Ajay, Anurag, et al. "Is Conditional Generative Modeling all you need for Decision Making?." The Eleventh International Conference on Learning Representations. 

[2] Li, Guanghe, et al. "DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching." Forty-first International Conference on Machine Learning.
limitations: limitations: - Please see the weaknesses and questions above.

Comments: title: Thanks for the feedback!
Thanks for the feedback!
Thanks for the feedback
Thank you for the clarifications
Thanks for the comments!
Thank you for your response
Thanks for the suggestion!
Thanks for the clarifications!
comment: Thanks for the clarification. I will raise my score.
Thanks for the suggestion. 

Fistly, comparing with MBRL, our setting centers on Trajectory Optimization, which assume model is available. In contrast, model-based RL baselines like Dream-v3 and TD-MPC2 are designed to concurrently learn the model and the policy. Therefore, a direct comparison is not entirely appropriate as it isn’t an apples-to-apples situation. However, examining the optimizer of TD-MPC2[4], one can observe that the planner module used for solving optimization tasks with learned dynamics is MPPI, which we also employ as a baseline across all our environments. In fact, MBD could be regarded as a subroutine within model-based RL algorithms, capable of optimizing the trajectory given a model. We believe this integration offers a promising avenue for future research.

Secondly, comparing with more advanced MFRL, we think PPO and SAC are still representative baselines for the asymptotic performance. 
Specifically, DIPO and QSM demonstrate better sample efficiency and multi-modal representation compared with PPO and SAC, but since we train RL util converge, we only focus on asymptotic performance, both of which does not show significant improvement over PPO/SAC.  
If we are proposing a diffusion-based RL algorithm, we would definitely compare with the diffusion-based RL algorithm like DIPO and QSM. 
However, since MBD is a trajectory optimization algorithm, we believe our choice of PPO and SAC as the aymptotic performance reference would be enough to demonstrate the effectiveness of our method, especially given the fact that there are no existing TO algorithm can solve those task well. 

If you have any further suggestions or concerns, please let us know. We appreciate your feedback!
We appreciate the reviewer's suggestion to empirically demonstrate the generalizability of MBD to various dynamics and cost functions. As shown in Sec.5 of the paper, we have conducted extensive experiments, which suggest that MBD, as a **Trajectory Optimization solver**  (aka planning, a very important problem in decision making and RL), inherently adapts to **ANY** given model and cost function. Alongside with its model-based nature as a Trajectory Optimization, we hope it's sufficient to showcast its capacity to handle diverse dynamics and costs.

However, we would like to seek further clarification on reviewer's specific concerns regarding the generalization abilities of MBD. If there are particular aspects or additional experiments you believe would more effectively demonstrate this capability, we would be delighted to explore these avenues to address your concerns.
Yes, you are right. 
The optimization variable should be $u_{0:T-1}, x_{1:T}$ instead of $u_{1:T}, x_{1:T}$.
Thank you so much for point it out! We will correct the notation in the revised version.
Thank you for the authors' detailed response. I understand that the settings of MFD and MBD are quite different, but it might be better to empirically demonstrate that MBD can generalize to various dynamics and cost functions. Therefore, I maintain the score.
Thank you for the clarifications. 

I have a question regarding the notation. Action $u_t$ at state $x_t$ leads to state $x_{t+1}$. Thus, the lhs of equation 1a should have $J_{x_{1:T}, u_{0:T-1}}$ instead of $J_{x_{1:T}, u_{1:T}}$? And I think $Y$ should also be modified to include $u_0$?
Thanks for your response. However, I still have some concerns about the comparison baselines. PPO and SAC can provide some reference for the performance, but they are not strong baselines. Recent model-free RL methods such as DIPO [1], QSM [2] show much better performance compared to PPO and SAC. Besides, since this method is model-based Diffusion, why not compare it to  model-based RL baselines, such as Dream-v3 [3], TD-MPC2 [4]. 

[1] Long Yang, et al, Policy Representation via Diffusion Probability Model for Reinforcement Learning, arXiv:2305.13122. 

[2] Michael Psenka, et al, Learning a Diffusion Model Policy from Rewards via Q-Score Matching, ICML 2024. 

[3] D Hafner, et al, Mastering Diverse Domains through World Models, arXiv:2301.04104. 

[4] Nicklas Hansen, et al, TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024.
Thank you for the comments. Here are our responses: 

> Regarding the extension to closed-loop control with a receding horizon, could you elaborate on how the current state x_t is incorporated into Algorithm 1?

We would like to refer to line 4 on Algorithm 2, where Y represents the entire control input sequence that $ Y = U = u_{1:T} $. Given the current state $x_t$, we can roll out the control sequences starting from $x_t$ and generate the predicted state sequence $x_{t+1:t+T}$. This procedure allows us to incorporate and current state $x_t$ into MBD's framework just by calculating the score function as shown in line 5 on Algorithm 2. The closed-loop process is effectively streamlined, as our non-receding-horizon MBD algorithm is capable of initiating from any initial state without requiring conditional training, as demonstrated in Algorithm 2, thanks to the model-based nature of MBD. We hope this would clarify the reviewer's inquiry. 

---

> In lines 506-507, you mention that for pushT, the reward is composed of the distance between the target and the object, as well as the orientation difference between them. Is the reward signal available at each state and computed based on the distance and orientation, or is it a sparse reward provided only when the distance/orientation are below certain thresholds?

Thanks for the clarification question. The reward signal is computed at each state without certain thereshold to the object. We consider this reward signal as a sparse since reward is not directly applied to the agent but to the object. The agent only gets reward when the object is moved. 
That's why PPO cannot solve this task effectively due to most of the movement of the agent is not rewarded and agent need to move a long distance to get the reward, especially when switching the contact point. 
However, we totally agree that the definition of sparse reward could be vague and we will clarify this as hard-to-explore reward in the revised version.
Thank you for your response to my questions and for conducting the additional experiments.

I have a few follow-up questions:

Regarding the extension to closed-loop control with a receding horizon, could you elaborate on how the current state x_t
is incorporated into Algorithm 1?

In lines 506-507, you mention that for pushT, the reward is composed of the distance between the target and the object, as well as the orientation difference between them. Is the reward signal available at each state and computed based on the distance and orientation, or is it a sparse reward provided only when the distance/orientation are below certain thresholds?
> Using the naive receding horizon MBD algorithm introduced by the authors, the overall compute time to obtain a new solution is probably much lower due to the use of single-step MBD and to the forward shifting of the trajectory to initialize the next optimization step. Reporting the possible close-loop rates of MBD would also be valuable for the readers.

Thanks for the suggestion. Showing the computation time of online MBD is indeed valuable for the audience who interested in real-time applications. 
We tested the online frequency on RTX 4070Ti GPU and the results are shown in Table 5.
We will also include the computation time of MBD in the updated manuscript.

Table 5: Online Running Frequency of receding horizon MBD
| Env | Frequency (Hz) |
| --- | --- |
| Hopper | 4.56 |
| HalfCheetah | 4.51 |
| Ant | 10.28 | 7.46 |
| Walker2d | 3.49 |
| Humanoid Standup | 6.82 |
| Humanoid Running | 4.03 |

Please note that the frequency is calculated under the assumption of solving the whole 50 steps TO problem without reduced model at each iteration, which quite different from the MPC approach proposed in [1] which only solve the reduced model. 
Besides, the mujoco code environment we used is not optimized for GPU, so the actual frequency could be higher with optimized environment. In our work we just use Brax as a simple and easy-to-use option. As the major computation time of MBD is spent on the forward dynamics simulation, it can be further improved by using more efficient physics engine. Moreover, MBD's online frequency can be further improved by introducing learning components including learned value function, policy, or dynamics model like TD-MPC [2]. 

In practice, by using GPU-optimized Unitree Go2 environment provided by [mujoco_menagerie](https://github.com/google-deepmind/mujoco_menagerie), we can achieve $52.3$ Hz control for Go2 robot with 12 DoF in walking task with 20 steps prediction horizon. 

Lastly, since MBD is designed to be a TO solver rather than a real-time controller, the computation time is not the main concern. But we totally agree that it is valuable to further explore the how to bring MBD to real-time which can solve the whole dynamic system without reduction model at each iteration and MBD has already shown its potential in this direction.

[1] D. Kim, J. Di Carlo, B. Katz, G. Bledt, and S. Kim, “Highly Dynamic Quadruped Locomotion via Whole-Body Impulse Control and Model Predictive Control,” Sep. 14, 2019, arXiv: arXiv:1909.06586.
[2] N. Hansen, X. Wang, and H. Su, “Temporal Difference Learning for Model Predictive Control,” Jul. 19, 2022, arXiv: arXiv:2203.04955. doi: 10.48550/arXiv.2203.04955.
I have read the author's response and appreciate the additional experiments and clarifications provided.

Regarding the extension to the online setup, it would be useful to provide the frequency at which the close-loop control can be executed. Table 3 reports a compute time of at least 17 seconds for MC-MBD, which is inadequate for close-loop control, especially for highly dynamic systems like the humanoids.

Using the naive receding horizon MBD algorithm introduced by the authors, the overall compute time to obtain a new solution is probably much lower due to the use of single-step MBD and to the forward shifting of the trajectory to initialize the next optimization step. Reporting the possible close-loop rates of MBD would also be valuable for the readers.

Thanks.