SlowGuess/NeurIPS_2023_OCR huggingface dataset:
of classical optimal algorithms, UCB and TS, which are specifically designed for exploration. This is notable because DPT was not explicitly trained to explore, but its emergent strategy is on par with some of the best. In Figure 2c, we show this property is robust to noise in the rewards not seen during pretraining by varying the standard deviation. In Appendix B, we show this generalization happens offline too and even with unseen Bernoulli rewards.
Leveraging latent structure from suboptimal data. We now investigate whether DPT can learn to leverage the inherent structure of a problem class, even without prior knowledge of this structure and even when learning from in-context datasets that do not explicitly utilize it. More precisely, we consider  $\mathcal{T}_{\mathrm{pre}}$  to be a distribution over linear bandits, where the reward function is given by  $\mathbb{E}[r|a,\tau] = \langle \theta_{\tau},\phi (a)\rangle$  and  $\theta_{\tau}\in \mathbb{R}^{d}$  is a task-specific parameter vector and  $\phi :\mathcal{A}\to \mathbb{R}^d$  is fixed feature vector that is the same for all tasks. Given the feature representation  $\phi$ , LinUCB [63], a UCB-style algorithm that leverages  $\phi$ , should achieve regret  $\tilde{\mathcal{O}} (d\sqrt{K})$  over  $K$  steps, a substantial gain over UCB and TS when  $d\ll |\mathcal{A}|$ . We do not assume  $\phi$  is known to DPT, besides whatever can be gleaned from the pretraining data. To evaluate the ability of DPT to uncover latent structure, we collect pretraining data with in-context datasets gathered by standard TS, which does not leverage the linear structure. Figures 3a and 3b show that DPT can exploit the unknown linear structure, implicitly learning a surrogate for  $\phi$ , allowing to do more informed exploration online and decision-making offline. It is nearly on par with LinUCB (which is given  $\phi$ ). We also see that DPT significantly outperforms the source of its pretraining data, TS which did not know or use the latent structure to gather the in-context datasets. These results present evidence that (1) DPT can automatically leverage structure, and (2) supervised learning-based approaches to RL can potentially learn novel exploration strategies that transcend the quality of their pretraining data.
Adapting to expert-biased datasets. A common assumption in offline RL is that datasets tend to be a mixture between optimal data (e.g. expert demonstrations) and suboptimal data (e.g. random interactions) [64]. Hence, LCB is generally effective in practice and the pretraining and testing distributions should be biased towards this setting. Motivated by this, we pretrain another DPT model where  $\mathcal{D}_{\mathrm{pre}}$  is generated by mixing the in-context datasets with varying fractions of expert data, biasing  $\mathcal{D}_{\mathrm{pre}}$  towards datasets that contain more examples of the optimal action. We denote this model by DPT-Exp. In Figure 3c, we plot the test-time performance of both pretrained models when evaluated on new offline datasets with varying percentages of expert data<sup>5</sup>. Our results suggest that when the pretraining distribution is also biased towards expert-suboptimal data, DPT-Exp behaves similarly to LCB, while DPT continues to resemble TS. This is quite interesting as, for other methods such as TS, it is less clear how to automatically incorporate the right amount of expert bias to yield the same effect, but DPT can leverage this from pretraining.
5 Learning in Markov Decision Processes
We next study how DPT can tackle Markov decision processes by testing its ability to perform exploration and credit assignment. In the following experiments, the DPT demonstrates generalization to new tasks, scalability to image-based observations, and capability to stitch in-context behaviors (Section 5.2). This section also examines whether DPT can be pretrained with datasets and action labels generated by a different RL algorithm, rather than the exact optimal policy (Section 5.3).
5That is,  $0\%$  is fully random while  $100\%$  has only optimal actions in the in-context dataset.
Figure 3: (a) Offline performance of DPT trained on linear bandits from TS source data. LinReg does linear regression and outputs the greedy action. (b) Online cumulative regret of the same model. The mean and standard error are computed over 200 test tasks. (c) Offline performance on expert-biased datasets. DPT pretrained on a different prior continues to match TS, but DPT-Exp trained from a more representative prior excels.
(a)
Figure 4: (a) Offline performance on held-out Dark Room goals, given random and expert datasets. (b) Online performance on held-out Dark Room goals. (c) Offline performance on Miniworld (images), given random and expert datasets. (d) Online performance on Miniworld (images) after 40 episodes. We report the average and standard error of the mean over 100 different offline datasets in (a) and (c) and 20 online trials in (b) and (d).
**Environments.** We consider environments that require targeted exploration to solve the task. The first is Dark Room [22, 7], a 2D discrete environment where the agent must locate the unknown goal location in a  $10 \times 10$  room, and only receives a reward of 1 when at the goal. We hold out a set of goals for generalization evaluation. The second is Miniworld [65], a 3D visual navigation problem to test the scalability of DPT to image observations. The agent is in a room with four boxes of different colors, and must find the target box. The target color is unknown initially. It receives a reward of 1 only when near the correct box. Environments and pretraining dataset details are in App. A.4 and A.5.
Comparisons. Our experiments aim to understand the effectiveness of DPT by comparing it to other context-based meta-RL algorithms based on supervised and RL objectives.
- Proximal Policy Optimization (PPO) [66]: We compare to this single-task RL algorithm, which trains from scratch without any pretraining data, to contextualize the performance of DPT and other meta-RL algorithms.
- Algorithm Distillation (AD) [7]: AD first generates a dataset of learning histories by running an RL algorithm in each training task. Then, given a sampled subsequence  $h_j = (s_j, a_j, r_j, \dots, s_{j + c})$  from a learning history, a transformer is trained to predict the next action  $a_{j + c}$  from the history.
-  $\mathrm{RL}^2$  [29]: This online meta-RL comparison uses a recurrent neural network to adapt the agent's policy from the given context. Unlike AD and DPT, which are trained with a supervised objective, the  $\mathrm{RL}^2$  agent is trained to maximize the expected return with PPO.
PPO and  $\mathbf{RL}^2$  are online algorithms, while AD is capable of learning both offline and online. Details on the implementation of these algorithms can be found in Appendix A.2.
Generalizing to new offline datasets and tasks. To study the generalization capabilities of DPT, we evaluate the model in Dark Room on a set of 20 held-out goals not in the pretraining dataset. When given an expert dataset, DPT achieves near-optimal performance. Even when given a random dataset, which has an average total reward of 1.1, DPT obtains a much higher average return of
(b)
Figure 5: (a) In Dark Room (Three Tasks), DPT stitches a new, optimal trajectory to the goal (blue) given two in-context demonstrations of other tasks (pink and orange). (b) Offline Dark Room performance of DPT trained on PPO data. (c) Online Dark Room performance of DPT trained on PPO data.
61.5 (see Fig. 4a). Qualitatively, we observe that when the in-context dataset contains a transition to the goal, DPT immediately exploits this and takes a direct path to the goal. In contrast, while AD demonstrates strong offline performance with expert data, it performs worse in-context learning with random data compared to DPT. The difference arises because AD is trained to infer a better policy than the in-context data, but not necessarily the optimal one.
We next evaluate DPT, AD,  $\mathrm{RL}^2$ , and PPO online without any prior data from the 20 test-time Dark Room tasks, shown in Fig. 4b. After 40 episodes, PPO does not make significant progress towards the goal, highlighting the difficulty of learning from such few interactions alone.  $\mathrm{RL}^2$  is trained to perform adaptation within four episodes each of length 100, and we report the performance after the four adaptation episodes. Notably, DPT on average solves each task faster than AD and reaches a higher final return than  $\mathrm{RL}^2$ , demonstrating its capability to explore effectively online even in MDPs. In Appendix B, we also present results on generalization to new dynamics.
Learning from image-based observations. In Miniworld, the agent receives RGB image observations of  $25 \times 25$  pixels. As shown in Fig. 4d, DPT can solve this high-dimensional task offline from both random and expert datasets. Compared to AD and  $\mathrm{RL}^2$ , DPT also learns online more efficiently.
Stitching novel trajectories from in-context subsequences. A desirable property of some offline RL algorithms is the ability to stitch suboptimal subsequences from the offline dataset into new trajectories with higher return. To test whether DPT exhibits stitching, we design the Dark Room (Three Tasks) environment in which there are three possible tasks. The pretraining data consists only of expert demonstrations of two of them. At test-time DPT is evaluated on third unseen task, but its offline dataset is only expert demonstrations of the original two. Despite this, it leverages the data to infer a path solving the third task (see Fig. 5a).
5.3 Learning from Algorithm-Generated Policies and Rollouts
So far, we have only considered action labels provided by an optimal policy. However, in some tasks, an optimal policy is not readily available even in pretraining. In this experiment, we use actions labeled by a policy learned via PPO and in-context datasets sampled from PPO replay buffers. We train PPO agents in each of the 80 train tasks for 1K episodes to generate 80K total rollouts, from which we sample the in-context datasets. This variant, DPT (PPO, PPO), performs on par with DPT and still better than AD, as shown in Figures 5b and 5c. DPT (PPO, PPO) can be viewed as a direct comparison between our pretraining objective and that of AD, given the same pretraining data but just used differently. We also evaluated a variant, DPT (Rand, PPO), which pretrains on random in-context datasets (like DPT), but still using PPO action labels. The performance is worse than other DPT variants in some settings, but only marginally so. In Appendix B, we analyze the sensitivity of DPT to other hyperparameters, such as context size and amount of pretraining data. These results give further evidence that DPT can learn an algorithm that transcends its pretraining data.
6 Theory
We now shed light on the observations of the previous empirical results through a theoretical analysis. Our main result shows that DPT (under a slight modification to pretraining) essentially is capable of executing posterior sampling (PS) reinforcement learning in-context. PS is a generalization of
Thompson Sampling for RL in MDPs [12]. It maintains and samples from a posterior over tasks  $\tau$  given historical data  $D$  and executes optimal policies  $\pi_{\tau}^{\star}$  (see Appendix C for a formal outline). It is provably sample-efficient with online Bayesian regret guarantees [13], but maintaining posteriors is generally computationally intractable. The ability for DPT to perform PS in-context suggests a path towards computation- and provably sample-efficient RL with priors learned from the data.
6.1 History-Dependent Pretraining and Assumptions
We start with a modification to the pretraining of DPT. Rather than conditioning only on  $s_{\mathrm{query}}$  and  $D$  to predict  $a^{\star} \sim \pi_{\tau}^{\star}(\cdot |s_{\mathrm{query}})$ , we propose also conditioning on a sequence  $\xi_h = (s_{1:h}, a_{1:h}^{\star})$  where  $s_{1:h} \sim \mathfrak{S}_h \in \Delta(S^h)$  is a distribution over sets of states, independent of  $\tau$ , and  $a_{h'}^{\star} \sim \pi_{\tau}^{\star}(\cdot |s_{h'})$  for  $h' \in [h]$ . Thus, we use  $\pi_{\tau}^{\star}$  to label both the query state (which is the prediction label) and the sequence of states sampled from  $\mathfrak{S}_h$ . Note that this does not require any environment interactions and hence no sampling from either  $T_{\tau}$  or  $R_{\tau}$ . At test-time at step  $h$ , this will allow us to condition on the history  $\xi_{h-1}$  of states that  $M_{\theta}$  visits and the actions that it takes in those states. Formally, the learned  $M_{\theta}$  is deployed as follows, given  $D$ . (1) At  $h = 0$ , initialize  $\xi_0 = (\cdot)$  to be empty. (2) At step  $h$ , visit  $s_h$  and find  $a_h$  by sampling from  $M_{\theta}(\cdot | s_{\mathrm{query}}, D, \xi_{h-1})$ . (3) Append  $(s_h, a_h)$  to  $\xi_{h-1}$  to get  $\xi_h$ . Note for bandits and contextual bandits  $(H = 1)$ , there is no difference between this and the original pretraining procedure of prior sections because  $\xi_0$  is empty. For MDPs, the original DPT can be viewed as a practically convenient approximation.
We now make several assumptions to simplify the analysis. First, assume  $\mathcal{D}_{\mathrm{query}}$ ,  $\mathcal{D}_{\mathrm{pre}}$ , and  $\mathfrak{S}$  have sufficient support such that all conditional probabilities of  $P_{pre}$  are well defined. Similar to other studies of in-context learning [67], we assume  $M_{\theta}$  fits the pretraining distribution exactly with enough coverage and data, so that the focus of the analysis is just the in-context learning abilities.
Assumption 1. (Learned model is consistent). Let  $M_{\theta}$  denote the pretrained model. For all  $(s_{query}, D, \xi_h)$ , we have  $P_{pre}(a | s_{query}, D, \xi_h) = M_{\theta}(a | s_{query}, D, \xi_h)$  for all  $a \in \mathcal{A}$ .
To provide some cursory justification, if  $M_{\theta}$  is the global minimizer of (2), then we might expect that  $\mathbb{E}_{P_{pre}}\| P_{pre}(\cdot |s_{query},D,\xi_h) - M_\theta (\cdot |s_{query},D,\xi_h)\| _1^2\to 0$  as the number of pretraining samples  $N\rightarrow \infty$  with high probability for transformer model classes of bounded complexity (see Proposition C.1). Approximate versions of the above assumptions are easily possible but obfuscate the key elements of the analysis. We also assume that the in-context dataset  $D\sim \mathcal{D}_{\mathrm{pre}}$  is compliant [62], meaning that the actions from  $D$  can depend only on the observed history and not additional confounders. Note that this still allows  $\mathcal{D}_{\mathrm{pre}}$  to be very general — it could be generated randomly or from adaptive algorithms like PPO or TS. Without compliance  $\mathcal{D}_{\mathrm{pre}}$  can influence  $M_{\theta}$ . In Proposition 6.4, we show that all compliant  $\mathcal{D}_{\mathrm{pre}}$  form a sort of equivalence class that generate the same  $M_{\theta}$ . For the remainder, we assume all  $\mathcal{D}_{\mathrm{pre}}$  are compliant.
Definition 6.1 (Compliance). The in-context dataset distribution  $\mathcal{D}_{pre}(\cdot ;\tau)$  is compliant if, for all  $i\in [n]$ , the  $i$ th action of the dataset,  $a_{i}$ , is conditionally independent of  $\tau$  given the  $i$ th state  $s_i$  and partial dataset,  $D_{i - 1}$ , so far. In other words, the distribution  $\mathcal{D}_{pre}(a_i|s_i,D_{i - 1};\tau)$  is invariant to  $\tau$ .
Equivalence of DPT and PS. We now state our main result which shows that the trajectories generated by a pretrained  $M_{\theta}$  will follow the same distribution as those from a well-specified PS algorithm. In particular, let PS use the well-specified prior  $\mathcal{T}_{\mathrm{pre}}$ . Let  $\tau_{c}$  be an arbitrary task. Let  $P_{ps}(\cdot |D,\tau_c)$  and  $P_{M_{\theta}}(\cdot |D,\tau_c)$  denote the distributions over trajectories  $\xi_H\in (S\times \mathcal{A})^H$  generated from running PS and  $M_{\theta}(\cdot |\cdot ,D,\cdot)$ , respectively, in task  $\tau_{c}$  given historical data  $D$ .
Theorem 1 (DPT  $\iff$  PS). Let the above assumptions hold. Then,  $P_{ps}(\xi_H\mid D,\tau_c) = P_{M_\theta}(\xi_H\mid D,\tau_c)$  for all trajectories  $\xi_H$ .
Regret implications. With almost no extra work besides verifying assumptions, we can see this result in action. Let us specialize to the finite MDP setting [13]. Suppose we pretrain  $M_{\theta}$  on a distribution  $\mathcal{T}_{\mathrm{pre}}$  over MDPs with  $S\coloneqq |\mathcal{S}|$  and  $A\coloneqq |\mathcal{A}|$ . Let  $\mathcal{D}_{\mathrm{pre}}$  be constructed by uniform sampling  $(s_i,a_i)$  and observing  $(r_i,s_i')$  for  $i\in [KH]$ . Let  $\mathbb{E}[r_h|s_h,a_h]\in [0,1]$ . And let  $\mathcal{D}_{\mathrm{query}}$  and  $\mathfrak{S}_h$  be uniform over  $\mathcal{S}$  and  $S^h$  (for all  $h$ ) respectively. Finally, let  $\mathcal{T}_{\mathrm{test}}$  be the distribution over test
It is still possible to implement this version, but it leaves some undesirable design choices to be made.
tasks with the same cardinalities. For a task  $\tau$ , define the online cumulative regret of DPT over  $K$  episodes as  $\mathrm{Reg}_{\tau}(M_{\theta}) := \sum_{k \in [K]} V_{\tau}(\pi_{\tau}^{\star}) - V_{\tau}(\hat{\pi}_k)$  where  $\hat{\pi}_k(\cdot | s_h) = M_{\theta}(\cdot | s_h, D_{(k-1)}, \xi_{h-1})$  and  $D_{(k)}$  contains the first  $k$  episodes collected from  $\hat{\pi}_{1:k}$ .
Corollary 6.2 (Finite MDPs). Suppose that  $\sup_{\tau} \mathcal{T}_{test}(\tau) / \mathcal{T}_{pre}(\tau) \leq \mathcal{C}$  for some  $\mathcal{C} > 0$ . For the above MDP setting, the pretrained model  $M_{\theta}$  satisfies  $\mathbb{E}_{\mathcal{T}_{test}}[Reg_{\tau}(M_{\theta})] \leq \widetilde{\mathcal{O}}(\mathcal{CH}^{3/2}S\sqrt{AK})$ .
A similar analysis due to [68] allows us to see theoretically why pretraining on (latently) linear bandits can lead to an in-context algorithm competitive with linear bandit algorithms, even when the in-context datasets are generated by algorithms unaware of this structure. We observed this empirically in Section 4. Consider a similar setup as there where  $S$  is a singleton,  $\mathcal{A}$  is finite but large,  $\theta_{\tau} \in \mathbb{R}^{d}$  is sampled as  $\theta_{\tau} \sim \mathcal{N}(0, I / d)$ ,  $\phi : \mathcal{A} \to \mathbb{R}^{d}$  is a fixed feature map with  $\sup_{a \in \mathcal{A}} \| \phi(a) \|_{2} \leq 1$ , and the reward of  $a \in \mathcal{A}$  in task  $\tau$  is distributed as  $\mathcal{N}(\langle \theta_{\tau}, \phi(a) \rangle, 1)$ . This time, we let  $D_{\mathrm{pre}}(\cdot; \tau)$  be given by running Thompson Sampling with Gaussian priors and likelihood functions on  $\tau$ .
Corollary 6.3 (Latent representation learning in linear bandits). For  $\mathcal{T}_{test} = \mathcal{T}_{pre}$  in the above linear bandit setting,  $M_{\theta}$  satisfies  $\mathbb{E}_{\mathcal{T}_{test}}[Reg_{\tau}(M_{\theta})] \leq \widetilde{\mathcal{O}}(d\sqrt{K})$ .
This significantly improves over the  $\tilde{\mathcal{O}} (\sqrt{|A|K})$  upper regret bound for TS that does not leverage the linear structure. This highlights how DPT can have provably tighter upper bounds on future bandit problems than the algorithms used to generate its (pretraining) data. Note that if there is additional structure in the tasks which yields a tighter regret bound (for example if there are only a small finite number of MDPs in the possible distribution), that may further improve performance, such as by removing the dependence on the state, action, or dimension in these cases.
Invariance of  $M_{\theta}$  to compliant  $\mathcal{D}_{\mathrm{pre}}$ . Our final result sheds light on how  $\mathcal{D}_{\mathrm{pre}}$  impacts the final DPT behavior  $M_{\theta}$ . Combined with Assumption 1,  $M_{\theta}$  is invariant to  $\mathcal{D}_{\mathrm{pre}}$  satisfying Definition 6.1.
Proposition 6.4. Let  $P_{pre}^{1}$  and  $P_{pre}^{2}$  be pretraining distributions that differ only by their in-context dataset distributions, denoted by  $\mathcal{D}_{pre}^{1}$  and  $\mathcal{D}_{pre}^{2}$ . If  $\mathcal{D}_{pre}^{1}$  and  $\mathcal{D}_{pre}^{2}$  are compliant with the same support, then  $P_{pre}^{1}(a^{\star}|s_{query}, D, \xi_{h}) = P_{pre}^{2}(a^{\star}|s_{query}, D, \xi_{h})$  for all  $a^{\star}, s_{query}, D, \xi_{h}$ .
That is, if we generate in-context datasets  $D$  by running various algorithms that depend only on the observed data in the current task, we will end up with the same  $M_{\theta}$ . For example, TS could be used for  $\mathcal{D}_{\mathrm{pre}}^{1}$  and PPO for  $\mathcal{D}_{\mathrm{pre}}^{2}$ . Expert-biased datasets discussed in Section 4 violate Definition 6.1, since privileged knowledge of  $\tau$  is being used. This helps explain our empirical results that pretraining on expert-biased datasets leads to a qualitatively different learned model at test-time (Fig. 3c).
We studied the problem of in-context RL. We introduced a new pretraining method and transformer model, DPT, which is trained via supervised learning to predict optimal actions given an in-context dataset of interactions. Through in-depth evaluations in classic decision problems in bandits and MDPs, we showed that this simple objective naturally gives rise to an in-context RL algorithm that is capable of online exploration and offline decision-making in a way that leverages structure in the pretraining data. In some cases, the RL algorithm learned can transcend the quality of the pretraining data. Our empirical and theoretical results provide first steps towards understanding these capabilities and what factors are important for success. Via the simplicity of pretraining, we can sidestep the complexities of hand-designing exploration or conservatism in RL algorithms while simultaneously allowing the transformer to derive novel few-shot strategies. These findings underscore the potential of supervised pretraining to equip transformer models with in-context RL abilities.
Limitations and future work. One limitation of DPT is the requirement of optimal actions at pretraining. Empirically, we find that this requirement can be relaxed by using actions generated by another RL-trained agent during pretraining, which only leads to a slight loss in performance. However, fully understanding this problem and how best to leverage multi-task decision-making datasets remains a key open problem. We also discussed that the practical implementation for MDPs differs from true posterior sampling. It would be interesting to further understand and bridge this empirical-theoretical gap in the future. Finally, further investigation is required to understand the implications of these findings for existing transformer-based large language models that are increasingly being deployed in decision-making settings [69].
We thank Yu Bai, Anirudh Goyal, Akshay Krishnamurthy, Evan Liu, Subhojyoti Mukherjee, Tung Nguyen, Lucy Shi, Ben Van Roy, and Sherry Yang for helpful discussions, feedback, and suggestions on related work. We also thank anonymous reviewers for their careful and constructive feedback. This work was supported in part by NSF grant 2112926 and ONR grant N00014-21-1-2685. JNL acknowledges support from the NSF GRFP. AP would like to thank the support of the Eric and Wendy Schmidt Center at the Broad Institute of MIT and Harvard. This work was supported in part by funding from the Eric and Wendy Schmidt Center at the Broad Institute of MIT and Harvard.