教師あり事前学習は文脈内強化学習を学習可能である

Supervised Pretraining Can Learn In-Context Reinforcement Learning

June 26, 2023

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

cs.AI

要旨

多様なデータセットで訓練された大規模なTransformerモデルは、文脈内学習において驚くべき能力を示し、明示的に訓練されていないタスクにおいても高いFew-shot性能を達成しています。本論文では、Transformerの意思決定問題における文脈内学習能力、すなわちバンディットやマルコフ決定過程における強化学習（RL）について研究します。これを行うために、我々はDecision-Pretrained Transformer（DPT）を導入し、その特性を調査します。DPTは、Transformerがクエリ状態と文脈内の相互作用データセットを基に最適な行動を予測するという、多様なタスクにわたる教師付き事前学習手法です。この手順は単純ながら、いくつかの驚くべき能力を持つモデルを生み出します。事前学習されたTransformerは、明示的に訓練されていないにもかかわらず、オンラインでの探索とオフラインでの保守性を示し、幅広いRL問題を文脈内で解決できることがわかりました。また、このモデルは事前学習分布を超えて新しいタスクに一般化し、未知の構造に対して自動的に意思決定戦略を適応させます。理論的には、DPTがベイズ事後サンプリングの効率的な実装と見なせることを示し、これは証明可能なサンプル効率の良いRLアルゴリズムです。さらに、この関連性を活用して、DPTによって生成される文脈内アルゴリズムの後悔に関する保証を提供し、事前学習データを生成するために使用されたアルゴリズムよりも速く学習できることを証明します。これらの結果は、Transformerに強力な文脈内意思決定能力を備えさせるための有望でシンプルな道筋を示唆しています。

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.