RAGDP: Retrieve-Augmented Generative Diffusion Policy

Thank you very much for your comments. Here are the answers to the discussion questions.

Q1:

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.

A1:

We have benchmarked LSTM-based models with Transformer-based models (BeT [1]) and Behaviour Cloning datasets, with the Diffusion Model as the best performing model as a starting point [2]. Another extension of the Diffusion Models is the Consistency Models [3]. It is shown in the diagram in the Appendix C that the proposed method can also improve the performance of Consistency Models.

Q2:

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.

A2:

Determining the hyperparameters is a very important point. We have made measurements on this at two additional points - check Sec 5.5.

Q3:

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.

A3:

Diffusion Policy uses only the diffusion model of DDPM for its implementation. There are two main types of diffusion models: VP-SDE (Variance Preserving Stochastic Differential Equations) and VE-SDE (Variance Exploding Stochastic Differential Equations). The difference between the properties of each is the difference in Eq 3. The differences in their properties can be explained using the parameters $\alpha(\tau), \sigma(\tau)$ in Eq 3. In VP-SDE, the parameters are constrained by $\alpha(\tau)^2 + \sigma(\tau)^2 = 1$. This means that the magnitude of the noise and the action are determined by a ratio. Therefore, the action retrieved from the database is used to calculate the final output from the ratio of action and noise corresponding to the starting diffusion step $\tau_0$. In the case of VE-SDE, $\alpha(\tau) = 1$ is fixed. Therefore, there is no limit on the size of the action and noise. Therefore, the action taken from the database adds noise of a magnitude corresponding to the starting diffusion step $\tau_0$, and the output is obtained where this noise becomes smaller. Section 4.2 has been modified.

Q4:

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.

A4:

Figure 3 has been updated to provide more detail. And code for a vector database using FAISS has been added to Appendix D. Information is stored in two files: a FAISS file for indexing and a Pickle file for storing action vectors.

[1] Shafiullah, Nur Muhammad (Mahi) et al. “Behavior Transformers: Cloning k modes with one stone.” ArXiv abs/2206.11251 (2022): n. pag.

[2] Chi, Cheng et al. “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion.” ArXiv abs/2303.04137 (2023): n. pag.

[3] Prasad, Aaditya et al. “Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation.” ArXiv abs/2405.07503 (2024): n. pag.

Thank you very much for your comments. Here are the answers to the discussion questions.

Q1:

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

A1:

The present method is proposed as a performance extension of the diffusion model. The key performance factors are accuracy, generation speed and robustness to the environment; Sec 5.2 and 5.3 show that speed can be increased without loss of accuracy. It has also been shown to improve accuracy when utilised with faster Consistency Models. Could you be more specific about the special requirements you indicate?

Q2:

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?

A2:

This is unchanged from the Diffusion Policy. If the time step at time $t$ is $a_t$, then $A_t$ is a chunked action as $a_{t-T_o:t-T_o+T_p}$. Section 3.2 has been modified.

Q3:

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?

A3:

We believe that the key here is in the design of the database. The data for PH is less noisy for experts and more noisy for MH. In this experiment, the models trained on SQUARE-PH were allowed to search the SQUARE-MH database and the models trained on SQUARE-MH were allowed to search the SQUARE-PH database. As a result, models trained on SQUARE-PH tend to be less accurate when searched on SQUARE-MH data. However, models trained on SQUARE-MH tend to go down when searched on SQUARE-PH. Therefore, it is desirable that the training data for the model and the data in the vector database are similar. We also investigated how the results changed as the size of the database changed. - Appendix B.4 added.

Q4:

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.

A4:

The training of the encoders for the observed data $O_t$ remains unchanged from the Diffusion Policy. In the case of images, the encoders of the CNN are used to compute the vectors, while in the case of states, they are converted to 1D vectors for retrieval. First, when training, the diffusion model and encoder are trained together. Then, when the database is created, the encoders used in training are used to calculate the vectors in the case of images, while in the case of states, the vectors are converted to 1D vectors. Then, when inferencing, the vectors calculated once with the encoder or vectors with the dimension changed are used to search and retrieve the behavioral data. Figure 3 has been updated with more detail. Code for searching has been added in Appendix D.

