Gesuperviseerde voorafgaande training kan in-context reinforcement learning aanleren.

Supervised Pretraining Can Learn In-Context Reinforcement Learning

June 26, 2023

Auteurs: Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, Emma Brunskill

cs.AI

Samenvatting

Grote transformermodellen die getraind zijn op diverse datasets hebben een opmerkelijke vaardigheid getoond om in-context te leren, waarbij ze hoge few-shot prestaties behalen op taken waar ze niet expliciet voor getraind zijn. In dit artikel bestuderen we de in-context leervaardigheden van transformers in besluitvormingsproblemen, namelijk reinforcement learning (RL) voor bandits en Markov-beslissingsprocessen. Hiervoor introduceren en bestuderen we de Decision-Pretrained Transformer (DPT), een supervised pretrainingsmethode waarbij de transformer een optimale actie voorspelt gegeven een querytoestand en een in-context dataset van interacties, over een diverse set van taken. Deze procedure, hoewel eenvoudig, produceert een model met verschillende verrassende capaciteiten. We ontdekken dat de voorgetrainde transformer gebruikt kan worden om een reeks RL-problemen in-context op te lossen, waarbij het zowel exploratie online als conservatisme offline vertoont, ondanks dat het hier niet expliciet voor getraind is. Het model generaliseert ook verder dan de pretrainingsdistributie naar nieuwe taken en past automatisch zijn besluitvormingsstrategieën aan aan onbekende structuren. Theoretisch laten we zien dat DPT kan worden gezien als een efficiënte implementatie van Bayesian posterior sampling, een bewezen sample-efficiënt RL-algoritme. We benutten deze verbinding verder om garanties te geven op de regret van het in-context algoritme dat door DPT wordt opgeleverd, en bewijzen dat het sneller kan leren dan de algoritmen die gebruikt worden om de pretrainingsdata te genereren. Deze resultaten suggereren een veelbelovende maar eenvoudige weg om sterke in-context besluitvormingsvaardigheden in transformers te ontwikkelen.

English

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.