Bootstrapping Imitation Learning for Long-horizon Manipulation via Hierarchical Data Collection Space

TL;DR

This paper introduces HD-Space, a hierarchical data collection method for imitation learning in robotic manipulation, significantly improving policy robustness and performance with less demonstration data in long-horizon tasks.

Contribution

The paper proposes a novel hierarchical data collection space (HD-Space) that enhances imitation learning by segmenting tasks and focusing on high-quality, atomic demonstrations.

Findings

HD-Space improves policy success rates in long-horizon tasks.

Using HD-Space reduces the amount of demonstration data needed.

Empirical results show superior performance over baseline methods.

Abstract

Imitation learning (IL) with human demonstrations is a promising method for robotic manipulation tasks. While minimal demonstrations enable robotic action execution, achieving high success rates and generalization requires high cost, e.g., continuously adding data or incrementally conducting human-in-loop processes with complex hardware/software systems. In this paper, we rethink the state/action space of the data collection pipeline as well as the underlying factors responsible for the prediction of non-robust actions. To this end, we introduce a Hierarchical Data Collection Space (HD-Space) for robotic imitation learning, a simple data collection scheme, endowing the model to train with proactive and high-quality data. Specifically, We segment the fine manipulation task into multiple key atomic tasks from a high-level perspective and design atomic state/action spaces for human demonstrations, aiming to generate robust IL data. We conduct empirical evaluations across two simulated and five real-world long-horizon manipulation tasks and demonstrate that IL policy training with HD-Space-based data can achieve significantly enhanced policy performance. HD-Space allows the use of a small amount of demonstration data to train a more powerful policy, particularly for long-horizon manipulation tasks. We aim for HD-Space to offer insights into optimizing data quality and guiding data scaling. project page: https://hd-space-robotics.github.io.