Thank you for your valuable points. Here are the answers to the discussion questions.

A1:

About real-time settings:

We appreciate this point as we want it to be clear that these models show potential in being used for real-world settings and that our method can further improve the performance. We modified the text to clarify that the method is more general and can help improve applications where the inference speed and accuracy are important.

A2:

1. About noisy or suboptimal expert demonstrations (Robustness for generating quality):

The experiments conducted in Appendix B.3 might be helpful to answer your question. Specifically, we evaluated the SQUARE-PH and SQUARE-MH datasets, where SQUARE-PH consists of proficient demonstrations, and SQUARE-MH includes a mix of proficient and suboptimal demonstrations. We perform two experiments: i) Training the base model on high quality demonstrations provided by experts and using a knowledge base on demonstrations of mixed quality and ii) Training the base model on a mixed quality of demonstrations with a knowledge base on high quality demonstrations. These experiments show both the effect of a distribution shift (as the model trained on mixed data and trained on expert data will behave differently) of the demonstrations as well as suboptimal demonstrations.

Table 5 (Base model trained on expert data and knowledge base trained on mixed data) shows that there is a decrease in performance with suboptimal demonstrations albeit still relatively high considering the harsh scenario of mismatched and suboptimal demonstrations. This can further be mitigated by setting a conservative $r$.

Table 6 (Base model trained on mixed data and knowledge base on expert data) shows a much smaller decrease in accuracy even with the difference in demonstrations. We hope that this experiment shows that the trade-off is acceptable even in sub-optimal conditions and that it can be proactively mitigated by picking the $r$ hyperparameter.

2. About vector database size:

We investigated the relationship between accuracy and retrieval distance as the size of the database varied. The results, shown in Figure 12, indicate that retrieval distance decreases as the database size increases. However, this change in retrieval distance does not significantly affect accuracy. This suggests that the accuracy of the method is robust to retrieval distance. Instead, the quality of the data stored in the database appears to play a more critical role than its quantity. Supporting this observation, prior studies [3] have demonstrated that retrieving high-quality data during training can improve accuracy, even when the dataset size is relatively small.

3. About embedding spaces:

The use of VAEs to embed observations and search in the embedding space, as in study [3], was also initially considered by us. However, the accuracy results did not change from the current method. We have also adopted the current retrieval method because the proposed method is easier to implement without the need to train additional models.

A3:

Additional experiments were conducted with Transport, which has twice the action dimension in the behavior cloning dataset. Similar results were obtained here - check Table 4, Figure 9.

[1] Prasad, Aaditya et al. “Consistency Policy: Accelerated Visuomotor Policies via Consistency Distillation.” ArXiv abs/2405.07503 (2024): n. pag.

[2] Tangkaratt, Voot et al. “VILD: Variational Imitation Learning with Diverse-quality Demonstrations.” International Conference on Machine Learning (2019).

[3] Du, Maximilian et al. “Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets.” ArXiv abs/2304.08742 (2023): n. pag.

We thank the reviewer for the valuable feedback which gives us the opportunity to improve the paper. We have made changes to the paper accordingly by clarifying the implementation as well as providing additional experiments.

1. Difficulty of adding RAG pipeline into existing method:

We appreciate the reviewer's concern about the engineering effort. To clarify, integrating RAGDP is relatively straightforward as it leverages existing diffusion models without requiring changes to their core architectures. Basically, to summarize the engineering process:

Stage 1: The training data is normalized and stored as a vector database using FAISS.

Stage 2: At inference, retrieved actions are incorporated into the diffusion denoising process using a few modifications to the sampling algorithm.

To make this clearer, we have updated Figure 3 to include more implementation details, and we added an explanation in Section 4.1 highlighting the modularity of the retrieval component. This modularity ensures that RAGDP can be easily adopted in conjunction with other diffusion-based approaches.

2. Choosing the optimal hyperparameter:

We agree with the reviewer that selecting appropriate hyperparameters is important. To give a more clear picture of the robustness of the hyperparameter choice, we have conducted additional experiments varying the hyperparameter r. The results, provided in Section 5.5, show that the method maintains robustness across a reasonable range of hyperparameter values, not showing a noticeable decrease until $r > 0.75$.