El preentrenamiento supervisado puede aprender aprendizaje por refuerzo en contexto.

Supervised Pretraining Can Learn In-Context Reinforcement Learning

June 26, 2023

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

cs.AI

Resumen

Los grandes modelos transformadores entrenados en conjuntos de datos diversos han demostrado una capacidad notable para aprender en contexto, logrando un alto rendimiento en pocas muestras en tareas para las que no fueron explícitamente entrenados. En este artículo, estudiamos las capacidades de aprendizaje en contexto de los transformadores en problemas de toma de decisiones, es decir, aprendizaje por refuerzo (RL) para bandidos y procesos de decisión de Markov. Para ello, introducimos y estudiamos el Transformador Preentrenado para Decisiones (DPT, por sus siglas en inglés), un método de preentrenamiento supervisado en el que el transformador predice una acción óptima dado un estado de consulta y un conjunto de datos en contexto de interacciones, a través de una variedad de tareas. Este procedimiento, aunque simple, produce un modelo con varias capacidades sorprendentes. Encontramos que el transformador preentrenado puede usarse para resolver una gama de problemas de RL en contexto, exhibiendo tanto exploración en línea como conservadurismo fuera de línea, a pesar de no haber sido explícitamente entrenado para ello. El modelo también generaliza más allá de la distribución de preentrenamiento a nuevas tareas y adapta automáticamente sus estrategias de toma de decisiones a estructuras desconocidas. Teóricamente, mostramos que DPT puede verse como una implementación eficiente del muestreo bayesiano posterior, un algoritmo de RL con eficiencia muestral demostrada. Además, aprovechamos esta conexión para proporcionar garantías sobre el arrepentimiento del algoritmo en contexto generado por DPT, y demostramos que puede aprender más rápido que los algoritmos utilizados para generar los datos de preentrenamiento. Estos resultados sugieren un camino prometedor y simple hacia la instilación de fuertes habilidades de toma de decisiones en contexto en los transformadores.

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.