One-Shot Imitation Learning

Yan Duan, Marcin Andrychowicz, Bradly C. Stadie, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, Wojciech Zaremba

Seminal March 2017

TL;DR One network is trained on pairs of demonstrations across 140 block-stacking tasks so that a single demonstration of a new task turns it into a policy for that task. With soft attention over the demonstration and the blocks, it succeeds 50.6% of the time on unseen 6-stage tasks, against 49.4% for the scripted demonstrator.

What they do

Instead of one policy per task, this paper trains a single network that takes one demonstration of a block-stacking task plus the current state and outputs an action. A demonstration of a new layout (e.g. ab cd ef gh: A on B, C on D, and so on) is then enough to build it from new block positions. Training is supervised learning on pairs of demonstrations of the same task.

How it works

Setup. A 7-DoF Fetch arm stacks 2-10 blocks. Observations are block (x, y, z) positions relative to the gripper plus gripper state. There are 140 training tasks and 43 held-out test tasks, with 1000 noisy scripted demonstrations per task, trained with behavioral cloning or DAgger.
Demonstration network. Temporal dropout discards 95% of the demonstration’s timesteps. Dilated temporal convolutions and neighborhood attention, in which each block queries the other blocks relative to itself, embed the rest.
Context network. A query from the current state attends over demonstration timesteps, then over the blocks in the current state. The result is a fixed-size embedding, whatever the demonstration length or block count.
Manipulation network. An MLP maps that embedding to the action. At test time, 10 downsampled demonstrations are ensembled.

Results

On unseen test tasks, full-demonstration DAgger succeeds 94.9% of the time at 2 stages, 65.9% at 5 and 50.6% at 6. The scripted demonstrator gets 95.8%, 65.9% and 49.4%. Conditioning only on the final state falls to 35.9% at 6 stages.
Behavioral cloning roughly matches DAgger (46.5% vs 50.6% on 6-stage test tasks), given noise injected into the demonstrations.
Trained only on abcd, the policy scores 82-95% on all 24 block permutations of it.

Limitations

Inputs are block positions, not images, and one demonstration must fully specify the task. Success falls as tasks get longer (18.0% on 8-stage test tasks), and most failures are irrecoverable manipulation errors, not misread demonstrations. Conditioning on full demonstrations trains more slowly than conditioning on final states or snapshots. No physical-robot experiments are reported.

Videos

Figures

From Figure 1: (a) Traditional imitation learning trains a separate policy per task from many demonstrations. (b) A o... — From Figure 1: (a) Traditional imitation learning trains a separate policy per task from many demonstrations. (b) A one-shot imitator is conditioned on a single demonstration of a new task and acts as its policy. (c) Training: the network sees one demonstration plus an observation from a second demonstration of the same task, and is supervised to output that demonstration's action.

From Figure 2: The network architecture. A demonstration network (temporal dropout, then neighborhood attention and t... — From Figure 2: The network architecture. A demonstration network (temporal dropout, then neighborhood attention and temporal convolution) embeds the demonstration; a context network attends over the demonstration and then over the blocks in the current state; a manipulation network maps the context embedding to an action.

From Figure 3: Average success rate by number of stages on (a) training tasks and (b) test tasks, for the hard-coded ... — From Figure 3: Average success rate by number of stages on (a) training tasks and (b) test tasks, for the hard-coded demonstrator (darkest bar), BC, DAgger, Snapshot and Final-state conditioning. Policies with temporal dropout average an ensemble of 10 downsampled demonstrations.

From Figure 4: Attention during one execution of task ab cde fg hij, one query head each. (a) Attention over blocks i... — From Figure 4: Attention during one execution of task ab cde fg hij, one query head each. (a) Attention over blocks in the current state stays on a small subset of blocks. (b) Attention over the downsampled demonstration is sparse, with roughly 6 high-weight regions matching the task's 6 stages.

Abstract

Imitation learning has been commonly applied to solve different tasks in isolation. This usually requires either careful feature engineering, or a significant number of samples. This is far from what we desire: ideally, robots should be able to learn from very few demonstrations of any given task, and instantly generalize to new situations of the same task, without requiring task-specific engineering. In this paper, we propose a meta-learning framework for achieving such capability, which we call one-shot imitation learning. Specifically, we consider the setting where there is a very large set of tasks, and each task has many instantiations. For example, a task could be to stack all blocks on a table into a single tower, another task could be to place all blocks on a table into two-block towers, etc. In each case, different instances of the task would consist of different sets of blocks with different initial states. At training time, our algorithm is presented with pairs of demonstrations for a subset of all tasks. A neural net is trained that takes as input one demonstration and the current state (which initially is the initial state of the other demonstration of the pair), and outputs an action with the goal that the resulting sequence of states and actions matches as closely as possible with the second demonstration. At test time, a demonstration of a single instance of a new task is presented, and the neural net is expected to perform well on new instances of this new task. The use of soft attention allows the model to generalize to conditions and tasks unseen in the training data. We anticipate that by training this model on a much greater variety of tasks and settings, we will obtain a general system that can turn any demonstrations into robust policies that can accomplish an overwhelming variety of tasks. Videos available at https://bit.ly/nips2017-oneshot .

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong