A new family of approaches to few-shot imitation

We want to make it easy for people to specify diverse tasks to artificial agents in diverse manners. Non-expert users may find it simpler to demonstrate a new task than to specify the task in other ways. For instance, it may be easier to teach a household robot how to clean a food blender by giving it a few demonstrations than to explain each step of the process verbally or to manually specify a reward function for reinforcement learning.

Existing approaches to learning from demonstrations tend to result in behaviours that are not robust. They fail to complete tasks when there are small changes to the environment. Moreover, they often make unrealistic assumptions. For instance, most approaches assume that the demonstrator and artificial agent are both robots with identical physical structure. But this is a showstopper if non-humanoid robots are to imitate humans.

We address the problem of few-shot imitation: we must design an agent, which when given a few (e.g. less than five) demonstrations of a new, previously unseen test task as input, outputs a policy (i.e. behaviour) that attains a high success rate or cumulative reward on that task. The design should cope in situations where: <

– the demonstrator’s behaviour is not fully optimal;
– the demonstrator’s state is available, but not the demonstrator’s actions;
– the demonstrator does not have the same physical structure as the agent; or
– the structure and dynamics of the environment vary from task to task.

Ideally, the agent’s policy should be fast enough for real-time control, and the agent should require little to no extra training or exploration of the test environment after receiving demonstrations of a test task.

Most existing work on imitation learning only addresses a single task at a time, requiring large numbers of demonstrations of that task, and the study of few-shot imitation is quite recent (1). There are two main pre-existing approaches to few-shot imitation: behaviour cloning and inverse reinforcement learning.

Behaviour cloning minimizes a loss that penalises inaccurate predictions of the demonstrator’s actions. Behaviour cloning assumes the demonstrator’s actions are available and appropriate for the agent. However, recent work has shown how to weaken such assumptions, for instance by inferring actions from state-only trajectories (i.e. demonstrations consisting of sequences of states) (2) or by training with pairs of human and robot demonstrations (3). The main limitation of behaviour cloning is error compounding (4). Here, small imitation errors lead the agent into parts of the state space that the demonstrations rarely visit and in which the learned policy performs poorly.

Inverse reinforcement learning (IRL) infers a reward function for a task from demonstrations. IRL can cope without actions in the demonstrations. It can also overcome error compounding, as its training involves interaction with the environment. However, existing IRL approaches to few-shot imitation (5) rely on the assumption that the structure and dynamics of the environment do not change from one task to another. Thus, they can address situations where tasks differ in terms of their objectives, such as the tasks ‘walk to the left’ and ‘walk to the right’. But they work less well in situations where tasks differ in terms of the structure and dynamics of their environments, such as ‘open a window’ and ‘close a drawer’. To address such changes, existing IRL algorithms mostly resort to additional exploration at test time. But this can be time-consuming, with 40 hours per test task reported (5); it can raise safety concerns; and it may not be viable in interactive use cases. In future, it would be interesting to devise IRL algorithms that avoid this assumption.

Our approach: demonstration-conditioned reinforcement learning

We propose a simple, new approach to learning few-shot imitation agents that we call demonstration-conditioned reinforcement learning or DCRL(6). For a set of training tasks, DCRL learns from a training set consisting of demonstrations, reward functions and transition distributions (a transition distribution describes how the state of an environment changes in response to an agent’s actions). The idea is to define a demonstration-conditioned policy—which takes demonstrations and the agent’s state or state history as input—and to train this policy to maximize the average of the cumulative reward over the training tasks. DCRL does not require access to the reward function for a new test task. Nor does it require any additional interaction with the test environment to learn a policy. Rather, you simply feed demonstrations of a new test task to the learned policy, along with the agent’s current state or history, and take actions according to the policy’s output. Figure 1 compares the learning objective of DCRL with those of pre-existing few-shot imitation methods.

