Quantization-Free Autoregressive Action Transformer

Claire Vernade (University of Tuebingen) · Ziyad Sheebaelhamd (Eberhard-Karls-Universität Tübingen) · Michael Tschannen (Google DeepMind) · Michael Muehlebach (Max-Planck Institute)

action spaceautoregressive transformer decodercontinuous policy parametrizationcontinuous structurediscrete action representationsgenerative infinite-vocabulary transformersgenerative modelimitation learning pipelinelatent codepolicy roll-outsquantization-free methodsampling algorithmssimulated robotics tasksstate-of-the-art performancetransformer-based imitation learning

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.