Retrieve-Augmented Generation for Speeding up Diffusion Policy without Additional Training

TL;DR

This paper introduces RAGDP, a retrieval-augmented framework that speeds up diffusion policies in imitation learning without extra training, maintaining high accuracy even with significant acceleration.

Contribution

RAGDP leverages a knowledge base and retrieval to reduce diffusion steps, eliminating additional training and improving speed-accuracy trade-offs in diffusion policies.

Findings

RAGDP improves inference speed by up to 20 times.

RAGDP achieves 7% higher accuracy than distillation methods.

RAGDP maintains performance without extra training.

Abstract

Diffusion Policies (DPs) have attracted attention for their ability to achieve significant accuracy improvements in various imitation learning tasks. However, DPs depend on Diffusion Models, which require multiple noise removal steps to generate a single action, resulting in long generation times. To solve this problem, knowledge distillation-based methods such as Consistency Policy (CP) have been proposed. However, these methods require a significant amount of training time, especially for difficult tasks. In this study, we propose RAGDP (Retrieve-Augmented Generation for Diffusion Policies) as a novel framework that eliminates the need for additional training using a knowledge base to expedite the inference of pre-trained DPs. In concrete, RAGDP encodes observation-action pairs through the DP encoder to construct a vector database of expert demonstrations. During inference, the current observation is embedded, and the most similar expert action is extracted. This extracted action is combined with an intermediate noise removal step to reduce the number of steps required compared to the original diffusion step. We show that by using RAGDP with the base model and existing acceleration methods, we improve the accuracy and speed trade-off with no additional training. Even when accelerating the models 20 times, RAGDP maintains an advantage in accuracy, with a 7% increase over distillation models such as CP.