Beyond the Majority: Long-tail Imitation Learning for Robotic Manipulation

TL;DR

This paper addresses the long-tail distribution challenge in imitation learning for robotic manipulation, proposing a novel transfer scheme called Approaching-Phase Augmentation (APA) that improves tail task performance without external data.

Contribution

The paper introduces APA, a new method that transfers knowledge from head to tail tasks, overcoming data scarcity and improving generalization in long-tail imitation learning.

Findings

APA significantly improves tail task performance in simulation and real-world tests.

Conventional re-sampling strategies are ineffective for long-tail policy learning.

Data scarcity impairs spatial reasoning, which APA effectively mitigates.

Abstract

While generalist robot policies hold significant promise for learning diverse manipulation skills through imitation, their performance is often hindered by the long-tail distribution of training demonstrations. Policies learned on such data, which is heavily skewed towards a few data-rich head tasks, frequently exhibit poor generalization when confronted with the vast number of data-scarce tail tasks. In this work, we conduct a comprehensive analysis of the pervasive long-tail challenge inherent in policy learning. Our analysis begins by demonstrating the inefficacy of conventional long-tail learning strategies (e.g., re-sampling) for improving the policy's performance on tail tasks. We then uncover the underlying mechanism for this failure, revealing that data scarcity on tail tasks directly impairs the policy's spatial reasoning capability. To overcome this, we introduce Approaching-Phase Augmentation (APA), a simple yet effective scheme that transfers knowledge from data-rich head tasks to data-scarce tail tasks without requiring external demonstrations. Extensive experiments in both simulation and real-world manipulation tasks demonstrate the effectiveness of APA. Our code and demos are publicly available at: https://mldxy.github.io/Project-VLA-long-tail/.