PaperlayerView exact revision on arXiv

Source-linked AI summary

One-Shot Imitation Learning

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

arXiv:1703.07326v3cs.AIcs.LGcs.NEcs.RO

TL;DR

Robots often need task-specific engineering or many samples to learn from demonstrations, motivating one-shot imitation learning. The paper trains a demonstration-conditioned policy with attention and reports strong generalization to novel block configurations, while proposing broader task scaling as future work.

  • Problem

    Existing imitation learning often requires careful feature engineering or substantial samples, but robots should generalize to new task instances from very few demonstrations.

  • Method

    The paper trains a policy over task distributions to map the current observation and one successful demonstration of the same task to controls.

  • Results

    Soft attention enabled strong generalization to novel block configurations not present in training data, and the approach was effective on block-stacking tasks.

  • Takeaways & Limitations

    A single successful demonstration can be mapped to an effective policy for solving the same task in a new situation.

Abstract

from arXiv · show

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 .

1 Introduction

The paper targets robots that learn complex manipulation tasks from very few demonstrations and generalize immediately to new instances. It introduces one-shot imitation learning, conditioning a policy on one demonstration to perform a previously unseen task.

  • Robotic systems need both dexterity for manipulating objects and communication mechanisms for conveying task intent across initial conditions.
  • Demonstrations communicate complex manipulation tasks while also providing clues about the motor skills required to execute them.
  • Training pairs one demonstration with an observation-action pair from another demonstration of the same task, teaching the network to output the corresponding action.
  • Existing demonstration-learning applications often require careful feature engineering or substantial system interaction, whereas the goal is rapid generalization from one or a few demonstrations.
  • One-shot imitation learning maximizes policy performance on a previously unseen task after receiving only one demonstration, without additional system interaction.
  • Soft attention processes demonstration sequences and block-location inputs, enabling strong generalization to novel block configurations absent from training data.

2 Related Work

Related imitation-learning methods commonly learn skills separately, while one-shot imitation learning addresses rapid acquisition of new skills from a single demonstration. The approach is positioned within meta-learning, transfer learning, and attention-based sequence modeling.

  • Behavioral cloning maps observations to actions, whereas inverse reinforcement learning estimates a reward function explaining demonstrated behavior.
  • Earlier imitation-learning work generally treats each skill separately, so learning one skill does not accelerate learning to imitate another.
  • Meta-learning studies systems that learn adaptation procedures, including one-shot and few-shot learning, fast reinforcement-learning agents, and fast weights.
  • Reinforcement learning typically requires many trials and a specified reward function, which can be more difficult than providing a demonstration.
  • Multi-task and transfer-learning methods reuse policies across tasks but do not provide immediate acquisition of a new skill from one demonstration.
  • The model relates to attention, interaction networks, and sequence-to-sequence models because it processes long demonstrations and emits action sequences.

3 One Shot Imitation Learning

The paper formalizes one-shot imitation over a distribution of tasks and successful demonstrations, then trains policies to map a demonstration and current observation to controls. Block stacking supplies a compositional testbed for generalization across configurations and initial states.

  • The framework defines a task distribution T, individual tasks t, demonstrations D(t), and a policy π_θ(a|o,d) conditioned on observation and demonstration.
  • The objective is to maximize expected policy performance over tasks and demonstrations, using an evaluation function that need not be available during training.
  • Block-stacking tasks control a 7-DOF Fetch arm to arrange varying numbers of cubes into user-specified towers, with each configuration defining a task.
  • Varying block starting positions require generalization within training tasks, while compositional configurations support evaluation on unseen tasks.
  • Behavioral cloning and DAGGER train the policy from demonstrations rather than requiring task reward functions.
  • Training samples tasks, demonstrations, and observation-action pairs, then regresses desired actions conditioned on the current observation and demonstration.

4 Architecture

