RAGDP: Retrieve-Augmented Generative Diffusion Policy

Authors: Sodtavilan Odonchimed, Tatsuya Matsushima, Simon Holk, Yusuke Iwasawa, Yutaka Matsuo
Primary area: applications to robotics, autonomy, planning
Keywords: Imitation Learning, Diffusion Models, Behavior Cloning

Abstract:
Diffusion Policy has attracted attention for its ability to achieve significant accuracy gains in a variety of imitation learning tasks. However, since Diffusion Policy relies on the Diffusion Model, it requires multiple denoising steps to generate a single action leading to long generation times. To address this issue, methods like DDIM and Consistency Models have been introduced to speed up the process. While these methods reduce computation time, this often comes at the cost of accuracy. In this paper, we propose RAGDP, a technique designed to improve the efficiency of learned Diffusion Policies without sacrificing accuracy. RAGDP builds upon the Retrieval-Augmented Generation (RAG) technique, which is commonly used in large language models to store and retrieve data from a vector database based on encoded embeddings. In RAGDP, pairs of expert observation and actions data are stored in a vector database. The system then searches the database using encoded observation data to retrieve expert action data with high similarity. This retrieved expert data is subsequently used by the RAGDP algorithm to generate actions tailored to the current environment. We introduce two action generation algorithms, RAGDP-VP and RAGDP-VE, which correspond to different types of Diffusion Models. Our results demonstrate that RAGDP can significantly improve the speed of Diffusion Policy without compromising accuracy. Furthermore, RAGDP can be integrated with existing speed-up methods to enhance their performance.

Decision: reject
Reviewer scores: 6, 3, 5, 3

Reviews and discussion:
<thread>
creation date:
1729997445332

thread_type:
official_review

<thread_object>
creator: reviewer_q22q

<content>
summary:
The paper proposes RAGDP, a new method aimed at enhancing the efficiency of diffusion policies in imitation learning by leveraging Retrieval-Augmented Generation (RAG). Traditional diffusion policies require multiple denoising steps, resulting in long generation times, which methods like DDIM and Consistency Models attempt to address at the expense of accuracy. RAGDP utilizes a vector database to store and retrieve pairs of expert observation and action data based on encoded embeddings, enabling faster action generation tailored to the current environment. The authors introduce two specific algorithms, RAGDP-VP and RAGDP-VE, corresponding to different types of diffusion models. Experimental results demonstrate that RAGDP significantly reduces generation time while maintaining accuracy and can be effectively integrated with existing acceleration methods.

strengths:
#### Originality
The paper presents a novel approach called RAGDP (Retrieve-Augmented Generative Diffusion Policy) that significantly enhances the efficiency of diffusion policies in imitation learning. By leveraging Retrieval-Augmented Generation (RAG), the authors introduce a fresh and creative solution to the problem of long generation times associated with traditional diffusion policies. The proposal of two distinct algorithms, RAGDP-VP and RAGDP-VE, further underscores the paper's originality by catering to different diffusion model types.

#### Quality
The paper demonstrates high-quality research through extensive experimentation and empirical validation. The authors provide comparative analyses with existing methods like DDIM and Consistency Models, showcasing the effectiveness of RAGDP in enhancing speed without sacrificing accuracy. 

#### Clarity
The paper is well-written and clearly structured, making it accessible to readers with a background in imitation learning and diffusion models. 

#### Significance
The significance of the proposed RAGDP method is evident in its potential to drastically reduce the computational cost of diffusion policies in imitation learning tasks. By maintaining high accuracy while improving speed, RAGDP addresses a critical bottleneck in practical applications, especially in real-time robotic tasks.

weaknesses:
- **Limited Novelty in Retrieval-Augmented Techniques:** While RAGDP is creative in combining retrieval-augmented generation with diffusion policies, the concept of retrieval-augmented generation itself is not entirely new. Prior works like RAG (Lewis et al., 2021) for text generation and ReDi (Zhang et al., 2023) for image generation are mentioned in the paper. The novelty mainly lies in its application to diffusion policies for imitation learning, which may not be sufficient to stand out unless significantly differentiated from these prior works.

