Towards More Generalizable One-shot Visual Imitation Learning

by   Zhao Mandi, et al.
berkeley college

A general-purpose robot should be able to master a wide range of tasks and quickly learn a novel one by leveraging past experiences. One-shot imitation learning (OSIL) approaches this goal by training an agent with (pairs of) expert demonstrations, such that at test time, it can directly execute a new task from just one demonstration. However, so far this framework has been limited to training on many variations of one task, and testing on other unseen but similar variations of the same task. In this work, we push for a higher level of generalization ability by investigating a more ambitious multi-task setup. We introduce a diverse suite of vision-based robot manipulation tasks, consisting of 7 tasks, a total of 61 variations, and a continuum of instances within each variation. For consistency and comparison purposes, we first train and evaluate single-task agents (as done in prior few-shot imitation work). We then study the multi-task setting, where multi-task training is followed by (i) one-shot imitation on variations within the training tasks, (ii) one-shot imitation on new tasks, and (iii) fine-tuning on new tasks. Prior state-of-the-art, while performing well within some single tasks, struggles in these harder multi-task settings. To address these limitations, we propose MOSAIC (Multi-task One-Shot Imitation with self-Attention and Contrastive learning), which integrates a self-attention model architecture and a temporal contrastive module to enable better task disambiguation and more robust representation learning. Our experiments show that MOSAIC outperforms prior state of the art in learning efficiency, final performance, and learns a multi-task policy with promising generalization ability via fine-tuning on novel tasks.


page 1

page 2

page 8

page 9

page 10


Transporters with Visual Foresight for Solving Unseen Rearrangement Tasks

Rearrangement tasks have been identified as a crucial challenge for inte...

Learning and Retrieval from Prior Data for Skill-based Imitation Learning

Imitation learning offers a promising path for robots to learn general-p...

CRIL: Continual Robot Imitation Learning via Generative and Prediction Model

Imitation learning (IL) algorithms have shown promising results for robo...

A System for Morphology-Task Generalization via Unified Representation and Behavior Distillation

The rise of generalist large-scale models in natural language and vision...

Comparing the Efficacy of Fine-Tuning and Meta-Learning for Few-Shot Policy Imitation

In this paper we explore few-shot imitation learning for control problem...

Stage Conscious Attention Network (SCAN) : A Demonstration-Conditioned Policy for Few-Shot Imitation

In few-shot imitation learning (FSIL), using behavioral cloning (BC) to ...

Multi-Task Imitation Learning for Linear Dynamical Systems

We study representation learning for efficient imitation learning over l...

Please sign up or login with your details

Forgot password? Click here to reset