Matt Strong

Matt Strong

CS PhD at Stanford University

← Paper Digest

Behavior Transformers: Cloning $k$ modes with one stone

Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, Lerrel Pinto

Seminal June 2022

Franka Kitchen rollout in which BeT chains tasks, starting with the microwave and the kettle.

TL;DR Behavior Transformer (BeT) clones unlabeled, multi-modal demonstrations by having a small transformer pick a k-means action bin and then add a learned continuous offset. On multimodal block pushing it gets both blocks into their squares in 71% of rollouts, where every baseline gets 0%, while still covering the demonstrated modes.

What they do

Standard behavior cloning fits a unimodal policy with an MSE loss, which averages across the different ways demonstrators solve a task. BeT learns from offline demonstrations with no reward or task labels that mix several behavior modes, and samples among those modes at rollout instead of collapsing to one. It is evaluated in simulation on point-mass toys, CARLA driving from RGB images, multimodal block pushing with an XArm, and Franka Kitchen with 566 human VR demonstrations.

How it works

  • Action discretization. k-means over the dataset’s actions gives k action centers. Each action becomes a bin index plus a residual, from which it is reconstructed exactly.
  • Transformer trunk. A small minGPT (10^4 to 10^6 parameters) reads the last h observations and predicts a categorical distribution over bins at each timestep, trained with focal loss.
  • Action offsets. A second head outputs a residual for each of the k bins. A masked multi-task loss trains only the residual of the ground-truth bin.
  • Sampling. At test time BeT samples a bin from the predicted probabilities and adds that bin’s residual.

Results

  • Block push (1,000 rollouts): BeT pushes one and two blocks into squares at 0.96 and 0.71. No baseline (MLP-MSE, nearest neighbor, LWR, VAE, Flow, IBC) exceeds 0.01 and 0.
  • Kitchen: BeT completes 4 tasks in 44% of rollouts, against 24% for IBC. CARLA: 0.98 success, narrowly behind LWR at 1.0.
  • Mode coverage: in CARLA BeT turns left/right 0.34/0.64 (demonstrations 0.50/0.50), while MLP, NN and LWR only turn right. Kitchen task-sequence entropy is 2.47, against 2.96 in the data and 2.41 for IBC.
  • Ablations: without binning, Block push falls to 0.25 of full BeT performance; swapping the transformer for an MLP drops it to 0 and Kitchen to 0.05. BeT trains on Block push in under an hour, against about 14 h for IBC.

Limitations

All experiments are in simulation. Rollouts are unconditional: there is no way yet to ask for one particular mode, which the authors leave to future work. Like other behavior cloning methods, BeT fails once observations drift out of distribution; in Block push its main failure is not noticing that a block has not fully entered the target. Mode frequencies are covered but not matched exactly, and some kitchen task sequences are over-visited.

Videos

Block push rollout pushing the green block to the red square, then the red block to the green square.
Overhead view of CARLA rollouts driving around the building block.

Figures

From Figure 1: Unconditional BeT rollouts in CARLA, Block push and Franka Kitchen. Successive rollouts in the same en...
From Figure 1: Unconditional BeT rollouts in CARLA, Block push and Franka Kitchen. Successive rollouts in the same environment reach different goals, or the same goals in different ways.
From Figure 3: Architecture. (A) k-means splits each action into a discrete bin and a continuous offset. (B) Offline ...
From Figure 3: Architecture. (A) k-means splits each action into a discrete bin and a continuous offset. (B) Offline training of the minGPT trunk with its bin and offset heads. (C) At test time BeT samples a bin, then adds that bin's offset to get a continuous action.
From Figure 4: Most frequent task sequences completed in Franka Kitchen, out of 1,000 unconditional rollouts per mode...
From Figure 4: Most frequent task sequences completed in Franka Kitchen, out of 1,000 unconditional rollouts per model, with each task in its own color.
Abstract

While behavior learning has made impressive progress in recent times, it lags behind computer vision and natural language processing due to its inability to leverage large, human-generated datasets. Human behaviors have wide variance, multiple modes, and human demonstrations typically do not come with reward labels. These properties limit the applicability of current methods in Offline RL and Behavioral Cloning to learn from large, pre-collected datasets. In this work, we present Behavior Transformer (BeT), a new technique to model unlabeled demonstration data with multiple modes. BeT retrofits standard transformer architectures with action discretization coupled with a multi-task action correction inspired by offset prediction in object detection. This allows us to leverage the multi-modal modeling ability of modern transformers to predict multi-modal continuous actions. We experimentally evaluate BeT on a variety of robotic manipulation and self-driving behavior datasets. We show that BeT significantly improves over prior state-of-the-art work on solving demonstrated tasks while capturing the major modes present in the pre-collected datasets. Finally, through an extensive ablation study, we analyze the importance of every crucial component in BeT. Videos of behavior generated by BeT are available at https://notmahi.github.io/bet

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.