PaperlayerView exact revision on arXiv

Source-linked AI summary

Behavior Transformers: Cloning $k$ modes with one stone

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

arXiv:2206.11251v2cs.LGcs.AIcs.CVcs.RO

TL;DR

Behavior learning needs methods that can learn from large, unlabeled, multimodal demonstration datasets without online interaction or reward labels. BeT retrofits transformers with discrete action-mode prediction and continuous offset correction, achieving higher rollout performance while covering major demonstration modes.

  • Problem

    Behavior cloning and offline RL are limited when pre-collected behavior is multimodal, noisy, and lacks task-specific reward labels.

  • Method

    BeT uses a transformer with k-means action discretization and a residual action corrector to predict continuous actions from multimodal demonstrations.

  • Results

    BeT achieves significantly higher online-rollout performance than prior behavior-modeling methods on multimodal datasets and covers their major behavior modes.

  • Takeaways & Limitations

    BeT models diverse continuous behavior without collapsing to a single mode, supporting multimodal imitation from pre-collected demonstrations.

Abstract

from arXiv · show

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

1 Introduction

Behavior learning from pre-collected demonstrations must handle multimodal, unlabeled human behavior without costly online interaction. BeT uses transformers with discretized action modes and continuous corrections, improving rollout performance while covering major dataset modes.

  • Motivation: RL often requires 1-10M reward-labeled online samples for benchmark control tasks, unlike data-driven vision and language learning.This motivates learning behavioral priors from pre-collected data.
  • Motivation: Behavior cloning learns actions from demonstrations without online interactions or reward labels, making it suitable for plentiful unlabeled behavior data.It has been applied to self-driving and robotic manipulation.
  • Problem: Existing behavior cloning methods often assume a unimodal expert, whereas natural and human demonstrations contain noisy, sub-optimal, multimodal behavior.The assumption is often embedded architecturally, such as through a Gaussian prior.
  • Approach: BeT combines transformer multimodal prediction, k-means action bins, and a residual action corrector to model continuous multimodal actions.The bins provide categorical mode modeling, while the corrector produces continuous actions for sampled bins.
  • Findings: BeT achieves significantly higher online-rollout performance than prior behavior-modeling methods on multimodal datasets.The paper evaluates five datasets spanning toy, robotic pushing, kitchen, and CARLA settings.
  • Findings: BeT covers major training-data modes rather than collapsing onto one mode.Unconditional rollouts can pursue different goals or achieve the same goals in different ways.

2 Behavior Transformers

Behavior Transformer models demonstration trajectories with history-dependent, multimodal action distributions by combining transformer sequence modeling with action-bin prediction and continuous corrections. Its factorization reconstructs continuous actions while retaining multimodal structure and reducing discretization loss.

  • Motivation: Regular behavior cloning can reduce to MSE when actions are modeled with unimodal isotropic Gaussians, encouraging a single central action.The figure contrasts this MSE-based behavior-cloning model with BeT's ability to capture multimodal distributions.
  • Modeling assumptions: BeT models action distributions as mixtures of Gaussians and conditions predictions on the current observation plus a history of previous observations.It models P(a_t | o_t, o_{t−1}, · · ·, o_{t−h+1}) rather than assuming purely Markovian behavior.
  • Action factorization: Each continuous action is encoded as a categorical action-bin index and a continuous residual relative to its nearest learned action center.The centers are optimized from dataset actions, and the action is reconstructed as its selected center plus residual.
  • Architecture: The minGPT-based transformer maps observation sequences to distributions over k action bins and uses an additional offset head to predict continuous corrections.The offset head produces proposed residuals for the action centers, restoring continuous-action fidelity after discretization.
  • Training: The offset loss is applied only to the ground-truth action class, preventing every bin-specific offset from fitting the same ground-truth action.This masked multi-task correction is inspired by object-detection offset prediction.
  • Inference: At test time, BeT samples an action bin from predicted probabilities, adds its corresponding residual, and reconstructs the continuous action.The procedure uses the latest h observations and computes the action as the selected center plus predicted residual.

3 Experiments

