Supervised Pretraining Can Learn In-Context Reinforcement Learning
Authors: Jonathan Lee
Annie Xie
Aldo Pacchiano
Yash Chandak
Chelsea Finn
Ofir Nachum
Emma Brunskill
Venue: NeurIPS 2023 spotlight (2023)
Keywords: decision making
reinforcement learning
in-context learning
bandits
transformers
offline reinforcement learning
exploration
reinforcement learning theory

Abstract: Large transformer models trained on diverse datasets have shown a remarkable ability to learn in-context, achieving high few-shot performance on tasks they were not explicitly trained to solve. In this paper, we study the in-context learning capabilities of transformers in decision-making problems, i.e., reinforcement learning (RL) for bandits and Markov decision processes. To do so, we introduce and study the Decision-Pretrained Transformer (DPT), a supervised pretraining method where a transformer predicts an optimal action given a query state and an in-context dataset of interactions from a diverse set of tasks. While simple, this procedure produces a model with several surprising capabilities. We find that the trained transformer can solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline, despite not being explicitly trained to do so. The model also generalizes beyond the pretraining distribution to new tasks and automatically adapts its decision-making strategies to unknown structure. Theoretically, we show DPT can be viewed as an efficient implementation of Bayesian posterior sampling, a provably sample-efficient RL algorithm. We further leverage this connection to provide guarantees on the regret of the in-context algorithm yielded by DPT, and prove that it can learn faster than algorithms used to generate the pretraining data. These results suggest a promising yet simple path towards instilling strong in-context decision-making abilities in transformers.

Supervised Pretraining Can Learn In-Context Reinforcement Learning
Jonathan N. Lee
Annie Xie
Aldo Pacchiano
Yash Chandak'
Chelsea Finn'
Ofir Nachum
Emma Brunskill'
Sttnnfor University 2 Broad Institute of MIT and Harvard Boston University
Abstract
Large transformer models trained diverse datasets have shown remarkable abil- ity to learn in-context, achieving high few-shot performance on tasks they were not explicitly trained to solve. In this paper, we study the in-context learning capabili- ties of transformers in ercision-makinng problems, i.e., reinforcement learning (RL) for bandits and Markov decision processes. To do so, we introduce and study the Decision-Pretrained Transformer (DPT), a supervised pretraining method where a transformer predicts an optimal action given query state and an in-context dataset of interactions from diverse set of tasks. While simple, this procedure produces model with several surprising capabilities. We find that the trained transformer can solve range of RL problems in-context, exhibiting both exploration online and conservatism offline, despite not being explicitly trained to do so. The model also generalizes beyond the pretraining distribution to new tasks and automatically adapts its decision- maaing strategies to unknown structure. Theoretically, we show DPT can be viewed as efficient implementation of Bayesian posterior sampling, a provably sample-efficient RL algorithm. We further leverage this connection provide guarantees on the regret of the in-context algorithm yielded by DPT, and prove that it can learn faster than algorithms used to generate the pretraining data. These results suggest promising yet simple path towards instilling strong in-context decision-making abilities transformers.
Introduction
For supervised learning, transformer-bast 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 [I. In this setting, pretrained model is presented with small number of supervised input- output examples in context, and then asked predict the most likely completion (i.e. output) of an unpaired input, without parameter updates. Over the last few years, in-context learning has been applied to solve a range of tasks [2) and a growing number works are beginning to understand and analyze in-context learning for supervised learning study and understand in-context learning applied to sequential decision-making specifically in the context of reinforcement learning (RL) settings. recisiin-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 agent's '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 recommendation systems.
3 4
5 6). In this work, focus is to
'Equal contrihution. Correspondence to jnl@stanford.ccommmm edu and anni anniexie@stanford edu. Code available stttps://github.... wtt.:o/ooooooeeeeeeeeeeeeeeeeeee trannfffrr
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

For in- in-context decision makiing 7 8 9 rather than input-output tuples, the context takes the form of state-action-reware tuples representing a dataset of interactions with an unknown environment. The agent must leverage these interactions to understand the dynamics of the world and what actions lead to good outcomes. hallmark of good decision- balance of selecting exploratory actions to gather information and selecting increasingly optimal actions by exploiting that information [0]. 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-make should exhibit similar behaviors.
in online RI. algorithms a judicious
To study in-contex decision beeiiinn-making. formally, we propose a new simple supervised pretraining objective, namely, to train (via supervised learning) a transformer to predict an optimal action labe given a query state and an in- in-context dataset of interactions, across a diverse set of tasks. We refer to the pretrained model as a Decision-Pretrained Pretrained Transformer (DPT). Once trained, DPT can be deployed as either an online or offline RL algorithm in new task by passing it an in-context dataset of interactions and querying 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- in-context decision-maker 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 [1 13
We summarize our main findings below.
Predicting optimal actions alone gives rise to near-optimal decision- decision-making algorithms. The DPT objective is solely based on predicting optimal actions from in-contex interactions. At the outset, it not immediately apparent that these predictions 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 rocision-makine problems, offline and online. We show DPT can handle reward distributions unseen in its pretraining data on bandit problems as well as unseen goals, dynamics, and datasets in simple MDPs. This suggests that the in-context strategies learned during pretraining are robust and generalizable without any parameter updates at test time.
DPT improves over the data used to pretrain it by exploiting latent structure. As an example, in parametric bandit problems, specialized algorithms can leverage structure (such as linear rewards) and offer provably better regret. However, 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), generalization of Thompson Sampling, can provably sample-efficiently solve online RL prob- lems 212 [3] but a common criticism is the lack of computationally efficient ways to update and sample from a posterior distribution. DPT can viewed as learning 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 priori. This suggests that in-context learning could be a possible path towards scalable and practical RL via posterior sampling.
2 Related Work
Meta-learning Algorithmically, in-context learning falls under the meta-learning framework [14] [3]. 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 mecksion-making and RL, there is often
'If not explicitly known, the optimal action can be determined by running any (potentially inefficient) minimax-optimal regret algorithm for each pretraining task.

---

a choice in what shared 'structure' is specifically learned such as the dynamics of the task [16 17
18 19,, task context identifier [20]
2 22
23 temporally extended skills and options 24 25 26
or initialization of neural network policy 27
28 ). In-context learning can be viewed as taking
more agnostic approach by learning the learning algorithm itself, more similar to [29 30 31 32
Algorithm Distillation (AD) 7 (3 falls under this category, using autoregressive supervised learning to distill traces of single-task RL algorithm into a task-agnostic model. While DPT also leverages autoregressive SL, it does not distill an existing RL algorithm in order to imitate how to learn. Instead, we train DPT to predict task-specific optimal actions, yielding potentially emergent strategies at test time that automatically leverage the task structure behave similarly to posterior sampling.
Autoregressive transformers for decision-making, In decision- makinn fields such as RL and imitation learning, transformer models trained using autoregressive supervised action prediction have proliferated [34), inspired by the successes of these techniques for large language models [55 36
uses a transformer to autoregressively model sequences of actions from offline experience data, conditioned on the achieved return. During inference, one can then query the model conditioned on a desired return value. This approach has been shown to scale favorably to large models and multi-task settings [39], at times exceeding the
For example, Decision Transformer (DT) 37 38 g
performance of large-scale multi-task imitation learning with transformers 40 41 42). However, DT is known to be provably (and unboundedly) sub-optimal in common scenarios 43 44 A common criticism of DT. and supervised learned transformers in general. is their inability to improve upon the dataset. For example, there is little reason for DT to output meaningful behavior if conditioned on return higher than any observed in training, without strong extrapolation assumptions [43]. In contrast, a major contribution of our work is theoretical and empirical evidence for the ability of DPT to improve over behaviors seen in the dataset in terms of regret.
Value and policy-based offline RL. Offline RL algorithms offer the opportunity to learn from existing datasets. To address distributional shift, many prior algorithms incorporate the principle of value pessimism 45
46 47 48), or policy regularization (49 50 51 53. To reduce the amount
of offline data required in a new task, methods for offline me can reuse interactions collected
in a set of related tasks [54] 55 56). However, they still must address distribution shift, requiring solutions such as policy regularization [54] or additional online interactions [57]. DPT follows the success of autoregressive models like DT and AD, avoiding these issues. With our pretraining objective, DPT also leverages offline datasets for new tasks more effectively than AD.
3 In-Context Learning Model
Basic decision models. The basic decision model of our study is the finite-horizon Markov decision process (MDP). An MDP is specified by the tuple = A, R H. p) to be solved, where S is the state space, A is the action space, T 8 x A + ASS) is the transition function, R: S x A - A(R) is the reward function, H E N is the horizon, and p E (SS) is the initial state distribution. A learner interacts with the environment through the following protocol: (1) an initial state s1 is sampled from p: (2) at time step h, the learner chooses an action an and transitions to state SA+1 ~ [-1s), ax). and receives a reward Th *R[Ist, an). The episode ends after H steps. A policy π maps states to distributions over actions and can be used to interact with the MDP. We denote the optimal policy as n², which maximizes the value function (iT') = max, V(T) == max, Ex Σn Th- When necessary, we use the subscript 7 to distinguish V- and म. for the specific MDP r. We assume the state space is partitioned by h E [H] so that 'r" is notationally independent of h. Note this framework encompasses multi-armed bandit settings where the state space is a single point, e.g. S = (1). H = 1, and the optimal policy is a* E argmax.e. E [r: (a
a].
Pretraining. We give pseudocode in Algorithm D and a visualization of the transformer in Figure I] Details on the precise structure of the transformer can be found in Section A.I Let Tpr be distribution over tasks at the time of pretraining. A task T ~ Tpe can be viewed as a specification of an MDP, functions and even different state and action spaces. We then sample a context (or prompt) which consists of a dataset D Dpe (:; T) of interactions between the learner and the MDP specified by T. D = {sj.aj, aj. sy is a collection of transition tuples taken in т. We refer to D as the in-context dataset because it provides the contextual information about T. D could be generated through variety of means, such as: (1) random interactions within T, (2) demonstrations from an expert, and (3) rollouts of an algorithm. Additionally, we independently sample a query state Squery from the distribution D query over states S and a label a" is sampled from the optimal policy π (- |8query)
AS,A, T, R, HI, p). The distribution Ty can span different reward and transition
3

---

