DiffuserLite: Towards Real-time Diffusion Planning

Thank you for your effort in reviewing and acknowledging our work. Based on the questions you have raised, it seems you are particularly interested in the implementation details of DiffuserLite and its potential applications. We have already released the initial version of the code on diffuserlite/diffuserlite.github.io to help you understand the algorithm better. And below is our response to your questions in the hope of addressing your concerns:

• Q1: Your question offers a different perspective on the planning procedure of DiffuserLite, particularly regarding treating the value critic as a classifier, which is indeed correct. This classifier plays two key roles in the planning framework: (a) guiding the Diffusion model to generate high-performance trajectories, a process that can be achieved either through gradient calculations (referred to as CG, commonly utilizing automatic differentiation in DL software packages) or by using the value as input to a neural network (referred to as CFG). Although mathematically equivalent, due to the high computational cost of gradient calculations that significantly reduce decision frequency, we have chosen to use CFG in DiffuserLite. Another reason for choosing CFG is that the Rectified flow backbone cannot use CG for guidance. Rectified flow is our alternative generative model backbone. To our knowledge, DiffuserLite is the first to validate its effectiveness in decision-making domains and leverage its reflow procedure further enhanced decision frequency. As Rectified flow cannot use CG, this is why we opted for CFG. To demonstrate the impact of using CG on decision frequency, we tested DiffuserLite-D on MuJoCo benchmarks, resulting in an average decision frequency of 27.2Hz (CG) versus 68.2Hz (CFG), suggesting that CFG should be preferred when high decision frequency is crucial.
• Q2: This question is highly relevant as real-world applications often encounter non-differentiable physical constraints. As a diffusion planning algorithm, DiffuserLite offers various ways to address this issue: (a) Explicitly filtering out trajectories that do not meet physical constraints during planning: DiffuserLite generates multiple candidate trajectories at each decision step and evaluates their quality to select the optimal one. Similar to search methods, this procedure allows for the explicit filtering of candidates that do not adhere to constraints. (b) Incorporating additional guidance: While many physical constraints are non-differentiable, we can soften these constraints and approximate them using a neural network classifier to guide the generation of trajectories that best satisfy these constraints. Previous research has shown significant improvements in the physical plausibility of generated trajectories by utilizing gradient guidance from the environment dynamic model [1]. Given DiffuserLite's flexible, plugin-based framework, this method can be incorporated with DiffuserLite to ensure compliance with physical constraints. In conclusion, we believe that a combined approach using (a) and (b) can lead to favorable performance in most real-world applications. As demonstrated in the paper, Robomimic benchmark offers a real-world robot arm dataset with inherent physical constraints. With method (a), DiffuserLite has already achieved very good performance on this benchmark.

I hope our response addresses your concerns. Please feel free to reach out if you need further clarification. Thank you once again for your insightful review.

[1] Ni, et al, "MetaDiffuser: diffusion model as conditional planner for offline meta-RL," in Proceedings of the 40th International Conference on Machine Learning, 2023.

Thank you for your thorough review. The concerns you have raised are insightful, and we aim to address them effectively in the following responses. To be concise, below, we use "Lite" to refer to "DiffuserLite".

Weakness: Limited Novelty

We aim to elucidate the differences between Lite and HDMI as well as HD-DA (It is actually a concurrent work and both it and Lite are published on arXiv in January 2024) from the perspectives of Motivation, Technology, and Reproducibility.

1. Motivation: The motivation behind Lite lies in reducing redundant information during planning to increase decision frequency, whereas HDMI and HD-DA focus on generating higher-quality plans through "subgoal planning" and "goal reaching". The primary contribution of Lite lies in speed enhancement. Experimental results indicate that in terms of D4RL score, Lite outperforms HD-DA and HDMI, with a significantly higher increase in decision frequency (Lite-R1: 52.2x $\gg$ HDMI: 1.8x $\approx$ HD-DA: 1.3x, compared to Diffuser) [1,2]. We argue this disparity is due to the failure of the other two works to eliminate redundant information in planning, unlike Lite. DiffuserLite addresses a distinct problem and is not merely an improvement upon HDMI.

2. Technology: While all three algorithms exhibit a hierarchical structure, Lite offers additional insights:

