Контролируемое предварительное обучение может освоить обучение с подкреплением в контексте.
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
Аннотация
Крупные трансформаторные модели, обученные на разнообразных наборах данных, продемонстрировали впечатляющую способность к обучению в контексте, достигая высокой производительности в условиях малого числа примеров на задачах, для которых они не были явно обучены. В данной работе мы исследуем способности трансформаторов к обучению в контексте в задачах принятия решений, а именно в обучении с подкреплением (RL) для задач бандитов и марковских процессов принятия решений. Для этого мы представляем и изучаем Decision-Pretrained Transformer (DPT) — метод предварительного обучения с учителем, в котором трансформатор предсказывает оптимальное действие для заданного состояния и набора данных взаимодействий в контексте, охватывающего разнообразные задачи. Этот подход, несмотря на свою простоту, создает модель с рядом удивительных возможностей. Мы обнаруживаем, что предварительно обученный трансформатор может использоваться для решения широкого спектра задач RL в контексте, демонстрируя как исследовательское поведение в онлайн-режиме, так и консерватизм в офлайн-режиме, несмотря на отсутствие явного обучения этим аспектам. Модель также обобщает свои знания за пределы распределения данных, использованных при предварительном обучении, на новые задачи и автоматически адаптирует свои стратегии принятия решений к неизвестной структуре. Теоретически мы показываем, что DPT можно рассматривать как эффективную реализацию байесовского сэмплирования апостериорного распределения — алгоритма RL с доказанной эффективностью по выборкам. Мы также используем эту связь для предоставления гарантий на сожаление (regret) алгоритма в контексте, полученного с помощью DPT, и доказываем, что он может обучаться быстрее, чем алгоритмы, использованные для генерации данных предварительного обучения. Эти результаты указывают на многообещающий, но простой путь к внедрению мощных способностей к принятию решений в контексте в трансформаторные модели.
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.