DCRL has several advantages over pre-existing approaches. It can improve upon suboptimal demonstrations, and operate if demonstrations contain state but not action information. DCRL successfully addresses situations where the demonstrator has a different physical structure to the agent, and where the structure and dynamics of the environment vary from task to task, even though DCRL requires no interaction with the environment to adapt its policy at test time. DCRL can address such situations as it interprets demonstrations considerably more flexibly than pre-existing methods. With DCLR, you can even use demonstrations to show the agent what to avoid doing (7)! Finally, in our experiments on robotic manipulation and navigation, we find that DCRL produces policies that outperform pre-existing methods by a large margin.

Paying attention to multiple demonstrations simultaneously with transformer-based policy architectures

Attention plays a major role in human task control (8), and attention-based architectures like transformers (9) currently deliver state-of-the-art performance in many natural language tasks (10). This has motivated several authors to explore the application of transformers to few-shot imitation (1, 11, 12, 13, 14,). We, too, experiment with a transformer-based policy architecture for DCRL. Our architecture explores two novel ideas.

1. Cross-demonstration attention. We process multiple input demonstrations simultaneously. This allows richer information integration across demonstrations than previous work, which encodes multiple demonstrations independently and then averages the encodings (12).
2. Axial attention. Our architecture attends to one dimension of its input at a time, resulting in a large reduction in its time and memory complexity on typical multidimensional inputs of shape

# demonstrations x # time steps x # sensor outputs

Axial attention has previously been applied in computer vision (15, 16), but our work is the first to apply it in a policy architecture.

Getting one robot to imitate another

We compared DCRL with a wide variety of baselines on a robotic manipulation benchmark, called Meta-World (17), and a navigation benchmark, as shown in Videos 1 and 2. Figure 2 gives success rates for these benchmarks.

First, we compared DCRL with a demonstration-conditioned behaviour cloning (DCBC) approach. DCBC has the same transformer-based architecture as DCRL, but we trained it to minimize an action-prediction loss. In all cases, DCRL had a significantly higher success rate than DCBC.

For the Meta-World benchmark, we found that the agents learned with our transformer-based architecture produce feedback controls at 370 Hz—fast enough to deploy on a real robot. Moreover, we used DCRL to train an agent that receives demonstrations from a robot that is physically different from the robot it must control. The results for this agent are very similar to those where the demonstrator and agent robots are identical, suggesting that DCRL can indeed cope when the demonstrator has a physically different structure! This intriguing capability relates to the fact that DCRL makes no attempt to ‘align’ the agent with the demonstrator. In contrast, IRL minimizes a divergence between the agent and demonstrator’s state distributions, and behaviour cloning minimizes an action-prediction loss.

For the navigation benchmark, to understand the benefit of cross-demonstration attention, we trained a version of our transformer-based policy that encodes multiple demonstrations independently and averages their encodings (12). The success rate of this agent when given five demonstrations is 10% worse than the success rate with our proposed architecture, showing that cross-demonstration attention can indeed help.

For both benchmarks, we fine-tuned the DCRL and DCBC policies using behaviour cloning, so as to imitate the actions in the test-task demonstrations. DCRL + fine-tuning then attained a 90% success rate on the Meta-World benchmark—a large improvement on the state of the art reported elsewhere (17). However, fine-tuning actually hurt the success rate on the navigation benchmark, particularly for DCBC. We are not sure why. Perhaps it’s because small shifts can move the goal through a wall, meaning that the mapping from the goal to an appropriate action is nearly discontinuous. In such a setting, fine-tuning could encourage the agent to take the wrong actions in some states.

The paper and supplementary material (6) present results for other baselines and task-by-task breakdowns of the results.

The future of few-shot imitation, and next steps

While DCRL does not require reward functions for test tasks, it does require reward functions for training tasks. Pre-existing methods do not require such training rewards. We argue that this extra cost is outweighed by DCRL’s higher success rate on new tasks, and its reliance on fewer assumptions about the nature of demonstrations. These advantages make DCRL more suitable for practical application.

