The Art of Imitation: Learning Long-Horizon Manipulation Tasks from Few Demonstrations

TL;DR

This paper introduces a novel approach for learning complex robot manipulation tasks from few demonstrations by modeling velocities with Riemannian GMMs, segmenting skills, and automatically detecting task parameters from visual data, achieving high efficiency and generalization.

Contribution

It proposes a Riemannian GMM-based method for modeling velocities, skill segmentation, and automatic task parameter detection from RGB-D data, enabling efficient learning from limited demonstrations.

Findings

Achieves state-of-the-art performance on RLBench tasks.

Learns from only five demonstrations with RGB-D observations.

Generalizes across environments, objects, and positions.

Abstract

Task Parametrized Gaussian Mixture Models (TP-GMM) are a sample-efficient method for learning object-centric robot manipulation tasks. However, there are several open challenges to applying TP-GMMs in the wild. In this work, we tackle three crucial challenges synergistically. First, end-effector velocities are non-Euclidean and thus hard to model using standard GMMs. We thus propose to factorize the robot's end-effector velocity into its direction and magnitude, and model them using Riemannian GMMs. Second, we leverage the factorized velocities to segment and sequence skills from complex demonstration trajectories. Through the segmentation, we further align skill trajectories and hence leverage time as a powerful inductive bias. Third, we present a method to automatically detect relevant task parameters per skill from visual observations. Our approach enables learning complex manipulation tasks from just five demonstrations while using only RGB-D observations. Extensive experimental evaluations on RLBench demonstrate that our approach achieves state-of-the-art performance with 20-fold improved sample efficiency. Our policies generalize across different environments, object instances, and object positions, while the learned skills are reusable.