基于潜在动作的行为生成

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