Generalizable Imitation Learning Through Pre-Trained Representations
This addresses generalization challenges in imitation learning for robotics and AI systems, representing an incremental improvement through hybrid methods.
The paper tackles the problem of generalization in imitation learning by using pre-trained vision transformer representations to create stable semantic keypoints, achieving improved performance across diverse object manipulation tasks with concrete evaluation metrics.
In this paper, we leverage self-supervised vision transformer models and their emergent semantic abilities to improve the generalization abilities of imitation learning policies. We introduce DVK, an imitation learning algorithm that leverages rich pre-trained Visual Transformer patch-level embeddings to obtain better generalization when learning through demonstrations. Our learner sees the world by clustering appearance features into groups associated with semantic concepts, forming stable keypoints that generalize across a wide range of appearance variations and object types. We demonstrate how this representation enables generalized behaviour by evaluating imitation learning across a diverse dataset of object manipulation tasks. To facilitate further study of generalization in Imitation Learning, all of our code for the method and evaluation, as well as the dataset, is made available.