Behavior Generation with Latent Actions
Seungjae Lee Yibin Wang Haritheja Etukuru H. Jin Kim Nur Muhammad Mahi Shafiullah Lerrel Pinto
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