Supervised Pretraining Can Learn In-Context Reinforcement Learning

Introduction

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 . 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 and a growing number works are beginning to understand and analyze in-context learning for supervised learning . 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.

For in-context decision-making , rather than input-output tuples, the context takes the form of state-action-reward tuples representing a dataset of interactions with an unknown environments. 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 . In contrast, an RL agent with access to only a suboptimal offline dataset should produce a policy that conservatively selects actions . 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 labelIf not explicitly known, the optimal action can be determined by running any (potentially inefficient) minimax-optimal regret algorithm for each pretraining task. 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 . 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, but 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 , 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 help unlock practical and efficient RL via posterior sampling.

Related Work

Meta-learning. Algorithmically, in-context learning falls under the meta-learning framework . 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 a often choice in what shared ‘structure’ is specifically learned such as the dynamics of the task , a task context identifier , temporally extended skills and options , or initialization of a neural network policy ). In-context learning can be viewed as taking a more agnostic approach by learning the learning algorithm itself, more similar to . Algorithm Distillation (AD) also falls under this category, applying autoregressive supervised learning to distill (sub-sampled) 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 pretrain DPT to predict optimal actions, yielding potentially emergent online and offline 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 , inspired by the successes of these techniques for large language models . For example, Decision Transformer (DT) 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 , at times exceeding the performance of large-scale multi-task imitation learning with transformers . However, DT is known to be provably (and unboundedly) sub-optimal in common scenarios . 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 . 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 , or policy regularization . 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 . However, they still must address distribution shift, requiring solutions such as policy regularization or additional online interactions . 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.

In-Context Learning Model

