Behavior Generation with Latent Actions

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

한 줄 요약

VQ-BeT는 행동 생성에서 다중 모드를 효과적으로 포착하고 추론 속도를 5배 향상시키는 행동 생성 모델이다.

핵심 기여도

VQ-BeT는 BeT의 k-means 기반 토큰화를 대체하여 계층적 Vector Quantization(VQ) 모듈을 도입함.
7개 환경에서 BeT와 Diffusion Policies를 초과하는 성능을 보임.
추론 속도가 Diffusion Policies 대비 5배 빠름.
실제 로봇 조작에서 장기적 작업 성능이 기존 방법 대비 73% 향상됨.

핵심 아이디어

기존 BeT는 행동 벡터를 k-means 클러스터링을 통해 이산화하여 다중 모드를 포착하려 했으나, 이는 고차원 행동 공간이나 긴 시퀀스에서는 확장성이 떨어지고, 기울기 정보가 없어 장기적 행동 모델링에 어려움이 있었다. VQ-BeT는 이 문제를 해결하기 위해 **계층적 VQ-VAE**(Vector Quantized Variational Autoencoder)를 도입하여 연속적인 행동 벡터를 이산 토큰으로 변환한다. 이는 **VQ-VAE의 이산 잠재 표현**이 이미지나 오디오 생성에서 입증된 성능을 행동 생성에도 적용한 것이다. 특히, VQ-VAE는 잠재 공간에서의 표현이 정확하면서도 계산 비용이 낮아, **단일 패스 모델**(single-pass model)로 구현되어 추론 속도를 크게 향상시킨다.

기술적 접근법

**모델 아키텍처**: VQ-BeT는 **Transformer** 기반 모델로, **계층적 VQ-VAE**를 사용하여 행동 벡터를 토큰화함.
**토큰화**: 연속 행동 벡터는 VQ-VAE의 잠재 공간에서 이산 토큰으로 변환됨.
**데이터셋**: 시뮬레이션 조작, 자율 주행, 로봇 환경 등 7개의 다양한 환경에서 실험됨.
**하이퍼파라미터**: 추론 속도는 Diffusion Policies 대비 **5× 가속**, 실제 로봇에서는 **25× 가속**됨.

주요 결과

**비조건적 생성**: 7개 환경 중 5개에서 BC, BeT, Diffusion Policies를 초과함.
**조건적 생성**: 7개 환경 중 6개에서 GCBC, C-BeT, BESO를 초과함.
**자율 주행**: nuScenes 데이터셋에서 기존 최고 성능과 유사한 결과를 보임.
**실제 로봇**: 장기적 작업에서 기존 방법 대비 **73% 성능 향상**됨.

의의 및 한계

VQ-BeT는 행동 생성에서 **다중 모드**를 효과적으로 포착하고, **단일 패스 모델**로 인해 추론 속도가 빠르며, 다양한 환경에서의 **확장성**을 입증했다. 특히, **VQ-VAE 기반 토큰화**는 행동 공간의 복잡성을 줄이며, 정확한 행동 재현을 가능하게 한다. 그러나, VQ-VAE는 잠재 공간의 표현이 제한적일 수 있으며, **복잡한 환경에서의 일반화 능력**은 추가 실험을 통해 검증이 필요하다. 또한, **다양한 로봇 사이의 정책 전이 가능성**은 아직 탐구 중이다.

실용적 활용

VQ-BeT는 자율 주행, 로봇 조작, 인간-로봇 정책 전이 등 다양한 산업 분야에서 활용 가능하다. 특히, **실시간 행동 생성**이 필요한 로봇 제어 시스템에서 빠른 추론 속도와 정확한 모델링 능력이 실용적 가치를 제공한다.