- **Hyperparameter Sensitivity:** The experiments indicate that RAGDP's performance is sensitive to the hyperparameter \( r \). However, the paper lacks a thorough investigation into how to optimally choose this parameter across different tasks and environments. More robust guidelines or heuristic methods for setting this parameter would strengthen the work.

- **Complexity of Algorithm Presentation:** While the algorithms RAGDP-VP and RAGDP-VE are explained, the presentation might be too complex for readers not deeply familiar with stochastic differential equations and diffusion models. Simplified explanations or additional illustrative examples could make the paper more approachable.

- **Elaboration on Vector Database Implementation:** The paper briefly mentions the use of FAISS for the vector database implementation but does not delve into the details of its configuration, indexing methods, and search efficiency. More information on these aspects would add clarity and allow for better reproducibility of the results.

questions:
### Questions and Suggestions for the Authors

1. **Comparison with More Baseline Methods:**
   - **Question:** Have you considered comparing RAGDP with other state-of-the-art imitation learning methods that do not use diffusion models?
   - **Suggestion:** Including such comparisons would provide a clearer picture of RAGDP's relative advantages and limitations.

2. **Hyperparameter Sensitivity:**
   - **Question:** Can you provide more insights into how the hyperparameter \( r \) was chosen for various tasks? Is there a heuristic or automated method to set this parameter?
   - **Suggestion:** A detailed analysis or guideline on setting \( r \) would enhance the robustness and reproducibility of the method.

3. **Algorithm Explanation:**
   - **Question:** Could you simplify the presentation of the RAGDP-VP and RAGDP-VE algorithms, perhaps with illustrative examples or a more intuitive walkthrough?
   - **Suggestion:** Simplified explanations and additional illustrations would make the algorithms more accessible to a broader audience.

4. **Details on Vector Database Implementation:**
   - **Question:** Can you elaborate on the configuration and optimization of the FAISS vector database, including indexing methods and search efficiency?
   - **Suggestion:** Providing these details would help in understanding the practical aspects of implementing the vector database and ensuring reproducibility.

In summary, it appears that RAGDP tends more towards being an engineering optimization and might not have enough originality.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730569390059

thread_type:
official_review

<thread_object>
creator: reviewer_fsfw

<content>
summary:
This paper proposes to improve the efficiency while preserve the effficacy of the diffusion-based policy generation, which is achieved by utilizing the retrieval.

strengths:
The logic is clear and easy to follow.

weaknesses:
There are several confusions that required to be answered:

1) Compared with other tasks, what are the special requriments for imitation learning on the utilization of retrieval-based diffusion?

2) Sec3.2 is confusing. Is mentioned prediction horizons T_p means steps within single diffusion? Or it means a sequence of actions in the imitation learning?

3) Given only the most similar are used for speeding up the diffusion process, what if the retrieval outputs a bad matching? Another thing is any possible that this retrieval operation degenerates the generalizability of the model?

4) It is not clear how the Encoder can be jointly trained with the policy generator, especially there is a matching algorithm involved. The details should be presented.

questions:
Please refer to the Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730599941806

thread_type:
official_review

<thread_object>
creator: reviewer_1zmy

<content>
summary:
This paper introduces RAGDP (Retrieve-Augmented Generative Diffusion Policy), which aims to improve the efficiency of diffusion policies without sacrificing accuracy. The main idea is to combine retrieval-based methods with diffusion models by storing expert observation-action pairs in a vector database. The paper presents two algorithms - RAGDP-VP and RAGDP-VE - which correspond to different types of diffusion models. The approach enables faster action generation while maintaining performance by leveraging retrieved expert demonstrations to inform the diffusion process.

strengths:
1. RAGDP does not necessarily require additional training and can be integrated with existing approaches.

