Model-Based Diffusion for Trajectory Optimization

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.

Strengths

- The motivation of the paper is valid. In the trajectory optimization problem, we already know the underlying dynamics of the environment. Therefore, it is natural to utilize such information when we use diffusion models to solve trajectory optimization problems. - The problem statement is clearly written and easy to follow, even for readers who are not familiar with the topic. - Experiment results demonstrate that MBD outperforms prior methods in various domains.

Weaknesses

- 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

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.

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

- Please see the weaknesses and questions above.

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.

Strengths

- The model-based diffusion framework enables an effective trajectory planner. - This algorithm shows good performance on various tasks compared to other state-of-the-art algorithms.

Weaknesses

- 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

Please refer to the Weaknesses.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are not discussed in the manuscript.

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.

Strengths

This paper introduces an innovative approach that leverages diffusion models directly as solvers for trajectory optimization. The key advantage is that if the cost function, dynamics, and constraints can be specified for a given problem, the proposed iterative refinement approach can obtain samples from the desired distribution without having to rely on any demonstration data. If demonstration data is available, the authors demonstrate how it could be used to augment their approach. The paper is also well-presented and easy to understand.

Weaknesses

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

- 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.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have not discussed the limitations of their work.

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.

Strengths

- The proposed method is evaluated on contact-rich, high-dimensional tasks and outperforms the baselines in terms of reward while keeping a computational time similar to the fastest baselines. - The paper is well-written and the main ideas are clear and well-motivated. Furthermore, the method shows interesting connections to multi-stage optimization and also to the Cross Entropy Method for sample-based optimization. - The method does not have many hyperparameters and they are robust across most tasks.

Weaknesses

- 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

- 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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations of the method are mentioned in the paper as potential lines of future work.

Thanks for the clarifications!

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.

Thank you for your response

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?