BEAST: Efficient Tokenization of B-Splines Encoded Action Sequences for Imitation Learning

TL;DR

BEAST introduces a B-spline based action tokenizer that encodes action sequences into uniform, smooth, and efficient tokens, enabling faster and more reliable imitation learning across various models and tasks.

Contribution

BEAST is a novel action tokenizer that requires no separate training, produces uniform tokens, and ensures smooth trajectories, improving efficiency and performance in imitation learning.

Findings

Reduces training and inference computational costs.

Generates smooth, high-frequency control signals.

Achieves competitive success rates on benchmark tasks.

Abstract

We present the B-spline Encoded Action Sequence Tokenizer (BEAST), a novel action tokenizer that encodes action sequences into compact discrete or continuous tokens using B-splines. In contrast to existing action tokenizers based on vector quantization or byte pair encoding, BEAST requires no separate tokenizer training and consistently produces tokens of uniform length, enabling fast action sequence generation via parallel decoding. Leveraging our B-spline formulation, BEAST inherently ensures generating smooth trajectories without discontinuities between adjacent segments. We extensively evaluate BEAST by integrating it with three distinct model architectures: a Variational Autoencoder (VAE) with continuous tokens, a decoder-only Transformer with discrete tokens, and Florence-2, a pretrained Vision-Language Model with an encoder-decoder architecture, demonstrating BEAST's compatibility and scalability with large pretrained models. We evaluate BEAST across three established benchmarks consisting of 166 simulated tasks and on three distinct robot settings with a total of 8 real-world tasks. Experimental results demonstrate that BEAST (i) significantly reduces both training and inference computational costs, and (ii) consistently generates smooth, high-frequency control signals suitable for continuous control tasks while (iii) reliably achieves competitive task success rates compared to state-of-the-art methods.