PASTA: Pretrained Action-State Transformer Agents

TL;DR

This paper introduces PASTA, a unified pre-trained transformer framework for reinforcement learning that systematically compares design choices across diverse downstream tasks, emphasizing simplicity, robustness, and broad applicability.

Contribution

It presents a comprehensive investigation of pre-trained action-state transformer agents (PASTA), highlighting a unified methodology and practical design choices for robust RL models.

Findings

Tokenization at component level improves model performance.

Pre-training with next token prediction or masked language modeling is effective.

Models trained across multiple domains demonstrate versatility.

Abstract

Self-supervised learning has brought about a revolutionary paradigm shift in various computing domains, including NLP, vision, and biology. Recent approaches involve pre-training transformer models on vast amounts of unlabeled data, serving as a starting point for efficiently solving downstream tasks. In reinforcement learning, researchers have recently adapted these approaches, developing models pre-trained on expert trajectories. This advancement enables the models to tackle a broad spectrum of tasks, ranging from robotics to recommendation systems. However, existing methods mostly rely on intricate pre-training objectives tailored to specific downstream applications. This paper conducts a comprehensive investigation of models, referred to as pre-trained action-state transformer agents (PASTA). Our study covers a unified methodology and covers an extensive set of general downstream tasks including behavioral cloning, offline RL, sensor failure robustness, and dynamics change adaptation. Our objective is to systematically compare various design choices and offer valuable insights that will aid practitioners in developing robust models. Key highlights of our study include tokenization at the component level for actions and states, the use of fundamental pre-training objectives such as next token prediction or masked language modeling, simultaneous training of models across multiple domains, and the application of various fine-tuning strategies. In this study, the developed models contain fewer than 7 million parameters allowing a broad community to use these models and reproduce our experiments. We hope that this study will encourage further research into the use of transformers with first principle design choices to represent RL trajectories and contribute to robust policy learning.