Imitation Learning Policy based on Multi-Step Consistent Integration Shortcut Model

TL;DR

This work tackles the high inference cost of diffusion-based robot imitation learning by proposing a one-step policy built on a multi-step integration framework. It introduces Adaptive Gradient Allocation (AGA) to balance the flow-matching loss $\mathcal{L}_{FM}$ and multi-step consistency loss $\mathcal{L}_{MC}$ across $n$ steps, with a closed-form solution for the gradient composition using $A=\|\nabla\mathcal{L}_{FM}\|_2$, $B=\|\nabla\mathcal{L}_{MC}\|_2$, and $\delta=\cos\angle(\nabla\mathcal{L}_{FM},\nabla\mathcal{L}_{MC})$, including an EMA update for the scale parameter $c$. The method unifies training and inference in a single framework, employs the UDiT1d velocity network with codebook-based sampling to stabilize inference, and demonstrates strong performance on two simulation benchmarks and five real-world tasks, including multimodal perception with visual and tactile data. Ablation studies show the effectiveness of 4-step configurations, the necessity of gradient adaptation, and the robustness of the approach to initialization and inference steps. Overall, the approach enables real-time diffusion-based robot policies with competitive accuracy and broader applicability to rich-contact manipulation tasks.

Abstract

The wide application of flow-matching methods has greatly promoted the development of robot imitation learning. However, these methods all face the problem of high inference time. To address this issue, researchers have proposed distillation methods and consistency methods, but the performance of these methods still struggles to compete with that of the original diffusion models and flow-matching models. In this article, we propose a one-step shortcut method with multi-step integration for robot imitation learning. To balance the inference speed and performance, we extend the multi-step consistency loss on the basis of the shortcut model, split the one-step loss into multi-step losses, and improve the performance of one-step inference. Secondly, to solve the problem of unstable optimization of the multi-step loss and the original flow-matching loss, we propose an adaptive gradient allocation method to enhance the stability of the learning process. Finally, we evaluate the proposed method in two simulation benchmarks and five real-world environment tasks. The experimental results verify the effectiveness of the proposed algorithm.

Figures (10)

• Figure 1: Overview of the differences among diffusion models, Shortcut models, and our approach.
• Figure 2: The difference between the two losses is too large in value.
• Figure 3: By adjusting the scaling factor $c$, we can control the projection intensity of the composite gradient in different directions. Red arrows denote the composite gradient; blue arrows denote its projections onto different directions.
• Figure 4: UDiT1d architecture.
• Figure 5: The experiments in the RoboTwin benchmark.