Few-shot Imitation Learning by Variable-Length Trajectory Retrieval from a Large and Diverse Dataset

Imitation learning offers an effective framework for enabling robots to acquire complex skills, but typically requires a large number of labeled demonstrations, making data collection costly. In contrast, large-scale unlabeled datasets containing diverse trajectories can be obtained more easily, motivating methods that leverage such data to enable learning from few demonstrations. A promising approach along this direction is to retrieve trajectories similar to demonstrations from unlabeled datasets. However, existing retrieval-based methods often rely on frame-level comparisons or fixed-length trajectory embeddings, limiting their ability to capture the temporal structure of behaviors and to generalize across different execution speeds. To address these limitations, this paper proposes a method that embeds variable-length trajectories using Self-Attention to capture sequence information and robustly retrieve semantically similar behaviors. Using a small number of demonstrations, relevant trajectories are extracted from unlabeled datasets and used alongside demonstrations to train the agent, thereby reducing data collection costs. Experiments on large-scale, diverse datasets demonstrate that the proposed method achieves higher retrieval accuracy and task success rates than existing retrieval-based imitation learning methods under realistic conditions.