2. The paper is well-written and largely clear.

weaknesses:
1. RAGDP seems sensitive to model selection and hyperparameters, but there's limited guidance on how this can be done without a reasonably exhaustive and task-specific search.

2. The effectiveness RAGDP depends heavily on the quality and coverage of the expert demonstration database, but this isn't thoroughly analyzed.

3. I don't find the vector database approach to be very scalable for very large datasets or high-dimensional action spaces.

questions:
1. I'm not very convinced by the problem setting described by the authors. A number of works have already demonstrated diffusion policy works in real-world , real-time settings, with limited real-world data and limited latency problems [1][2][3]. Furthermore, the evaluation suite for RAGDP strictly relies on simulation tasks, whereas the authors describe their method in lines 76 and 77 as being "ideal for real-time robotic tasks." Finally, retrieval for control seems ill-advised, where small perturbations in the real world can deviate significantly from expert demonstrations, which often times are themselves suboptimal. Can the authors provide a more convincing portrayal for why one would want to use RAG in the diffusion policy context?

2. Related to the above question, but also more concretely:

i) How does the method perform with noisy or suboptimal expert demonstrations?

ii) What is the impact of the vector database size on performance?

iii) How sensitive is the method to the choice of embedding space for retrieval?

Without answering these questions, I remain unconvinced that RAG for DP is sensible.

3. The action dimensions for the tasks (as presented in Table 1) are all quite low (i.e. less than 10). How well does RAGDP do on tasks with much higher action dimensions (even just on simulation-based control tasks like those in Mujoco)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703289946

thread_type:
official_review

<thread_object>
creator: reviewer_kpzd

<content>
summary:
RAGDP addresses the slow performance of sampling from diffusion policy model by incorporating a vector database containing pairs of expert observations and actions. This database is used to retrieve relevant actions during the generation process, which helps maintain high accuracy while speeding up policy inference.

strengths:
The paper introduces an innovative method, RAGDP (Retrieve-Augmented Generative Diffusion Policy), which integrates retrieval-augmented generation into the realm of Diffusion Policies for imitation learning. This approach is original in that it applies the concept of retrieval augmentation, typically used in large language models, to enhance the efficiency of Diffusion Policies. By leveraging a vector database containing expert observation-action pairs, the authors create a system that retrieves high-similarity actions during inference to accelerate the generation process without sacrificing accuracy

The paper is generally clear and well-structured, providing readers with sufficient background on Diffusion Policies, stochastic differential equations, and retrieval-augmented methods

The paper evaluates this approach on four common and widely used samplers and show performance gains across all of them.

weaknesses:
* Adding RAG pipelines to existing frameworks is comparatively engineering heavy compared to simpler methods, but the results are strong enough that they might warrant the extra effort.
* The method is sensitive to hyper parameters as the authors admit in the conclusion.

questions:
* A lot of white space is wasted in Figure 3
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a new approach called Retrieve-Augmented Generative Diffusion Policy (RAGDP), which combines retrieval-augmented generation with diffusion policies in the context of imitation learning. The authors claim this integration can speed up the learning process without compromising on accuracy.

The strengths of the paper lie in its innovative approach to enhancing the efficiency of diffusion policies, which is a notable contribution to the field. The paper is also commended for its clarity and thorough explanation of the methods involved.

However, there are several weaknesses that need to be addressed. The novelty of the approach is limited, as retrieval-augmented techniques have been explored in other domains. Reviewers also raised concerns about the paper's sensitivity to hyperparameters and its scalability, especially with larger datasets or high-dimensional action spaces. There's also an apparent lack of analysis on how the quality of the expert demonstration database affects the results.

The decision to reject is primarily due to these weaknesses, particularly the limited novelty, the issues with hyperparameters and scalability, and the insufficient analysis on the influence of expert demonstration quality. These concerns were highlighted by multiple reviewers and were not satisfactorily resolved during the rebuttal phase.