SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[2] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837, 2022.
[3] Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878-18891, 2022.
[4] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598, 2022.
[5] Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot reasoning. arXiv preprint arXiv:2202.07206, 2022.
[6] Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
[7] Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14215, 2022.
[8] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In International Conference on Machine Learning, pages 24631-24645. PMLR, 2022.
[9] Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyperdecision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487, 2023.
[10] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[11] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[12] Malcolm Strens. A bayesian framework for reinforcement learning. In ICML, volume 2000, pages 943-950, 2000.
[13] Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
[14] Tom Schaul and Jürgen Schmidhuber. Metalearning. *Scholarpedia*, 5(6):4650, 2010.
[82] Xiuyuan Lu and Benjamin Van Roy. Ensemble sampling. Advances in neural information processing systems, 30, 2017.
[83] Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
[84] Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems, 31, 2018.
[85] Ian Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy. Approximate thompson sampling via epistemic neural networks. arXiv preprint arXiv:2302.09205, 2023.
[86] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[87] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
[88] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. The Journal of Machine Learning Research, 22(1):12348-12355, 2021.
[89] Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095-20107, 2020.
Additional Related Work
In-context learning. Beyond decision-making and reinforcement learning, our approach takes inspiration from general in-context learning, a phenomenon observed most prominently in large language models in which large-scale autoregressive modelling can surprisingly lead to a model that exhibits meta-learning capabilities [1]. Recently, there has been great interest in understanding the capabilities and properties of in-context learning in various settings [4, 70, 5, 71, 72, 73, 74, 75, 6, 76, 77, 78]. While a common hypothesis suggests that this phenomenon is due to properties of the data used to train large language models [3], our work suggests that this phenomenon can also be encouraged in general settings via adjustments to the pre-training objective. In fact, DPT could be interpreted as explicitly encouraging the ability to perform Bayesian inference, which is a popular explanation for the mechanism behind in-context learning for large language models [67]. For meta-learning in-context, work by [79] proposed Transformer Neural Processes to address problems traditionally addressed with neural processes. Their work focuses on supervised learning-style input-output pairs for meta-learning, but they also applied this to contextual bandits to learn rewards.
For decision-making, in addition to the works discussed in the main paper, several works have applied transformers for in-context RL. AdA [32] meta-learns a transformer online (like  $\mathrm{RL}^2$ ). A base online RL algorithm is used to optimize the transformer, which uses online access to the simulator. The authors also use a curriculum learning method to improve the learning process. DPT is learned from offline supervised pretraining and it is agnostic to how the pretraining data is acquired. We believe many of the ideas in AdA are complementary and can be ported to DPT. Prompt-DT [8] is designed for offline RL and conditions on expert demos and uses return-to-go in the context much like the Decision Transformer [37]. This necessitates expert demos as input at test-time. DPT is meant for both online and offline in-context RL and does not require test-time demos (but it can benefit from them e.g. Fig. 4a). DPT does not use return-to-go, and instead predicts optimal actions from arbitrary in-context interaction datasets.
Posterior Sampling. Posterior sampling originates from the seminal work of [59], and has been popularized and thoroughly investigated in recent years by a number of authors [58, 80, 12, 13, 81, 68]. For bandits, it is often referred to as Thompson Sampling, but the framework is easily generalizable to RL, as shown in [12, 13]. The principle is as follows: begin with a prior over possible models (i.e. reward and transition functions), and maintain a posterior distribution over models by updating as new interactions are made. At decision-time, sample a model from the posterior and execute its optimal policy. The aforementioned prior works have developed strong theoretical guarantees on Bayesian and frequentist regret for posterior sampling. Despite its desirable theoretical characteristics, a major limitation is that computing the posterior is often computationally intractable. Recently, the community has made efforts towards scalable posterior sampling often by leveraging approximations [82, 83, 84, 85]. In Section 6, we show that a version of the DPT model learned from pretraining can be viewed as implementing posterior sampling without specifying complicated priors or posterior updates. Instead, the posterior update is implicitly learned through pretraining to predict the optimal action. This suggests that in-context learning via supervised pretraining (or meta-learning more generally) could be a promising alternative approach towards scalable posterior sampling.
On the theoretical side, [19] show it is possible to meta-learn prior distribution parameters for Thompson Sampling (TS) using a meta-exploration policy. The resulting TS algorithm can then be deployed on new problems. This is potentially consistent with the observation that DPT can approximate the right prior from offline data.
A Implementation and Experiment Details
A.1 DPT Architecture: Formal Description
In this section, we provide a detailed description of the architecture alluded to in Section 3 and Figure 1. See hyperparameter details for models in their respective sections. The model is implemented in Python with PyTorch [86]. The backbone of the transformer architecture we use is an autoregressive GPT-2 model from the HuggingFace transformers library.
For the sake of exposition, we suppose that  $S$  and  $\mathcal{A}$  are subsets of  $\mathbb{R}^{d_S}$  and  $\mathbb{R}^{d_A}$  respectively. We handle discrete state and action spaces with one-hot encoding. Consider a single training datapoint
Algorithm 2 Decision-Pretrained Transformer (detailed)
1: // Collecting pretraining dataset  
2: Initialize empty dataset  $\mathcal{B}$   
3: for  $i$  in  $[N]$  do  
4: Sample task  $\tau \sim \mathcal{T}_{\mathrm{pre}}$   
5: Sample interaction dataset  $D \sim \mathcal{D}_{\mathrm{pre}}(:, \tau)$  of length  $n$   
6: Sample  $s_{\mathrm{query}} \sim \mathcal{D}_{\mathrm{query}}$  and  $a^{\star} \sim \pi_{\tau}^{\star}(\cdot | s_{\mathrm{query}})$   
7: Add  $(s_{\mathrm{query}}, D, a^{\star})$  to  $\mathcal{B}$   
8: end for  
9: // Training model on dataset  
10: Initialize model  $M_{\theta}$  with parameters  $\theta$   
11: while not converged do  
12: Sample  $(s_{\mathrm{query}}, D, a^{\star})$  from  $\mathcal{B}$   
13: Predict  $\hat{p}_j(\cdot) = M_\theta (\cdot | s_{\mathrm{query}}, D_j)$  for all  $j \in [n]$ .  
14: Compute loss in (5) with respect to  $a^{\star}$  and backpropagate to update  $\theta$ .  
15: end while
Algorithm 3 Offline test-time deployment (detailed)
1: // Task and offline dataset are generated without learner's control  
2: Sample unknown task  $\tau \sim \mathcal{T}_{\mathrm{test}}$   
3: Sample dataset  $D \sim \mathcal{D}_{\mathrm{test}}(\cdot; \tau)$   
4: // Deploying offline policy  $M_{\theta}(\cdot|\cdot, D)$   
5:  $s_1 = \text{reset}(\tau)$   
6: for  $h$  in  $[H]$  do  
7:  $a_h = \operatorname{argmax}_{a \in \mathcal{A}} M_{\theta}(\cdot|s_h, D) // \text{Most likely action}$   
8:  $s_{h+1}, r_h = \text{step}(\tau, a_h)$   
9: end for
Algorithm 4 Online test-time deployment (detailed)
1: // Online, dataset is empty as learning is from scratch  
2: Initialize  $D = \{\}$   
3: Sample unknown task  $\tau \sim \mathcal{T}_{\mathrm{test}}$   
4: for ep in max_epochs do  
5:  $s_1 = \text{reset}(\tau)$   
6: for h in [H] do  
7:  $a_h \sim M_\theta(\cdot | s_h, D)$  // Sample action from predicted distribution  
8:  $s_{h+1}, r_h = \text{step}(\tau, a_h)$   
9: end for  
10: // Experience from previous episode added to dataset  
11: Add  $(s_1, a_1, r_1, \ldots)$  to  $D$   
12: end for
derived from an (potentially unknown) task  $\tau$ : we have a dataset  $D$  of interactions within  $\tau$ , a query state  $s_{\mathrm{query}}$ , and its corresponding optimal action  $a^{\star} = \pi_{\tau}^{\star}(s_{\mathrm{query}})$ . We construct the embeddings to be passed to the GPT-2 backbone in the following way. From the dataset  $D = \{(s_j, a_j, s_j', r_j)\}_{j \in [n]}$ , we construct vectors  $\xi_j = (s_j, a_j, s_j', r_j)$  by stacking the elements of the transition tuple into dimension  $d_{\xi} := 2d_S + d_A + 1$  for each  $j$  in the sequence. This sequence of  $n$  elements is concatenated with another vector  $v := (s_{\mathrm{query}}, \mathbf{0})$  where the  $\mathbf{0}$  vector is a vector of zeros of sufficient length to make the entire element dimension  $d_{\xi}$ . The  $(n + 1)$ -length sequence is given by  $X = (v, \xi_1, \ldots, \xi_n)$ . As order does not often matter for the dataset  $D^7$ , we do not use positional encoding in order to take advantage of this invariance. We first apply a linear layer Linear  $(X)$  and pass the result to the transformer, which outputs the sequence  $Y = (\hat{y}_0, \hat{y}_1, \ldots, \hat{y}_n)$ . In the continuous action case, these can be used as is for predictions of  $a^{\star}$ . For the discrete action case, we use them as logits to be converted to either a distribution over actions in  $\mathcal{A}$  or one-hot vector predictions of  $a^{\star}$ . Here, we compute action probabilities
$$
\hat {p} _ {j} = \operatorname {s o f t m a x} \left(\hat {y} _ {j}\right) \in \Delta (\mathcal {A}) \tag {3}
$$
Because of the GPT-2 causal architecture (we defer details to the original papers [87, 1]), we note that  $\hat{p}_j$  depends only on  $s_{\mathrm{query}}$  and the partial dataset  $D_{j} = \{(s_{k},a_{k},s_{k}^{\prime},r_{k})\}_{k\in [j]}$ , which is why we write the model notation,
$$
M _ {\theta} (\cdot | s _ {\text {q u e r y}}, D _ {j}) = \hat {p} _ {j} (\cdot), \tag {4}
$$
to denote that the predicted probabilities of the  $j$ th element only depend on  $D_{j}$  and not the entire  $D$  for the model  $M$  with parameters  $\theta \in \Theta$ . For example, with  $j = 0$ , the prediction of  $a^{\star}$  is made without any contextual information about the task  $\tau$  except for  $s_{\mathrm{query}}$ , which can be interpreted as the prior over  $a^{\star}$ . We measure loss of this training example via the cross entropy for each  $j \in [n]$ :
$$
- \sum_ {j \in [ n ]} \log \hat {p} _ {j} \left(a ^ {\star}\right) \tag {5}
$$
Intuition. Elements of the inputs sequence  $X$  represent transitions in the environment. When passed through the GPT-2 transformer, the model learns to associate elements of the sequence via the standard query-key-value mechanism of the attention model. The query state  $s_{\mathrm{query}}$  is demarcated by its zeros vector (which also acts as padding). Unlike other examples of transformers used for decision-making such as the Decision Transformer [37] and Algorithm Distillation [7], DPT does not separate the individual  $(s, a, s', r)$  into their own embeddings to be made into one long sequence. This is because we view the transition tuples in the dataset as their own singletons, to be related with other singletons in the dataset through the attention mechanism. We note that there are various other implementation variations one could take, but we found success and robustness with this one.
A.2.1 Bandit algorithms
First, we describe the comparisons from the bandit experiments with hyperparameters.
Empirical Mean (Emp). Emp has no hyperparameters, but we give it some mechanism to avoid degenerate scenarios. In the offline setting, Emp will only choose from actions that have at least one example in the dataset. This gives Emp and LCB-style effect when actions are missing. Similarly, online, Emp will sample each action at least once before defaulting to its real strategy. These changes only improve Emp.
Upper Confidence Bound (UCB). According to the Hoeffding bound, we choose actions as  $\hat{a} \in \operatorname{argmax}_{a \in \mathcal{A}} \left\{\hat{\mu}_a + \sqrt{1 / n_a}\right\}$  where  $\hat{\mu}_a$  is the empirical mean so far for action  $a$  and  $n_a$  is the number of times  $a$  has been chosen so far. To arrive at this constant for the bonus, we coarsely tried a set of plausible values given the noise and found this to perform the best.
Lower Confidence Bound (LCB). We choose actions as  $\hat{a} \in \operatorname{argmax}_{a \in \mathcal{A}} \left\{\hat{\mu}_a - \sqrt{1 / n_a}\right\}$  where  $\hat{\mu}_a$  is the empirical mean so far for action  $a$  and  $n_a$  is the number of times  $a$  has been chosen so far.
This is not always true such as when data comes from an algorithm such as PPO or Thompson Sampling.
Thompson Sampling (TS). Since the means are sampled uniformly from  $[0,1]$ , Gaussian TS is partially misspecified; however, we set prior mean and variance to  $\frac{1}{2}$  and  $\frac{1}{12}$  to match the true ones. The noise model was well-specified with the correct variance. In the linear experiments of Figure 3a and Figure 3b, we set the prior mean and variance to 0 and 1 to fit the true ones better.
LinUCB. We choose  $\hat{a}_t \in \operatorname{argmax}_{a \in \mathcal{A}} \langle \hat{\theta}_t, \phi(a) \rangle + \beta \| \phi(a) \|_{\hat{\Sigma}_t^{-1}}$  where  $\beta = 1$  and  $\hat{\Sigma}_t = I + \sum_{s \in [t-1]} \phi(a_s) \phi(a_s)^\top$  and  $\hat{\theta}_t = \hat{\Sigma}_t^{-1} \sum_{s \in [t-1]} r_s \phi(a_s)$ . Here,  $r_s$  and  $a_s$  are the reward and action observed at time  $s$ .
LinReg. LinReg (offline) is the same as LinUCB except we set  $\beta = 0$  to greedily choose actions.
DPT. The transformer for DPT has an embedding size of 32, context length of 500 for basic bandits and 200 for linear bandits, 4 hidden layers, and 4 attention heads per attention layer for all bandits. We use the AdamW optimizer with weight decay 1e-4, learning rate 1e-4, and batch-size 64. For all experiments, we shuffle the in-context dataset  $D$  since order does not matter except in the linear bandit.
A.2.2 RL Algorithms
Below, we describe the comparisons from the MDP experiments and their hyperparameters.
Proximal Policy Optimization (PPO). The reported results for PPO use the Stable Baselines3 implementation [88] with the default hyperparameters, which successfully learns each task given 100K environment steps in Dark Room and 125K environment steps in Miniworld. In Dark Room, the policy is implemented as a multi-layer perceptron with two hidden layers of 64 units each. In Miniworld, the policy is a convolutional neural network with two convolutional layers with  $16 \times 3$  kernels each, followed by a linear layer with output dimension of 8.
Algorithm Distillation (AD). We first collect learning histories with PPO for each of the training tasks. Then, given a cross-episodic context of length  $H$ , where  $H$  is the task horizon, the model is trained to predict the actions taken  $K$  episodes later (given the states visited in that episode). This was shown to lead to faster algorithms in [7]. We evaluated AD across different values of  $K$ . Between  $K = 10, 50, 100$ , we found  $K = 100$  to be most performant in the Dark Room environment. In Miniworld, we also subsampled with  $K = 100$ . In Dark Room, the transformer has similar hyperparameters as DPT: an embedding size of 32, context length of 100 steps, 4 hidden layers, and 4 attention heads per attention layer. In Miniworld, as with DPT, we first encode the image with a convolutional network with two convolutional layers with  $16 \times 3$  kernels each, followed by a linear layer with output dimension of 8.
$\mathbf{RL}^2$ . The reported results for  $\mathrm{RL}^2$  use an open-sourced implementation from [22]. The implementation uses PPO as the RL algorithm and defines a single trial as four consecutive episodes. The policy is implemented with one hidden layer of 32 units in Dark Room. In Miniworld, the policy is parameterized with a convolutional neural network with two convolutional layers with  $163\times 3$  kernels each, followed by a linear layer with output dimension of 8.
DPT. The transformer for DPT has an embedding size of 32, context length of 100 steps, 4 hidden layers, and 4 attention heads per attention layer in Dark Room. In Miniworld, the image is first passed through a convolutional network with two convolutional layers  $16 \times 3$  kernels each, followed by a linear layer with output dimension of 8. The transformer model that processes these image embeddings otherwise has the same hyperparameters as in Dark Room. We use the AdamW optimizer with weight decay 1e-4, learning rate 1e-3, and batch-size 128.
A.3 Bandit Pretraining and Testing
Basic Bandit. Offline, to generate the in-context datasets for pretraining, we used a Dirichlet distribution to sample action frequencies in order to generate datasets with diverse compositions (i.e. some more uniform, some that only choose a few actions, etc.):  $p_1 \sim \mathrm{Dir}(\mathbb{1})$  where  $p_1 \in \Delta(\mathcal{A})$  and  $\mathbb{1} \in \mathbb{R}^{|\mathcal{A}|}$ . We also mixed this with a distribution that has all mass on one action:  $\hat{a} \sim \mathrm{Unif}(\mathcal{A})$  and  $p_2(\hat{a}) = 1$  and  $p_2(a) = 0$  for all  $a \neq \hat{a}$ . The final action distribution is  $p = (1 - \omega)p_1 + \omega p_2$  where  $\omega \sim \mathrm{Unif}(0.1[10])$ . We train on 100,000 pretraining samples for 300 epochs with an 80/20 train/validation split. In Figure 2a,  $D_{\mathrm{test}}$  is generated in the same way.
Expert-Biased Bandit. To generate expert-biased datasets for pretraining, we compute the action frequencies to bias the dataset towards the optimal action. Let  $a^{\star}$  be the optimal one. As before, we take  $p_1 \sim \mathrm{Dir}(\mathbb{1})$ . Then,  $p_2(a^\star) = 1$  and  $p_2(a) = 0$  for all  $a \neq a^{\star}$ . For of bias of  $\omega$ , we take  $p = (1 - \omega)p_1 + \omega p_2$  with  $\omega \sim \mathrm{Unif}(0.1[10])$ . We use the same pretraining sample size and epochs as before. For testing,  $D_{\mathrm{test}}$  is generated the same way except we fix a particular  $\omega \in \{0, 0.5, 1\}$  to test on.
Linear Bandit. We consider the case where  $|\mathcal{A}| = 10$  and  $d = 2$ . To generate environments from  $\mathcal{T}_{\mathrm{pre}}$ , we first sampled a fixed set of actions from  $\mathcal{N}(\mathbf{0}, I_d / d)$  in  $\mathbb{R}^d$  to represent the features. Then, for each  $\tau$ , we sampled  $\theta_{\tau} \sim \mathcal{N}(\mathbf{0}, I_d / d)$  to produce the means  $\mu_a = \langle \theta_{\tau}, \phi(a) \rangle$  for  $a \in \mathcal{A}$ . To generate the in-context dataset, we ran Gaussian TS (which does not leverage  $\phi$ ) over  $n = 200$  steps (see hyperparameters in previous section). Because order matters, we did not shuffle and used 1,000,000 pretraining samples over 200 epochs with an 80/20 train/validation split. At test time, we set  $\mathcal{T}_{\mathrm{test}} = \mathcal{T}_{\mathrm{pre}}$  and  $\mathcal{D}_{\mathrm{test}} = \mathcal{D}_{\mathrm{pre}}$ . Note that  $\phi$  is fixed over all  $\tau$ , as is standard for a linear bandit.
A.4 MDP Environment Details
Dark Room. The agent must navigate a  $10 \times 10$  grid to find the goal within  $H = 100$  steps. The agent's observation is its  $xy$ -position, the allowed actions are left, right, up, down, and stay, and the reward is only  $r = 1$  when the agent is at the goal, and  $r = 0$  otherwise. At test time, the agent begins at the  $(0,0)$  position. We randomly designate 80 of the 100 grid squares to be goals for the training tasks, and hold out the remaining 20 for evaluation.
Miniworld. The agent must navigate to the correct box, which is initially unknown, from  $25 \times 25$  RGB image observations. The agent is additionally conditioned on its own direction vector. In each episode, the environment is initialized with four boxes of different colors, one in each corner of the square room. The agent can turn left, turn right, or move forward. The reward is only  $r = 1$  when the agent is near the correct box and  $r = 0$  otherwise, and each episode is 50 time-steps long. At test time, the agent begins in the middle of the room.
A.5 MDP Pretraining Datasets
Dark Room. In Dark Room, we collect 100K in-context datasets, each of length  $H = 100$  steps, with a uniform-random policy. The 100K datasets are evenly collected across the 100 goals. The query states are uniformly sampled from the state space, and the optimal actions are computed as follows: move up/down until the agent is on the same  $y$ -position as the goal, then move left/right until the agent is on the  $x$ -position as the goal. Of the 100K collections of datasets, query states, and optimal actions, we use the first 80K (corresponding to the first 80 goals) for training and the remaining 20K for validation.