Implicit Behavioral Cloning

Pete Florence, Corey Lynch, Andy Zeng, Oscar Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, Jonathan Tompson

Seminal September 2021

TL;DR Behavioral cloning with an energy-based model, where the policy outputs the action that minimizes a learned energy of observation and action. On four real-robot pushing tasks from images it beats an otherwise identical MSE policy, with 83% vs 7% success on a 1 mm-tolerance insertion.

What they do

Behavioral cloning policies are usually explicit networks that map an observation straight to an action. This paper changes only the form of the policy: an energy-based model scores observation-action pairs, and the policy returns the lowest-energy action. Implicit and explicit policies with near-identical architectures are compared on D4RL human-expert tasks, an N-D particle task, simulated pushing and sweeping, and real xArm6 pushing from RGB images.

How it works

Training. For each demonstration, negative actions are sampled uniformly within the action bounds, and an InfoNCE-style loss pushes the demonstrated action to the lowest energy.
Inference. Three optimizers: derivative-free sampling with resampling and shrinking noise (2-D actions), an autoregressive per-dimension version (sweeping, up to 12-D), and Langevin MCMC with a gradient penalty (D4RL, up to 30-D).
Images. A CNN encodes the image and the action is concatenated with the encoding before the MLP.
Why it helps. In 1D fits, implicit models keep discontinuities sharp and represent multi-valued functions, where MSE models interpolate. They prove that the argmin of a continuous function can represent any set-valued function with a closed graph.

Results

Real robot, implicit vs MSE (60 trials each): Insert-Blue 83.3% vs 6.7%, sorting 48.3% vs 19.6%, Push-Red-then-Green 85.0% vs 35.0%, multimodal pushing 88.3% vs 55.0%.
Simulated image-based pushing: EBM 100%, MSE 87.0%, MDN 10.0%. Bi-manual sweeping: EBM beats MSE by 14%.
N-D particle: 95% success up to 16 dimensions for implicit, 8 for MSE.
D4RL human-expert tasks: competitive with CQL and S4RL without rewards, e.g. pen-human 2586 vs 1419.6 for S4RL, though worse on hammer-human.

Limitations

More compute: on a D4RL task, train+eval takes 3.4 h vs 0.66 h for MSE, and real-robot inference takes 7.22 ms vs 3.49 ms. Implementation is also more complex, and the plain sampling optimizer fails past 5 dimensions on the particle task.

Figures

From Figure 1: (a) Explicit policies map observations to actions; implicit policies take both observations and action... — From Figure 1: (a) Explicit policies map observations to actions; implicit policies take both observations and actions as input and pick the action that minimizes the energy (b). Implicit policies do substantially better on precise block insertion (c) and sorting (d) learned from human demonstrations.

From Figure 2: Implicit (top) vs explicit (bottom) fits of 1D functions: a single discontinuity, piecewise sections w... — From Figure 2: Implicit (top) vs explicit (bottom) fits of 1D functions: a single discontinuity, piecewise sections with different slopes, and random noise. Red insets zoom in on the discontinuities, which the implicit model keeps sharp and the explicit model interpolates across.

From Figure 5: Implicit vs explicit policies across 6 simulated and real domains, with author-reported baselines on t... — From Figure 5: Implicit vs explicit policies across 6 simulated and real domains, with author-reported baselines on the human-expert D4RL tasks.

From Figure 9: The real-world setup (a) and tasks: multi-modal block pushing (b), oriented insertion needing 1 mm pre... — From Figure 9: The real-world setup (a) and tasks: multi-modal block pushing (b), oriented insertion needing 1 mm precision (c), and sorting blue from yellow blocks (d).

Abstract

We find that across a wide range of robot policy learning scenarios, treating supervised policy learning with an implicit model generally performs better, on average, than commonly used explicit models. We present extensive experiments on this finding, and we provide both intuitive insight and theoretical arguments distinguishing the properties of implicit models compared to their explicit counterparts, particularly with respect to approximating complex, potentially discontinuous and multi-valued (set-valued) functions. On robotic policy learning tasks we show that implicit behavioral cloning policies with energy-based models (EBM) often outperform common explicit (Mean Square Error, or Mixture Density) behavioral cloning policies, including on tasks with high-dimensional action spaces and visual image inputs. We find these policies provide competitive results or outperform state-of-the-art offline reinforcement learning methods on the challenging human-expert tasks from the D4RL benchmark suite, despite using no reward information. In the real world, robots with implicit policies can learn complex and remarkably subtle behaviors on contact-rich tasks from human demonstrations, including tasks with high combinatorial complexity and tasks requiring 1mm precision.

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