title: Supervised Pretraining Can Learn In-Context Reinforcement Learning
summary: 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.
paper_authors._id: 649a84b0a37f2e731852670e
paper_authors.name: Jonathan N. Lee
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2023-06-29T20:55:56.608Z
paper_authors._id: 649de792e685e40238a07f65
paper_authors.avatarUrl: /avatars/96c3ac0b410b85719bffa8a1a9b035b1.svg
paper_authors.fullname: Jonathan Lee
paper_authors.type: user
paper_authors.user: jnlee
paper_authors._id: 649a84b0a37f2e731852670f
paper_authors.name: Annie Xie
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2023-06-30T10:23:10.030Z
paper_authors._id: 649e23bb3e5e75047640b3b5
paper_authors.avatarUrl: /avatars/019d008a5bd7057b8665e1ef856f13a7.svg
paper_authors.fullname: Annie Xie
paper_authors.type: user
paper_authors.user: anniexie
paper_authors._id: 649a84b0a37f2e7318526710
paper_authors.name: Aldo Pacchiano
paper_authors._id: 649a84b0a37f2e7318526711
paper_authors.name: Yash Chandak
paper_authors._id: 649a84b0a37f2e7318526712
paper_authors.name: Chelsea Finn
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2023-07-11T07:56:18.636Z
paper_authors._id: 64ac22a9193f0a807deb673d
paper_authors.avatarUrl: /avatars/fcac4912678ad3cb6e817d40bdee9aea.svg
paper_authors.fullname: Chelsea Finn
paper_authors.type: user
paper_authors.user: cbfinn
paper_authors._id: 649a84b0a37f2e7318526713
paper_authors.name: Ofir Nachum
paper_authors._id: 649a84b0a37f2e7318526714
paper_authors.name: Emma Brunskill
paper_title: Supervised Pretraining Can Learn In-Context Reinforcement Learning
paper_summary: 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.
paper_ai_keywords: transformers reinforcement learning (RL) bandits Markov decision processes Decision-Pretrained Transformer (DPT) supervised pretraining query state in-context dataset interactions exploration online conservatism offline Bayesian posterior sampling sample-efficient RL algorithm regret decision-making abilities