Pretraining. We give pseudocode in Algorithm 1 and a visualization in Figure 1. Let ${\mathcal{T}}_{\text{pre}}$ be a distribution over tasks at the time of pretraining. A task $\tau\sim{\mathcal{T}}_{\text{pre}}$ can be viewed as a specification of an MDP, $\tau=\langle{\mathcal{S}},\mathcal{A},T,R,H,\rho\rangle$. The distribution ${\mathcal{T}}_{\text{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}_{\text{pre}}(\cdot;\tau)$ of interactions between the learner and the MDP specified by $\tau$. $D=\{s_{j},a_{j},s^{\prime}_{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_{\text{query}}$ from the distribution $\mathcal{D}_{\text{query}}$ over states ${\mathcal{S}}$ and a label $a^{\star}$ is sampled from the optimal policy $\pi^{\star}_{\tau}(\cdot|s_{\text{query}})$ 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}$:

Given the in-context dataset $D$ and a query state $s_{\text{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^{\prime}_{1},r_{1}),\ldots,(s_{j},a_{j},s^{\prime}_{j},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:

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}$.

Sources of distribution mismatch. Inherent to pretraining, like nearly all foundation models, is distribution mismatch on downstream test-time tasks. DPT pretrained on sufficiently diverse data should ideally be robust (to some extent) to these mismatches. (1) When deployed, $M_{\theta}$ will execute its learned policy which invariably induces a distribution over states different from $\mathcal{D}_{\text{query}}$. (2) Pretraining ${\mathcal{T}}_{\text{pre}}$ likely differs from the downstream ${\mathcal{T}}_{\text{test}}$. (3) Similarly, the test-time datasets prompts can also differ, especially online where they are collected by $M_{\theta}$ itself.

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 ${\mathcal{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.

Comparisons. We compare to several well-known algorithms for banditsSee Appendix A.2 for additional details such as hyperparameters.. 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 can both be used for offline or online learning; UCB is known to be provably optimal online by ensuring exploration through optimism under uncertainty; and LCB is used to minimize suboptimality given an offline dataset by selecting actions pessimistically. It is the opposite of UCB. 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 2(a), 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 distributionNote our randomly generated environments are equally likely to have expert-biased datasets and adversarial datasets, so LCB is not expected to outperform here ..

Interestingly, the same transformer produces an extremely effective online bandit algorithm when sampling actions instead of taking an argmax. As shown in Figure 2(b), DPT matches the performance 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 2(c), 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.

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) . Hence, LCB is generally effective in practice and the pretraining and testing distributions should be biased towards this setting. Motivated by this, we pretrain a second DPT model where $\mathcal{D}_{\text{pre}}$ is generated by mixing the in-context datasets with varying fractions of expert data, biasing $\mathcal{D}_{\text{pre}}$ towards datasets that contain more examples of the optimal action. We denote this model by DPT-Exp. In Figure 3(c), we plot the test-time performance of both pretrained models when evaluated on new offline datasets with varying percentages of expert dataThat is, $0\%$ is fully random while $100\%$ has only optimal actions in the in-context dataset.. 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.

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).

Environments. We consider environments that require targeted exploration to solve the task. The first is Dark Room , 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. Our second environment is Miniworld , 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 color of which is unknown to the agent initially. It receives a reward of $1$ only when near the correct box. Details on these environments and the pre-training datasets are in App. A.4 and A.5.

Comparisons. Our experiments aim to understand the effectiveness of DPT in comparison to that of other context-based meta-RL algorithms. To that end, we compare to meta-RL algorithms based on supervised and RL objectives.

Proximal Policy Optimization (PPO) : 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) : 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 tranformer is trained to predict the next action $a_{j+c}$ from the learning history.

$\text{RL}^{2}$ : 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 $\text{RL}^{2}$ agent is trained to maximize the expected return with PPO.

PPO and RL2 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.

2 Main Results

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 $61.5$ (see Fig. 4(a)). 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, RL2, and PPO online without any prior data from the $20$ test-time Dark Room tasks, shown in Fig. 4(b). After $40$ episodes, PPO does not make significant progress towards the goal, highlighting the difficulty of learning from such few interactions alone. RL2 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 RL2, 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. 4(d), DPT can solve this high-dimensional task offline from both random and expert datasets. Compared to AD and RL2, 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. 5(a)).

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 $1$K episodes to generate $80$K 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 5(b) and 5(c). 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 the other DPT variants in some settings, but only marginally so. In Appendix B, we analyze the sensitivity of DPT to other hyperparameters, such as the context size and amount of pretraining data.

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 performs in-context posterior sampling (PS). PS is a generalization of Thompson Sampling for RL in MDPs. It maintains and samples from a posterior over tasks $\tau$ given historical data $D$ and executes optimal policies $\pi^{\star}_{\tau}$ (see Appendix C for a formal outline). It is provably sample-efficient with online Bayesian regret guarantees , 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.

We start with a modification to the pretraining of DPT. Rather than conditioning only on $s_{\text{query}}$ and $D$ to predict $a^{\star}\sim\pi^{\star}_{\tau}(\cdot|s_{\text{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({\mathcal{S}}^{h})$ is a distribution over sets of states, independent of $\tau$, and $a^{\star}_{h^{\prime}}\sim\pi^{\star}_{\tau}(\cdot|s_{h^{\prime}})$ for $h^{\prime}\in[h]$. Thus, we use $\pi^{\star}_{\tau}$ 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}=()$ to be empty. (2) At step $h$, visit $s_{h}$ and find $a_{h}$ by sampling from $M_{\theta}(\cdot|s_{\text{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 convenient approximation.

We now make several assumptions to simplify the analysis. First, assume $\mathcal{D}_{\text{query}}$, $\mathcal{D}_{\text{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 , 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.

(Learned model is consistent). Let $M_{\theta}$ denote the pretrained model. For all $(s_{\text{query}},D,\xi_{h})$, we have $P_{pre}(a|s_{\text{query}},D,\xi_{h})=M_{\theta}(a|s_{\text{query}},D,\xi_{h})$ for all $a\in\mathcal{A}$.

The in-context dataset distribution $\mathcal{D}_{\text{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}_{\text{pre}}(a_{i}|s_{i},D_{i-1};\tau)$ is invariant to $\tau$.

Generally, $\mathcal{D}_{\text{pre}}$ can influence $M_{\theta}$. In Proposition 6.4, we show that all compliant $\mathcal{D}_{\text{pre}}$ form a sort of equivalence class that generate the same $M_{\theta}$. For the remainder, we assume all $\mathcal{D}_{\text{pre}}$ are compliant.

2 Main Results

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}}_{\text{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({\mathcal{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$.

Let the above assumptions hold. Then, $P_{ps}(\xi_{H}\ |\ D,\tau_{c})=P_{M_{\theta}}(\xi_{H}\ |\ D,\tau_{c})$ for all trajectories $\xi_{H}$.

This significantly improves over the $\widetilde{\mathcal{O}}(\sqrt{|\mathcal{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 known MDPs in the possible distribution), that may further improve performance, such as by removing the dependence on the problem finite state, action or full d-dimensional representation.

Our final result sheds light on how $\mathcal{D}_{\text{pre}}$ impacts the final DPT behavior $M_{\theta}$. Combined with Assumption 1, $M_{\theta}$ is invariant to $\mathcal{D}_{\text{pre}}$ satisfying Definition 6.1.

Let $P_{pre}^{1}$ and $P_{pre}^{2}$ be pretraining distributions that differ only by their in-context dataset distributions, denoted by $\mathcal{D}_{\text{pre}}^{1}$ and $\mathcal{D}_{\text{pre}}^{2}$. If $\mathcal{D}_{\text{pre}}^{1}$ and $\mathcal{D}_{\text{pre}}^{2}$ are compliant with the same support, then $P_{pre}^{1}(a^{\star}|s_{\text{query}},D,\xi_{h})=P_{pre}^{2}(a^{\star}|s_{\text{query}},D,\xi_{h})$ for all $a^{\star},s_{\text{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}_{\text{pre}}^{1}$ and PPO for $\mathcal{D}_{\text{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.

Discussion

In this paper, we studied the problem of in-context decision-making. 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, unlike other algorithms that are explicitly trained or designed to do these. Our empirical and theoretical results provide first steps towards understanding these capabilities that arise from DPT and what factors are important for it to succeed. The inherent strength of pretraining lies in its simplicity–we can sidestep the complexities of hand-designing exploration or conservatism in RL algorithms and while simultaneously allowing the transformer to derive novel strategies that best leverage problem structure. These findings underscore the potential of supervised pretraining in equipping transformer models with in-context decision-making 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. We also remark that our preliminary analysis shows promise for DPT to generalize to new tasks beyond its pretraining distribution. This suggests that diversifying the task distributions during pretraining could significantly enhance the model’s ability to generalize to new tasks. This possibility holds an exciting avenue for future work. Finally, further investigation is required to understand the implications of these findings for existing foundation models, such as instruction-finetuned models, that are increasingly being deployed in decision-making settings .

Acknowledgments and Disclosure of Funding

We thank Evan Liu, Sherry Yang, and Lucy Shi for helpful discussions and 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.

References

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 . Recently, there has been great interest in understanding the capabilities and properties of in-context learning . While a common hypothesis suggests that this phenomenon is due to properties of the data used to train large language models , 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 .

Posterior Sampling. Posterior sampling originates from the seminal work of , and has been popularized and thoroughly investigated in recent years by a number of authors . For bandits, it is often referred to as Thompson Sampling, but the framework is easily generalizable to RL. 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, leading practitioners to rely on approximation-based solutions . In Section 6, we show that a version of the DPT model learned from pretraining can be viewed as implementing posterior sampling as it should be without resorting to approximations or deriving complicated posterior updates. Instead, the posterior update is implicitly learned through pretraining to predict the optimal action. This suggests that in-context learning (or meta-learning more generally) could be a key in unlocking practically applicable posterior sampling for RL.

Appendix A Implementation and Experiment Details

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 . The backbone of the transformer architecture we use is an autoregressive GPT-2 model from the HuggingFace transformers library.

Because of the GPT-2 causal architecture (we defer details to the original papers ), we note that $\hat{p}_{j}$ depends only on $s_{\text{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,

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_{\text{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]$:

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_{\text{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 and Algorithm Distillation , DPT does not separate the individual $(s,a,s^{\prime},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 Implementation Details

First, we describe the comparisons from the bandit experiments with hyperparameters.

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.

Since the means are sampled uniformly from $$, 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 3(a) and Figure 3(b), we set the prior mean and variance to and$1$ to fit the true ones better.

LinReg (offline) is the same as LinUCB except we set $\beta=0$ to greedily choose actions.

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.

The reported results for PPO use the Stable Baselines3 implementation with the default hyperparameters, which successfully learns each task given $100$K environment steps in Dark Room and $125$K 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$ $3\times 3$ kernels each, followed by a linear layer with output dimension of $8$.

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 . 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$ $3\times 3$ kernels each, followed by a linear layer with output dimension of $8$.

The reported results for $\text{RL}^{2}$ use an open-sourced implementation from . 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 $16$ $3\times 3$ kernels each, followed by a linear layer with output dimension of $8$.

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$ $3\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

A.4 MDP Environment Details

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.

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

In Dark Room, we collect $100$K in-context datasets, each of length $H=100$ steps, with a uniform-random policy. The $100$K 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 $100$K collections of datasets, query states, and optimal actions, we use the first $80$K (corresponding to the first $80$ goals) for training and the remaining $20$K for validation.

While this task is solved from image-based observations, we also note that there are only four distinct tasks (one for each colored box), and the agent does not need to handle new tasks at test time. Hence, the number of in-context datasets required in pretraining is fewer – we use $40$K datasets each of length $H=50$ steps. So as to reduce computation, the in-context datasets only have only $(s,a,r)$ tuples. The query states, which consist of image and direction are sampled uniformly from the entire state space, i.e., the agent is place uniformly at random in the environment, pointing in a random direction. The optimal actions are computed as follows: turn towards the correct box if the agent is not yet facing it (within $\pm 15$ degrees), otherwise move forward. Of the $40$K collections of datasets, query states, and optimal actions, we use $32$K for training and the remaining $8$K for validation.

Appendix B Additional Experimental Results

This section reports additional experimental results in bandit environments.

In Figures 2(c) and 6(a), we demonstrate the robustness of the basic pretrained model under shifts in the reward distribution at test time by varying the amount of noise observed in the rewards. DPT maintains robustness to these shifts similar to TS.

We test the out-of-distribution ability of DPT further by completely changing the reward distribution from Gaussian to Bernoulli bandits. Despite being trained only on Gaussian tasks during pretraining, DPT maintains strong performance both offline and online in Figures 6(b) and 6(c).

B.2 Markov Decision Processes

This section reports additional experimental results in the Dark Room and Miniworld environments.

In Fig. 8, we show the performance of each method on the training tasks in Dark Room. Offline, DPT and AD demonstrate comparable performance as on the training tasks, indicating a minimal generalization gap to new goals. Online, DPT, AD, and RL2 also achieve performance on the training tasks similar to that on the test tasks.

In this experiment, we study generalization to variations in a different aspect of the MDP, namely the dynamics. We design Dark Room (Permuted), a variant of Dark Room in which the goal is fixed to a corner but the action space is randomly permuted. Hence, the agent must leverage its historical context to infer the effect of each action. On a held-out set of $20$ permutations, DPT infers the optimal policy correctly every time offline, given only $100$ offline samples, matching the optimal policy at $83$ return. Similarly, the online performance immediately snaps to a near optimal policy in one episode once it identifies the novel permutation in Figure 7.

B.3 Sensitivity Analysis

We next seek to understand the sensitivity of DPT to different hyperparameters, including the model size and size of the pretraining dataset. These experiments are performed in the Dark Room environment. As shown in Fig. 9, the performance of DPT is robust to the model size; it is the same across different embedding sizes, number of layers, and number of attention heads. Notably, the performance is slightly worse with $8$ attention heads, which may be attributed to slight overfitting. We do see that when the pretraining dataset is reduced to $10\%$ of its original size ($10000$ samples) the performance degrades, but otherwise has similar performance with larger pretraining datasets.

Appendix C Additional Theory and Omitted Proofs

We start with a well-known concentration inequality for the maximum-likelihood estimate (MLE) to provide some more justification for the approximation made in Assumption 1. We state a version from . Let $\mathcal{F}$ be a finite function class used to model a conditional distribution $p_{Y|X}(y|x)$ for $x\in\mathcal{X}$ and $y\in\mathcal{Y}$. Assume there is $f^{\star}\in\mathcal{F}$ such that $p(y|x)=f^{\star}(y|x)$ (realizable), and $f(\cdot|x)\in\Delta(\mathcal{Y})$ for all $x\in\mathcal{X}$ and $f\in\mathcal{F}$ (proper). Let $D=\{x_{i},y_{i}\}_{i\in[N]}$ denote a dataset of i.i.d samples where $x_{i}\sim p_{X}$ and $y_{i}\sim p_{Y|X}(\cdot|x_{i})$. Let

Let $D$ and $\hat{f}$ be given as above under the aforementioned conditions. Then, with probability at least $1-\delta$,

Posterior sampling is most generally described with the following procedure . Initialize a prior distribution ${\mathcal{T}}_{1}={\mathcal{T}}_{\text{pre}}$ and dataset $D=\{\}$. For $k\in[K]$

Sample $\tau_{k}\sim{\mathcal{T}}_{k}$ and compute $\hat{\pi}_{\tau_{k}}$

Execute $\pi^{\star}_{\tau_{k}}$ and add interactions to $D$

Update posterior distribution ${\mathcal{T}}_{k+1}(\tau)=P(\tau|D)$.

The prior and posteriors are typically over models such as reward functions in bandits or transition dynamics in MDPs.

C.2 Proof of Theorem 1

Without loss of generality, for a task $\tau$, we take $\pi^{\star}_{\tau}(\cdot|s)$ to be deterministic and denote the optimal action in state $s$ as $\pi^{\star}_{\tau}(s)$. Recall that we consider a fixed current task $\tau_{c}$ and a fixed in-context dataset $D$. Define $\xi_{h}=(s_{1},a_{1},\ldots,s_{h},a_{h})$.

We now formally state the variant of the full joint distribution from which we sample during pretraining. Let $\tau$ and $D^{\prime}$ be an arbitrary task and dataset and let $a^{\star}\in\mathcal{A}$, $s_{\text{query}}\in{\mathcal{S}}$, $\xi_{H-1}\in({\mathcal{S}}\times\mathcal{A})^{H-1}$, and $h\in[0,H-1]$ be arbitrary.

The $\text{Unif}[0,H-1]$ is due to the fact that we sample $h\sim\text{Unif}[0,H-1]$ and then truncate $\xi_{h}$ from $\xi_{H-1}$ (or, equivalently, sample $\xi_{h}\sim{\mathfrak{S}}_{h}$ directly), marginalizing out the other variables. For $h^{\prime}\leq h-1$, recall that we also use the notation ${\mathfrak{S}}_{h^{\prime}}(s_{1:h^{\prime}})$ to denote the marginalization of the full joint ${\mathfrak{S}}_{H}$. We will eventually work with the posterior of this distribution given the data $D$ and history $\xi_{h}$:

We define the following random sequences and subsequences:

where the variables are generated according to the following conditional process: $\tau_{ps}\sim P(\cdot|D)$, $S^{ps}_{1}\sim\rho_{\tau_{c}}$, $A^{ps}_{h}\sim\pi^{\star}_{\tau_{ps}}(\cdot|S^{ps}_{h})$, and $S_{h+1}^{ps}\sim T_{\tau_{c}}(\cdot|S_{h}^{ps},A^{ps}_{h})$. We also define $\Xi_{ps}(h^{\prime}:h;D)$ to be the last $h-h^{\prime}$ elements of $\Xi_{ps}(h;D)$. Analogously, we define

where the variables are from the process: $S_{1}^{pre}\sim\rho_{\tau_{c}}$, $A_{h}^{pre}\sim P_{pre}(\cdot|S_{h}^{pre},D,\Xi_{pre}(h-1;D))$, and $S_{h+1}^{pre}\sim T_{\tau_{c}}(\cdot|S_{h}^{pre},A_{h}^{pre})$. Note that $A_{h}^{pre}$ is sampled conditioned on the sequence $\Xi_{pre}(h;D)$ so far.

We will show that $\Xi_{ps}(h;D)$ and $\Xi_{pre}(h;D)$ follow the same distribution for all $h\in[H]$. For convenience, we will drop notational dependence on $D$, except where it resolves ambiguity. Also, because of Assumption 1, we have that $P_{pre}(\cdot|S^{pre}_{h},D,\Xi_{pre}(h-1))=M_{\theta}(\cdot|S^{pre}_{h},D,\Xi_{pre}(h-1))$, so we will just work with $P_{pre}$ for the remainder of the proof. We will also make use of the following lemma.

If $\mathcal{D}_{\text{pre}}$ is complaint, then $P_{pre}(\tau|D)=P(\tau_{ps}=\tau|D)$.

From the definition of posterior sampling (using the same prior, ${\mathcal{T}}_{\text{pre}}$), we have that

where the second line crucially uses the fact that posterior sampling chooses actions based only on the prior and history so far. Similarly, the third line uses the fact that $\mathcal{D}_{\text{pre}}$ is compliant. Since the two sides are proportional in $\tau$, they are equivalent. ∎

We will prove Theorem 1 via induction for each $h\in[H]$. First, consider the base case for a sequence of length $h=1$. Recall that $\rho_{\tau_{c}}$ denotes the initial state distribution of $\tau_{c}$. We have that the densities can be written as

where the second line uses the sampling process of $S^{pre}_{1}$; the third marginalizes over $\tau_{ps}$, which is the task that posterior sampling samples to find the optimal policy; the fourth decomposes this into the optimal policy and the posterior over $\tau_{ps}$ given $D$ and $S^{ps}_{1}$. Since $S^{ps}_{1}$ is independent of sampling of $\tau_{ps}$ this dependence goes away in the next line. The sixth line applies Lemma C.2 and then, for $h=1$, there is no history to condition on.

Now, we leverage the inductive hypothesis to prove the full statement. Suppose that the hypothesis holds for $h-1$. Then,

By the hypothesis, we have that $P(\Xi_{ps}(h-1)=\xi_{h-1})=P(\Xi_{pre}(h-1)=\xi_{h-1})$. For the second factor,

As before, we can further rewrite the last factor as

where $\propto$ denotes that the two sides are equal up to multiplicative factors independent of $\tau$. In the first line, we used Bayes rule. In the second line, given that $\tau_{ps}=\tau$ (i.e. posterior sampling selected $\tau$ to deploy), we decompose the probability of observing that sequence of states of actions. We also used Lemma C.2. The denominator does not depend on $\tau$. Similarly, for the third and fourth lines, $T_{\tau_{c}}$ and ${\mathfrak{S}}$ do not depend on $\tau$. The final line follows from the definition of the joint pretraining distribution in this regime.

Therefore, we conclude that the posterior over the value of $\tau_{ps}$ is the same as the posterior over the task in the pretraining distribution, given $s_{h},D,\xi_{h-1}$. Substituting back through all the previous equations, we have

C.3 Proof of Corollary 6.2

Note that $\mathcal{D}_{\text{pre}}$ is clearly compliant since it is generated by random sampling. We use the equivalence between $M_{\theta}$ and posterior sampling established in Theorem 1. The proof then follows immediately from Theorem 1 of to guarantee that

where the notation $\widetilde{\mathcal{O}}$ omits polylogarithmic dependence. The bound on the test task distribution follows from the assumed bound on the likelihood ratio under the priors:

C.4 Proof of Corollary 6.3

The distribution $\mathcal{D}_{\text{pre}}$ satisfies compliance by definition because it is generated by an adaptive algorithm TS. The proof once again follows by immediately deferring to the established result of (Proposition 3) for linear bandits by the posterior sampling equivalence of Theorem 1. This ensures that posterior sampling achieves regret $\widetilde{\mathcal{O}}(d\sqrt{K})$. It remains, however, to justify that $P_{pre}(\cdot|D_{k})$ will be covered by Gaussian Thompson Sampling for all $D_{k}$ with $k\in[K]$. This is verified by noting that $P_{ps}(a|D_{k})>0$ for non-degenerate Gaussian Thompson Sampling (positive variances of the prior and likelihood functions) and finite $K$. This guarantees that any $D_{k}$ will have support. ∎

C.5 Proof of Proposition 6.4

The proof follows by direct inspection of the pretraining distributions. For $P^{1}_{pre}$, we have

The posterior distribution over tasks is simply

Then, the distribution over the in-context dataset can be decomposed as

where the second equality holds because $\mathcal{D}_{\text{pre}}^{1}(a_{j}|s_{j},D_{j};\tau)$ is assumed to be invariant to $\tau$ by compliance, and the fifth equality holds because $\mathcal{D}_{\text{pre}}^{2}(a_{j}|s_{j},D_{j};\tau)$ is assumed to be invariant to $\tau$.

Therefore, we conclude that, for any $s,D,\xi$,

Since also $\int_{\tau}P^{1}_{pre}(\tau|s,D,\xi)=1=\int_{\tau}P^{2}_{pre}(\tau|s,D,\xi)$, then

Substituting this back into Equation45 yields $P^{1}_{pre}(a^{\star}|s,D,\xi)=P^{1}_{pre}(a^{\star}|s,D,\xi)$. ∎