The experiments test whether BeT imitates and preserves multimodal behaviors across diagnostic, robotic, kitchen, and self-driving environments, while ablations examine its components and computational cost.

  • Experimental setup: Experiments span five environments, including point-mass diagnostics, CARLA self-driving, block pushing, and Franka Kitchen task sequences.The study evaluates imitation, multimodality, component importance, and computational considerations across these settings.
  • Imitation performance: BeT outperforms other methods in every evaluated environment except CARLA, where LWR narrowly performs better.All behavioral cloning methods can fail after observations leave the training distribution, but their tolerance differs across environments.
  • Multimodality: BeT covers all demonstration modes in CARLA and Block push, while Kitchen rollouts preserve diversity and produce longer task strings more often than strong baselines.Kitchen task frequencies do not always exactly match demonstrations, but BeT avoids collapsing to a single mode.
  • Ablations: Removing discrete action binning significantly reduces performance and limits Kitchen rollouts to one observed task subsequence after 100 random rollouts.The experiments also find k-means binning performs better overall than mixture density networks and uniform quantization.
  • Ablations: Removing residual action offsets degrades performance, with the largest loss occurring in the highest-dimensional Kitchen action space.The authors relate this pattern to greater discretization fidelity loss in higher-dimensional actions.
  • Ablations: Historical context improves BeT performance across the second point-mass, CARLA, Block push, and Kitchen environments.Replacing the transformer with MLP, temporal convolution, or LSTM backbones lowers performance and makes training less stable.
  • Computation: BeT uses 10^4-10^6 parameters, trains on the largest dataset within an hour on one desktop GPU, and evaluates much faster than IBC.The reported Block push rollout times are 1.65 seconds for BeT versus 17.70 seconds for IBC.

4 Related Work

The related work situates BeT within offline behavior learning, generative imitation models, transformer control, and datasets containing unlabeled multimodal behavior.

  • Learning from offline data: Offline behavior learning includes Offline RL with reward-labeled mixed-quality data and imitation learning from expert behavior without reward labels.BeT belongs to the imitation-learning category as a behavior cloning model.
  • Learning from offline data: Behavior cloning models expert actions from observations using supervised learning and are generally faster and simpler than RL or Offline RL methods.The passage notes that such methods can sometimes achieve competitive results.
  • Generative models for behavior learning: Generative imitation approaches include inverse reinforcement learning, which constructs a reward function, and GAIL, which combines generative adversarial models with imitation learning.These approaches aim to generate desirable or expert-like behavior from demonstrations.
  • Transformers for control: Transformer-based methods have been applied to reinforcement learning, Offline RL, and imitation learning, with some using transformers to summarize historical visual context.BeT extends transformer use toward multimodal behavior modeling.
  • Datasets for distributionally multi-modal data: Behavior datasets vary in whether they include agent goals or rewards; open-ended datasets often lack these labels and contain greater action-distribution multimodality.The absence of labeling requirements also makes unlabeled datasets cheaper to collect.

5 Discussions

BeT models continuous action sequences from open-ended, multi-modal demonstrations using a transformer-decoder backbone, a discrete action mode predictor, and a continuous offset corrector. The authors identify diverse real-world human demonstrations and extracting unimodal policies during online interactions as promising future uses.

  • BeT combines a transformer-decoder backbone with discrete action mode prediction and continuous action offset correction.The method targets continuous action sequences from open-ended, multi-modal demonstrations.
  • Future applications include learning diverse behavior from human demonstrations or real-world interactions.
  • Extracting a unimodal behavior policy during online interactions could make BeT useful as a prior for online Reinforcement Learning.The paper suggests distillation or generating suitable prompts as possible routes.

Checklist

The checklist records that the paper addresses claims, limitations, reproducibility, assets, and ethics requirements. Most items are marked Yes or N/A, with linked code, appendix training details, and a stated limitations section cited where relevant.

  • The paper states that code and dataset licensing is covered and directs readers to videos of diverse model behaviors.
  • The authors state that the main claims accurately reflect the paper’s contributions and scope and that the work’s limitations are described in Section 5.
  • The paper reports reproducibility materials, training details, and error-bar information through the abstract, appendix B, and the checklist responses.
  • The checklist marks theoretical-result questions, human-subject procedures, participant compensation, and personally identifiable information as not applicable.
  • Existing assets are cited and their licenses are addressed through linked code, while no new assets are reported in the checklist.

A Environment and Dataset Details

The experiments use simple point-mass environments, simulated robotic pushing and kitchen manipulation, and photo-realistic self-driving. These settings introduce controlled path, target, task, noise, and observation variations for evaluating multi-modal behavior cloning.

  • Point mass environments: Point-mass environments use two-dimensional observations and actions, with actions specifying immediate (Δx, Δy) displacement.
  • Point mass environments: Multipath datasets contain two or three trajectory modes distinguished by fork paths or first-step direction.The supplied passages describe two modes for one dataset and three modes for another.
  • CARLA environment: CARLA uses 224 × 224 × 3 RGB observations processed by a frozen ImageNet-pretrained ResNet-18 into 512-dimensional vectors.The environment is a photo-realistic driving simulation on the Town04 map, with 100 demonstrations collected by a tuned PID agent and noisy execution.
  • Block-push environment: Block-push demonstrations use a hard-coded controller, and the models are trained on 1,000 demonstrations generated under its randomized modes.
  • Block-push environment: The block-push environment randomizes block and target positions and rotations while permitting four controller behavior modes.Modes arise from independent choices of the first block and its target, followed by pushing the remaining block to the remaining target.
  • Franka kitchen environment: The kitchen environment uses a Franka Panda with a 9-dimensional action space and zeroes goal dimensions because demonstrations lack task labels.The dataset contains 566 human VR demonstrations from the Relay Policy Learning environment.