The architecture combines demonstration, context, and manipulation modules to convert variable-length demonstrations and variable-sized block states into actions. Soft and neighborhood attention preserve task-relevant relationships while producing representations usable for control.

  • 4 Architecture: The proposed architecture contains demonstration, context, and manipulation modules, with a full specification provided in the appendix.
  • 4.1 Demonstration Network: The demonstration network embeds a trajectory, while temporal dropout reduces hundreds-to-thousands-step demonstrations during training by randomly discarding time steps.
  • 4.1 Demonstration Network: Using p = 0.95 discards 95% of demonstration time steps, reducing demonstration length by a factor of 20; test-time trajectory ensembles can improve performance.
  • 4.1 Demonstration Network: Neighborhood attention preserves variable-dimensional information by producing one output per input, with each output attending to other inputs relative to its corresponding input.
  • 4.1 Demonstration Network: Soft attention computes weighted combinations of memory vectors from query and context inputs, and multiple query heads produce multiple output vectors.
  • 4.2 Context network: The context network attends from the current state over demonstration time steps, propagates information across block embeddings, and repeats this process with an LSTM.
  • 4.2 Context network: Attention over the current state produces fixed-dimensional context vectors from block positions and robot state for the manipulation network.
  • 4.3 Manipulation network: The manipulation network uses source and target block information to compute the action for stacking one block on another through an MLP.

5 Experiments

The experiments evaluate conditioning strategies, training procedures, and generalization across block-stacking tasks. Full demonstrations outperform final-state and snapshot conditioning on harder tasks, while behavioral cloning matches DAGGER and attention visualizations suggest selective use of blocks and key frames.

  • Experimental design: The evaluation compares behavioral cloning with DAGGER and conditions policies on entire demonstrations, final states, or informative trajectory snapshots.It also evaluates generalization across training tasks, unseen configurations, and entirely new tasks.
  • Experimental design: The benchmark contains 140 training tasks and 43 test tasks with layouts involving 2 to 10 blocks.Training uses 1000 trajectories per task, with separate trajectories and initial configurations for evaluation.
  • Performance evaluation: As task difficulty increases, conditioning on the entire demonstration outperforms conditioning on the final state.For easier tasks with fewer stages, the conditioning strategies perform almost perfectly and similarly; demonstrations may additionally indicate which block to stack first.
  • Performance evaluation: Entire-demonstration conditioning also outperforms snapshot conditioning, contrary to the original expectation.The authors suspect temporal dropout regularizes training by effectively augmenting the demonstrations available to the policy.
  • Performance evaluation: Behavioral cloning reaches the same performance level as DAGGER in the reported experiments.Noise injection into trajectory collection was important for behavioral cloning to work well, and was used throughout these experiments.
  • Visualization: Attention visualizations show focus on a small subset of current-state blocks and sparse demonstration time steps corresponding to important stages.The visualized temporal pattern contains roughly six high-attention regions for a task requiring six stages.

6 Conclusions

The paper presents a model that maps one successful demonstration to a policy for solving the same task in a new situation. It demonstrates this approach on block-stacking tasks and identifies scaling to broader task distributions as future work.

  • Conclusion: The model maps a single successful demonstration of a task to an effective policy for solving that task in a new situation.The demonstrated setting is a family of block-stacking tasks.
  • Future work: Future work includes scaling the method to a much larger and broader distribution of tasks.The authors frame this as a step toward a general robotics imitation-learning system covering many tasks.

A Illustrative Example: Particle Reaching

Particle reaching tests whether a demonstration identifies which landmark to reach in a new scene. Attention-based, increasingly specialized architectures improve generalization, with final-state conditioning performing best on this simple task.

  • Task setup: Each task specifies a landmark target, while the agent and landmark positions vary across task instances.Without a demonstration, the robot cannot know which landmark to reach; after learning, it should identify the demonstrated target in a new instance.
  • Architectures: The experiment compares three increasingly specialized architectures for particle reaching.The architectures differ in how they process the demonstration and landmark information.
  • Results: More specialized architectures achieve much better test generalization, while conditioning on the entire demonstration hurts performance in this simple task.Conditioning only on the final state performs best because it sufficiently characterizes the target landmark.
  • Results: When the number of landmarks increases, LSTM-based architectures overfit, whereas final-state attention achieves perfect generalization across scenarios.The policies are evaluated over 100 configurations using training and test success rates.
  • Learning dynamics: Learning exhibits a phase transition that may correspond to inferring the task from the demonstration before control learning becomes nearly trivial.The supplied results describe this interpretation as intuitive rather than definitive.

B.1 Full Description of Architecture

The architecture section specifies the implementation context for particle-reaching policies and reports training through learning curves. Policies are trained for up to 1000 epochs, with most experiments finishing sooner.

  • Implementation: The architecture uses weight normalization with data-dependent initialization for dense and convolution operations.This implementation choice is stated alongside omitted minibatch, padding, and masking details.
  • Training procedure: Learning curves report moving-average success rates over the preceding 10 epochs.Each policy is trained for up to 1000 epochs, taking up to an hour on a Titan X Pascal GPU.
  • Training procedure: Most experiments finish before the maximum 1000-epoch training budget.The reported upper bound is one hour using a Titan X Pascal GPU.

