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

翻译：从标注数据集中对复杂行为进行生成建模一直是决策领域的一个长期问题。与语言或图像生成不同，决策需要建模动作——这些连续值向量分布呈多模态特性，可能来自未经过滤的数据源，且其生成误差会在序列预测中累积。近期一类名为行为变换器（BeT）的模型通过使用k-means聚类对动作进行离散化来捕捉不同模式。然而，k-means难以扩展到高维动作空间或长序列，且缺乏梯度信息，因此BeT在长程动作建模方面表现不佳。本文提出向量量化行为变换器（VQ-BeT），这是一种用于行为生成的通用模型，可处理多模态动作预测、条件生成和部分观测。VQ-BeT通过引入分层向量量化模块对连续动作进行分词化，对BeT进行了增强。在包括模拟操控、自动驾驶和机器人操作在内的七个环境中，VQ-BeT相较于BeT和扩散策略等最先进模型表现更优。重要的是，我们证明了VQ-BeT在捕捉行为模式方面的能力提升，同时其推理速度较扩散策略提升5倍。视频和代码请见https://sjlee.cc/vq-bet