B.1 Baselines

The baseline suite includes supervised, nearest-neighbor, continuous generative, energy-based, and trajectory-transformer behavior models. The paper reports that continuous approaches and some baselines face context, extrapolation, overfitting, or rollout limitations in the evaluated settings.

  • Multi-layer Perceptron with MSE: MLP-MSE baselines use fully connected networks with optional BatchNorm and select depth and width using a bias-variance trade-off.They train on 95% of each dataset and evaluate MSE on the remaining 5%.
  • Nearest Neighbor: Nearest Neighbor stores observation-action pairs and executes the action associated with the closest test observation in representation space.
  • Nearest Neighbor: Nearest Neighbor performs well in simple CARLA and dense Kitchen settings but fails more when Block-push requires interpolation or extrapolation.
  • Continuous Generative Model: VAE with Gaussian Prior: The VAE compresses sequences of T = 10 actions into a 10-dimensional latent variable and reconstructs action sequences from a state-conditioned Gaussian prior.
  • Continuous Generative Model: VAE with Gaussian Prior: The VAE’s unconditional rollouts are not very successful, which the authors attribute to latent sampling that ignores recently executed actions and a single-mode Gaussian prior.
  • Continuous generative models: Continuous generative approaches underperform partly because they omit historical context and can return less likely actions that drive rollouts out of distribution.The authors contrast this with BeT’s strong action prior and offset prediction.
  • Implicit Behavioral Cloning: IBC is described as the strongest baseline but is easy to overfit, requiring early stopping for CARLA and Block-pushing.
  • Trajectory Transformers: Trajectory Transformer failed to complete any Kitchen tasks in unconditioned, greedy, or beam-search rollouts.

B.2 Algorithm Details

BeT trains a transformer to predict discrete action classes and class-specific continuous offsets, combining focal and multi-task losses for behavior modeling.

  • Loss functions: Focal loss modifies cross entropy to emphasize errors on low-probability classes, helping predict imbalanced k-means action classes.Its gradient is steeper for smaller pt and flatter for larger pt.
  • Loss functions: The multi-task loss penalizes the action offset only for the ground-truth class rather than across all classes.This directs offset learning through the correct action bin.
  • Training pipeline: Training optimizes the combined focal and multi-task losses, with α scaling them to comparable magnitude at initialization.The implementation samples trajectory subsequences, predicts class probabilities and offsets, and backpropagates the normalized combined loss.
  • Runtime: In the Kitchen environment, BeT requires 2.8 ms per action, compared with 52 ms for IBC, 0.5 ms for MLP, and 867.86 ms for Trajectory Transformer.Wall-clock measurements were averaged over three runs.
  • Training pipeline: BeT encodes continuous actions with k-means cluster centers and trains MinGPT to predict categorical classes plus residual action offsets.The encoder assigns each action to its nearest center and represents the residual relative to that center; decoding adds the offset back to the selected center.
  • Rollout: At deployment, BeT samples an action class from the final predicted distribution, decodes its offset into a continuous action, and executes it.The model uses the last h observations as context before producing the action.

C Ablation studies

The ablations test historical context and alternative sequence-modeling trunks, finding that context improves behavior cloning and transformer models outperform the tested alternatives.

  • Historical context: Providing historical context improves BeT behavior cloning compared with training and testing the same model on single observations.The experiment uses identical models while varying sequence length.

C.2 Ablating the number of discrete bin centers, k

The number of discrete action bins creates a loss trade-off: increasing k raises log-likelihood loss while reducing MSE, and k = 1 lowers BeT performance.

  • Bin-count trade-off: Increasing the number of bins raises log-likelihood loss while lowering MSE loss.BeT combines focal loss for the binning head with MSE loss for the offset head.
  • Performance versus k: Using only one bin (k = 1) decreases BeT's performance level.The ablation varies k and reports reward normalized to the best-performing model.

C.3 Ablating the core model in the architecture

The core-model ablation replaces MinGPT with MLP, temporal convolution, or LSTM trunks; the alternatives show weaker performance or slower fitting.

  • Ablation design: The ablation replaces MinGPT with a fully connected MLP, temporal convolution network, or LSTM-based recurrent neural network.These alternatives serve as architecture comparisons for the core transformer trunk.
  • MLP comparison: The MLP ablation stacks the last t observation frames and zero-pads the stack near trajectory beginnings.Its intermediate layers retain the corresponding MinGPT width and layer count.
  • Temporal convolution: Temporal convolution models perform consistently worse than MinGPT models but are easier to fit than RNNs.Both model types receive the same history period used by the transformer.
  • LSTM comparison: LSTM models take significantly longer to fit, produce worse training and test fits, and achieve lower open-ended rollout success than MinGPT.The reported differences remain even with sufficient model capacity.
Loading 2206.11251v2…