• Lite demonstrates high-performance diffusion planning WITHOUT meticulous dataset handling or the inclusion of prior knowledge, unlike HDMI.
• Lite proves that disregarding distant redundant information and using a significantly shorter planning horizon can still yield excellent results and greatly enhance speed compared to HDMI and HD-DA.
• Lite conducts extensive ablation experiments discussing the impact of the number of hierarchical levels $L$ and interval at each level $I_l$ on performance, summarizing best practices in hierarchical structure design, a discussion lacking in HDMI and HD-DA (and they only use $L=2$).
• Lite validates its algorithm on real-world datasets like Robomimic and FinRL, showcasing its practical applicability, which is not extensively covered in HDMI and HD-DA.
• Lite boasts a simpler structure that can serve as a flexible plugin for other diffusion planners, as validated on AlignDiff, setting it apart from HDMI and HD-DA.

3. Reproducibility: Currently, only DiffuserLite has released code on Lite/Lite.github.io, making its results the most reliable.

Questions

Q1: Thank you for your reminder. The website content is updated. : )

Q2: Planning-based methods suffer from compounding errors, leading agents to deviate from the plan over time [3]. Hence, agents execute only a few steps based on the plan (typically 1 step in planning-based RL algorithms). Therefore, we believe that detailing every aspect of plan generation is unnecessary. We believe successful long-term planning requires: (a) Sparser details for more distant parts (as long as they can accurately reflect the plan's quality). (b) More details for nearer parts to ensure decision consistency with the plan.

Q3: In lines 176-178, the described critic used in Diffuser/DD predicts the cumulative return within the horizon, leading to greedy plans, lacking of long-term judgment. For example, Hopper in Table 2 is a single-legged hopping robot, faster movement yields higher rewards. A common failure mode in Diffuser and DD is that the agent jumps abruptly and then falls, attributing this to the short-sightedness induced by greedy planning within the horizon. Critics with values can alleviate this issue by providing long-term judgment.

Q4: In line 186, a 'uniform critic' refers to assigning the same judgment to all plans, assuming equal quality. DD generates one high-performance plan a priori for execution without a plan selection process. To align with our unified diffusion planning framework, DD can be viewed as utilizing a uniform critic for plan selection.

Q5: The metric used in Table 2 is the normalized D4RL score, which is the episodic return normalized against expert performance (100 means expert-level [4]). The metric used in Table 3 is the success rate of manipulation tasks in Robomimic [5]. The metric used in Table 4 is the episodic return in the realistic stock market simulation [6]. Across all tables, higher values indicate better performance.

I hope our response addresses your concerns. Please feel free to reach out if you need further clarification. Thank you once again for your insightful review. Moreover, we have released the code at https://github.com/diffuserlite/diffuserlite.github.io to facilitate a better understanding of the implementation details. I hope it can help you : )

References

[1] Chen, et al, "Simple Hierarchical Planning with Diffusion," in The Twelfth International Conference on Learning Representations, ICLR, 2024.

[2] Li, et al, "Hierarchical Diffusion for Offline Decision Making," in Proceedings of the 40th International Conference on Machine Learning, ICML, 2023.

[3] Xiao, et al, "Learning to Combat Compounding-Error in Model-Based Reinforcement Learning," in arXiv,1912.11206, 2019.

[4] Fu, et al, "D4RL: Datasets for Deep Data-Driven Reinforcement Learning," in arXiv,2004.07219, 2021.

[5] Mandlekar, et al, "What Matters in Learning from Offline Human Demonstrations for Robot Manipulation," in Conference on Robot Learning (CoRL), 2021.

[6] Liu, et al, "FinRL: A Deep Reinforcement Learning Library for Automated Stock Trading in Quantitative Finance," in arXiv,2011.09607, 2022.

Reviewer 7Mwp,

As the author-reviewer discussion phase comes to a close, have the authors satisfactorily addressed your concerns? Please engage in the discussion if you still have any unsolved concerns.

Best,

AC

Thank you for your insightful suggestions. I am deeply touched by your diligent review and appreciate your dedication. I summarize the questions you raised and respond below:

• Implementation Details: We have released the code at diffuserlite/diffuserlite.github.io to facilitate a better understanding of the implementation details. Furthermore, in response to your suggestion, we will incorporate more details on implementation and design choices into main text during revision to avoid the necessity of referring to Appendix for complete comprehension.

• Clearer Presentations: We commit to double-checking all figures and tables during revision to ensure they are presented more clearly and address any potential ambiguity issues.

• Real-World Tests: Robomimic and FinRL benchmarks provide real-world datasets that significantly differ from D4RL. The strong performance of DiffuserLite on these benchmarks demonstrates its potential for real-world applications. Besides, our released code allows for easy adjustments of model size and training datasets to evaluate scalability in more complex real-world scenarios. We will explore this in future work.

• PRP-Introduced Complexity: While the methodology may appear complex, the implementation of our algorithm is not significantly more intricate than one-shot generation. After excluding components such as neural networks, generative backbones, and the reflow procedure, the core algorithm comprises approximately 100 lines of code, alleviating concerns regarding complexity.

• Design Choice Recommendations: More PRP levels allow for more redundant information to be discarded but may also accumulate errors across levels. In the revision, we will provide theoretical justifications from this perspective. However, due to numerical calculation and many estimations in PRP, theoretical justifications may not suggest an optimal configuration. As such, we conduct additional experiments to explore various design choices, as detailed in the following table. This table is an expanded version of Table 6 in the paper. You can see that the third column is the right part of Table 6. The results align with the recommendations in Appendix C, emphasizing that 3 levels are a practical choice balancing performance and decision frequency, while too many levels can lead to error accumulation and reduced decision frequency, and too few levels may complicate trajectory distribution simplification, resulting in lower performance. It is also important to balance the horizon length at each level, avoiding overweighting at the top or bottom. Additional design choice recommendations can be found in Appendix C, and we will integrate these details into the main text based on your feedback for the revision.

• Performance with $L=5$: To address your query, we include an experiment with $L=5$ in the table above (column 4). Increasing the number of levels beyond a certain point may lead to error accumulation and reduced decision frequency. Considering these factors, we still recommend using $L=3$.

I hope our response addresses your concerns. Please feel free to reach out if you need further clarification. Thank you once again for your insightful review.

Thank you for your insightful suggestions. I am deeply touched by the diligent review and appreciate your dedication. I summarize the issues and respond below to address your concerns:

• Paper revision: Following your suggestion, during revision, we will further introduce Reflow in main text, address PRP’s multiple abbreviation issue, update the D4RL score report for Diffuser/DD, and shift the design choices from Appendix to main text.
• About PRP configuration: More PRP levels allow more redundant information discarded but may also accumulate errors across levels. We will provide theoretical justifications from this perspective during revision. However, due to numerical calculation and many estimations in PRP, theoretical justifications may not suggest an optimal configuration. As such, we conduct additional experiments to explore various design choices, as detailed in the following table. This table is an expanded version of Table 6 in the paper. You can see that the third column is the right part of Table 6. The results align with the recommendations in Appendix C, emphasizing that 3 levels are a practical choice balancing performance and decision frequency. It is also important to balance the horizon length at each level, avoiding overweighting at the top or bottom. Additional design choice recommendations can be found in Appendix C, and we will integrate these details into the main text based on your feedback for the revision.

• About CG and CFG: We do not use CG primarily for two considerations: (a) CG requires NN gradients computation, which is very expensive. A comparison of decision frequency for DiffuserLite-D on MuJoCo, 27.2Hz (CG) vs. 68.2Hz (CFG), suggests that CFG should be preferred when high decision frequency is crucial. (b) Rectified flow cannot use CG for gudiance. Therefore, we solely employ CFG in Lite. Although CFG requires specifying a target condition for each level, various methods can address this: (a) using an additional NN to predict the max reachable return for the current state as the target or employing more sophisticated methods in [1], (b) we also observe that Level 0 has the most significant impact on decision performance, with most cases requiring only the tuning of target condition for Level 0.
• About critic design: (Why is it helpful when sparse rewards?) Assuming $x_s$ as the planned trajectory and $\tau$ as $x_s$ and its future trajectory, $R$ estimates the trajectory's discounted return. A critic with only reward terms optimizes $\mathbb E_{q_0(x_0),q(\epsilon),s,\tau}\left[\Vert\epsilon_\theta(x_s,s,R(\tau))-\epsilon\Vert^2\right]$, whereas a critic with value terms optimizes $\mathbb E_{q_0(x_0),q(\epsilon),s}\left[\Vert\epsilon_\theta(x_s,s,\mathbb E_\tau[R(\tau)])-\epsilon\Vert^2\right]$. When rewards are sparse, a critic with value terms can receive more stable conditions during training, leading to better performance in sparse reward tasks. Additionally, DiffuserLite does not restrict the use of any specific critic; as demonstrated in AlignDiff-Lite, we utilized an attribute strength model [2] as a critic for preference aligning. You can also deploy DiffuserLite in online tasks using curiosity-based rewards or value-based methods.
• In Section 5.5, Lite w/ only last level refers to one level with a horizon of 9, while Lite w/o PRP refers to one level with a horizon of 129. I apologize for any lack of clarity.
• Low Performance of Diffuser and DD in Antmaze: In Antmaze, we run official code for Diffuser and referenced the report in [3] for DD. We did not intentionally report low scores. The higher scores in LDGC [4] may be due to (a) their use of an inpainting version of Diffuser as the original version yielded poor results (page 9, line 2), and (b) careful tuning of DD hyperparameters. To ensure a fair comparison, we will update the score report during revision. However, Diffuser and DD still exhibit significantly lower scores than Lite. Regarding Kitchen tasks, we have double-checked to confirm that both LDGC and our reports align with official reports.
• About DiffuserLite-R1: We attribute its best performance to the straight flow property of Rectified Flow, which confers an advantage in very few sampling steps. Compared to Lite-D, we only change the generative backbone from diffusion model to rectified flow, without any other tricks.
• Practical Aspects: Robomimic and FinRL provide real-world datasets with substantial differences compared to D4RL. The strong performance of Lite on these datasets underscores its potential for real-world applications.

I hope our response addresses your concerns. Please feel free to reach out if you need further clarification. Thank you once again for your insightful review.

[1] Yu, et al, "Regularized Conditional Diffusion Model for Multi-Task Preference Alignment," in arXiv, 2404.04920, 2024.

[2] Dong, et al, "AlignDiff: Aligning Diverse Human Preferences via Behavior-Customisable Diffusion Model," in The Twelfth International Conference on Learning Representations, 2023.

[3] Hong, et al, "Diffused Task-Agnostic Milestone Planner," in arXiv, 2312.03395, 2023.

[4] Venkatraman, et al, "Reasoning with Latent Diffusion in Offline Reinforcement Learning," in arXiv, 2309.06599, 2023.

I have carefully reviewed your detailed responses to my comments. Your commitment to addressing the points raised is appreciated. Given the promised revisions and the forthcoming code release, I believe the simple yet effective idea of PRP could indeed be valuable in reducing the notorious planning time of diffusion models for planning.

While I find the empirical results compelling, I would have been even more supportive if there were stronger theoretical justifications for the approach. Nonetheless, the practical benefits of the method are clear.

Considering these factors, I am revising my overall assessment from 5 to 6.

Here are a few remaining questions.

1. Regarding CFG and CG: In your response, you mentioned that CG was not used for several reasons, including decision frequency. Is there a difference in performance (score) between using CG and using CFG without considering rectified flow? In Section 4, you stated that you used CFG for decision frequency. However, in the conclusion, you mentioned the limitations of DiffuserLite with CFG being the main cause of these limitations. Could using CG address this issue? It would be helpful to explain why it was necessary to improve decision frequency by opting for CFG instead of CG. As it currently reads, your writing suggests, "We use CFG because CG is not good, but CFG also has limitations!" If CG can address this issue, it would be better to explain why decision frequency is important despite those limitations. If not, it would be clearer to explain that the issue is not about the limitations between CG and CFG.

2. Additional question about the plugin: Have you tried using this method as a plugin for other backbones besides AlignDiff? It would be beneficial to include any tests that show whether it can be applied to more major methods like Diffuser or HDMI, rather than AlignDiff.