L'apprentissage supervisé par prétraitement peut acquérir l'apprentissage par renforcement en contexte.

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

Résumé

Les grands modèles de transformateurs entraînés sur des ensembles de données diversifiés ont démontré une capacité remarquable à apprendre en contexte, atteignant des performances élevées en few-shot sur des tâches pour lesquelles ils n'ont pas été explicitement entraînés. Dans cet article, nous étudions les capacités d'apprentissage en contexte des transformateurs dans des problèmes de prise de décision, c'est-à-dire l'apprentissage par renforcement (RL) pour les bandits et les processus de décision markoviens. Pour ce faire, nous introduisons et étudions le Decision-Pretrained Transformer (DPT), une méthode de pré-entraînement supervisée où le transformateur prédit une action optimale étant donné un état de requête et un ensemble de données d'interactions en contexte, à travers une diversité de tâches. Cette procédure, bien que simple, produit un modèle doté de plusieurs capacités surprenantes. Nous constatons que le transformateur pré-entraîné peut être utilisé pour résoudre une gamme de problèmes RL en contexte, manifestant à la fois une exploration en ligne et un conservatisme hors ligne, bien qu'il n'ait pas été explicitement entraîné pour cela. Le modèle généralise également au-delà de la distribution de pré-entraînement à de nouvelles tâches et adapte automatiquement ses stratégies de prise de décision à des structures inconnues. Théoriquement, nous montrons que le DPT peut être vu comme une implémentation efficace de l'échantillonnage bayésien a posteriori, un algorithme RL prouvé comme étant efficace en termes d'échantillons. Nous exploitons en outre cette connexion pour fournir des garanties sur le regret de l'algorithme en contexte produit par le DPT, et prouvons qu'il peut apprendre plus rapidement que les algorithmes utilisés pour générer les données de pré-entraînement. Ces résultats suggèrent une voie prometteuse et simple pour instiller de solides capacités de prise de décision en contexte dans les transformateurs.

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.