B.1.1 Demonstration Network

The demonstration network converts a variable-length, block-structured demonstration into an embedding that can be processed with the current scene. Temporal dropout, block-wise processing, convolutions, and attention support this representation.

  • Inputs: Each demonstration observation contains 3D coordinates for B blocks and a 2-dimensional gripper state, giving dimension 3B+2.The architecture uses observations rather than actions from the input demonstration.
  • Preprocessing: Temporal dropout and splitting separate block information from robot information, producing tensors organized by time and block.The robot state is broadcast across blocks, yielding shapes approximately T̃×B×3 and T̃×B×2.
  • Block processing: A 1×1 convolution projects block states to the per-block embedding dimension before neighborhood attention and further convolutions.Residual connections are used while combining attention results with block and robot states.
  • Block processing: Neighborhood processing lets each block query other blocks and compare those results with its own and the robot’s state.This provides relational information across block positions during demonstration encoding.
  • Context integration: The context network repeatedly applies attention over the demonstration and current state, then produces a fixed-dimensional embedding independent of sequence length and block count.The resulting context is passed to a multilayer perceptron that predicts actions.
  • Permutation generalization: Training one representative task per block-permutation equivalence class tests whether the architecture can transfer across reordered block identities.For example, training on abcd should support dcba after receiving a demonstration of dcba.

B.3 Effect of Ensembling

The ensemble evaluation varies the number of downsampled demonstrations used at evaluation time. More ensembles help especially on easier tasks and consistently improve harder tasks, motivating 10 ensembles for the main evaluation.

  • Evaluation: Evaluation compares ensemble counts from 1 to 20 across training and test tasks.The ensembles consist of multiple downsampled demonstrations evaluated together.
  • Results: More ensembles help most on tasks with fewer stages, while harder tasks improve consistently with a smaller performance gap.The authors suspect the policy attends to frames where blocks are already stacked together.
  • Architecture: The context network uses temporal attention over demonstration steps and spatial attention over blocks before action prediction.The resulting context embedding is passed to a manipulation network that outputs an action distribution.
  • Architecture: The manipulation network maps the context embedding through dense layers to a distribution over available actions.Its output dimension is |A|, the number of possible actions.
  • Results: Ten ensembles are used in the main evaluation because additional ensembles provide no significant improvement in the reported setting.This choice follows the observed ensemble-performance analysis.

B.4 Breakdown of Failure Cases

The failure analysis separates wrong moves, manipulation failures, and recoverable timeouts, finding that manipulation failures account for most failures. Conditioning only on the final state produces more wrong moves than other architectures.

  • Most failures are manipulation failures that are mostly irrecoverable, indicating that stronger manipulation skills are needed for reliability.Examples include blocks shaken off the table, which the current hard-coded policy cannot handle.
  • Final-state conditioning produces more wrong moves than the other architectures.
  • Figure 6 represents each success or failure scenario by the area occupied by its color.

B.5 Learning Curves

Learning curves reveal relative training trends rather than final performance because evaluations use partial, training-sampled trajectories. Full-trajectory conditioning performs best but trains longer, while intermediate information helps learning and later stages progress more slowly.

  • Learning Curves: The learning curves do not measure final performance because evaluations sample training tasks and demonstrations, reset to a stage, and run for at most one stage.The figures remain useful for comparing relative trends.
  • Learning Curves: Full-trajectory conditioning gives the best performance but requires much longer training because processing the entire demonstration is more computationally expensive.Training may also vary because demonstrations are downsampled and the network must learn to segment them.
  • Learning Curves: Snapshot conditioning learns faster than final-state conditioning, suggesting that intermediate demonstration information helps facilitate training.
  • Learning Curves: Learning progresses most rapidly in initial stages and more slowly later, when manipulation is harder and later stages are sampled less frequently.
  • Exact Performance Numbers: Tables 3–11 report success rates across training and test tasks, ensemble counts, equivalent tasks, and success or failure scenarios.
  • More Visualizations: Attention visualizations show that the policy mostly attends to one block at a time during stacking operations.This matches the need to target one block during grasping and placing.
Loading 1703.07326v3…