Supervised Pretraining Can Learn In-Context Reinforcement Learning

TL;DR

This paper demonstrates that supervised pretraining of transformers on decision-making tasks enables in-context reinforcement learning, allowing models to adapt to new tasks, perform exploration, and achieve sample-efficient learning with theoretical guarantees.

Contribution

Introducing Decision-Pretrained Transformer (DPT), a supervised pretraining method that enables transformers to perform in-context RL with theoretical and empirical benefits.

Findings

DPT can solve RL problems in-context, showing exploration and offline conservatism.

The model generalizes to new tasks beyond pretraining.

Theoretical connection to Bayesian posterior sampling with regret guarantees.

Abstract

Large transformer models trained on diverse datasets have shown a remarkable ability to learn in-context, achieving high few-shot performance on tasks they were not explicitly trained to solve. In this paper, we study the in-context learning capabilities of transformers in decision-making problems, i.e., reinforcement learning (RL) for bandits and Markov decision processes. To do so, we introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action given a query state and an in-context dataset of interactions, across a diverse set of tasks. This procedure, while simple, produces a model with several surprising capabilities. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline, despite not being explicitly trained to do so. The model also generalizes beyond the pretraining distribution to new tasks and automatically adapts its decision-making strategies to unknown structure. Theoretically, we show DPT can be viewed as an efficient implementation of Bayesian posterior sampling, a provably sample-efficient RL algorithm. We further leverage this connection to provide guarantees on the regret of the in-context algorithm yielded by DPT, and prove that it can learn faster than algorithms used to generate the pretraining data. These results suggest a promising yet simple path towards instilling strong in-context decision-making abilities in transformers.