Generalizable Imitation Learning Through Pre-Trained Representations

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.

Authors

• Wei-Di Chang
• Francois Robert Hogan
• Scott Fujimoto
• David Meger
• Gregory Dudek

Keywords

• Visualization
• Computer vision
• Codes
• Imitation learning
• Semantics
• Clustering algorithms
• Transformers
• Robotics and automation
• Semantic Knowledge
• Policy Learning
• Appearance Variations
• Coding Method
• Pre-trained Embeddings
• Vision Transformer
• Training Data
• Training Dataset
• Visual Representation
• Recurrent Neural Network
• Large-scale Datasets
• ImageNet
• Multilayer Perceptron
• Object Classification
• Representation Learning
• Training Environment
• Training Objective
• Robot Manipulator
• Image Coordinates
• Unseen Objects
• Expert Demonstrations
• Inverse Reinforcement Learning
• Candidate Clusters
• Multi-head Self-attention
• Robot Behavior
• Object Dataset
• Object Shape