# Towards More Generalizable One-shot Visual Imitation Learning

arXiv: https://arxiv.org/abs/2110.13423 · HTML: https://mlsanity.com/paper/2110.13423/

## Key facts

- Cited by 49 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 34 of its 45 references that are in the map.
- First posted on arXiv 26 Oct 2021.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Non-local Neural Networks](https://mlsanity.com/paper/1711.07971/) (2017)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [End to End Learning for Self-Driving Cars](https://mlsanity.com/paper/1604.07316/) (2016)
- [AI2-THOR: An Interactive 3D Environment for Visual AI](https://mlsanity.com/paper/1712.05474/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Data-Efficient Image Recognition with Contrastive Predictive Coding](https://mlsanity.com/paper/1905.09272/) (2019)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications](https://mlsanity.com/paper/1701.05517/) (2017)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Learning Latent Plans from Play](https://mlsanity.com/paper/1903.01973/) (2019)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [Learning Agile Robotic Locomotion Skills by Imitating Animals](https://mlsanity.com/paper/2004.00784/) (2020)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [One-Shot Visual Imitation Learning via Meta-Learning](https://mlsanity.com/paper/1709.04905/) (2017)
- [Language Conditioned Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2005.07648/) (2020)
- [One-Shot Imitation from Observing Humans via Domain-Adaptive Meta-Learning](https://mlsanity.com/paper/1802.01557/) (2018)
- [Language-Conditioned Imitation Learning for Robot Manipulation Tasks](https://mlsanity.com/paper/2010.12083/) (2020)
- [Visual Imitation Made Easy](https://mlsanity.com/paper/2008.04899/) (2020)
- [Playing hard exploration games by watching YouTube](https://mlsanity.com/paper/1805.11592/) (2018)
- [Transformers for One-Shot Visual Imitation](https://mlsanity.com/paper/2011.05970/) (2020)
- [Neural Task Programming: Learning to Generalize Across Hierarchical Tasks](https://mlsanity.com/paper/1710.01813/) (2017)
- [Task-Embedded Control Networks for Few-Shot Imitation Learning](https://mlsanity.com/paper/1810.03237/) (2018)

## Built on by (most cited first, 30 of 49)

- [VIMA: General Robot Manipulation with Multimodal Prompts](https://mlsanity.com/paper/2210.03094/) (2022)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [RH20T: A Comprehensive Robotic Dataset for Learning Diverse Skills in One-Shot](https://mlsanity.com/paper/2307.00595/) (2023)
- [Large Language Models for Robotics: A Survey](https://mlsanity.com/paper/2311.07226/) (2023)
- [From Play to Policy: Conditional Behavior Generation from Uncurated Robot Data](https://mlsanity.com/paper/2210.10047/) (2022)
- [Toward General-Purpose Robots via Foundation Models: A Survey and Meta-Analysis](https://mlsanity.com/paper/2312.08782/) (2023)
- [Preference Transformer: Modeling Human Preferences using Transformers for RL](https://mlsanity.com/paper/2303.00957/) (2023)
- [Teach a Robot to FISH: Versatile Imitation from One Minute of Demonstrations](https://mlsanity.com/paper/2303.01497/) (2023)
- [XSkill: Cross Embodiment Skill Discovery](https://mlsanity.com/paper/2307.09955/) (2023)
- [EquiBot: SIM(3)-Equivariant Diffusion Policy for Generalizable and Data Efficient Learning](https://mlsanity.com/paper/2407.01479/) (2024)
- [In-Context Imitation Learning via Next-Token Prediction](https://mlsanity.com/paper/2408.15980/) (2024)
- [Behavior Retrieval: Few-Shot Imitation Learning by Querying Unlabeled Datasets](https://mlsanity.com/paper/2304.08742/) (2023)
- [RAM: Retrieval-Based Affordance Transfer for Generalizable Zero-Shot Robotic Manipulation](https://mlsanity.com/paper/2407.04689/) (2024)
- [Vid2Robot: End-to-end Video-conditioned Policy Learning with Cross-Attention Transformers](https://mlsanity.com/paper/2403.12943/) (2024)
- [Generalization to New Sequential Decision Making Tasks with In-Context Learning](https://mlsanity.com/paper/2312.03801/) (2023)
- [VLM See, Robot Do: Human Demo Video to Robot Action Plan via Vision Language Model](https://mlsanity.com/paper/2410.08792/) (2024)
- [On Transforming Reinforcement Learning by Transformer: The Development Trajectory](https://mlsanity.com/paper/2212.14164/) (2022)
- [SEIL: Simulation-augmented Equivariant Imitation Learning](https://mlsanity.com/paper/2211.00194/) (2022)
- [One-Shot Imitation Learning: A Pose Estimation Perspective](https://mlsanity.com/paper/2310.12077/) (2023)
- [RH20T-P: A Primitive-Level Robotic Dataset Towards Composable Generalization Agents](https://mlsanity.com/paper/2403.19622/) (2024)
- [On the Effectiveness of Fine-tuning Versus Meta-reinforcement Learning](https://mlsanity.com/paper/2206.03271/) (2022)
- [Chain-of-Thought Predictive Control](https://mlsanity.com/paper/2304.00776/) (2023)
- [One-Shot Imitation Learning with Invariance Matching for Robotic Manipulation](https://mlsanity.com/paper/2405.13178/) (2024)
- [Do You Need Proprioceptive States in Visuomotor Policies?](https://mlsanity.com/paper/2509.18644/) (2025)
- [Few-Shot In-Context Imitation Learning via Implicit Graph Alignment](https://mlsanity.com/paper/2310.12238/) (2023)
- [One-Shot Imitation under Mismatched Execution](https://mlsanity.com/paper/2409.06615/) (2024)
- [NetHack is Hard to Hack](https://mlsanity.com/paper/2305.19240/) (2023)
- [Learning by Watching: A Review of Video-based Learning Approaches for Robot Manipulation](https://mlsanity.com/paper/2402.07127/) (2024)
- [Learning to Drive Anywhere](https://mlsanity.com/paper/2309.12295/) (2023)
- [MILES: Making Imitation Learning Easy with Self-Supervision](https://mlsanity.com/paper/2410.19693/) (2024)

## Abstract (arXiv)

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.