SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jonathan N. Lee*1
Annie Xie\*1
Aldo Pacchiano $^{2}$
Yash Chandak<sup>1</sup>
Chelsea Finn
Ofir Nachum
$^{1}$ Stanford University  $^{2}$ Broad Institute of MIT and Harvard  $^{3}$ Boston University
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 the Decision-Pretrained Transformer (DPT), a supervised pretraining method where a transformer predicts an optimal action given a query state and an in-context dataset of interactions from a diverse set of tasks. While simple, this procedure produces a model with several surprising capabilities. We find that the trained transformer can 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.
For supervised learning, transformer-based models trained at scale have shown impressive abilities to perform tasks given an input context, often referred to as few-shot prompting or in-context learning [1]. In this setting, a pretrained model is presented with a small number of supervised input-output examples in its context, and is then asked to predict the most likely completion (i.e. output) of an unpaired input, without parameter updates. Over the last few years, in-context learning has been applied to solve a range of tasks [2] and a growing number works are beginning to understand and analyze in-context learning for supervised learning [3, 4, 5, 6]. In this work, our focus is to study and understand in-context learning applied to sequential decision-making, specifically in the context of reinforcement learning (RL) settings. Decision-making (e.g. RL) is considerably more dynamic and complex than supervised learning. Understanding and leveraging in-context learning here could potentially unlock significant improvements in an agent's ability to adapt and make few-shot decisions in response to observations from the world. Such capabilities are instrumental for practical applications ranging from robotics to recommendation systems.
*Equal contribution. Correspondence to jnl@stanford.edu and anniexie@stanford.edu. Code is available at https://github.com/jon-lee/decision-pretrained-transformer
For in-context decision-making [7, 8, 9], rather than input-output tuples, the context takes the form of state-action-reward tuples representing a dataset of interactions with an unknown environment. The agent must leverage these interactions to understand the dynamics of the world and what actions lead to good outcomes. A hallmark of good decision-making in online RL algorithms is a judicious balance of selecting exploratory actions to gather information and selecting increasingly optimal actions by exploiting that information [10]. In contrast, an RL agent with access to only a suboptimal offline dataset should produce a policy that conservatively selects actions [11]. An ideal in-context decision-maker should exhibit similar behaviors.
To study in-context decision-making formally, we propose a new simple supervised pretraining objective, namely, to train (via supervised learning) a transformer to predict an optimal action label<sup>2</sup> given a query state and an in-context dataset of interactions, across a diverse set of tasks. We refer to the pretrained model as a Decision-Pretrained Transformer (DPT). Once trained, DPT can be deployed as either an online or offline RL algorithm in a new task by passing it an in-context dataset of interactions and querying it for predictions of the optimal action in different states. For example, online, the in-context dataset is initially empty and DPT's predictions are uncertain because the new task is unknown, but it fills the dataset with its interactions as it learns and becomes more confident about the optimal action. We show empirically and theoretically that DPT yields a surprisingly effective in-context decision-maker with regret guarantees. As it turns out, DPT effectively performs posterior sampling — a provably sample-efficient Bayesian RL algorithm that has historically been limited by its computational burden [12, 13]. We summarize our main findings below.
- Predicting optimal actions alone gives rise to near-optimal decision-making algorithms. The DPT objective is solely based on predicting optimal actions from in-context interactions. At the outset, it is not immediately apparent that these predictions at test-time would yield good decision-making behavior when the task is unknown and behaviors such as online exploration are necessary to solve it. Intriguingly, DPT as an algorithm is capable of dealing with this uncertainty in-context. For example, despite not being explicitly trained to explore, DPT exhibits an exploration strategy on par with hand-designed algorithms, as a means to discover the optimal actions.
- DPT generalizes to new decision-making problems, offline and online. We show DPT can handle reward distributions unseen in its pretraining data on bandit problems as well as unseen goals, dynamics, and datasets in simple MDPs. This suggests that the in-context strategies learned during pretraining are robust and generalizable without any parameter updates at test time.
- DPT improves over the data used to pretrain it by exploiting latent structure. As an example, in parametric bandit problems, specialized algorithms can leverage structure (such as linear rewards) and offer provably better regret. However, a representation must be known in advance. Perhaps surprisingly, we find that pretraining on linear bandit problems, even with unknown representations, leads DPT to select actions and explore in a way that matches an efficient linear bandit algorithm. This holds even when the source pretraining data comes from a suboptimal algorithm (i.e., one that does not take advantage of any latent structure), demonstrating the ability to learn improved in-context strategies beyond what it was trained on.
- Posterior sampling can be implemented via in-context learning. Posterior sampling (PS), a generalization of Thompson Sampling, can provably sample-efficiently solve online RL problems [12, 13], but a common criticism is the lack of computationally efficient ways to update and sample from a posterior distribution. DPT can be viewed as learning a posterior distribution over optimal actions, shortcutting the PS procedure. Under some conditions, we show theoretically that DPT in-context is equivalent to PS. Furthermore, DPT's prior and posterior updates are grounded in data rather than needing to be specified a priori. This suggests that in-context learning could be a possible path towards scalable and practical RL via posterior sampling.
Meta-learning. Algorithmically, in-context learning falls under the meta-learning framework [14, 15]. At a high-level, these methods attempt to learn some underlying shared structure of the training distribution of tasks to accelerate learning of new tasks. For decision-making and RL, there is often
2If not explicitly known, the optimal action can be determined by running any (potentially inefficient) minimax-optimal regret algorithm for each pretraining task.
a choice in what shared 'structure' is specifically learned such as the dynamics of the task [16, 17, 18, 19], a task context identifier [20, 21, 22, 23], temporally extended skills and options [24, 25, 26], or initialization of a neural network policy [27, 28]). In-context learning can be viewed as taking a more agnostic approach by learning the learning algorithm itself, more similar to [29, 30, 31, 32, 8]. Algorithm Distillation (AD) [7, 33] falls under this category, using autoregressive supervised learning to distill traces of a single-task RL algorithm into a task-agnostic model. While DPT also leverages autoregressive SL, it does not distill an existing RL algorithm in order to imitate how to learn. Instead, we train DPT to predict task-specific optimal actions, yielding potentially emergent strategies at test time that automatically leverage the task structure to behave similarly to posterior sampling.
Autoregressive transformers for decision-making. In decision-making fields such as RL and imitation learning, transformer models trained using autoregressive supervised action prediction have proliferated [34], inspired by the successes of these techniques for large language models [35, 36, 1]. For example, Decision Transformer (DT) [37, 38, 8] uses a transformer to autoregressively model sequences of actions from offline experience data, conditioned on the achieved return. During inference, one can then query the model conditioned on a desired return value. This approach has been shown to scale favorably to large models and multi-task settings [39], at times exceeding the performance of large-scale multi-task imitation learning with transformers [40, 41, 42]. However, DT is known to be provably (and unboundedly) sub-optimal in common scenarios [43, 44]. A common criticism of DT, and supervised learned transformers in general, is their inability to improve upon the dataset. For example, there is little reason for DT to output meaningful behavior if conditioned on return higher than any observed in training, without strong extrapolation assumptions [43]. In contrast, a major contribution of our work is theoretical and empirical evidence for the ability of DPT to improve over behaviors seen in the dataset in terms of regret.
Value and policy-based offline RL. Offline RL algorithms offer the opportunity to learn from existing datasets. To address distributional shift, many prior algorithms incorporate the principle of value pessimism [45, 46, 47, 48], or policy regularization [49, 50, 51, 52, 53]. To reduce the amount of offline data required in a new task, methods for offline meta-RL can reuse interactions collected in a set of related tasks [54, 55, 56]. However, they still must address distribution shift, requiring solutions such as policy regularization [54] or additional online interactions [57]. DPT follows the success of autoregressive models like DT and AD, avoiding these issues. With our pretraining objective, DPT also leverages offline datasets for new tasks more effectively than AD.
3 In-Context Learning Model
Basic decision models. The basic decision model of our study is the finite-horizon Markov decision process (MDP). An MDP is specified by the tuple  $\tau = \langle S, \mathcal{A}, T, R, H, \rho \rangle$  to be solved, where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space,  $T: \mathcal{S} \times \mathcal{A} \to \Delta(\mathcal{S})$  is the transition function,  $R: \mathcal{S} \times \mathcal{A} \to \Delta(\mathbb{R})$  is the reward function,  $H \in \mathbb{N}$  is the horizon, and  $\rho \in \Delta(\mathcal{S})$  is the initial state distribution. A learner interacts with the environment through the following protocol: (1) an initial state  $s_1$  is sampled from  $\rho$ ; (2) at time step  $h$ , the learner chooses an action  $a_h$  and transitions to state  $s_{h+1} \sim T(\cdot | s_h, a_h)$ , and receives a reward  $r_h \sim R(\cdot | s_h, a_h)$ . The episode ends after  $H$  steps. A policy  $\pi$  maps states to distributions over actions and can be used to interact with the MDP. We denote the optimal policy as  $\pi^{\star}$ , which maximizes the value function  $V(\pi^{\star}) = \max_{\pi} V(\pi) := \max_{\pi} \mathbb{E}_{\pi} \sum_{h} r_h$ . When necessary, we use the subscript  $\tau$  to distinguish  $V_{\tau}$  and  $\pi_{\tau}^{\star}$  for the specific MDP  $\tau$ . We assume the state space is partitioned by  $h \in [H]$  so that  $\pi^{\star}$  is notationally independent of  $h$ . Note this framework encompasses multi-armed bandit settings where the state space is a single point, e.g.  $S = \{1\}$ ,  $H = 1$ , and the optimal policy is  $a^{\star} \in \operatorname{argmax}_{a \in \mathcal{A}} \mathbb{E}[r_1 | a_1 = a]$ .
**Pretraining.** We give pseudocode in Algorithm 1 and a visualization of the transformer in Figure 1. Details on the precise structure of the transformer can be found in Section A.1. Let  $\mathcal{T}_{\mathrm{pre}}$  be a distribution over tasks at the time of pretraining. A task  $\tau \sim \mathcal{T}_{\mathrm{pre}}$  can be viewed as a specification of an MDP,  $\tau = \langle S, \mathcal{A}, T, R, H, \rho \rangle$ . The distribution  $\mathcal{T}_{\mathrm{pre}}$  can span different reward and transition functions and even different state and action spaces. We then sample a context (or a prompt) which consists of a dataset  $D \sim \mathcal{D}_{\mathrm{pre}}(\cdot; \tau)$  of interactions between the learner and the MDP specified by  $\tau$ .  $D = \{s_j, a_j, s_j', r_j\}_{j \in [n]}$  is a collection of transition tuples taken in  $\tau$ . We refer to  $D$  as the in-context dataset because it provides the contextual information about  $\tau$ .  $D$  could be generated through variety of means, such as: (1) random interactions within  $\tau$ , (2) demonstrations from an expert, and (3) rollouts of an algorithm. Additionally, we independently sample a query state  $s_{\mathrm{query}}$  from the distribution  $\mathcal{D}_{\mathrm{query}}$  over states  $S$  and a label  $a^{\star}$  is sampled from the optimal policy  $\pi_{\tau}^{\star}(\cdot | s_{\mathrm{query}})$ .
Figure 1: A transformer model  $M_{\theta}$  is pretrained to predict an optimal action  $a_{\mathrm{query}}^{\star}$  from a state  $s_{\mathrm{query}}$  in a task, given a dataset of interactions from that task. The resulting Decision-Pretrained Transformer (DPT) learns a distribution over the optimal action conditioned on an in-context dataset.  $M_{\theta}$  can be deployed in new tasks online by collecting data on the fly or offline by immediately conditioning on a static dataset.
Algorithm 1 Decision-Pretrained Transformer (DPT): Training and Deployment
1: // Collecting pretraining dataset  
2: Initialize empty pretraining dataset  $\mathcal{B}$   
3: for  $i$  in  $[N]$  do  
4: Sample task  $\tau \sim T_{\mathrm{pre}}$ , in-context dataset  $D \sim D_{\mathrm{pre}}(\cdot; \tau)$ , query state  $s_{\mathrm{query}} \sim D_{\mathrm{query}}$   
5: Sample label  $a^{\star} \sim \pi_{\tau}^{\star}(\cdot|s_{\mathrm{query}})$  and add  $(s_{\mathrm{query}}, D, a^{\star})$  to  $\mathcal{B}$   
6: end for  
7: // Pretraining model on dataset  
8: Initialize model  $M_{\theta}$  with parameters  $\theta$   
9: while not converged do  
10: Sample  $(s_{\mathrm{query}}, D, a^{\star})$  from  $\mathcal{B}$  and predict  $\hat{p}_j(\cdot) = M_{\theta}(\cdot|s_{\mathrm{query}}, D_j)$  for all  $j \in [n]$   
11: Compute loss in (2) with respect to  $a^{\star}$  and backpropagate to update  $\theta$ .  
12: end while  
13: // Offline test-time deployment  
14: Sample unknown task  $\tau \sim T_{\mathrm{test}}$ , sample dataset  $D \sim D_{\mathrm{test}}(\cdot; \tau)$   
15: Deploy  $M_{\theta}$  in  $\tau$  by choosing  $a_h \in \operatorname{argmax}_{a \in \mathcal{A}} M_{\theta}(a|s_h, D)$  at step  $h$   
16: // Online test-time deployment  
17: Sample unknown task  $\tau \sim T_{\mathrm{test}}$  and initialize empty  $D = \{\}$   
18: for ep in max_eps do  
19: Deploy  $M_{\theta}$  by sampling  $a_h \sim M_{\theta}(\cdot|s_h, D)$  at step  $h$   
20: Add  $(s_1, a_1, r_1, \ldots)$  to  $D$   
21: end for
for task  $\tau$  (see Section 5.3 for how to implement this in common practical scenarios). We denote the joint pretraining distribution over tasks, in-context datasets, query states, and action labels as  $P_{pre}$ :
$$
P _ {p r e} (\tau , D, s _ {\text {q u e r y}}, a ^ {\star}) = \mathcal {T} _ {\text {p r e}} (\tau) \mathcal {D} _ {\text {p r e}} (D; \tau) \mathcal {D} _ {\text {q u e r y}} (s _ {\text {q u e r y}}) \pi_ {\tau} ^ {\star} (a ^ {\star} | s _ {\text {q u e r y}}). \tag {1}
$$
Given the in-context dataset  $D$  and a query state  $s_{\mathrm{query}}$ , we can train a model to predict the optimal action  $a^{\star}$  in response simply via supervised learning. Let  $D_{j} = \{(s_{1},a_{1},s_{1}^{\prime},r_{1}),\ldots ,(s_{j},a_{j},s_{j}^{\prime},r_{j})\}$  denote the partial dataset up to  $j$  samples. Formally, we aim to train a causal GPT-2 transformer model  $M$  parameterized by  $\theta$ , which outputs a distribution over actions  $\mathcal{A}$ , to minimize the expected loss over samples from the pretraining distribution:
$$
\min  _ {\theta} \mathbb {E} _ {P _ {p r e}} \sum_ {j \in [ n ]} \ell \left(M _ {\theta} (\cdot \mid s _ {\text {q u e r y}}, D _ {j}), a ^ {\star}\right). \tag {2}
$$
Throughout, we set the loss to be the negative log-likelihood with  $\ell(M_{\theta}(\cdot | s_{\mathrm{query}}, D_j), a^{\star}) := -\log M_{\theta}(a^{\star} | s_{\mathrm{query}}, D_j)$ . This framework can work for both discrete and continuous  $\mathcal{A}$ . For our experiments with discrete  $\mathcal{A}$ , we use a softmax parameterization for the distribution of  $M_{\theta}$  essentially treating this as a classification problem. The resulting output model  $M_{\theta}$  can be viewed as an algorithm that takes in a dataset of interactions  $D$  and can be queried with a forward pass for predictions of the optimal action via inputting a query state  $s_{\mathrm{query}}$ . We refer to the trained model  $M_{\theta}$  as a Decision-Pretrained Transformer (DPT).
Testing. After pretraining, a new task (MDP)  $\tau$  is sampled from a test-task distribution  $\mathcal{T}_{\mathrm{test}}$ . If the DPT is to be tested offline, then a dataset (prompt) is a sampled  $D \sim \mathcal{D}_{\mathrm{test}}(\cdot ;\tau)$  and the policy that the model in-context learns is given conditionally as  $M_{\theta}(\cdot |\cdot ,D)$ . Namely, we evaluate the policy by selecting action  $a_{h} \in \operatorname{argmax}_{a} M_{\theta}(a|s_{h},D)$  when the learner visits state  $s_h$ . If the model is to be tested online through multiple episodes of interaction, then the dataset is initialized as empty  $D = \{\}$ . At each episode,  $M_{\theta}(\cdot |\cdot ,D)$  is deployed where the model samples  $a_{h} \sim M_{\theta}(\cdot |s_{h},D)$  upon observing state  $s_h$ . Throughout a full episode, it collects interactions  $\{s_1,a_1,r_1,\dots,s_H,a_H,r_H\}$  which are subsequently appended to  $D$ . The model then repeats the process with another episode, and so on until a specified number of episodes has been reached. A key distinction of the testing phase is that there are no updates to the parameters of  $M_{\theta}$ . This is in contrast to hand-designed RL algorithms that would perform parameter updates or maintain statistics using  $D$  to learn from scratch. Instead, the model  $M_{\theta}$  performs a computation through its forward pass to generate a distribution over actions conditioned on the in-context  $D$  and query state  $s_h$ .
4 Learning in Bandits
We begin with an empirical investigation of DPT in a multi-armed bandit, a well-studied special case of the MDP where the state space  $S$  is a singleton and the horizon  $H = 1$  is a single step. We will examine the performance of DPT both when aiming to select a good action from offline historical data and for online learning where the goal is to maximize cumulative reward from scratch. Offline, it is critical to account for uncertainty due to noise as certain actions may not be sampled well enough. Online, it is critical to judiciously balance exploration and exploitation to minimize overall regret. For detailed descriptions of the experiment setups, see Appendix A.
Pretraining distribution. For the pretraining task distribution  $\mathcal{T}_{\mathrm{pre}}$ , we sample 5-armed bandits  $(|\mathcal{A}| = 5)$ . The reward function for arm  $a$  is a normal distribution  $R(\cdot |s,a) = \mathcal{N}(\mu_a,\sigma^2)$  where  $\mu_{a}\sim \mathrm{Unif}[0,1]$  independently and  $\sigma = 0.3$ . To generate in-context datasets  $\mathcal{D}_{\mathrm{pre}}$ , we randomly generate action frequencies by sampling probabilities from a Dirichlet distribution and mixing them with a point-mass distribution on one random arm (see details in Appendix A.3). Then we sample the actions accordingly from this distribution. This encourages diversity of the in-context datasets. The optimal policy  $\pi_{\tau}^{\star}$  for bandit  $\tau$  is  $\arg \max_{a}\mu_{a}$ , which we can easily compute during pretraining. We pretrain the model  $M_{\theta}$  to predict  $a^{\star}$  from  $D$  as described in Section 3 for datasets up to size  $n = 500$ .
Comparisons. We compare to several well-known algorithms for bandits<sup>3</sup>. All of the algorithms are designed to reason in a particular way about uncertainty based on their observations.
- Empirical mean algorithm (Emp) selects the action with the highest empirical mean reward naively.
- Upper Confidence Bound (UCB) selects the action with the highest upper confidence bound.
- Lower Confidence Bound (LCB) selects the action with the highest lower confidence bound.
- Thompson Sampling (TS) selects the action with the highest sampled mean from a posterior distribution over reward models. The prior and likelihood functions are Gaussian.
Emp and TS [58, 59] can both be used for offline or online learning; UCB [60] is an online algorithm that achieves exploration via optimism under uncertainty; and LCB [61, 62] is used to minimize suboptimality given an offline dataset by selecting actions pessimistically. We evaluate algorithms with standard bandit metrics. Offline, we use the suboptimality  $\mu_{a^{\star}} - \mu_{\hat{a}}$  where  $\hat{a}$  is the chosen action. Online, we use cumulative regret:  $\sum_{k} \mu_{a^{\star}} - \mu_{\hat{a}_{k}}$  where  $\hat{a}_{k}$  is the  $k$ th action chosen.
DPT learns to reason through uncertainty. As shown in Figure 2a, in the offline setting, DPT significantly exceeds the performance of Emp and LCB while matching the performance of TS, when the in-context datasets are sampled from the same distribution as during pretraining. The results suggest that the transformer is capable of reasoning through uncertainty caused by the noisy rewards in the dataset. Unlike Emp which can be fooled by noisy, undersampled actions, the transformer has learned to hedge to a degree. However, it also suggests that this hedging is fundamentally different from what LCB does, at least on this specific distribution<sup>4</sup>.
Interestingly, the same transformer produces an extremely effective online bandit algorithm when sampling actions instead of taking an argmax. As shown in Figure 2b, DPT matches the performance
See Appendix A.2 for additional details such as hyperparameters.
Note our randomly generated environments are equally likely to have expert-biased datasets and adversarial datasets, so LCB is not expected to outperform here [61].
(c)
Figure 2: (a) Offline performance on in-distribution bandits, given random in-context datasets. (b) Online cumulative regret on bandits. (c) Final (after 500 steps) cumulative regret on out-of-distribution bandits with different Gaussian noise standard deviations. The mean and standard error are computed over 200 test tasks.