Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining

Licong Lin, Yu Bai, Song Mei

TL;DR

This work analyzes the theoretical capacity of pretrained Transformers to perform in-context reinforcement learning via supervised pretraining. It proves that under realizability, a supervised-pretrained Transformer imitates the conditional expectation of a given expert algorithm, with generalization error governed by the transformer's capacity and a distribution-shift factor between offline data and expert behavior. It further shows that Transformers with ReLU attention can efficiently approximate near-optimal online RL algorithms (LinUCB, Thompson sampling for linear bandits, and UCB-VI for tabular MDPs), yielding quantitative regret bounds. The framework encompasses Algorithm Distillation and Decision-Pretrained Transformers, offers concrete constructions for in-context RL, and provides preliminary experiments validating the theory. The results illuminate how offline trajectories can endow Transformers with powerful ICRL capabilities while highlighting practical considerations like distribution ratio and approximation errors.

Abstract

Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabilities, where they can make good decisions when prompted with interaction trajectories from unseen environments. However, when and how transformers can be trained to perform ICRL have not been theoretically well-understood. In particular, it is unclear which reinforcement-learning algorithms transformers can perform in context, and how distribution mismatch in offline training data affects the learned algorithms. This paper provides a theoretical framework that analyzes supervised pretraining for ICRL. This includes two recently proposed training methods -- algorithm distillation and decision-pretrained transformers. First, assuming model realizability, we prove the supervised-pretrained transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory. The generalization error will scale with model capacity and a distribution divergence factor between the expert and offline algorithms. Second, we show transformers with ReLU attention can efficiently approximate near-optimal online reinforcement learning algorithms like LinUCB and Thompson sampling for stochastic linear bandits, and UCB-VI for tabular Markov decision processes. This provides the first quantitative analysis of the ICRL capabilities of transformers pretrained from offline trajectories.

Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining

TL;DR

This work analyzes the theoretical capacity of pretrained Transformers to perform in-context reinforcement learning via supervised pretraining. It proves that under realizability, a supervised-pretrained Transformer imitates the conditional expectation of a given expert algorithm, with generalization error governed by the transformer's capacity and a distribution-shift factor between offline data and expert behavior. It further shows that Transformers with ReLU attention can efficiently approximate near-optimal online RL algorithms (LinUCB, Thompson sampling for linear bandits, and UCB-VI for tabular MDPs), yielding quantitative regret bounds. The framework encompasses Algorithm Distillation and Decision-Pretrained Transformers, offers concrete constructions for in-context RL, and provides preliminary experiments validating the theory. The results illuminate how offline trajectories can endow Transformers with powerful ICRL capabilities while highlighting practical considerations like distribution ratio and approximation errors.

Abstract

Large transformer models pretrained on offline reinforcement learning datasets have demonstrated remarkable in-context reinforcement learning (ICRL) capabilities, where they can make good decisions when prompted with interaction trajectories from unseen environments. However, when and how transformers can be trained to perform ICRL have not been theoretically well-understood. In particular, it is unclear which reinforcement-learning algorithms transformers can perform in context, and how distribution mismatch in offline training data affects the learned algorithms. This paper provides a theoretical framework that analyzes supervised pretraining for ICRL. This includes two recently proposed training methods -- algorithm distillation and decision-pretrained transformers. First, assuming model realizability, we prove the supervised-pretrained transformer will imitate the conditional expectation of the expert algorithm given the observed trajectory. The generalization error will scale with model capacity and a distribution divergence factor between the expert and offline algorithms. Second, we show transformers with ReLU attention can efficiently approximate near-optimal online reinforcement learning algorithms like LinUCB and Thompson sampling for stochastic linear bandits, and UCB-VI for tabular Markov decision processes. This provides the first quantitative analysis of the ICRL capabilities of transformers pretrained from offline trajectories.
Paper Structure (86 sections, 18 theorems, 228 equations, 4 figures)

This paper contains 86 sections, 18 theorems, 228 equations, 4 figures.