N
Mo(.1 D)
Sy. A1, Fi. x
Syr uzy 1s, $
Sax Aer Fax 8n
Online Exploration
Mg (R"1squrry. D)
Transformer
Ma D)
8 query
Aquery
Offline Learning
Figure I: A transformer model Me is pretrained to predict an optimal action aieery from state Squery a task, given a dataset of interactions from that task. The resulting Decision-Pretrained Transformer (DPT) learns a distribution over the optimal action conditioned an in-context dataset. Ma can be deployed in new tasks online by collecting data the fly offline by immediately conditioning static dataset.
Algorithm Decision- Prettiiine Transformer (DPT): Training and Deployment
// Collecting pretraining dataset 2: Initialize empty pretraining dataset B
3: for in [N] do
4: Sample task Tpe. in-context dataset D ~ (pee (),, query state Squery ~ D
5: Sample label a* ~ n, (qqurr) and add (Squery: D, a to B
6: end for
7: // Pretraining model on dataset
8: Initialize model Ma with parameters 0
9: while not converged do
10: Sample (Squery: D, a²) from B and predict p-()) Mo(-Isqwayy Dj) for all [n]
11: Compute loss in (2) with respect to a" and backpropagate to update e.
12: end while
13: // Dffline test-time deployment
14: Sample unknown task r Tiea sample dataset D) ~ (::T)
15: Deploy Mo in r by choosing dA E argmax Ma(a) 8h. D) at step
Online test-time deployment
17: Sample unknown task T ~ Trest and initialize empty D = 0
18: for ep in max_epa do
19: Deploy Mg by sampling en ~ Mol-Ish, D) at step h
Add (s1, (1.......) to
21: end for
for task (see Section 5.3 for how implement this in common practical scenarios). We denote the joint pretraining distribution over tasks, in-contex datasets, query states, and action labels as Ppre:
PTre D, Squery, a*) = Tpe (T) D; T T)D query (Squery ) (a" [Squery).
(1)
Given the in-context dataset D and query state Squery. we can train a model to predict the optimal action a" in response simply via supervised learning. Let Dj = {(1 01, $1. ri)...., (s), aj. s r,)} denote the partial dataset up to j samples. Formally, we aim to train a causal GPT-2 transformer model M parameterized by which outputs a distribution over actions A, minimize the expected loss over samples from the pretraining distribution:
ming Ep,r. Σ}e(n) (Me(- Squery, Dj),a")
(2)
Throughout, we set the loss to be the negative log- ligllikl with {(Mo(- |Squery: Dj), a*)
log Mo(a" Squery, D,). This framework can work for both discrete and continuous A. For our experiments with discrete A. we use a softmax parameterization for the distribution of Mg. essentially treating this as a classification problem. The resulting output model M, can viewed as algorithm that takes in dataset of interactions D and can be queried with forward pass for predictions of the optimal action via inputting query state Squery. We refer the trained model Mo Decision-Pretrainee Transformer (DPT).

---

Testing, After pretraining, a new task (MDP) is sampled from a test-task distribution Twst If the DPT be tested offfine, then dataset (prompt) is a sampled D Dieu( :77 and the policy that the model in-context learns is given conditionally as Mo D). Namely, we evaluate the policy by selecting action 04 E a Mo(a) Sh. D) when the learner visits state 86- If the model is to be tested online through multiple episodes of interaction, then the dataset is initialized as empty D {}. At each episode, Mp(- D) is deployed where the model samples ag ~ Ma(-|sa, D) upon observing state sΛ. Throughout a full episode, it collects interactions {s1, "1;"",, ........ 3H aH, rH) which are subsequently appended to D. The model then repeats the process with another episode, and until specified number of episodes has been reached. A key distinction of the testing phase is that there are no updates to the parameters of Mg. 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 Mo performs a computation through its forward pass to generate a distribution over actions conditioned on the in-context D and query state 84-
4 Learning in Bandits
We begin with an empirical investigation of DPT in multi-arme' bandit, well-studied special case of the MDP where the state space S is a singleton and the horizon H = 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. Offiiee, it is critical to account for uncertainty due noise certain actions may not be sampled well enough. Online, it is critical to judiciously balance exploration and exploitation to minimize overall regret. For detailed descriptions of the experiment setups, see Appendix A
Pretraining distribution. For the pretraining task distribution Tpres we sample 5-armed bandits (IA| 5). The reward function for arm is a normal distribution R(-|s, a) = [N(,a where Ma ~ Unif(0, 1] independently and a = 0.3. To generate in in-context datasets D we randomly generate action frequencies by sampling probabilities from Dirichlet distribution and mixing them with point-mass distribution one random arm (see details in Appendix [A.3]. Then we sample the actions accordingly from this distribution. This encourages diversity of the in-context datasets. The optimal policy म. for bandit is argmax, Ma, which we can easily compute during pretraining. We pretrain the model Mg to predict a" from D as described in Section 3 for datasets up to size n = 500.
Comparisons. We compare to several well-known algorithms for bandity" 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 (L.B) selects the action with the highest lower confidence bound.
Thompson Sampling (TS) selects the action with the highest sampled mean from a posterior distribution over reward models. The prior and likelihood functions are Gaussian.
Emp and TS 58 59 can both be used for offline online learning: UCB 60 is an online algorithm that achieves exploration via optimism under uncertainty; and LCB 6 62 is used to minimize suboptimality given offline dataset by selecting actions pessimistically. We evaluate algorithms with standard bandit metrics. Offline, we use the suboptimality Ha (4a where à is the chosen action. Online, we cumulative regret: Σσ Pa- Ma, where âR is the kehh action chosen.
DPT learns to reason through uncertainty. As shown in Figure 2a 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 degree. However, it also suggests that this hedging is fundamentally different from what LCB does, at least on this specific distribution
Interestingly, the same transformer produces an extremely effective online bandit algorithm when sampling actions instead of taking an argmax. As shown in Figure 26 DPT matches the performance
P See Appendix A.2 for additional details such as hyperparameters. 'Note our randomly generated euvironments are equally likely to have expert-biased datasets and adversarial datasets, LCB is not expected to outperform here (61).
5

---

Offline Bandit
20.0
Bandit
Online Rosustiess to Rewardd Shint
Emp
LGB
Curulltivve Regret
17.5
UCB
Brol Ounulittie Regnst
Suboptmaalty
15
15
DPT (ours)
DPT (ours)
7.5
5.0
2.5
0.02
400
300
400
Nos: stanoord deviztion
(a)
(b)
(c)
Figure 2: (a) Offline performance on in-distribution bandits, given random in-context datasets. (b) Online cumulative regret on bandits. (c) Final (after 500 steps) cumulative regret on cut-of-distribution bandits with different Gaussian noise standard deviations. The mean and standard error are computed over 200 test tasks.
of classical optimal algorithms, UCB and TS, which are specifically designed for exploration. This is notable because DPT not explicitly trained to explore, but its emergent strategy is on par with some of the best. In Figure [c] we show this property is robust to noise in the rewards not seen during we show this generalization happens
pretraining by varying the standard deviation. In Appendix B 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 preciscly, we consider Tpre to be a distribution over linear bandits, where the reward function is given by [r a, = (0,, v(a)) and 0, E Rd is a task-specific parameter vector and Φ : A - Rd is fixed feature vector that is the same for all tasks. Given the feature representation o, LinUCB [63)
UCB-style algorithm that leverages d, should achieve regret Ödv K) over steps, a substantial gain over UCB and TS when d < |AA. We do not assume & is known to DPT, besides whatever can be gleaned from the pretraining data. To evaluate the ability of DPT uncover latent structure, we collect pretraining data with in-context datasets gathered by standard TS, which does not leverage the linear structure. Figures [b sho show that DPT can exploit the unknown linear structure, implicitly learning a surrogate for Ø, allowing to do more informed exploration online and decision-making offline. nearly on par with LinUCB (which given Φ). 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 Dpr is generated by mixing the in-context datasets with varying fractions of expert data, biasing Dpe towards datasets that contain more examples of the optimal action. We denote this model by DPT-Exp. In Figure 33] we plot the test-time performance of both pretrained models when evaluated on new offline datasets with varying percentages of expert dat=F) Our results suggest that when the pretraining distribution is also biased towards expert-suboptimal data, DPT-Ext behaves similarly to LCB, while DPT continues to resemble TS. This is quite interesting as, for other methods such 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 different RL algorithm, rather than the exact optimal policy (Section 5.3
'That is. 0% is fully random while 100% has only optimal actions in the in-context dataset.

---

Bandit
Adaptation Expert Datasels
(ssurte)
UnPog
Cumulaativ Regrch
17.5
Suboptimaiity
Suboptinallyg
10"4
Dff lours)
15.0
[ours)
0.04
DP7 (ours)
0.03
DPPHEap joun)
1..00
10-2
7.5
5.0
175
0.0
23
150
160
Expert
(a)
(b)
(c)
Figure 3: (a) Offiine 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 етror are computed over 200 test tasks. (c) Offline performance on expert- biased datascts, DPT pretrained on different prior continues match TS, but DPT- trained from more representative prior excels.
Ortineg Durk Room Flezt Bocks
Offine
On ins Wiinorr
DIT CUIS
H.tum
Redum
fastin Smen
Dabasit Composttion
Fan-do Deen
Dacaser Compostien
Essodes
(a)
(b)
(c)
(d)
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) Offlie performance on Miniworld (images), given random and expert datasets. (d) Online performance Miniworld (images) after 40 episodes. We report the average and standard error of the
100 different offline datasets (a) and (c) and 20 online trials in (b) and (d).
5.1 Experimental Setup
Environments. We consider environments that require targeted exploration to solve the task. The first is Dark Room 22 7) 2D discrete environment where the agent must locate the unknown goal location in a 10 x 10 room, and only receives reward of 1 when at the goal. We hold out set of goals for generalization evaluation. The second is Miniworld [65] a visual navigation problem to test the scalability of DPT to image observations. The agent is in room with four boxes of different colors, and must find the target box. The target color is unknown initially. receives reward of 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-tas 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 dataset of learning histories by running RL algorithm each training task. Then, given a sampled subsequence h = (*j. T Ty....,.- 8jte) from learning history, transformer is trained predict the next action aj+c from the history.
This online ne comparison uses recurrent neural network adapt the agent's policy from the given context. Unlike AD and DPT, which are trained with supervised objective, the RL2 agent trained maximize the expected return with PPO.
RL2
29
PPO and RL are online algorithms, while AD is capable of learning both offline and online. Details on the implementation these algorithms can found Appendix A.2
5.2 Main Results
Generalizing to new offline datasets and tasks. To study the generalization capabilities of DPT, we evaluate the model in Dark Room 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 average total reward of 1.1, DPT obtains much higher average return of
7

---

.14
Offine PPO Ablition (lss: Tases)
Ablations Tasks)
hand
06-
60
Returr
Return
40
DPT
DPT (Sana, (90)
(PPO, PPO1
Dataset Composition
Episodes
(a)
(b)
(c)
Figure 5: (a) In Dark Room (Three Tasks), DPT stitches a new, optimal trajectory the goal (blue) given two in-context demonstrations of other tasks (pink and orange). (b) Offline Dark Room performance 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 transition to the goal, DPT immediately exploits this and takes direct path to the goal. In contrast, while AD demonstrates strong offline performance with expert data, performs worse in-context learning with random data compared to DPT. The difference arises because AD is trained to infer better policy than the in-context data, but not necessarily the optimal one.
We next evaluate DPT, AD, RL?, 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 diffiultyy of learning from such few interactions alone. RL is trained to perform adaptation within four episodes each length 100, and we report the performance after the four adaptation episodes. Notably, DPT on average solves each task faster than AD and reaches 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 observa- tions of 25 x 25 pixels. As shown in Fig. 4d DPT can solve this nigh-dimensional task offline from both random and expert datasets. Compared AD and RL 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. Sa).
5.3 Learning from Algorithm-Generated Policies and Rollouts
So far, 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 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 55) 5c DPT (PPO, PPO) can be viewed as direct comparison between our pretraining objective and that of AD, given the same pretraining data but just used differently. We also evaluated variant, DPT (Rand, PPO), which pretrains on random in-context datasets (like DPT), but still using PPO action labels. The performance 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 theoretical analysis. Our main result shows that DPT (under a slight modification to pretraining) essentially capable of executing posterior sampling (PS) reinforcement learning in-context. PS is generalization of

---

Thompson Sampling for RL in MDPs [1] It maintains and samples from a posterior over tasks given historical data D and executes optimal policies मद (see Appendix [ a formal outline). is provably sample-efficient with online Bayesian regret guarantees [3]. but maintaining posteriors is generally computationally intractable. The ability for DPT to perform PS in-context suggests a path towards computation- and provably sample-efficint RL with priors learned from the data.
6.1
Pretraining and Assumptions
We start with a modification to the pretraining of DPT. Rather than conditioning only on Squery and D to predict at ~ (+ (querr).. we propose also conditioning on a sequence En. = ($1:h: (i,A) where $1:h ~ 6₁ E (SS)) is a distribution over sets of states, independent of t, and d (+19)))) for h' E [h] Thus, we use ; to label both the query state (which is the prediction label) and the sequence of states sampled from 6₁, Note that this does not require any environment interactions and hence no sampling from either T- or R,. At test-time at step h, this will allow us condition on the history - of states that Mo visits and the actions that it takes in those states. Formally, the learned Ma is deployed as follows, given D. (1) At = 0, initialize Eo = () to be empty. (2) At step h, visit 8,, and find dn by sampling from Mø(- |squery D, (11) (3) Append (sh, as.) to E1-1 to Mouels
get En. Note for bandits and contextual bandits (H = 1). there is no difference between this and the original pretraining procedure of prior sections because 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 Dquery, D and 6 have sufficient support such that all conditional probabilities of Ppre are well defined. Similar other studies of in-context learning [67]. we assume Mg 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 Me denote the pretrained model. For all (8 (quuery, ),), we have P Ppre(a) Squery; D, Ex) = Ma (a) Squery, D2E,) for all E A.
To provide some cursory justification, if Mo is the global minimizer of (2. then we might expect that Ma[-ISquer,,
Ep,,.. [..Ilprre Spre (-ISquery. D, (b. Mp Squuery D, Ex. ) - 0 as the number of pretraining samples N 0o with high probability for transformer model classes of bounded complexity (see Proposition C. 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 Ppre 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 Pare be very general it could be generated randomly or from adaptive algorithms like PPO or TS. Without compliance Dree can influence Mg. In Proposition 6.44 we show that all compliant Dree form sort of equivalence class that generate the same Mg. For the remainder, we assume all D
pre are compliant.
Definition 6.1 (Compliance). The in-context dataset distribution (:re(;;)) T) is compliant if, for all i E [n), the ith action of the dataset, og is conditionally independent of given the ith state s; and partial dataset, D;-1, so far. In other words, the distribution Dpre (as ( Di- 1;T) is invariant to T.
6.2 Main Results
Equivalence of DPT and PS. We now state main result which shows that the trajectories generated by a pretrained Mg will follow the same distribution as those from a well-specified PS algorithm. In particular, let PS use the well-specified prior Tpr. Let To be arbitrary task. Let
Ppe Te) and (M D, T2) denote the distributions over trajectories &H (S A)" generated from running PS and Mo(-I, D, .), respectively, in task T. given historical data D.
Theorem (DPT = PS), Let the above assumptions hold. Then, Ppo(EH D,Te) = PM, (EH D, Te) all trajectories {и-
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 Ma on a distribution Tpme over MDPs with S == [S] and := AA.. Let Dpre constructed by uniform sampling (s; a;) and observing (r,, 8() for i E [KH] Let E (ra (7,8,, an] E [0, 1). And let Dquery and G₁ be uniform over S and Sh (for all h) respectively. Finally, let Test be the distribution over test
II is still possible to implement this version, but it leaves some undesirable design choices be made.

---

tasks with the same cardinalities. For a task T, define the online cumulative regret of DPT over K episodes as Reg, (Ma) := ExE[K] (=(x)) - V++(+) where (-(s))))) Mo((Isn, (a- Eh-1) and (a)) contains the first is episodes collected from 1i:k-
Corollary 6.2 (Finite MDPs), Suppose that sup, Tppr S C for some C > 0. For the above MDP setting, the pretrained model Me satisfies ET.. [Reg, (Mo)) - O(CH" AK).
A similar analysis due to 68 allows
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 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, AA is finite but large, 0, E TR is sampled 02 ~ ((,, r/d), é : A TR is fixed feature map with suPaE. [|ó(a)] S 1, and the reward of a E A in task T is distributed as ((0r,, (a)), 1). This time, we let D (:r be given by running Thompson Sampling with Gaussian priors and likelihood functions on T.
Corollary 6.3 (Latent representation learning in linear bandits). For Trea Tpre in the above linear bandit setting, Me satisfies T [Reg, (Ma) S Odd/K).
This significantly improves over the Ö(v ( 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 additional structure in the tasks which yields tighter regret bound (for example if there are only 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 Ma to compliant D Prre- Our final result sheds light on how Dine impacts the final DPT behavior Mg. Combined with Assumption []
Mg is invariant to Dpre satisfying Definition 6.1 Proposition 6.4. Let Pere and pare be pretraining distributions that differ only by their in-context dataset distributions, denoted by Dir and D2 If D and D2, are compliant with the same support, then P pire (a") Squrry D, ) = -P Ppre(a" (Squery: D, Eh) for all a", Squery Eh-
That is, if we generate in 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 Ma. For example, TS could be used for D1 pre and PPO for D2 pre. Expert-biased datasets discussed Section violate Definition 6.1 since privileged knowledge of 7 is being used. This helps explain our empirical results that pretraiming on expert-biase biased datasets leads to a qualitatively different learned model at test-time (Fig. 3c
7 Discussion
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- 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-maie makinn 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 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 abilitics.
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-makee datasets remains 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
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 makiing settings [69]
10

---

Acknowledgments and Disclosure of Funding
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 part by NSF grant 21 12926 and ONR grant N00014-21-1-2666 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 Schmidr Center at the Broad Institute of MIT and Harvard.
JNL
References
[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-19001 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. 2022.
[3] Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emer- gent in-context learning in transformers. Advances in Neural Information Processing Systems, 35:18878-1889,, 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 Neural Information Processing Systems, 35:30583-3088 -05988 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 Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma. and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.156611 2022.
[7] Michael Laskin, Luyu Wang. Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprini urXiv: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-24644 PMLR, 2022.
[9] Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyper- decision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.0848) 2023.
[10] Richard Sutton and Andrew 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.0164: 2020.
[12] Malcolm Strens. 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):450) 2010.

---

[15] Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule. Citeseer 1990.
[16] Justin Fu, Sergey Levine, and Pieter Abbeel. One-shot learning of manipulation skills with online dynamics adaptation and neural network priors. In 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 4019-4026. IEEE, 2016.
[17] Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald $ Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta- reinforcement learning. arXiv preprint arXiv: 1347,2018, 2018.
[18] Nicholas C Landolfi, Garrett Thomas, and Tengyu Ma. A model-based approach for sample- efficient multi-task reinforcement learning. arXiv preprint arXiv:1907.04980 2019.
[19] Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel J Hsu, Thodoris Lykouris, Miro Dudik, and Robert E Schapire. Bayesian decision-makes under misspecified priors with applications to meta-learning. Advances in Neural Information Processing Systems, 34:26382- 26394, 2021.
[20] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference machine learning, pages 5331-5344 PMLR, 2019.
[21] Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: 06444, 2019.
[22] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. In International Conference on Learning Representation (ICLR), 2020.
[23] Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and exploitation for mettareiinforcement learning without sacrifices. In International conference on machine learning, pages 6925-6935. PMLR, 2021.
[24] Theodore J Perkins, Doina Precup, al. Using options for knowledge transfer in reinforcement learning. Technical report, Citeseer, 1999.
[25] Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. Advances in neural information processing systems, 31, 2018.
[26] Yiding Jiang, Evan Liu, Benjamin Eysenbach. Zico Kolter, and Chelsea Finn. Learning options via compression. Advances in Neural Information Processing Systems, 35:21 1821184-211 118421117 199, 2022.
[27] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostie meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126 1135. PMLR, 2017.
[28] Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprint arXiv:1810.06784, 2018.
[29] Yan Duan, John Schulman. Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 161 2016.
[30] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv: 1.05663, 2016.
[31] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. simple neural attentive meta-learner. arXiv preprint 1xXiv:1707.03141 2017.
12

---

[32] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani. Avishkar Bhoopchand, Nathalie Gradley-Schmiee Michael Chang, Natalie Clay, Adrian Collister, et al. Human-timescale adaptation in an open-ended task space. arXiv preprint wwxi:23301.06000 2023.
[33] Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. arXiv preprint (xXXx::33033398822 2023.
[34] Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.041299 2023.
[35] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1))5485-555) 2020.
[37] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34: 150844177797,
2021.
[38] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[39] Kuang-Huce Lee, Ofrr Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, Ian Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances in Neural Information Processing Systems, 35:27921-279366. 2022.
[40] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.061111 2022.
[41] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:22 12.066/72.. 2022.
[42] Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning h modes with one stone. Advances in neural information processing systems, 11-22955-22963 2022.
[43] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return return-ondiitooned supervised learning work for offline reinforcement learning? arXiv preprint arXiv:2206.01077. 2022.
[44] Mengjiao Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of control: Separating what you can control from what you cannot. arXiv preprint arXiv:2210.13435,, 2022.
[45] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179- 1191, 2020.
[46] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances neural information processing systems, 34:28954-28966 2021.
13

---

[47] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off-policy reinforcement learning without great exploration. Advances neural information processing systems, 15:1264-12744 2020.
[48] Kamyar Chasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? estimating uncertainties for offline rl through ensembles, and why their independence matters. Advances in Neural Information Processing Systems, 15:18267-18281, 2022.
[49] Scott Fujimoto, David Meger, and Doina Precup. Off Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062.
PMLR, 2019.
[50] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[51] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.1136,, 2019.
[52] Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmalcki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.088900 2020.
[53] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. UAI, 2019.
[54] Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforrcement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.01112, 2020.
[55] Eric Mitchell, Rafael Rafailov, Xue Bin Peng. Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021.
[56] Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning- identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-444 4618, 2021.
[57] Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline metu-reinforrement learning with online self-supervision. In International Conference on Machine Learning, pages 17811-17829 PMLR, 2022.
[58] Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, lan Osband, Zheng Wen, et tutorial on thompson sampling. Foundations and Trends@ in Machine Learning, 11(1):1-96, 2018.
[59] William Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-290 -944, 1933.
[60] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47:235-225 256, 2002.
[61] Chenjun Xiao, Yifan Wu, Jincheng Mei, Bo Dai, Tor Lattimore. Lihong Li, Csaba Szepes- vari, and Dale Schuurmans. On the optimality of batch policy optimization algorithms. In International Conference on Machine Learning, pages 11362-11371, PMIR, 2021.
[62] Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline ri? In International Conference on Machine Learning, pages 5084-5096. PMIRR 2021.
[63] Yasin Abbasi-Yadkori, Dávid Pâl, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. In Advances Neural Information Processing Systems, pages 2312-2320, 2011.
14

---

[64] Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline reinforcement learning and imitation learning: tale of pessimism. Advances in Neural Information Processing Systems, 34:1 34:11702-11716, 2021.
[65] Maxime Chevalier-Boisvert. Miniworld: Minimalistic 3d environment for rl and robotics research, 2018
[66] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint (xXv:1707.06347 2017.
[67] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 2021.
[68] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research, 39(4):1221-12: 2014.
[69] Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint 2xXv:2305.1629/1 2023.
[70] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mord- vintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.0766 07667777, 2022.
[71] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895. 2022.
[72] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in- context learning by meta-learning transformers. arXiv preprint arXiv:2212.044558 2022.
[73] Seongjin Shin, Sang- Woo Lee, Hwijeen Ahn, Sungdong Kim, HyoungSeok Kim, Boseop Kim, Kyunghyun Cho, Gichang Lee, Woomyoung Park, Jung-Woo Ha, et al. On the effect of pretraining corpora on in-context learning by a large-scale language model. arXiv preprint arXiv:2204. 13509, 2022.
[74) Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and implicit model selection in in-context learning. arXiv preprint arXiv:2301.07067, 2023.
[75] Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. arXiv preprint 2.03.0.:3000.789... 2023.
[76] Jacob Abernethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. mech- anism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305. /7040, 2023.
[77] Nan Rosemary Ke, Silvia Chiappa, Jane Wang, Anirudh Goyal, Jorg Bornschein, Melanie Rey, Theophane Weber, Matthew Botvinic, Michael Mozer, and Danilo Jimenez Rezende. Learning induce causal structure, arXiv preprint arXiv:2204.0 04485,, 2022.
[78] Anirudh Goyal, Abram Friesen, Andrea Banino, Theophane Weber, Nan Rosemary Ke, Adria Puigdomenech Badia, Arthur Guez, Mehdi Mirza, Peter C Humphreys, Ksenia Konyushova, et al. Retrieval-augmented reinforcement learning. In International Confer- ence on Machine Learning, pages 7740-7765. PMLR, 2022.
[79] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04179, 2022.
[80] Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM), 64(5): 2017.
[81] Shipra Agrawal and Randy Jia. Optimistic posterior sampling for reinforcement learning: worst-case regret bounds. Advances in Neural Information Processing Systems, 30, 2017.
15

---

[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 rein- forcement 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 S.xiv:2302.09203. 2023.
[86] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, 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. OpenAl 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, 1231)11338-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-2001 201077, 2020.
16

---

Additional Related Work
In-context learning. Beyond Decision-making and reinforcement learning, our approach takes inspiration from general in-context learning, 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 M.. 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] [4]
[75] [. (6 777 [78). 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 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- weta-learning, but they also applied this to contextual bandits to learn rewards.
For decision-making. making, in addition to the works discussed in the main paper, several works have applied transformers for in-context RL. AdA 32 meta-leams a transformer online (like RL A base online RI. algorithm is used 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 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. 4ap. DPT does not use return-to-pe 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 [3). The principle as follows: begin with prior over possible models (i.e. reward and transition functions), and maintain posterior distribution over models by updating new interactions are made. At decision-time, sample 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,
In Section 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 promising alternative approach towards scalable posterior sampling.
83
84
85.
On the theoretical side, [19] show it is possible 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
and Figure In this section, we provide a detailed description of the architecture alluded to in Section and See hyperparameter details for models in their respective sections. The model is implemented in Python with Py PyTorch [86]. The backbone of the transformer architecture we use is an autoregressive GPT-2 model from the HuggingFace transf library.
For the sake of exposition, we suppose that S and A are subsets of Rds and Idda respectively. We handle discrete state and action spaces with one-hot encoding, Consider single training datapoint
17

---

Algorithm 2 Decision-Prretranned Transformer (detailed)
1: // Collecting pretraining dataset
2: Initialize empty dataset B
3: for i in [N] do
4: Sample task T
Tp
5: Sample interaction dataset ~ Ppre(::7 T) of length n
6: Sample Squery D query and a* TT Squery)
7: Add (Aquery, D, a." to B
8: end for
9: // Training model on dataset
10: Initialize model Mg with parameters A
11: while not converged do
Sample (Squery, D,a*) from B
13:
Predict ()) Ma(-Isquery Dj) for all j E |n].
Compute loss in (5) with respect to and backpropagate update 0.
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 Test
3: Sample dataset D ~ (:7)
4: Deploying offline policy M)
5: = reset(T)
6: for in [H] do
arrgmx Mo(-|sh, D) // Most likely action 8: S/+1.7% step an)
7:
9: end for
Algorithm 4 Online test-time deployment (detailed)
1: // Online, dataset is empty as learning is from scratch
2: Initialize =
3: Sample unknown task F Tist
4: for ep in max_eps do reset(T)
5:
6: for h in [H] do
7: dh Ma(-Ish, D) // Sample action from predicted distribution
8: 8h41, T4 step (T)
9: end for
10: // Experience from previous episode added to dataset
11: Add (I,..), to
12:
end for
18

---

derived from an (potentially unknown) task T: we have dataset D of interactions within T, query state Squery and its corresponding optimal action a" " (Squery) We construct the embeddings be passed to the GPT-2 backbone in the following way. From the dataset D> = {(8j, Gj. k/, ry)) je(n), we construct vectors { = ($j. aj y ri) by stacking the elements of the transition tuple into dimension dg :- 2ds + d + for each j in the sequence. This sequence of n elements is concatenated with another vector := (5query, 0) where the O vector is a vector of zeros of sufficient length to make the entire element dimension dg. The (n + Jlenngth sequence is given by X = (v:
As order does not often matter for the dataset 7 we do not use positional encoding in order to take advantage of this invariance. We first apply linear layer Linear(X (X) and pass the result to the transformer, which outputs the sequence Y (jo. ÚI.... Dn).
the continuous action case, these can be used as is for predictions of a*. For the discrete action case, we use them as logits to be converted to either a distribution over actions in A or one-hot vector predictions of a". Here, we compute action probabilities
β3 softmax(0s) E ((A)
Because of the GPT-2 causal architecture (we defer details to the original papers 88 that p. depends only on Squery and the partial dataset Dj- {(S,s ags S/ rk (EEs)) which is why we write the model notation,
17, we note
= ( Squery, D,)= (.),
to denote that the predicted probabilities of the jth element only depend on D, and not the entire D for the model M with parameters 0 E O. For example, with j O. the prediction of a" is made without any contextual information about the task t except for Squery, which can be interpreted as the prior over a We measure loss of this training example via the cross entropy for each j E [nt]
Σ JEIm]
log paa)
(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 Squery is demarcated by its zeros vector (which also acts padding). Unlike other examples of transformers used for becision-making such as the Decision Transformer (37 and Algorithm Distillation (7], DPT does not separate the individual (s, 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
Implementation Details
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 E argmax [a + v1/na where ia is the empirical mean so far for action a and ni, the number of times has been chosen far. To arrive at this constant for the bonus, coarsely tried set of plausible values given the noise and found this to perform the best.
Lower Confidence Bound (LCB). We choose actions as å E argmax {
/1/na, } where A. is the empirical mean so far for action a and n. is the number of times a has been chosen so far.
This is not always true such when data comes from an algorithm such PPO Thompson Sampling.
19

---

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 2 and tz to match the true ones. The noise model was well-specified with the correct variance. In the linear experiments of Figure 3a and Figure (bb]
we set the prior mean and variance to and to fit the true ones better.
We choose ār E argmax (0e, ((a)) _1 where 3 and Σ. I Ej ((a.)o(a,)" and 0. = E 1] (a,)). Here, rs and a, are the reward and action observed at time 8.
LinUCB,
LinReg. LinReg (offline) is the same as LinUCB except we set B = 0 to greedily choose actions.
DPT. The transformer for DPT has embedding size of 32, context length of 500 for basic bandits and 200 for linear bandits, hidden layers, and attention heads per attention layer for all bandits. We use the AdamW optimizer with weight decay le-4, learning rate le-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 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 x 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 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 lead to faster algorithms in 77. 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 x 3 kernels each, followed by a linear layer with output dimension of 8.
RL. The reported results for RL 2 use an open-sourced implementation from 22 The implemen- tation 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 x 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 x kernels each, followed by 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 W optimizer with weight decay le-4, learning rate 1e-3, and batch-size 128.
DPT.
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 few actions, etc.): p1 Dir(1) where P1 E ((A)
and 1 E TR IR, We also mixed this with a distribution that has all mass on one action: ~ Unif(A)
and p2 (a) = 1 and p2(a) 0 for all a &f à. The final action distribution is (I - -w)pt + wp2
where w Unif(0. [10]). We train on 100,000 pretraining samples for 300 epochs with an 80/20 train/validation split. In Figure 2a Drest is generated in the same way.
20

---

Expert-Biasse Biased Bandit. To generate expert-biased datasets for pretraining, we compute the action frequencies to bias the dataset towards the optimal action. Let a* be the optimal one. As before, we take P1 ~ Dir(1). Then, (2 = 1 and p2(a) = 0 for all * a². For of bias of w, we take p = (1 w)p + wP2 with w ~ Unif (10)). We use the same pretraining sample size and epochs as before. For testing, Dest is generated the same way except we fix particular E {0, 0.5, 1) to test on.
Linear Bandit. We consider the case where |A| = 10 and d = 2. To generate environments from Tpre. we first sampled a fixed set of actions from ( Ia/d) in R to represent the features. Then, for each 7, we sampled 0, (0, Ia/d) to produce the means Pa generate the in-context dataset, we ran Gaussian TS (which does not leverage )) over n = 200 steps (see hyperparameters in previous section). Because order matters, we did not shuffle and used 1, 1000, 000 pretraining samples over 200 epochs with an 80/20 train/validation split. At test time, we Pre. Note that is fixed over all т, as standard for linear bandit.
((a)) for E A. To
set Trest Tpe and Drest
A.4 MDP Environment Details
Dark Room. The agent must navigate a 10 10 grid to find the goal within H 100 steps. The agent's 's observation is its zyyposition, the allowed actions are left, right, up, down, and stay, and the reward is only r when the agent is at the goal, and 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 x 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 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 y-position as the goal, then move left/right until the agent is on the z-position 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.
Miniworld. 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 60K datasets each of length /l = 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 the agent is not yet facing it (within 15 degrees), otherwise move forward. Of the 60K collections of datasets, query states, and optimal actions, we use 80% for training and the remaining 20% for validation.
B Additional Experimental Results
B.1 Bandits
This section reports additional experimental results in bandit environments.
reward variances. In Figures 2c and 6a 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.
21

---

Offinn Nobustness Roward
(Bernoulli)
(Bernouli))
Emp UCB
8:00
Suhoptimalin
0.3
tours)
25
DPT tours)
Cumui
20
400
500
Noise standard desfation
(a)
(b)
(c)
Figure 6: Final (after 500 steps) offline suboptimality out-of-distribution bandits with different Gaussian noise standard deviations. (b) Offline performance on -fstriibbiiiiiooonnnn Bernoulli bandits, given random in- context datasets. (c) Online cumulative regret on Bernoulli bandits. The mean and standard error are computed over 200 test tasks.
Offine Dark Beom Ihui Taoks)
Offine Bar flom (Test Tasks)
Onsiiee Dank Boom (Tain Tader)
Online
142
10
15
20
25
30
15
Datanet Compoction
Detavie Compaction
5p scdcs
Ep
sodos
(d)
Figure 7: All comparisons in Dark Room evaluated on the tasks that were seen during pretraining, displayed next their evaluations task counterparts from the main text.
Bernoulli rewards. 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 6b
and 6c
B.2 Markoy Decision Processes
This section reports additional experimental results in the Dark Room and Miniworld environments.
Performance on training tasks. In Fig. 7 we show the performance each method on the training tasks in Dark Room. Offline, DPT and AD demonstrate comparable performance tasks, indicating a minimal generalization gap to new goals. Online, DPT, AD, and RL² also achieve performance on the training tasks similar to that the test tasks.
the training
Generalization new dynamics. In this experiment, we study generalization to variations in different aspect of the MDP, namely the dynamics. We design Dark Room (Permuted), a variant of Dark Room in which the goal is fixed randomly permuted. Hence, the agent must leverage its historical context to infer the effect of each action. On 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 near optimal policy in one episode once it identifies the novel permutation in Figure g
Online Dark Room Permuted
corner but the action space is
770
20
Figure 8: Online evaluation of DPT on Dark Room when tested novel actions set permutations.
B.3 Sensitivity Analysis
We next seek to understand the sensitivity of DPT to different hy- perparameters, 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,
22

---

Layers cOMfiee
Afterion Hesos roffine)
Enbeodan Omeesan
Freraiiigg Samsles coftine!
90.
50
10c
121
1144
(a)
(b)
(c)
(d)
Figure 9: Sensitivity analysis of the offiine Dark Rook task over the GPT-2 transformer's hyperparameters: (a)
layers (b) attention heads (c) embedding dimensions (d) pretraining samples.
number layers, and number attention heads. Notably, the performance slightly worse with attention heads, which may be attributed to slight overfitting. We do 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.
C Additional Theory and Omitted Proofs
start with a well-known concentration inequality for the maximum-likelihood estimate (MLE) to provide some more justification for the approximation made in Assumption [ We state version from [89]. F be finite function class used to model a conditional distribution x (yla) for x and y E y. Assume there is f' E F such that p(y(z) (*(aa>) (realizable), and f(-lar) E A()) for all x E & and f E F (proper). Let D = {Ia, LE[M] denote a dataset of i.i.d samples where Px and 7/s ~ Pr|x (+12g). Let
f argmax
Σ [EIN]
log fy)(22))
(6)
FEF
Proposition C.1 (Theorem 21 of (81). Let D and f be given as above under the aforementioned conditions. Then, with probability at least 1 5,
- ( 8oog (1F1/8)
E Exppx
N
The finiteness of F done for simplicity, but we see that this yields dependence the cog-cardinality, common measure of complexity. Extensions to infinite F of bounded statis- tical complexity can be readily made replace this. For our setting, the bound suggests that Squery. D. D,Ex) - 0 as N + oo with high probability, provided the function class of Mo has bounded statistical complexity.
Ep,.. Pree (- D. (r)) Mp M
C.1 Posterior Sampling
Posterior sampling is most generally described with the following procedure [3 distribution T Tpre and dataset = {}. For
Initialize prior
[K]
1. Sample Th Th and compute मं.
2. Execute nī and add interactions D
3. Update posterior distribution (T+1(T) = P(r[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
Theorem (DPT > PS). Let the above assumptions hold. Then, Ppo(En D, PMo (EH D, Te) for all trajectories Ey-
23

---

Proof. Without loss of generality, for a task T, we take (-[s) to be deterministic and denote the optimal action in state as n (s). Recall that we consider fixed current task T. and fixed in-context dataset D. Define EΛ (81, 01, 8h; an).
We now formally state the variant of the full joint distribution from which sample during pretrain- ing. Let T and D be an arbitrary task and dataset and let a" E A, Squery E S, EH-1 E (S x A)"-1
and h E [0, 1] be arbitrary.
Ppre (T, a", Squery; D', hH- = (T) D'; TT) query (8query) (
(query)
(8)
x Unif(0, H - 1] II (a:ls))
KE[H]
The Onif(0, H 1] due to the fact that we sample h Unif(0, 1] and then truncate &n- from f (or, equivalently, sample &a. ~ 6, directly), marginalizing out the other variables. For I S recall that we also use the notation G5 (#1:5) to denote the marginalization of the full joint Gn. We will eventually work with the posterior of this distribution given the data 1 and history En:
Ppre (T)( x TT (T) (D:T)
r) II (a:ls)
(10)
iE(M)
P Prre (T|
II
क. (asls;)
(11)
AE[A]
We define the following random sequences and subsequences:
Eps (h: D) = ,
S% 47
(12)
where the variables are generated according to the following conditional process: Tps P()
ST ~ Pro. Ab. ~ Thm (-15)% and Sh+1 ~ T,, (-15% A%"). We also define Eps (N' h; D) to be the last h - ht' elements of Epa (h; D). Analogously, we define
Epre(h: D) = (STre Aprr
STTE APPC)
(13)
where the variables are from the process: Spre Pr., A ~ Pere( D 11; D)), and Sh+1 Tr. (-15%% ).pee. Note that Ag is sampled conditioned on the sequence Epre (h: D) so far.
We will show that Eps (h; D) and (pre(h; D) follow the same distribution for all h E [H]. For convenience, we will drop notational dependence on D, except where resolves ambiguity. Also, because of Assumption]] I we have that (-1S% D, Epre(h- 1)) = Me(-|Spre, D. Epre (h- 1)), so we will just work with pree for the remainder of the proof. We will also make use of the following lemma.
Lemma C.2. If Df is complaint, then (pre(T/DD = P(Tpe D).
Proof. From the definition of posterior sampling (using the same prior, Tpre), we have that
P(Tp= = D) x P(D T (p)
(14)
Tpp
II T, T,(6//ss,,))R,((735....,
(15)
jE[m]
α Tp
II ((( D)-1)
(16)
5E(m]
(T) (D;T) (T
(17)
(18)
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 Dpr is compliant, Since the two sides are proportional in they are equivalent.
24

---

We will prove Theorem I via induction for each h E [H] First, consider the base case for sequence
of length h 1. Recall that pr. denotes the initial state distribution of Tc. We have that the densities can be written as
(1) (1) P(S 81, Aft =01 S4 =81)
(19)
=P(c(s))P(AP =1
(20)
=pro(s)) P(AT" Tps =Tss $1)dT
(21)
= por( 1. a (1 (popps = 7 r|D, S9e = 81 $1)
(22)
-p) /. TT(ai1s1)PPs (Tps = D(D)
(23)
Ppre = a1] D)
(24)
P(EEpre = E1)
(25)
where the second line uses the sampling process of S the third marginalizes over Tps. 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 Tpe given D and Sp Since S is independent of sampling of Tys this dependence goes away in the next line. The sixth line applies Lemma C.2 and then, for 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,
p(Epa(h) (h) = En) = (h 1) = -n- P(S = 8h, Ah = ap/Epa 1) = En-1)
(26)
(27)
By the hypothesis, we have that P(Eps(h - 1) = En-11 = P(Epre (n- -1) = En- For the second factor,
P(Sh = 8h, Ahe = Epa 11 = En-1)
(28)
Tr. (8) (h-1)) P(Ap- = aap an1shh = Sh: =ps (n- - 1) = Ek-1)
(29)
=Tr, TT. 88-1 (sh18)h111 1n))
P(AAh ah Tps
-15h S/: Epo (h- 1) 1
(30)
As before, we can further rewrite the last factor as
P(A ahe Tps = -/5%% Sh: Epe (h- 1) Eh-
(31)
(anls))) P(Tps -1Sh% 8h: Spo(h 1) En-1)
(32)
where
P(S sh. Eps - 1) En- Tps PT r(D)
P(Tps = T/500 = Sh. Ep - 1) = Eh-1)
P(S = sh, Eps (h- - 1) = En-1)
(33)
x P Prre (D
II iE(h-1]
Th (8141181, a:);
(34)
x (T)
II kE(K-1)
क (a;ls))
(35)
α Pere (T D) DDuerryyyaaaaaaaa
query ( (sh) G (3
II (arls:)
iE 1]
(36)
x Ppre(T)app D, En-11)
(37)
(38)
where ix denotes that the two sides are equal up to multiplicative factors independent of r. In the first line, we used Bayes rule. In the second line, given that Tys = 7 (i.c. posterior sampling selected 7 to
25

---

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 т. Similarly, for the third and fourth lines, T,, and G do not depend on т. The final line follows from the definition of the joint pretraining distribution this regime.
Therefore, we conclude that the posterior over the value of Tpe is the same as the posterior over the task in the pretraining distribution, given Sh, E Substituting back through all the previous equations, we have
(h) En)
(39)
P(Epre (h- 1) = 1) Tr, (sh18h-1, 0n-1) /. a (an (sh) Ppre((Tr, Eh-11)d
(40)
P(Epre = 1)- TT (shlsn- au Pr D. En-1)
(41)
(h) En)
(42)
This concludes the proof.
C.3 Proof of Corollary 6.2
Corollary C.3 (Finite MDPs). Suppose that sup, Tper(T)/Tpre(T S C for some C > O. For the above MDP setting, the pretrained model Ma satisfies ET [Reg, (Mo)l S Ö(CH3/2SVAK).
Proof. Note that D is clearly compliant since it is generated by random sampling. We use the equivalence between My and posterior sampling established Theorem The proof then follows immediately from Theorem of [13] to guarantee that
T Reg, (Ma)) s H ( 3/32
(43)
where the notation Ô omits polylogarithmic dependence. The bound on the test task distribution follows from the assumed bound on the likelihood ratio under the priors:
(Ma) / (og, (Ma)dr s (Teg (Mo)
(44)
C.4 Proof of Corollary 6.3
Corollary C.4 (Latent representation learning in linear bandits). For Trett Tpre in the above linear bandit setting, Mg satisfies Tr,, [Reg, (Ma)) S (dv/K).
Proof. The distribution Pree satisfies compliance by definition because it is generated by adaptive algorithm TS. The proof once again follows by immediately deferring to the established result of (68) (Proposition 3) for linear bandits by the posterior sampling equivalence of Theorem [ This ensures that posterior sampling achieves regret Odd K). It remains, however, to justify that (-re(-1Da) will be covered by Gaussian Thompson Sampling for all D₁ with & E [K]. This is verified by noting that Pps (al D2)) for non-degenerate Gaussian Thompson Sampling (positive variances of the prior and likelihood functions) and finite K. This guarantees that any D₃ will have support.
C.5 Proof of Proposition 6.4
Proposition C.5. Let pre and phr he pretraining distributions that differ only by their in-context dataset distributions, denoted by Din and D If D1 and D2 are compliant with the same support, then p1 (a" Squery D, En. -2 (a D. End for all a² Squery D. En-
26

---

Proof. The proof follows by direct inspection of the pretraining distributions. For Pp pre- we have
Pp pre (a" Squary D,E) = J. J. (Squeey)) (pr( dd
(T|squery:
(45)
The posterior distribution over tasks is simply
Pll pre Squery: D2C) Ppre(T, Squery D.E)
(46)
Piree (Squeyy D)E)
x P pre (T) prc ( (Squery) DD prc. (D;
(47)
2 (T) P2 pre (EIr))2uucc querry (Aquurr) )D) pre (D:))
(48)
Then, the distribution over the in-context dataset can be decomposed as
D proe (D;
II iE(m]
(a:|s;; D:-1:T)
(49)
II (r(ri (T) (ails: Di-1) (E(m]
(50)
x
II iE(n]
R2 (r. u, (T, (5) (a, (al1s D
(51)
DD (D;T),
(52)
where the second equality holds because D (aylaj, D);T) is assumed to be invariant to 7 by compliance, and the fifth equality holds because D2 (a Dy; is assumed be invariant to
Therefore, conclude that, for any s, D.
Pire (rls, D2C) P P2 ( (Dlr) query (s) (D;7)
(53)
x p2 prc (T|8, E)
(54)
Since also Jr (T|s, D,E)= Jr P (T|s, D) then
pl (T D2E) = pere (r|s, D,
(55)
Substituting this back into Equation yields pare( (a ss D,E) Pare(a")) D,E).
27


Reviews:
[{"review_id": "EiZSLJlrNT", "reviewer": "Reviewer_CX5S", "summary": "\nThis work leverages transformers to solve decision making problems in a few-shot manner. In particular, they train transformers to predict optimal actions from a task given a query state and a few-shot dataset of state, action, reward transitions. The approach is tested in both multi-armed-bandit problems and POMDPs, showing superior performance than other RL and Meta-RL approaches. It also provides theoretical analysis of the proposed policy, providing regret bounds and guarantees of performance when different algorithms are used to generate the pre-training data.", "questions": "It is not intuitive to me that a model trained with expert data (DPT-Exp) behaves poorly when the % of expert data is high (Figure 3.a) since that is what is seen during training, I would expect it to b heave the best. Could authors elaborate on why this is the case?\n\n\nThe setting in lines 312-317 makes a lot of sense. Why not using that in all the experiments, if as claimed \"the original pretraining method can be seen as a simpler approximation\n316 of this modified method\"?\n\nI don't understand assumption 1.1, if we are following that assumption, why not directly using P_{pre} as out policy? Some intuition would be helpful.\n\nProposition 6.3 is very unitnuitive to me, if I understand it correctly. Is it stating that a pretraining dataset coming from PPO and one coming from a random policy would perform the same? I don't see anything in the proposition that implies otherwise, but this doesnt match in the empirical experiments.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "Method:\nThe work conducts a thorough study of the effectiveness of transformer for in-context learning in decision making problems. While the domains chosen are quite simple, there are some valuable learnings (especially if they transfer to more complex domains), such as: \n- The value of sampling actions from the transformer's likelihood in multi-armed bandits \n- The effect of using expert datasets for in-context learning \n- Better generalization than competing baselines, when tested on new tasks\n\nThe work also includes examples of POMDPS with higher dimensional observations (25x25 pixels) showing the proposed pretraining paradigm achieves better performance than competing baselines.\n\nValuable analysis of the performance of in-context learning when having access to PPO-trained policies, bringing the setting closer to AD. I particularly value how authors evaluate the effect of different sources for the context states and actions. I am wondering whether PPO, PPO is directly comparable to AD though since in their case they provide in-context actions and trajectories for different stages of training. \n\nBy chosing simple domains, the work can thoroughly explore the behavior of in-context learning for different properties of the training data, particularly in the multi-armed bandit section. To my knowledge, this is the first work to build connections of in-context learning in decision making problems and posterior sampling, which allows to:\n1. Provide an upper bound on the regret of the proposed transformer-based policy.\n2. Show that if the pretraining in-context dataset comes from policies that were only trained on the data and task present in the dataset, the final policy will be the same.  \n\nClarity:\nGenerally clear writing and background. The experiments are simple but informative, and the work is easily reproducible. More notation and clarity should be given in the theory section, such as stating what is H, and providing more intuition in the propositions.\n\n", "weaknesses": "Novelty: while this paper conducts a reasonable study on the advantages of in-context learning in decision making problems, this is a feature that has been recently studied in more complex domains and tasks (see related works next). My main concern in this work is that there is not much novelty or learnings compared to those works, and that despite making a reasonable study of the capabilities of transformers for few-shot decision making problems, they are studied in very simple scenarios. \n\nRelated work: Several works have studied in context-learning for decision making problems, in more complex scenarios [1, 2]. These works should be cited. I am also not clear on what novelty this work provides with respect to the aforementioned papers. I would like a clarification on this as well.\n\n[1] https://arxiv.org/pdf/2206.13499.pdf\n[2] https://arxiv.org/abs/2301.07608\n\nBaselines: While RL2 is a fair baseline, I think authors should look at other meta-RL algorithms, or tune the existing baselines to be more comparable. I understand that PPO is used as a reference, but it should at least be finetuned with the few-shot examples given - it is hard to compare methods when some are actually seeing less data. AD is designed to learn to do RL and therefore works best when having a curriculum of trajectories, which is not the case here.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback!\n\n**Novelty**: We apologize this was not clear. DPT indeed offers significant novel insights, distinct from prior work. \n1. It is not necessary to start with a good RL algorithm to learn a good RL algorithm. Supervised pretraining can produce one. Prior work has assumed access to good data through either a good online RL algorithm or expert demos at test-time. Instead we make the much more flexible assumption that one can compute task-specific (approximately) optimal actions for the histories *offline during pretraining*. This can be done under different assumptions, subsuming prior work, including:\n   - data collected from a (inefficient or efficient) single-task RL algorithm (Figs 5.bc, 3.bc, Cor 6.2). The final learned policy can be used to relabel the training trajectories (Figs 5.bc)\n   - data gathered in a way that includes the task-specific optimal policy distribution, which can then be used to compute the optimal policy via offline RL\n   - privileged information (e.g. sim-to-real, hindsight info) or task-specific expert labelers\n2. DPT shows provably sample-efficient posterior sampling (an algorithm that has long been studied primarily theoretically) can be scaled up via transformers and *supervised* pretraining (Thm 1, Cor 6.1, Figs 2.ab)\n3. DPT offers a new way of training a transformer, showing that learning distributions over task-specific optimal actions (rather than imitating existing algorithms) better leverages data-driven priors (Figs 3.abc, 4.b, 5.bc) \n\nWe do not believe these important and sometimes surprising implications have been empirically or theoretically elucidated before, in either simple or complex settings.\n\n**Comparison to [1] and [2]**: Thanks for recommending these papers. In addition to the aforementioned new insights, there are technical differences in operation.\n\n- Prompt-DT [1] is designed for offline RL and conditions on expert demos and return-to-go in the context. 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. Figs 4.ac). DPT does not use return-to-go, which saves a test-time hyperparameter. Pretraining is also different as DPT predicts optimal actions from arbitrary datasets.\n- AdA [2] meta learns online (like RL^2). Major differences are (1) a base online RL alg to optimize and online access to the simulator, and (2) a curriculum learning method. DPT is learned from offline supervised pretraining and it is agnostic to how the pretraining data is acquired. The curriculum is important but orthogonal to our contributions. We believe many of the ideas in AdA are complementary and can be ported to DPT.\n\n**Baselines**: Thanks for your suggestions. We are not aware of an immediate extension of PPO to few-shot finetuning. Training on different tasks does not work because the policy cannot know its task only from state. Using history effectively would make it RL^2. We presume the reviewer is suggesting to do something like MAML. We have included MAML in the response PDF (Fig 1.a). MAML is notably worse and this has been observed before in exploration problems [1]. We inspected the pre-adaptation episodes (Fig 1.b) and found that they fail to explore the task space.\n\nWe also included a new comparison with Prompt-DT (Fig 1.c). Prompt-DT requires expert demos and only works offline, whereas DPT works with expert or random data and works online.\n\n[1] Gupta et al. Meta-RL of Structured Exploration Strategies. 2018.\n\n**AD implementation**: We apologize for any confusion. We believe our implementation of AD is already consistent with both your suggestion and the recipe in the original paper. We distill PPO histories by predicting the next actions. As the authors suggest, we apply skipping to improve adaptation.\n\n**DPT-Exp performance in Fig 3.a**: We believe there may be a misunderstanding. Fig 3.a shows that the suboptimality of DPT-Exp goes to zero as the expert data % is increased. Lower is better. This appears to be consistent with what the reviewer has expected.\n\n**Why not do modified DPT (presented in the theory) in the experiments?**: While theoretically elegant, the modified version is not the most practical for a transformer. It is possible to implement, but leaves open some ambiguous design choices. E.g. one should differentiate the in-context dataset D from the history $\\xi$ and $\\xi$ should have positional encoding. In practice, it is easier to implement the original, which still has strong empirical performance (and is still equivalent to PS for bandits).\n\n**If we are following Asmp 1, why not directly use P_{pre}?**: \nYour understanding is correct. In Section 6, after Asmp 1, you may assume that any reference to $M_{\\theta}$ can be replaced with $P_{pre}$. Asmp 1 helps us answer: “if we had perfect pretraining conditions and exactly learned the target distribution, what are the theoretical characteristics of DPT?” We will clarify this in the text.\n\nHowever, the proof of Thm 1 is not simply proof-by-definition! PS and DPT are fundamentally different in their operation. The key technical challenge of the proof is showing that the posterior distribution over actions collapses in just the right way to make the two trajectory distributions equivalent.\n\n**Prop 6.3: Is it stating that a pretraining dataset from PPO and a random policy would perform the same?**: Thanks for highlighting this. Prop 6.3 helps explain the following observation. Whether we pretrain with datasets sampled randomly or from an algorithm, DPT has similar behavior. However, some pretraining datasets produce different results (DPT-Exp in Fig 3.a). Prop 6.3 shows pretraining with adaptively collected datasets will *distributionally* lead to the same model. Statistical differences may arise due to finite samples and coverage in practice. This is qualitatively consistent with what we observe in Figs 3.bc and 5.bc."}, {"review_id": "xbPij8Lvzt", "reviewer": "Reviewer_pHmJ", "summary": "This paper studies in-context learning for sequential decision-making problems. It proposes an approach, Decision Pretrained Transformer (DPT), which takes as input a (state, context) pair and produces an action. It is pretrained on a large number of tasks drawn from some task distribution and evaluated on held-out tasks. This approach can be applied in both offline settings (in which case the context is a given set of trajectories from the task at hand) and online settings (in which case the context is the agent’s experience on the current task so far). \n\nThe approach is evaluated in two settings: toy bandits (including linear contextual bandits) and simple MDPs. In the online bandit setting, it is shown to learn behaviors which efficiently balance exploration and exploitation, similarly to hand-coded bandit algorithms like UCB/LinUCB. This hold both in the stateless bandit setting and the linear contextual bandit setting. In the offline setting, it also matches Thompson sampling. \n\nThe second setting the algorithm is evaluated on is simple MDPs, namely a small partially observed grid world and MiniWorld (a pixel-based version of Minigrid). Here again, DPT is able to perform well both online and offline, even if the dataset quality is poor. \n\nFinally, the paper shows that under certain assumptions, DPT can be shown to perform posterior sampling, which is consistent with the empirical results in the bandit settings. \n\nOverall, this is pretty good paper. The idea, though simple, has not to my knowledge been explored before. The experiments are well-designed, illustrate the claims well and the writing is clear. The main downside is that the tasks feel pretty toy - the bandit examples are just Gaussians and the MDPs are grid worlds (possibly dressed up with pixels, as in MiniWorld). However, as a first proof of concept I think this still meets the bar for acceptance at NeurIPS. ", "questions": "- How does the computational cost of the proposed method compare to that of the baselines? Due to the need to perform attention over the context each step, this might be expensive. Please add a short discussion of this. \n\n- Will the code to reproduce experiments be open-sourced?", "limitations": "Yes. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The problem is definitely interesting, since large-scale, diverse datasets are often available for pretraining and zero-shot adaptation to new tasks is often desirable\n- The writing is very clear, and the experiments are well presented\n- The experiments are well-designed to show the properties of the proposed approach", "weaknesses": "- The experiments are on toy environments\n- There is no mention of open sourcing the code. I think this is important since the setting is new and future work should be able to compare to this approach on the same tasks. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive comments and detailed feedback!\n\n**How does the computational cost of the proposed method compare to that of the baselines? Due to the need to perform attention… this might be expensive**: AD also uses a GPT-like architecture, so its computational cost is approximately the same as DPT. RL^2 uses an RNN which will do inference faster (although progress in transformers is rapidly trying to close this gap). PPO is perhaps the fastest with a simple MLP, but of course it does not benefit from pretraining.\n\nFortunately, driven by the surge of interest in language models, there is considerable effort in the community to develop fast computational methods for inference with either transformers or other architectures. Known examples include FlashAttention and S4. There is already work beginning to address this for in-context RL [Lu et al, 2023]. We expect the general framework of DPT to immediately inherit benefits from these advances.\n\n**Will the code to reproduce experiments be open-sourced**: Yes, we intend to open-source the code to reproduce the experiments in a convenient and easily extensible manner. In the meantime, original code may be found in the supplementary material for the reviewers’ convenience.\n\n**Experiments on toy environments**: The environments are indeed simple, but also we believe thoroughly studying the problem empirically and theoretically with simple models is key to developing generalizable and still non-trivial insights about it. We look forward to scaling these insights in future work!"}, {"review_id": "WHVHqPTZlI", "reviewer": "Reviewer_R3a7", "summary": "This paper proposes an in-context learning based algorithm for meta-learning decision-making algorithms, combined with a novel analysis which shows that unlike prior works, the proposed algorithm is guaranteed to work. Moreover, the authors also show empirically and theoretically that their algorithm is able to improve upon its training demonstrations.", "questions": "Could you please explain why in Figure 2, UCB is outperformed by Thompson Sampling even though you stated that UCB is an optimal online multi-armed bandit algorithm?", "limitations": "Since the analysis for MDP assumes that the training demonstrations are generated with an optimal policy, the claim that the proposed method can theoretically improve upon the training demonstrations should explicitly specify that the improvement is only for the bandit and contextual bandit cases, not the more general MDP case.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper conducts a novel theoretical analysis on the ability to generalize well when meta-learning with the in-context learning paradigm applied to decision-making problems.\n\n2. The latent representation learning results in linear bandits which enable the improvement over the training demonstrations is a novel motivation for in-context learning based decision-making algorithms.\n\n3. The paper is clearly written and accessible to readers with varying backgrounds.", "weaknesses": "The formulation of a consistent learned model assumes the meta-learning algorithm is perfect. Can you relax this assumption to provide the fine-grained pre-training sample complexity?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments and detailed feedback! \n\n**Can you relax the consistency assumption on the learned model for the theory?**: Thank you for the great question. In Appendix C, we provide some justification for this via a generic MLE concentration inequality (cf. Zhang, 2006). This will be true for realizable function classes of bounded statistical complexity. Asymptotically (with N, the number of pretraining samples), the distribution can be learned through this MLE procedure. While finite sample versions can be derived [2], the calculations are long and can obfuscate the key theoretical takeaways of the theorem and proof, which is primarily about how the pretraining can lead to in-context posterior sampling behavior. While indeed extremely important, the study of the generalizability and complexity of transformers is in early stages and beyond the scope of this paper. This assumption also allows the theory to be somewhat agnostic to the underlying model used to learn the distribution.\n\n[1] Zhang. From ɛ-entropy to KL-entropy: Analysis of minimum information complexity density estimation. 2006.\n\n[2] Du et al. Few-shot learning via learning the representation, provably. 2020.\n\n**Why is UCB outperformed by Thompson Sampling when it is said to be optimal?**: We apologize we were unclear and will update the draft to clarify this statement. UCB is theoretically *minimax* optimal for regret in the sense that its proven regret *rates* match an information-theoretic lower bound, up to log factors and constants. TS is also minimax optimal. In specific instances TS or UCB may outperform each other.\n\n **Since the analysis for MDP assumes that the training demonstrations are generated with an optimal policy, the claim that the proposed method can theoretically improve upon the training demonstrations should explicitly specify that the improvement is only for the bandit and contextual bandit cases, not the more general MDP case**: \nTo clarify, we do not assume the in-context datasets are necessarily generated with an optimal policy or optimal RL algorithm. For example, a random algorithm could be executed in a well-connected tabular MDP and, if there is sufficient data, one can extract the task-specific optimal policy from collected data. In such settings, DPT can learn a RL strategy that is better than the random strategy used to generate data for the pretraining process, so that it can solve new RL problems more efficiently than the random algorithm. When we refer to it improving over pretraining data, we mean it can learn an RL strategy that may be better than the process used to generate the pretraining data (which could be an inefficient RL algorithm). We do not require demonstrations of an optimal RL algorithm. For MDPs, this is evidenced by Corollary 6.1 (pretraining in-context datasets generated by random interactions) and Figs 5.bc (pretraining data generated by PPO). We will be sure to articulate this point better in the revision."}, {"review_id": "QVIGLPk4IZ", "reviewer": "Reviewer_qnXx", "summary": "This paper introduces a new pretraining objective for decision-making tasks. Leveraging the in-context learning capabilities of transformers, the authors propose to pretrain a model that predicts the optimal action given the current state and a set of history interactions as context. Empirically, the pretrained model is capable of learning various algorithms implicitly and can transfer to unseen tasks. Theoretically, the pretraining objective can be viewed as training the model to perform posterior sampling.\n\n**After rebuttal**: The rebuttal addressed my concerns sufficiently. The new experimental results highlight the importance of both in-context datasets and optimal actions. I have increased my score to 6.", "questions": "1. In Appendix A.5, you mention that the in-context datasets for Miniworld do not include images as observations. Does this imply that in-context datasets are not important as expected for the final performance?\n2. Can you please explain the difference between DPT(PPO, PPO) and AD, and why DPT(PPO, PPO) is superior to AD?\n3. How are the datasets constructed for Dark Room (Three Tasks)?", "limitations": "The authors acknowledge certain limitations of the proposed method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This work proposes a new pretraining framework for decision-making tasks. It is one of the pioneering efforts in the problem of in-context decision-making.\n2. The paper is generally sound and presents adequate empirical results. The connection to posterior sampling is also interesting.\n3. The paper is well-written.", "weaknesses": "1. My main concern is about the requirement of optimal actions at pretraining. This assumption appears to be overly strong and I doubt its scalability to large-scale pretraining and more complex tasks.\n2. I find it unclear whether the pretrained model benefits more from in-context datasets or optimal actions. It is suggested to conduct an ablation study on the context size, which appears to be missing from the appendix. I am curious about the performance when the context size is set to 0.\n3. The modification of history-dependent pretraining in Section 6 needs further clarification. The importance of this modification is unclear to me.\n4. The authors should provide a more comprehensive comparison between their proposed pretraining objective and previous research, such as retrieval-augmented RL [1].\n\n[1] Goyal A, Friesen A, Banino A, et al. Retrieval-augmented reinforcement learning[C]//International Conference on Machine Learning. PMLR, 2022: 7740-7765.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for detailed feedback and questions!\n\n**Optimal actions at pretraining.**: Indeed, the theoretical formulation of DPT requires that task-specific optimal actions can be obtained for pretraining. While restrictive in some cases, we actually view this as being potentially quite flexible. This is because these labels can be obtained *offline* through many different means: \n- If task-specific expert labelers or demonstrations are available, those can be used directly\n- If the data is collected via a (inefficient or efficient) single-task RL algorithm, we can retroactively use labels from the learned policy (Fig 5.bc, also Fig 3.bc, Cor 6.2)\n- If privileged information is only available during pretraining but not testing (e.g. synthetic tasks, sim-to-real, hindsight [1]), we can obtain action labels that would otherwise be too hard to obtain\n- If the data covers the optimal policy distribution, offline RL can compute a near-optimal policy (same setting as [2]).\n\nIn all cases, one can “relabel” the data with the optimal action in a query state. A key feature of DPT is that it is agnostic to how the training tasks are solved as long as the labels are found. The theory of DPT does not require demonstrations from an optimal RL algorithm, only task-specific optimal actions.\n\nMany meta-RL strategies also implicitly assume the training tasks can be “solved,” but they often require that they are solved in a specific way (e.g. online like PEARL [3] or by using a realizable source algorithm like AD). \n\nPerhaps more importantly, DPT can work with approximately optimal actions. The procedure is agnostic to whether the provided labels are actually optimal or not. We showed empirically that if we only know “good” actions such as those from a final PPO policy, DPT still can work well (see Figs 5.bc).\n\n[1] Sinclair et al. Hindsight learning for mdps with exogenous inputs. 2023.\n\n[2] Mendonca et al. Guided meta-policy search. 2019.\n\n[3] Rakelly et al. Efficient off-policy meta-RL via probabilistic context variables. 2019.\n\n**Does pretraining benefit more from in-context datasets or optimal actions? What if the context size is set to 0?**: Both are critical. The in-context dataset provides information about the task through rewards and transitions. It serves the same purpose as an offline dataset for an offline RL or a history for online RL. With no data in a new task, DPT makes a guess from its learned prior. Information-theoretically, no algorithm can do better than this. The prediction becomes more refined with more data or online interactions. See Fig 2.a at $n = 0$ for offline bandits and Fig 2.b for online. For MDPs, see Figs 4.bd. We also included a new figure for offline Dark Room with DPT in the response PDF (Fig 2.c).\n\nOptimal actions are important to connect with posterior sampling (PS) on new RL problems. Again, approximate ones suffice in practice. Pretraining to predict actions only in the in-context datasets would be equivalent to AD.\n\n\n**The modification of history-dependent pretraining in Section 6 needs further clarification.**: Thanks for pointing this out. The context is augmented with a history of optimal actions and states up to the current timestep during pretraining. During deployment, $D$ behaves the same, and the learner’s in-episode history is recorded in $\\xi$. This modification is important to establish the precise connection between DPT and PS. In particular, conditioning on the history allows us to ‘collapse’ the posterior distribution over actions that are only consistent with policies that could have been generated by PS.\n\n\n**Comparison to retrieval-augmented RL [1]**: Thank you for pointing out this relevant work. [1] augments a standard RL algorithm with a dataset retrieval process to supply additional immediate information for the agent to use. They demonstrate this can also be used for multitask settings. DPT learns directly to map datasets and states to actions via offline supervised pretraining, and we also provide a theoretical analysis.\n\n**The in-context datasets for Miniworld do not include images as observations. Does this imply that in-context datasets are not important as expected for the final performance?**: We apologize for the confusion. To make this experiment more clear, we have included a new version in the response PDF (Fig 2.ab) which uses images and completely removes the xy-position information. To answer your question, in the original version, the in-context datasets provide contextual information about the task through the xy-position + rewards. Observations are indeed necessary to identify and solve the task.\n\n**Difference between DPT(PPO, PPO) and AD, and why DPT(PPO, PPO) is superior to AD?**: Both use data from PPO. A key difference is that AD is directly trained on the trajectories generated by PPO (i.e. to predict the next action from the learning algorithm), but DPT(PPO, PPO) is trained to predict the action from the *final PPO policy*, given an offline history and a state. This action acts as an approximation of the optimal policy.  AD will only be as good as the algorithm it imitates, but DPT may be better because it can act as a good approximation of posterior sampling with a data-driven prior. [Note that the AD authors consider skipping trajectories to predict slightly further ahead, which we implemented.]\n\n\n**How are the datasets constructed for Dark Room (Three Tasks)?**: The three tasks in Dark Room (Three Tasks) differ in the start and goal location of the agent: (1) start at (0, 5) and goal at (9, 5); (2) start at (5, 0) and goal at (5, 9); (3) start at (0, 5) and goal at (5, 9). The pretraining in-context dataset here is collected by a random policy for the two goals. The in-context dataset consists of demonstrations of the first two tasks, where the reward is 1 only for transitions to (5, 9), and our evaluation is in the third task. This demonstrates the ability to do offline stitching of suboptimal demonstrations."}]

Author rebuttal:
Dear Reviewers and AC,

Thank you for your detailed and thoughtful reviews. We are delighted to hear that the reviewers found the problem and results interesting and novel (R3a7, pHmJ), and the analyses well-designed and thorough (pHmJ, CX5S). We share excitement about the insights such as the connection to posterior sampling (qnXx, CX5S) and the ability to exploit latent structure from pretraining (R3a7). We also greatly appreciate the feedback and will incorporate it in the revision. Please see the individual responses for answers to specific questions and clarifications. We have also attached a PDF response with new figures addressing certain questions that warranted new experiments. Here are a few key points:
- Reviewers CX5S and qnXx: We have included discussion of how (approximately) optimal actions during offline pretraining can be acquired in many settings, that are a superset of related meta-RL problem settings.
- Reviewer CX5S: We have highlighted major contributions/implications of the work that we believe have not been elucidated in prior work, including how
   - it is not necessary to start with a good RL algorithm to learn one, via supervised pretraining.
   - provably sample-efficient posterior sampling (studied primarily theoretically) can be scaled up via transformers and *supervised* pretraining.
   - learning distributions over task-specific optimal actions (rather than imitating existing algorithms) better leverages data-driven priors.
- Reviewer CX5S: In the attached PDF, we have included additional MAML and Prompt-DT baselines in their respective settings, as well as a MAML visualization.
- Reviewer qnXx: In the attached PDF, we have included a new version of the MiniWorld results with position completely removed and instead images used in the in-context dataset. The pretraining was also scaled up for all methods. We have also included an ablation varying the context size for offline Dark Room (in addition to existing ones in the original submission for offline/online bandits and online Dark Room).

We are happy to answer any additional questions.

Meta-review: This paper presents DPT, a new approach for in-context reinforcement learning (ICRL) with transformers. The approach differs from existing ICRL approaches such as Algorithm Distillation in several aspects, such as using the optimal actions as training labels, as well as the flexibility of the in-context dataset. The paper demonstrates that DPT works well experimentally for solving new tasks in-context in several bandits and small RL environments. Theoretically, DPT is shown to implement the posterior sampling algorithm at the minimizer of the population objective.

All reviewers are positive about the contributions of the paper after rebuttal. I believe DPT could be of broad interest to the community as an approach for decision making with transformers. Therefore, I recommend acceptance with spotlight, and congratulate the authors for the nice work.
Decision: Accept (spotlight)