Currently, we are exploring the use of DCRL with real robots, using demonstrations from a human-operated haptic device. At this stage, the major challenge is the gap between the simulated robot and its real-world cousin. In future, we would like to train DCRL on many more tasks, to improve its generalization to new tasks. We are also looking forward to extending our transformer-based policy to accept videos of humans as demonstrations.

References

1: One-shot imitation learning. Yan Duan, Marcin Andrychowicz, Bradly Stadie, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel et al. 31st Conference on Neural Information Processing Systems (NIPS), Long Beach, California, USA, 4–9 December 2017.

2: Hybrid reinforcement learning with expert state sequences. Xiaoxiao Guo, Shiyu Chang, Mo Yu, Gerald Tesauro and Murray Campbell. 33rd AAAI Conference on Artificial Intelligence (AAAI), Honolulu, Hawaii, USA, 27 January – 1 February 2019.

3: Learning one-shot imitation from humans without humans. Alessandro Bonardi, Stephen James and Andrew J. Davison. IEEE Robotics and Automation Letters, vol. 5, no. 2, 2020, pp. 3533–3539.

4: Efficient reductions for imitation learning. Stéphane Ross and J. Andrew Bagnell. 13th International Conference on Artificial Intelligence and Statistics (AISTATS), Sardinia, Italy, 2010.

5: Meta-inverse reinforcement learning with probabilistic context variables. Lantao Yu, Tianhe Yu, Chelsea Finn and Stefano Ermon. 33rd Conference on Neural Information Processing Systems (NeurIPS), Vancouver, Canada, 3–14 December 2019.

6: Demonstration-conditioned reinforcement learning for few-shot imitation. Theo Cachet, Julien Perez and Christopher Dance. International Conference on Machine Learning (ICML), virtual event, 18–24 July 2021.

7: Avoidance learning using observational reinforcement learning. David Venuto, Leonard Boussioux, Junhao Wang, Rola Dali, Jhelum Chakravorty, Yoshua Bengio and Doina Precup. arXiv:1909.11228, 2019.

8: Task and context determine where you look. Constantin A. Rothkopf, Dana H. Ballard and Mary M. Hayhoe. Journal of Vision, vol. 7, no. 14, 2007, pp. 1–20.

9: Attention is all you need. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser et al. 31st Conference on Neural Information Processing Systems (NIPS), Long Beach, California, USA, 4–9 December 2017.

10: Language models are few-shot learners. Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan et al. arXiv:2005.14165, 2020.

11: A simple neural attentive meta-learner. Nikhil Mishra, Mostafa Rohaninejad, Xi Chen and Pieter Abbeel. Sixth International Conference on Learning Representations (ICLR), Vancouver, Canada, 30 April – 3 May 2018.

12: Task-embedded control networks for few-shot imitation learning. Stephen James, Michael Bloesch and Andrew J. Davison. Conference on Robot Learning (CoRL), Zurich, Switzerland, 29–31 October 2018.

13: Transformers for one-shot visual imitation. Sudeep Dasari and Abhinav Gupta. Conference on Robot Learning (CoRL), virtual conference, 16–18 November 2020.

14: Transformer-based meta-imitation learning for robotic manipulation. Theo Cachet, Julien Perez and Seungsu Kim. 3rd Robot Learning Workshop at NeurIPS, virtual conference, 6–12 December 2020.

15: Axial attention in multidimensional transformers. Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn and Tim Salimans. arXiv:1912.12180, 2019.

16: Translating math formula images to LaTeX sequences using deep neural networks with sequence-level training. Zelun Wang and Jyh-Charn Liu. International Journal on Document Analysis and Recognition (IJDAR), vol. 24, pp. 63–75, 2021.

17: Meta-World: a benchmark and evaluation for multi-task and meta reinforcement learning. Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur et al. Conference on Robot Learning (CoRL), London, UK, and virtual, 8–22 November 2019.