Table of Contents

  1. Introduction
  2. Summary of contributions and paper outline
  3. Related work
  4. Meta-learning and meta-reinforcement learning
  5. In-context learning
  6. Transformers for decision making
  7. Expressivity of transformers
  8. Statistical theories of imitation learning
  9. Statistical theories of imitation learning
  10. Framework for In-Context Reinforcement Learning
  11. Distributions of offline trajectories
  12. Three special cases of expert algorithms
  13. Transformer architecture
  14. Algorithm induced by Transformers
  15. Statistical analysis of supervised pretraining
  16. ...and 71 more sections

Key Result

Theorem 6

Let Assumption asp:realizability hold and let ${\widehat{\theta}}$ be a solution to Eq. eq:general_mle. Take ${\mathcal{R}} = {\mathcal{R}}_{\overline{\mathsf{Alg}}_{E},{\mathsf{Alg}}_0}$ as defined in Definition def:dist_ratio, and $\cN_{{\Theta}} = \cN_{{\Theta}}(({n}{T})^{-2})$ as defined in Defi Further assume that $|{r}_t| \leq 1$ almost surely. Then with probability at least $1-\delta$, the

Figures (4)

  • Figure 1: Regrets of transformer (TF), empirical average (Emp), Thompson sampling (TS) and LinUCB or UCB (LinUCB reduces to UCB for Bernoulli bandits). Left: linear bandit with $d=5$, $A=10$, $\sigma=1.5$, ${\mathsf{Alg}}_0={\mathsf{Alg}}_{E}={\mathrm{LinUCB}}$. Right: Bernoulli bandit with $d=5$, ${\mathsf{Alg}}_0=({\mathsf{Alg}}_{\mathrm{unif}}+{\mathsf{Alg}}_{{\mathrm{TS}}})/2$ and ${\mathsf{Alg}}_{E}={a}^*$. The simulation is repeated 500 times. Shading displays the standard deviation of the regret estimates.
  • Figure 2: Suboptimalities of transformer (TF), empirical average (Emp), Thompson sampling (TS), and LinUCB (or UCB). Left: linear bandit with $d=5$, $A=10$, $\sigma=1.5$, ${\mathsf{Alg}}_0={\mathsf{Alg}}_{E}={\mathrm{LinUCB}}$. Right: Bernoulli bandit with $d=5$, ${\mathsf{Alg}}_0=({\mathsf{Alg}}_{\mathrm{unif}}+{\mathsf{Alg}}_{{\mathrm{TS}}})/2$, and ${\mathsf{Alg}}_{E}={a}_t^*$. The simulation is repeated 500 times. Shading displays the standard deviation of the sub-optimality estimates.
  • Figure 3: Regrets and suboptimalities of transformer (TF), empirical average (Emp), Thompson sampling (TS), and UCB. Settings: Bernoulli bandit with $d=5$, and ${\mathsf{Alg}}_0={\mathsf{Alg}}_{E}={\mathrm{LinUCB}}$. The simulation is repeated 500 times. Shading displays the standard deviation of the estimates.
  • Figure 4: Regrets and difference of regrets between transformers and Thompson sampling, for different context algorithms. Settings: Bernoulli bandit with $d=5$, ${\mathsf{Alg}}_{E}={a}_t^*$ and ${\mathsf{Alg}}_0=\alpha{\mathsf{Alg}}_{\mathrm{TS}}+(1-\alpha){\mathsf{Alg}}_{\mathrm{unif}}$ with $\alpha \in \{0,0.1,0.5,1\}$. The simulation is repeated 500 times. Shading displays the standard deviation of the estimates.

Theorems & Definitions (29)

  • Definition 1: Masked attention layer
  • Definition 2: MLP layer
  • Definition 3: Decoder-based Transformer
  • Definition 4: Covering number
  • Definition 5: Distribution ratio
  • Theorem 6: Performance gap between expected cumulative rewards
  • Proposition 7
  • Theorem 8: Approximating the soft LinUCB
  • Theorem 9: Regret of LinUCB and ICRL
  • Theorem 10: Approximating the Thompson sampling, Informal
  • ...and 19 more