Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets

摘要

Enabling robots to learn novel visuomotor skills in a data-efficient manner remains an unsolved problem with myriad challenges.A popular paradigm for tackling this problem is through leveraging large unlabeled datasets that have many behaviors in them and then adapting a policy to a specific task using a small amount of task-specific human supervision (i.e.interventions or demonstrations).However, how best to leverage the narrow task-specific supervision and balance it with offline data remains an open question.Our key insight in this work is that task-specific data not only provides new data for an agent to train on but can also inform the type of prior data the agent should use for learning.Concretely, we propose a simple approach that uses a small amount of downstream expert data to selectively query relevant behaviors from an offline, unlabeled dataset (including many sub-optimal behaviors).The agent is then jointly trained on the expert and queried data.We observe that our method learns to query only the relevant transitions to the task, filtering out sub-optimal or task-irrelevant data.By doing so, it is able to learn more effectively from the mix of task-specific and offline data compared to naïvely mixing the data or only using the task-specific data.Furthermore, we find that our simple querying approach outperforms more complex goal-conditioned methods by 20% across simulated and real robotic manipulation tasks from images.See https://sites.google.com/view/behaviorretrievalfor videos and code.