Quantization-Free Autoregressive Action Transformer

Ziyad Sheebaelhamd, Michael Tschannen, Michael Muehlebach, Claire Vernade

Keywords: Imitation Learning, Reinforcement Learning, Transformers

TL;DR: We introduce a quantization-free way to train autoregressive transformers for continuous action decision making, improving on discretized action methods.

Abstract: Current transformer-based imitation learning approaches introduce discrete action representations and train an autoregressive transformer decoder on the resulting latent code. However, the initial quantization breaks the continuous structure of the action space thereby limiting the capabilities of the generative model. We propose a quantization-free method instead that leverages Generative Infinite-Vocabulary Transformers (GIVT) as a direct, continuous policy parametrization for autoregressive transformers. This simplifies the imitation learning pipeline while achieving state-of-the-art performance on a variety of popular simulated robotics tasks. We enhance our policy roll-outs by carefully studying sampling algorithms, further improving the results.

Primary Area: Reinforcement learning (e.g., decision and control, planning, hierarchical RL, robotics)

Submission Number: 12363