Behavior Generation with Latent Actions

TL;DR

VQ-BeT introduces hierarchical vector quantization into behavior transformers, enabling scalable, multimodal action generation with faster inference across diverse environments, surpassing previous models like BeT and Diffusion Policies.

Contribution

The paper presents VQ-BeT, a novel hierarchical vector quantization approach that enhances behavior generation models by improving scalability, mode capturing, and inference speed.

Findings

VQ-BeT outperforms BeT and Diffusion Policies in multiple environments.

VQ-BeT captures behavior modes more effectively.

Inference speed is increased 5x over Diffusion Policies.

Abstract

Generative modeling of complex behaviors from labeled datasets has been a longstanding problem in decision making. Unlike language or image generation, decision making requires modeling actions - continuous-valued vectors that are multimodal in their distribution, potentially drawn from uncurated sources, where generation errors can compound in sequential prediction. A recent class of models called Behavior Transformers (BeT) addresses this by discretizing actions using k-means clustering to capture different modes. However, k-means struggles to scale for high-dimensional action spaces or long sequences, and lacks gradient information, and thus BeT suffers in modeling long-range actions. In this work, we present Vector-Quantized Behavior Transformer (VQ-BeT), a versatile model for behavior generation that handles multimodal action prediction, conditional generation, and partial observations. VQ-BeT augments BeT by tokenizing continuous actions with a hierarchical vector quantization module. Across seven environments including simulated manipulation, autonomous driving, and robotics, VQ-BeT improves on state-of-the-art models such as BeT and Diffusion Policies. Importantly, we demonstrate VQ-BeT's improved ability to capture behavior modes while accelerating inference speed 5x over Diffusion Policies. Videos and code can be found https://sjlee.cc/vq-bet