Home

Awesome

Behavior Transformers: Cloning k modes with one stone

[Arxiv] [Code] [miniBET] [Data] [Project page and videos]

Created by Nur Muhammad (Mahi) Shafiullah, Zichen Jeff Cui, Ariuntuya (Arty) Altanzaya, and Lerrel Pinto at New York University.

Update: Also look into the miniBeT github repo if you are trying to use (Conditional) Behavior Transformers in your own project.

https://user-images.githubusercontent.com/3000253/174540589-35c2d51b-3a5b-42d8-9f80-483b61df91a7.mp4

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 behavior has a wide variance, multiple modes, and human demonstrations naturally don’t 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 further analyze the importance of every crucial component in BeT.

Code release

In this repository, you can find the code to reproduce Behavior Transformer (BeT). The following assumes our current working directory is the root folder of this project repository; tested on Ubuntu 20.04 LTS (amd64).

Getting started

Setting up the project environments

Getting the training datasets

Datasets used for training will be available at this OSF link: https://osf.io/983qz/.

Reproducing experiments

The following assumes our current working directory is the root folder of this project repository.

To reproduce the experiment results, the overall steps are:

  1. Activate the conda environment with
    conda activate behavior-transformer
    
  2. Train with python3 train.py. A model snapshot will be saved to ./exp_local/...;
  3. In the corresponding environment config, set the load_dir to the absolute path of the snapshot directory above;
  4. Eval with python3 run_on_env.py.

See below for detailed steps for each environment.

CARLA

Franka kitchen

Block push

Update (11/17/22): There was a small error in the published hyperparameters for Block push (namely window_size and batch_size). We have updated the hyperparameters in the config files to the parameters that replicate the results in the paper. We apologize for any inconvenience.

</details>

Speeding up evaluation

Acknowledgements

We are indebted to the following codebases and tools for making our lives significantly easier.