Title: Supervised Pretraining Can Learn In-Context Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

For supervised learning, transformer-based models trained at scale have shown impressive abilities to perform tasks given an input context, often referred to as few-shot prompting or in-context learning [1]. In this setting, a pretrained model is presented with a small number of supervised inputoutput examples in its context, and is then asked to predict the most likely completion (i.e. output) of an unpaired input, without parameter updates. Over the last few years, in-context learning has been applied to solve a range of tasks [2] and a growing number works are beginning to understand and analyze in-context learning for supervised learning [3, 4, 5, 6]. In this work, our focus is to study and understand in-context learning applied to sequential decision-making, specifically in the context of reinforcement learning (RL) settings. Decision-making (e.g. RL) is considerably more dynamic and complex than supervised learning. Understanding and leveraging in-context learning here could potentially unlock significant improvements in an agent's ability to adapt and make few-shot decisions in response to observations from the world. Such capabilities are instrumental for practical applications ranging from robotics to recommendation systems.

For in-context decision-making [7, 8, 9], rather than input-output tuples, the context takes the form of state-action-reward tuples representing a dataset of interactions with an unknown environment. The agent must leverage these interactions to understand the dynamics of the world and what actions lead to good outcomes. A hallmark of good decision-making in online RL algorithms is a judicious balance of selecting exploratory actions to gather information and selecting increasingly optimal actions by exploiting that information [10]. In contrast, an RL agent with access to only a suboptimal offline dataset should produce a policy that conservatively selects actions [11]. An ideal in-context decision-maker should exhibit similar behaviors.

To study in-context decision-making formally, we propose a new simple supervised pretraining objective, namely, to train (via supervised learning) a transformer to predict an optimal action label 2 given a query state and an in-context dataset of interactions, across a diverse set of tasks. We refer to the pretrained model as a Decision-Pretrained Transformer (DPT). Once trained, DPT can be deployed as either an online or offline RL algorithm in a new task by passing it an in-context dataset of interactions and querying it for predictions of the optimal action in different states. For example, online, the in-context dataset is initially empty and DPT's predictions are uncertain because the new task is unknown, but it fills the dataset with its interactions as it learns and becomes more confident about the optimal action. We show empirically and theoretically that DPT yields a surprisingly effective in-context decision-maker with regret guarantees. As it turns out, DPT effectively performs posterior sampling - a provably sample-efficient Bayesian RL algorithm that has historically been limited by its computational burden [12, 13]. We summarize our main findings below.

Predicting optimal actions alone gives rise to near-optimal decision-making algorithms. The DPT objective is solely based on predicting optimal actions from in-context interactions. At the outset, it is not immediately apparent that these predictions at test-time would yield good decisionmaking behavior when the task is unknown and behaviors such as online exploration are necessary to solve it. Intriguingly, DPT as an algorithm is capable of dealing with this uncertainty in-context. For example, despite not being explicitly trained to explore, DPT exhibits an exploration strategy on par with hand-designed algorithms, as a means to discover the optimal actions.

DPT generalizes to new decision-making problems, offline and online. We show DPT can handle reward distributions unseen in its pretraining data on bandit problems as well as unseen goals, dynamics, and datasets in simple MDPs. This suggests that the in-context strategies learned during pretraining are robust and generalizable without any parameter updates at test time.

DPT improves over the data used to pretrain it by exploiting latent structure. As an example, in parametric bandit problems, specialized algorithms can leverage structure (such as linear rewards) and offer provably better regret. However, a representation must be known in advance. Perhaps surprisingly, we find that pretraining on linear bandit problems, even with unknown representations, leads DPT to select actions and explore in a way that matches an efficient linear bandit algorithm. This holds even when the source pretraining data comes from a suboptimal algorithm (i.e., one that does not take advantage of any latent structure), demonstrating the ability to learn improved in-context strategies beyond what it was trained on.

Posterior sampling can be implemented via in-context learning. Posterior sampling (PS), a generalization of Thompson Sampling, can provably sample-efficiently solve online RL problems [12, 13], but a common criticism is the lack of computationally efficient ways to update and sample from a posterior distribution. DPT can be viewed as learning a posterior distribution over optimal actions, shortcutting the PS procedure. Under some conditions, we show theoretically that DPT in-context is equivalent to PS. Furthermore, DPT's prior and posterior updates are grounded in data rather than needing to be specified a priori . This suggests that in-context learning could be a possible path towards scalable and practical RL via posterior sampling.

2 Related Work

Meta-learning. Algorithmically, in-context learning falls under the meta-learning framework [14, 15]. At a high-level, these methods attempt to learn some underlying shared structure of the training distribution of tasks to accelerate learning of new tasks. For decision-making and RL, there is often

a choice in what shared 'structure' is specifically learned such as the dynamics of the task [16, 17, 18, 19], a task context identifier [20, 21, 22, 23], temporally extended skills and options [24, 25, 26], or initialization of a neural network policy [27, 28]). In-context learning can be viewed as taking a more agnostic approach by learning the learning algorithm itself, more similar to [29, 30, 31, 32, 8]. Algorithm Distillation (AD) [7, 33] falls under this category, using autoregressive supervised learning to distill traces of a single-task RL algorithm into a task-agnostic model. While DPT also leverages autoregressive SL, it does not distill an existing RL algorithm in order to imitate how to learn. Instead, we train DPT to predict task-specific optimal actions, yielding potentially emergent strategies at test time that automatically leverage the task structure to behave similarly to posterior sampling.

Autoregressive transformers for decision-making. In decision-making fields such as RL and imitation learning, transformer models trained using autoregressive supervised action prediction have proliferated [34], inspired by the successes of these techniques for large language models [35, 36, 1]. For example, Decision Transformer (DT) [37, 38, 8] uses a transformer to autoregressively model sequences of actions from offline experience data, conditioned on the achieved return. During inference, one can then query the model conditioned on a desired return value. This approach has been shown to scale favorably to large models and multi-task settings [39], at times exceeding the performance of large-scale multi-task imitation learning with transformers [40, 41, 42]. However, DT is known to be provably (and unboundedly) sub-optimal in common scenarios [43, 44]. A common criticism of DT, and supervised learned transformers in general, is their inability to improve upon the dataset. For example, there is little reason for DT to output meaningful behavior if conditioned on return higher than any observed in training, without strong extrapolation assumptions [43]. In contrast, a major contribution of our work is theoretical and empirical evidence for the ability of DPT to improve over behaviors seen in the dataset in terms of regret.

Value and policy-based offline RL. Offline RL algorithms offer the opportunity to learn from existing datasets. To address distributional shift, many prior algorithms incorporate the principle of value pessimism [45, 46, 47, 48], or policy regularization [49, 50, 51, 52, 53]. To reduce the amount of offline data required in a new task, methods for offline meta-RL can reuse interactions collected in a set of related tasks [54, 55, 56]. However, they still must address distribution shift, requiring solutions such as policy regularization [54] or additional online interactions [57]. DPT follows the success of autoregressive models like DT and AD, avoiding these issues. With our pretraining objective, DPT also leverages offline datasets for new tasks more effectively than AD.

3 In-Context Learning Model

Basic decision models. The basic decision model of our study is the finite-horizon Markov decision process (MDP). An MDP is specified by the tuple τ = ⟨S , A , T, R, H, ρ ⟩ to be solved, where S is the state space, A is the action space, T : S × A → ∆( S ) is the transition function, R : S × A → ∆( R ) is the reward function, H ∈ N is the horizon, and ρ ∈ ∆( S ) is the initial state distribution. A learner interacts with the environment through the following protocol: (1) an initial state s 1 is sampled from ρ ; (2) at time step h , the learner chooses an action a h and transitions to state s h +1 ∼ T ( ·| s h , a h ) , and receives a reward r h ∼ R ( ·| s h , a h ) . 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 π ⋆ , which maximizes the value function V ( π ⋆ ) = max π V ( π ) := max π E π ∑ h r h . When necessary, we use the subscript τ to distinguish V τ and π ⋆ τ for the specific MDP τ . We assume the state space is partitioned by h ∈ [ H ] so that π ⋆ 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 ⋆ ∈ argmax a ∈A E [ r 1 | a 1 = a ] .

Pretraining. We give pseudocode in Algorithm 1 and a visualization of the transformer in Figure 1. Details on the precise structure of the transformer can be found in Section A.1. Let T pre be a distribution over tasks at the time of pretraining. A task τ ∼ T pre can be viewed as a specification of an MDP, τ = ⟨S , A , T, R, H, ρ ⟩ . The distribution T pre can span different reward and transition functions and even different state and action spaces. We then sample a context (or a prompt) which consists of a dataset D ∼ D pre ( · ; τ ) of interactions between the learner and the MDP specified by τ . D = { s j , a j , s ′ j , r j } j ∈ [ n ] is a collection of transition tuples taken in τ . We refer to D as the in-context dataset because it provides the contextual information about τ . D could be generated through variety of means, such as: (1) random interactions within τ , (2) demonstrations from an expert, and (3) rollouts of an algorithm. Additionally, we independently sample a query state s query from the distribution D query over states S and a label a ⋆ is sampled from the optimal policy π ⋆ τ ( ·| s query )

Algorithm 1 Decision-Pretrained Transformer (DPT): Training and Deployment

for task τ (see Section 5.3 for how to implement this in common practical scenarios). We denote the joint pretraining distribution over tasks, in-context datasets, query states, and action labels as P pre :

Given the in-context dataset D and a query state s query, we can train a model to predict the optimal action a ⋆ in response simply via supervised learning. Let D j = { ( s 1 , a 1 , s ′ 1 , r 1 ) , . . . , ( s j , a j , s ′ j , r j ) } denote the partial dataset up to j samples. Formally, we aim to train a causal GPT-2 transformer model M parameterized by θ , which outputs a distribution over actions A , to minimize the expected loss over samples from the pretraining distribution:

Throughout, we set the loss to be the negative log-likelihood with ℓ ( M θ ( · | s query , D j ) , a ⋆ ) := -log M θ ( a ⋆ | s query , D j ) . 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 M θ , essentially treating this as a classification problem. The resulting output model M θ can be viewed as an algorithm that takes in a dataset of interactions D and can be queried with a forward pass for predictions of the optimal action via inputting a query state s query. We refer to the trained model M θ as a Decision-Pretrained Transformer (DPT).

Testing. After pretraining, a new task (MDP) τ is sampled from a test-task distribution T test . If the DPT is to be tested offline , then a dataset (prompt) is a sampled D ∼ D test ( · ; τ ) and the policy that the model in-context learns is given conditionally as M θ ( · | · , D ) . Namely, we evaluate the policy by selecting action a h ∈ argmax a M θ ( a | s h , D ) when the learner visits state s h . If the model is to be tested online through multiple episodes of interaction, then the dataset is initialized as empty D = {} . At each episode, M θ ( · | · , D ) is deployed where the model samples a h ∼ M θ ( ·| s h , D ) upon observing state s h . Throughout a full episode, it collects interactions { s 1 , a 1 , r 1 , . . . , s H , a H , r H } which are subsequently appended to D . The model then repeats the process with another episode, and so on until a specified number of episodes has been reached. A key distinction of the testing phase is that there are no updates to the parameters of M θ . This is in contrast to hand-designed RL algorithms that would perform parameter updates or maintain statistics using D to learn from scratch. Instead, the model M θ performs a computation through its forward pass to generate a distribution over actions conditioned on the in-context D and query state s h .

4 Learning in Bandits

We begin with an empirical investigation of DPT in a multi-armed bandit, a well-studied special case of the MDP where the state space S is a singleton and the horizon H = 1 is a single step. We will examine the performance of DPT both when aiming to select a good action from offline historical data and for online learning where the goal is to maximize cumulative reward from scratch. Offline, it is critical to account for uncertainty due to noise as certain actions may not be sampled well enough. Online, it is critical to judiciously balance exploration and exploitation to minimize overall regret. For detailed descriptions of the experiment setups, see Appendix A.

Pretraining distribution. For the pretraining task distribution T pre, we sample 5 -armed bandits ( |A| = 5 ). The reward function for arm a is a normal distribution R ( ·| s, a ) = N ( µ a , σ 2 ) where µ a ∼ Unif [0 , 1] independently and σ = 0 . 3 . To generate in-context datasets D pre, we randomly generate action frequencies by sampling probabilities from a Dirichlet distribution and mixing them with a point-mass distribution on one random arm (see details in Appendix A.3). Then we sample the actions accordingly from this distribution. This encourages diversity of the in-context datasets. The optimal policy π ⋆ τ for bandit τ is argmax a µ a , which we can easily compute during pretraining. We pretrain the model M θ 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 bandits 3 . All of the algorithms are designed to reason in a particular way about uncertainty based on their observations.

Empirical mean algorithm (Emp) selects the action with the highest empirical mean reward naively.

Upper Confidence Bound (UCB) selects the action with the highest upper confidence bound.

Lower Confidence Bound (LCB) selects the action with the highest lower confidence bound.

Thompson Sampling (TS) selects the action with the highest sampled mean from a posterior distribution over reward models. The prior and likelihood functions are Gaussian.

Emp and TS [58, 59] can both be used for offline or online learning; UCB [60] is an online algorithm that achieves exploration via optimism under uncertainty; and LCB [61, 62] is used to minimize suboptimality given an offline dataset by selecting actions pessimistically. We evaluate algorithms with standard bandit metrics. Offline, we use the suboptimality µ a ⋆ -µ ˆ a where ˆ a is the chosen action. Online, we use cumulative regret: ∑ k µ a ⋆ -µ ˆ a k where ˆ a k is the k th action chosen.

DPT learns to reason through uncertainty. As shown in Figure 2a, in the offline setting, DPT significantly exceeds the performance of Emp and LCB while matching the performance of TS, when the in-context datasets are sampled from the same distribution as during pretraining. The results suggest that the transformer is capable of reasoning through uncertainty caused by the noisy rewards in the dataset. Unlike Emp which can be fooled by noisy, undersampled actions, the transformer has learned to hedge to a degree. However, it also suggests that this hedging is fundamentally different from what LCB does, at least on this specific distribution 4 .

Interestingly, the same transformer produces an extremely effective online bandit algorithm when sampling actions instead of taking an argmax. As shown in Figure 2b, DPT matches the performance

0

100

200

300

400

500

Data

0.1

0.3

0.02

Suboptimality

Offline Bandit

Emp

LCB

TS

DPT (ours)

(a)

0

100

200

300

400

500

Data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

Cumulative Regret

Online Bandit

Emp

UCB

TS

DPT (ours)

(b)

0.0

0.3

0.5

Noise standard deviation

0

5

10

15

20

25

30

Final Cumulative Regret

Online Robustness to Reward Shift

Emp

UCB

TS

DPT (ours)

(c)

of classical optimal algorithms, UCB and TS, which are specifically designed for exploration. This is notable because DPT was not explicitly trained to explore, but its emergent strategy is on par with some of the best. In Figure 2c, we show this property is robust to noise in the rewards not seen during pretraining by varying the standard deviation. In Appendix B, we show this generalization happens offline too and even with unseen Bernoulli rewards.

Leveraging latent structure from suboptimal data. We now investigate whether DPT can learn to leverage the inherent structure of a problem class, even without prior knowledge of this structure and even when learning from in-context datasets that do not explicitly utilize it. More precisely, we consider T pre to be a distribution over linear bandits, where the reward function is given by E [ r | a, τ ] = ⟨ θ τ , ϕ ( a ) ⟩ and θ τ ∈ R d is a task-specific parameter vector and ϕ : A → R d is fixed feature vector that is the same for all tasks. Given the feature representation ϕ , LinUCB [63], a UCB-style algorithm that leverages ϕ , should achieve regret ˜ O ( d √ K ) over K steps, a substantial gain over UCB and TS when d ≪|A| . We do not assume ϕ is known to DPT, besides whatever can be gleaned from the pretraining data. To evaluate the ability of DPT to uncover latent structure, we collect pretraining data with in-context datasets gathered by standard TS, which does not leverage the linear structure. Figures 3a and 3b show that DPT can exploit the unknown linear structure, implicitly learning a surrogate for ϕ , allowing to do more informed exploration online and decision-making offline. It is nearly on par with LinUCB (which is 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 D pre is generated by mixing the in-context datasets with varying fractions of expert data, biasing D pre towards datasets that contain more examples of the optimal action. We denote this model by DPT-Exp. In Figure 3c, we plot the test-time performance of both pretrained models when evaluated on new offline datasets with varying percentages of expert data 5 . Our results suggest that when the pretraining distribution is also biased towards expert-suboptimal data, DPT-Exp behaves similarly to LCB, while DPT continues to resemble TS. This is quite interesting as, for other methods such as TS, it is less clear how to automatically incorporate the right amount of expert bias to yield the same effect, but DPT can leverage this from pretraining.

5 Learning in Markov Decision Processes

We next study how DPT can tackle Markov decision processes by testing its ability to perform exploration and credit assignment. In the following experiments, the DPT demonstrates generalization to new tasks, scalability to image-based observations, and capability to stitch in-context behaviors (Section 5.2). This section also examines whether DPT can be pretrained with datasets and action labels generated by a different RL algorithm, rather than the exact optimal policy (Section 5.3).

0

25

50

75

100

125

150

175

200

Data

10

2

10

1

Suboptimality

Offline Linear Bandit

TS (source)

LinReg

DPT (ours)

(a)

0

25

50

75

100

125

150

175

200

Data

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

20.0

Cumulative Regret

Online Linear Bandit

TS (source)

LinUCB

DPT (ours)

(b)

0

50

100

% of Expert Data

0.00

0.01

0.02

0.03

0.04

0.05

Final Suboptimality

Offline Adaptation to Expert Datasets

Emp

LCB

TS

DPT (ours)

DPT-Exp (ours)

(c)

Random

Expert

Dataset Composition

0

20

40

60

80

Return

Offline Dark Room (Test Tasks)

Dataset

AD

DPT (ours)

(a)

0

5

10

15

20

25

30

35

40

Episodes

0

20

40

60

80

Return

DPT (ours)

AD

PPO

RL2

(b)

Random

Expert

Dataset Composition

0

10

20

30

40

Return

Offline Miniworld

Dataset

AD

DPT (ours)

(c)

0

5

10

15

20

25

30

35

40

Episodes

0

10

20

30

40

50

Return

Online Miniworld

DPT (ours)

AD

PPO

RL2

(d)

5.1 Experimental Setup

Environments. We consider environments that require targeted exploration to solve the task. The first is Dark Room [22, 7], a 2D discrete environment where the agent must locate the unknown goal location in a 10 × 10 room, and only receives a reward of 1 when at the goal. We hold out a set of goals for generalization evaluation. The second is Miniworld [65], a 3D visual navigation problem to test the scalability of DPT to image observations. The agent is in a room with four boxes of different colors, and must find the target box. The target color is unknown initially. It receives a reward of 1 only when near the correct box. Environments and pretraining dataset details are in App. A.4 and A.5.

Comparisons. Our experiments aim to understand the effectiveness of DPT by comparing it to other context-based meta-RL algorithms based on supervised and RL objectives.

Proximal Policy Optimization (PPO) [66]: We compare to this single-task RL algorithm, which trains from scratch without any pretraining data, to contextualize the performance of DPT and other meta-RL algorithms.

Algorithm Distillation (AD) [7]: AD first generates a dataset of learning histories by running an RL algorithm in each training task. Then, given a sampled subsequence h j = ( s j , a j , r j , . . . , s j + c ) from a learning history, a transformer is trained to predict the next action a j + c from the history.

RL 2 [29]: This online meta-RL comparison uses a recurrent neural network to adapt the agent's policy from the given context. Unlike AD and DPT, which are trained with a supervised objective, the RL 2 agent is trained to maximize the expected return with PPO.

PPO and RL 2 are online algorithms, while AD is capable of learning both offline and online. Details on the implementation of these algorithms can be found in Appendix A.2.

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 on a set of 20 held-out goals not in the pretraining dataset. When given an expert dataset, DPT achieves near-optimal performance. Even when given a random dataset, which has an average total reward of 1 . 1 , DPT obtains a much higher average return of

y

10

8

6

4

2

0

2

0.0

2.5

5.0

7.5

x

(a)

Random

Expert

Dataset Composition

0

20

40

60

80

100

Return

Offline PPO Ablations (T est T asks)

DPT

DPT (Rand, PPO)

DPT (PPO, PPO)

AD

(b)

0

5

10

15

20

25

30

35

40

Episodes

0

20

40

60

80

Return

Online PPO Ablations (Test Tasks)

DPT

DPT (Rand, PPO)

DPT (PPO, PPO)

AD

(c)

61 . 5 (see Fig. 4a). Qualitatively, we observe that when the in-context dataset contains a transition to the goal, DPT immediately exploits this and takes a direct path to the goal. In contrast, while AD demonstrates strong offline performance with expert data, it performs worse in-context learning with random data compared to DPT. The difference arises because AD is trained to infer a better policy than the in-context data, but not necessarily the optimal one.

We next evaluate DPT, AD, RL 2 , and PPO online without any prior data from the 20 test-time Dark Room tasks, shown in Fig. 4b. After 40 episodes, PPO does not make significant progress towards the goal, highlighting the difficulty of learning from such few interactions alone. RL 2 is trained to perform adaptation within four episodes each of length 100 , and we report the performance after the four adaptation episodes. Notably, DPT on average solves each task faster than AD and reaches a higher final return than RL 2 , demonstrating its capability to explore effectively online even in MDPs. In Appendix B, we also present results on generalization to new dynamics.

Learning from image-based observations. In Miniworld, the agent receives RGB image observations of 25 × 25 pixels. As shown in Fig. 4d, DPT can solve this high-dimensional task offline from both random and expert datasets. Compared to AD and RL 2 , DPT also learns online more efficiently.

Stitching novel trajectories from in-context subsequences. A desirable property of some offline RL algorithms is the ability to stitch suboptimal subsequences from the offline dataset into new trajectories with higher return. To test whether DPT exhibits stitching, we design the Dark Room (Three Tasks) environment in which there are three possible tasks. The pretraining data consists only of expert demonstrations of two of them. At test-time DPT is evaluated on third unseen task, but its offline dataset is only expert demonstrations of the original two. Despite this, it leverages the data to infer a path solving the third task (see Fig. 5a).

5.3 Learning from Algorithm-Generated Policies and Rollouts

So far, we have only considered action labels provided by an optimal policy. However, in some tasks, an optimal policy is not readily available even in pretraining. In this experiment, we use actions labeled by a policy learned via PPO and in-context datasets sampled from PPO replay buffers. We train PPO agents in each of the 80 train tasks for 1 K episodes to generate 80 K total rollouts, from which we sample the in-context datasets. This variant, DPT (PPO, PPO), performs on par with DPT and still better than AD, as shown in Figures 5b and 5c. DPT (PPO, PPO) can be viewed as a direct comparison between our pretraining objective and that of AD, given the same pretraining data but just used differently. We also evaluated a variant, DPT (Rand, PPO), which pretrains on random in-context datasets (like DPT), but still using PPO action labels. The performance is worse than other DPT variants in some settings, but only marginally so. In Appendix B, we analyze the sensitivity of DPT to other hyperparameters, such as context size and amount of pretraining data. These results give further evidence that DPT can learn an algorithm that transcends its pretraining data.

6 Theory

We now shed light on the observations of the previous empirical results through a theoretical analysis. Our main result shows that DPT (under a slight modification to pretraining) essentially is capable of executing posterior sampling (PS) reinforcement learning in-context. PS is a generalization of

Thompson Sampling for RL in MDPs [12]. It maintains and samples from a posterior over tasks τ given historical data D and executes optimal policies π ⋆ τ (see Appendix C for a formal outline). It is provably sample-efficient with online Bayesian regret guarantees [13], but maintaining posteriors is generally computationally intractable. The ability for DPT to perform PS in-context suggests a path towards computation- and provably sample-efficient RL with priors learned from the data.

6.1 History-Dependent Pretraining and Assumptions

We start with a modification to the pretraining of DPT. Rather than conditioning only on s query and D to predict a ⋆ ∼ π ⋆ τ ( ·| s query ) , we propose also conditioning on a sequence ξ h = ( s 1: h , a ⋆ 1: h ) where s 1: h ∼ S h ∈ ∆( S h ) is a distribution over sets of states, independent of τ , and a ⋆ h ′ ∼ π ⋆ τ ( ·| s h ′ ) for h ′ ∈ [ h ] . Thus, we use π ⋆ τ to label both the query state (which is the prediction label) and the sequence of states sampled from S h . 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 to condition on the history ξ h -1 of states that M θ visits and the actions that it takes in those states. Formally, the learned M θ is deployed as follows, given D . (1) At h = 0 , initialize ξ 0 = () to be empty. (2) At step h , visit s h and find a h by sampling from M θ ( ·| s query , D, ξ h -1 ) . (3) Append ( s h , a h ) to ξ h -1 to get ξ h . Note for bandits and contextual bandits ( H = 1 ), there is no difference between this and the original pretraining procedure of prior sections because ξ 0 is empty. For MDPs, the original DPT can be viewed as a practically convenient approximation. 6

We now make several assumptions to simplify the analysis. First, assume D query , D pre, and S have sufficient support such that all conditional probabilities of P pre are well defined. Similar to other studies of in-context learning [67], we assume M θ fits the pretraining distribution exactly with enough coverage and data, so that the focus of the analysis is just the in-context learning abilities.

Assumption 1. (Learned model is consistent). Let M θ denote the pretrained model. For all ( s query , D, ξ h ) , we have P pre ( a | s query , D, ξ h ) = M θ ( a | s query , D, ξ h ) for all a ∈ A .

To provide some cursory justification, if M θ is the global minimizer of (2), then we might expect that E P pre ∥ P pre ( ·| s query , D, ξ h ) -M θ ( ·| s query , D, ξ h ) ∥ 2 1 → 0 as the number of pretraining samples N → ∞ with high probability for transformer model classes of bounded complexity (see Proposition C.1). Approximate versions of the above assumptions are easily possible but obfuscate the key elements of the analysis. We also assume that the in-context dataset D ∼ D pre is compliant [62], meaning that the actions from D can depend only on the observed history and not additional confounders. Note that this still allows D pre to be very general - it could be generated randomly or from adaptive algorithms like PPO or TS. Without compliance D pre can influence M θ . In Proposition 6.4, we show that all compliant D pre form a sort of equivalence class that generate the same M θ . For the remainder, we assume all D pre are compliant.

Definition 6.1 (Compliance) . The in-context dataset distribution D pre ( · ; τ ) is compliant if, for all i ∈ [ n ] , the i th action of the dataset, a i , is conditionally independent of τ given the i th state s i and partial dataset, D i -1 , so far. In other words, the distribution D pre ( a i | s i , D i -1 ; τ ) is invariant to τ .

6.2 Main Results

Equivalence of DPT and PS. We now state our main result which shows that the trajectories generated by a pretrained M θ will follow the same distribution as those from a well-specified PS algorithm. In particular, let PS use the well-specified prior T pre . Let τ c be an arbitrary task. Let P ps ( · | D,τ c ) and P M θ ( · | D,τ c ) denote the distributions over trajectories ξ H ∈ ( S ×A ) H generated from running PS and M θ ( ·|· , D, · ) , respectively, in task τ c given historical data D .

Theorem 1 (DPT ⇐⇒ PS) . Let the above assumptions hold. Then, P ps ( ξ H | D,τ c ) = P M θ ( ξ H | D,τ c ) for all trajectories ξ H .

Regret implications. With almost no extra work besides verifying assumptions, we can see this result in action. Let us specialize to the finite MDP setting [13]. Suppose we pretrain M θ on a distribution T pre over MDPs with S := |S| and A := |A| . Let D pre be constructed by uniform sampling ( s i , a i ) and observing ( r i , s ′ i ) for i ∈ [ KH ] . Let E [ r h | s h , a h ] ∈ [0 , 1] . And let D query and S h be uniform over S and S h (for all h ) respectively. Finally, let T test be the distribution over test

tasks with the same cardinalities. For a task τ , define the online cumulative regret of DPT over K episodes as Reg τ ( M θ ) := ∑ k ∈ [ K ] V τ ( π ⋆ τ ) -V τ (ˆ π k ) where ˆ π k ( ·| s h ) = M θ ( ·| s h , D ( k -1) , ξ h -1 ) and D ( k ) contains the first k episodes collected from ˆ π 1: k .

Corollary 6.2 (Finite MDPs) . Suppose that sup τ T test ( τ ) / T pre ( τ ) ≤ C for some C > 0 . For the above MDP setting, the pretrained model M θ satisfies E T test [ Reg τ ( M θ )] ≤ ˜ O ( C H 3 / 2 S √ AK ) .

Asimilar analysis due to [68] allows us to see theoretically why pretraining on (latently) linear bandits can lead to an in-context algorithm competitive with linear bandit algorithms, even when the incontext datasets are generated by algorithms unaware of this structure. We observed this empirically in Section 4. Consider a similar setup as there where S is a singleton, A is finite but large, θ τ ∈ R d is sampled as θ τ ∼ N (0 , I/d ) , ϕ : A → R d is a fixed feature map with sup a ∈A ∥ ϕ ( a ) ∥ 2 ≤ 1 , and the reward of a ∈ A in task τ is distributed as N ( ⟨ θ τ , ϕ ( a ) ⟩ , 1) . This time, we let D pre ( · ; τ ) be given by running Thompson Sampling with Gaussian priors and likelihood functions on τ .

Corollary 6.3 (Latent representation learning in linear bandits) . For T test = T pre in the above linear bandit setting, M θ satisfies E T test [ Reg τ ( M θ )] ≤ ˜ O ( d √ K ) .

This significantly improves over the ˜ O ( √ |A| K ) upper regret bound for TS that does not leverage the linear structure. This highlights how DPT can have provably tighter upper bounds on future bandit problems than the algorithms used to generate its (pretraining) data. Note that if there is additional structure in the tasks which yields a tighter regret bound (for example if there are only a small finite number of MDPs in the possible distribution), that may further improve performance, such as by removing the dependence on the state, action, or dimension in these cases.

Invariance of M θ to compliant D pre . Our final result sheds light on how D pre impacts the final DPT behavior M θ . Combined with Assumption 1, M θ is invariant to D pre satisfying Definition 6.1.

Proposition 6.4. Let P 1 pre and P 2 pre be pretraining distributions that differ only by their in-context dataset distributions, denoted by D 1 pre and D 2 pre . If D 1 pre and D 2 pre are compliant with the same support, then P 1 pre ( a ⋆ | s query , D, ξ h ) = P 2 pre ( a ⋆ | s query , D, ξ h ) for all a ⋆ , s query , D, ξ h .

That is, if we generate in-context datasets D by running various algorithms that depend only on the observed data in the current task, we will end up with the same M θ . For example, TS could be used for D 1 pre and PPO for D 2 pre . Expert-biased datasets discussed in Section 4 violate Definition 6.1, since privileged knowledge of τ is being used. This helps explain our empirical results that pretraining on expert-biased datasets leads to a qualitatively different learned model at test-time (Fig. 3c).

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-context dataset of interactions. Through in-depth evaluations in classic decision problems in bandits and MDPs, we showed that this simple objective naturally gives rise to an in-context RL algorithm that is capable of online exploration and offline decision-making in a way that leverages structure in the pretraining data. In some cases, the RL algorithm learned can transcend the quality of the pretraining data. Our empirical and theoretical results provide first steps towards understanding these capabilities and what factors are important for success. Via the simplicity of pretraining, we can sidestep the complexities of hand-designing exploration or conservatism in RL algorithms while simultaneously allowing the transformer to derive novel few-shot strategies. These findings underscore the potential of supervised pretraining to equip transformer models with in-context RL abilities.

Limitations and future work. One limitation of DPT is the requirement of optimal actions at pretraining. Empirically, we find that this requirement can be relaxed by using actions generated by another RL-trained agent during pretraining, which only leads to a slight loss in performance. However, fully understanding this problem and how best to leverage multi-task decision-making datasets remains a key open problem. We also discussed that the practical implementation for MDPs differs from true posterior sampling. It would be interesting to further understand and bridge this empirical-theoretical gap in the future. Finally, further investigation is required to understand the implications of these findings for existing transformer-based large language models that are increasingly being deployed in decision-making settings [69].

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 in part by NSF grant 2112926 and ONR grant N00014-21-1-2685. JNL acknowledges support from the NSF GRFP. AP would like to thank the support of the Eric and Wendy Schmidt Center at the Broad Institute of MIT and Harvard. This work was supported in part by funding from the Eric and Wendy Schmidt Center at the Broad Institute of MIT and Harvard.

References

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint arXiv:2202.12837 , 2022.

Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Neural Information Processing Systems , 35:18878-18891, 2022.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

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.

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.15661 , 2022.

Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14215 , 2022.

Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In International Conference on Machine Learning , pages 24631-24645. PMLR, 2022.

Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyperdecision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487 , 2023.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Malcolm Strens. A bayesian framework for reinforcement learning. In ICML , volume 2000, pages 943-950, 2000.

Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems , 26, 2013.

Tom Schaul and Jürgen Schmidhuber. Metalearning. Scholarpedia , 5(6):4650, 2010.

Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule . Citeseer, 1990.

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.

Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through metareinforcement learning. arXiv preprint arXiv:1803.11347 , 2018.

Nicholas C Landolfi, Garrett Thomas, and Tengyu Ma. A model-based approach for sampleefficient multi-task reinforcement learning. arXiv preprint arXiv:1907.04964 , 2019.

Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel J Hsu, Thodoris Lykouris, Miro Dudik, and Robert E Schapire. Bayesian decision-making under misspecified priors with applications to meta-learning. Advances in Neural Information Processing Systems , 34:2638226394, 2021.

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning , pages 5331-5340. PMLR, 2019.

Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:1905.06424 , 2019.

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 metalearning. In International Conference on Learning Representation (ICLR) , 2020.

Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices. In International conference on machine learning , pages 6925-6935. PMLR, 2021.

Theodore J Perkins, Doina Precup, et al. Using options for knowledge transfer in reinforcement learning. Technical report, Citeseer, 1999.

Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Metareinforcement learning of structured exploration strategies. Advances in neural information processing systems , 31, 2018.

Yiding Jiang, Evan Liu, Benjamin Eysenbach, J Zico Kolter, and Chelsea Finn. Learning options via compression. Advances in Neural Information Processing Systems , 35:21184-21199, 2022.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprint arXiv:1810.06784 , 2018.

Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. Rl2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

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:1611.05763 , 2016.

Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learner. arXiv preprint arXiv:1707.03141 , 2017.

Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, et al. Human-timescale adaptation in an open-ended task space. arXiv preprint arXiv:2301.07608 , 2023.

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 arXiv:2303.03982 , 2023.

Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foundation models for decision making: Problems, methods, and opportunities. arXiv preprint arXiv:2303.04129 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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 a unified text-to-text transformer. The Journal of Machine Learning Research , 21(1):5485-5551, 2020.

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:15084-15097, 2021.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Kuang-Huei Lee, Ofir 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-27936, 2022.

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.06175 , 2022.

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:2212.06817 , 2022.

Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing systems , 35:22955-22968, 2022.

David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-conditioned supervised learning work for offline reinforcement learning? arXiv preprint arXiv:2206.01079 , 2022.

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.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:11791191, 2020.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems , 34:28954-28967, 2021.

Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off-policy reinforcement learning without great exploration. Advances in neural information processing systems , 33:1264-1274, 2020.

Kamyar Ghasemipour, 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 , 35:18267-18281, 2022.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, 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.08396 , 2020.

Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. UAI , 2019.

Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.01112 , 2020.

Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline metareinforcement learning with advantage weighting. In International Conference on Machine Learning , pages 7780-7791. PMLR, 2021.

Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learningidentifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems , 34:4607-4618, 2021.

Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reinforcement learning with online self-supervision. In International Conference on Machine Learning , pages 17811-17829. PMLR, 2022.

Daniel J Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on thompson sampling. Foundations and Trends® in Machine Learning , 11(1):1-96, 2018.

William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika , 25(3-4):285-294, 1933.

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning , 47:235-256, 2002.

Chenjun Xiao, Yifan Wu, Jincheng Mei, Bo Dai, Tor Lattimore, Lihong Li, Csaba Szepesvari, and Dale Schuurmans. On the optimality of batch policy optimization algorithms. In International Conference on Machine Learning , pages 11362-11371. PMLR, 2021.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári. Improved algorithms for linear stochastic bandits. In Advances in Neural Information Processing Systems , pages 2312-2320, 2011.

Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline reinforcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems , 34:11702-11716, 2021.

Maxime Chevalier-Boisvert. Miniworld: Minimalistic 3d environment for rl and robotics research, 2018.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

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.

Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research , 39(4):1221-1243, 2014.

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 arXiv:2305.16291 , 2023.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677 , 2022.

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.

Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose incontext learning by meta-learning transformers. arXiv preprint arXiv:2212.04458 , 2022.

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.

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.

Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. arXiv preprint arXiv:2303.07895 , 2023.

Jacob Abernethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305.17040 , 2023.

Nan Rosemary Ke, Silvia Chiappa, Jane Wang, Anirudh Goyal, Jorg Bornschein, Melanie Rey, Theophane Weber, Matthew Botvinic, Michael Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. arXiv preprint arXiv:2204.04875 , 2022.

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 Conference on Machine Learning , pages 7740-7765. PMLR, 2022.

Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04179 , 2022.

Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM) , 64(5):1-24, 2017.

Shipra Agrawal and Randy Jia. Optimistic posterior sampling for reinforcement learning: worst-case regret bounds. Advances in Neural Information Processing Systems , 30, 2017.

Xiuyuan Lu and Benjamin Van Roy. Ensemble sampling. Advances in neural information processing systems , 30, 2017.

Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems , 29, 2016.

Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems , 31, 2018.

Ian Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Morteza Ibrahimi, Xiuyuan Lu, and Benjamin Van Roy. Approximate thompson sampling via epistemic neural networks. arXiv preprint arXiv:2302.09205 , 2023.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. The Journal of Machine Learning Research , 22(1):12348-12355, 2021.

Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems , 33:20095-20107, 2020.

Additional Related Work

In-context learning. Beyond decision-making and reinforcement learning, our approach takes inspiration from general in-context learning, a phenomenon observed most prominently in large language models in which large-scale autoregressive modelling can surprisingly lead to a model that exhibits meta-learning capabilities [1]. Recently, there has been great interest in understanding the capabilities and properties of in-context learning in various settings [4, 70, 5, 71, 72, 73, 74, 75, 6, 76, 77, 78]. While a common hypothesis suggests that this phenomenon is due to properties of the data used to train large language models [3], our work suggests that this phenomenon can also be encouraged in general settings via adjustments to the pre-training objective. In fact, DPT could be interpreted as explicitly encouraging the ability to perform Bayesian inference, which is a popular explanation for the mechanism behind in-context learning for large language models [67]. For meta-learning in-context, work by [79] proposed Transformer Neural Processes to address problems traditionally addressed with neural processes. Their work focuses on supervised learning-style input-output pairs for meta-learning, but they also applied this to contextual bandits to learn rewards.

For decision-making, in addition to the works discussed in the main paper, several works have applied transformers for in-context RL. AdA [32] meta-learns a transformer online (like RL 2 ). A base online RL algorithm is used to optimize the transformer, which uses online access to the simulator. The authors also use a curriculum learning method to improve the learning process. DPT is learned from offline supervised pretraining and it is agnostic to how the pretraining data is acquired. We believe many of the ideas in AdA are complementary and can be ported to DPT. Prompt-DT [8] is designed for offline RL and conditions on expert demos and uses return-to-go in the context much like the Decision Transformer [37]. This necessitates expert demos as input at test-time. DPT is meant for both online and offline in-context RL and does not require test-time demos (but it can benefit from them e.g. Fig. 4a). DPT does not use return-to-go, and instead predicts optimal actions from arbitrary in-context interaction datasets.

Posterior Sampling. Posterior sampling originates from the seminal work of [59], and has been popularized and thoroughly investigated in recent years by a number of authors [58, 80, 12, 13, 81, 68]. For bandits, it is often referred to as Thompson Sampling, but the framework is easily generalizable to RL, as shown in [12, 13]. The principle is as follows: begin with a prior over possible models (i.e. reward and transition functions), and maintain a posterior distribution over models by updating as new interactions are made. At decision-time, sample a model from the posterior and execute its optimal policy. The aforementioned prior works have developed strong theoretical guarantees on Bayesian and frequentist regret for posterior sampling. Despite its desirable theoretical characteristics, a major limitation is that computing the posterior is often computationally intractable. Recently, the community has made efforts towards scalable posterior sampling often by leveraging approximations [82, 83, 84, 85]. In Section 6, we show that a version of the DPT model learned from pretraining can be viewed as implementing posterior sampling without specifying complicated priors or posterior updates. Instead, the posterior update is implicitly learned through pretraining to predict the optimal action. This suggests that in-context learning via supervised pretraining (or meta-learning more generally) could be a promising alternative approach towards scalable posterior sampling.

On the theoretical side, [19] show it is possible to meta-learn prior distribution parameters for Thompson Sampling (TS) using a meta-exploration policy. The resulting TS algorithm can then be deployed on new problems. This is potentially consistent with the observation that DPT can approximate the right prior from offline data.

A Implementation and Experiment Details

A.1 DPT Architecture: Formal Description

In this section, we provide a detailed description of the architecture alluded to in Section 3 and Figure 1. See hyperparameter details for models in their respective sections. The model is implemented in Python with PyTorch [86]. The backbone of the transformer architecture we use is an autoregressive GPT-2 model from the HuggingFace transformers library.

For the sake of exposition, we suppose that S and A are subsets of R d S and R d A respectively. We handle discrete state and action spaces with one-hot encoding. Consider a single training datapoint

Algorithm 2 Decision-Pretrained Transformer (detailed)

Algorithm 3 Offline test-time deployment (detailed)

Algorithm 4 Online test-time deployment (detailed)

derived from an (potentially unknown) task τ : we have a dataset D of interactions within τ , a query state s query, and its corresponding optimal action a ⋆ = π ⋆ τ ( s query ) . We construct the embeddings to be passed to the GPT-2 backbone in the following way. From the dataset D = { ( s j , a j , s ′ j , r j ) } j ∈ [ n ] , we construct vectors ξ j = ( s j , a j , s ′ j , r j ) by stacking the elements of the transition tuple into dimension d ξ := 2 d S + d A +1 for each j in the sequence. This sequence of n elements is concatenated with another vector v := ( s query , 0 ) where the 0 vector is a vector of zeros of sufficient length to make the entire element dimension d ξ . The ( n +1) -length sequence is given by X = ( v, ξ 1 , . . . , ξ n ) . As order does not often matter for the dataset D 7 , we do not use positional encoding in order to take advantage of this invariance. We first apply a linear layer Linear ( X ) and pass the result to the transformer, which outputs the sequence Y = (ˆ y 0 , ˆ y 1 , . . . , ˆ y n ) . In 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

Because of the GPT-2 causal architecture (we defer details to the original papers [87, 1]), we note that ˆ p j depends only on s query and the partial dataset D j = { ( s k , a k , s ′ k , r k ) } k ∈ [ j ] , which is why we write the model notation,

to denote that the predicted probabilities of the j th element only depend on D j and not the entire D for the model M with parameters θ ∈ Θ . For example, with j = 0 , the prediction of a ⋆ is made without any contextual information about the task τ except for s query, which can be interpreted as the prior over a ⋆ . We measure loss of this training example via the cross entropy for each j ∈ [ n ] :

Intuition. Elements of the inputs sequence X represent transitions in the environment. When passed through the GPT-2 transformer, the model learns to associate elements of the sequence via the standard query-key-value mechanism of the attention model. The query state s query is demarcated by its zeros vector (which also acts as padding). Unlike other examples of transformers used for decision-making such as the Decision Transformer [37] and Algorithm Distillation [7], DPT does not separate the individual ( s, a, s ′ , r ) into their own embeddings to be made into one long sequence. This is because we view the transition tuples in the dataset as their own singletons, to be related with other singletons in the dataset through the attention mechanism. We note that there are various other implementation variations one could take, but we found success and robustness with this one.

A.2 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 ˆ a ∈ argmax a ∈A { ˆ µ a + √ 1 /n a } where ˆ µ a is the empirical mean so far for action a and n a is the number of times a has been chosen so far. To arrive at this constant for the bonus, we coarsely tried a set of plausible values given the noise and found this to perform the best.

Lower Confidence Bound (LCB). We choose actions as ˆ a ∈ argmax a ∈A { ˆ µ a -√ 1 /n a } where ˆ µ a is the empirical mean so far for action a and n a is the number of times a has been chosen so far.

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 1 2 and 1 12 to match the true ones. The noise model was well-specified with the correct variance. In the linear experiments of Figure 3a and Figure 3b, we set the prior mean and variance to 0 and 1 to fit the true ones better.

LinUCB. We choose ˆ a t ∈ argmax a ∈A ⟨ ˆ θ t , ϕ ( a ) ⟩ + β ∥ ϕ ( a ) ∥ ˆ Σ -1 t where β = 1 and ˆ Σ t = I + ∑ s ∈ [ t -1] ϕ ( a s ) ϕ ( a s ) ⊤ and ˆ θ t = ˆ Σ -1 t ∑ s ∈ [ t -1] r s ϕ ( a s ) . Here, r s and a s are the reward and action observed at time s .

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

DPT. The transformer for DPT has an embedding size of 32 , context length of 500 for basic bandits and 200 for linear bandits, 4 hidden layers, and 4 attention heads per attention layer for all bandits. We use the AdamW optimizer with weight decay 1e-4, learning rate 1e-4, and batch-size 64. For all experiments, we shuffle the in-context dataset D since order does not matter except in the linear bandit.

A.2.2 RL Algorithms

Below, we describe the comparisons from the MDP experiments and their hyperparameters.

Proximal Policy Optimization (PPO). The reported results for PPO use the Stable Baselines3 implementation [88] with the default hyperparameters, which successfully learns each task given 100 K environment steps in Dark Room and 125 K environment steps in Miniworld. In Dark Room, the policy is implemented as a multi-layer perceptron with two hidden layers of 64 units each. In Miniworld, the policy is a convolutional neural network with two convolutional layers with 16 3 × 3 kernels each, followed by a linear layer with output dimension of 8 .

Algorithm Distillation (AD). We first collect learning histories with PPO for each of the training tasks. Then, given a cross-episodic context of length H , where H is the task horizon, the model is trained to predict the actions taken K episodes later (given the states visited in that episode). This was shown to lead to faster algorithms in [7]. We evaluated AD across different values of K . Between K = 10 , 50 , 100 , we found K = 100 to be most performant in the Dark Room environment. In Miniworld, we also subsampled with K = 100 . In Dark Room, the transformer has similar hyperparameters as DPT: an embedding size of 32 , context length of 100 steps, 4 hidden layers, and 4 attention heads per attention layer. In Miniworld, as with DPT, we first encode the image with a convolutional network with two convolutional layers with 16 3 × 3 kernels each, followed by a linear layer with output dimension of 8 .

RL 2 . The reported results for RL 2 use an open-sourced implementation from [22]. The implementation uses PPO as the RL algorithm and defines a single trial as four consecutive episodes. The policy is implemented with one hidden layer of 32 units in Dark Room. In Miniworld, the policy is parameterized with a convolutional neural network with two convolutional layers with 16 3 × 3 kernels each, followed by a linear layer with output dimension of 8 .

DPT. The transformer for DPT has an embedding size of 32 , context length of 100 steps, 4 hidden layers, and 4 attention heads per attention layer in Dark Room. In Miniworld, the image is first passed through a convolutional network with two convolutional layers 16 3 × 3 kernels each, followed by a linear layer with output dimension of 8 . The transformer model that processes these image embeddings otherwise has the same hyperparameters as in Dark Room. We use the AdamW optimizer with weight decay 1e-4, learning rate 1e-3, and batch-size 128 .

A.3 Bandit Pretraining and Testing

Basic Bandit. Offline, to generate the in-context datasets for pretraining, we used a Dirichlet distribution to sample action frequencies in order to generate datasets with diverse compositions (i.e. some more uniform, some that only choose a few actions, etc.): p 1 ∼ Dir ( ✶ ) where p 1 ∈ ∆( A ) and ✶ ∈ R |A| . We also mixed this with a distribution that has all mass on one action: ˆ a ∼ Unif ( A ) and p 2 (ˆ a ) = 1 and p 2 ( a ) = 0 for all a = ˆ a . The final action distribution is p = (1 -ω ) p 1 + ωp 2 where ω ∼ Unif (0 . 1[10]) . We train on 100,000 pretraining samples for 300 epochs with an 80/20 train/validation split. In Figure 2a, D test is generated in the same way.

̸

̸

Expert-Biased Bandit. To generate expert-biased datasets for pretraining, we compute the action frequencies to bias the dataset towards the optimal action. Let a ⋆ be the optimal one. As before, we take p 1 ∼ Dir ( ✶ ) . Then, p 2 ( a ⋆ ) = 1 and p 2 ( a ) = 0 for all a = a ⋆ . For of bias of ω , we take p = (1 -ω ) p 1 + ωp 2 with ω ∼ Unif (0 . 1[10]) . We use the same pretraining sample size and epochs as before. For testing, D test is generated the same way except we fix a particular ω ∈ { 0 , 0 . 5 , 1 } to test on.

Linear Bandit. We consider the case where |A| = 10 and d = 2 . To generate environments from T pre, we first sampled a fixed set of actions from N ( 0 , I d /d ) in R d to represent the features. Then, for each τ , we sampled θ τ ∼ N ( 0 , I d /d ) to produce the means µ a = ⟨ θ τ , ϕ ( a ) ⟩ for a ∈ A . To 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 , 000 , 000 pretraining samples over 200 epochs with an 80/20 train/validation split. At test time, we set T test = T pre and D test = D pre. Note that ϕ is fixed over all τ , as is standard for a linear bandit.

A.4 MDPEnvironment Details

Dark Room. The agent must navigate a 10 × 10 grid to find the goal within H = 100 steps. The agent's observation is its xy -position, the allowed actions are left, right, up, down, and stay, and the reward is only r = 1 when the agent is at the goal, and r = 0 otherwise. At test time, the agent begins at the (0 , 0) position. We randomly designate 80 of the 100 grid squares to be goals for the training tasks, and hold out the remaining 20 for evaluation.

Miniworld. The agent must navigate to the correct box, which is initially unknown, from 25 × 25 RGB image observations. The agent is additionally conditioned on its own direction vector. In each episode, the environment is initialized with four boxes of different colors, one in each corner of the square room. The agent can turn left, turn right, or move forward. The reward is only r = 1 when the agent is near the correct box and r = 0 otherwise, and each episode is 50 time-steps long. At test time, the agent begins in the middle of the room.

A.5 MDPPretraining Datasets

Dark Room. In Dark Room, we collect 100 K in-context datasets, each of length H = 100 steps, with a uniform-random policy. The 100 K datasets are evenly collected across the 100 goals. The query states are uniformly sampled from the state space, and the optimal actions are computed as follows: move up/down until the agent is on the same y -position as the goal, then move left/right until the agent is on the x -position as the goal. Of the 100 K collections of datasets, query states, and optimal actions, we use the first 80 K (corresponding to the first 80 goals) for training and the remaining 20 K for validation.

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 60 K datasets each of length H = 50 steps. So as to reduce computation, the in-context datasets only have only ( s, a, r ) tuples. The query states, which consist of image and direction are sampled uniformly from the entire state space, i.e., the agent is place uniformly at random in the environment, pointing in a random direction. The optimal actions are computed as follows: turn towards the correct box if the agent is not yet facing it (within ± 15 degrees), otherwise move forward. Of the 60 K 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.

Out-of-distribution 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.

0.0

0.3

0.5

Noise standard deviation

0.00

0.01

0.02

0.03

0.04

0.05

0.06

0.07

0.08

Final Suboptimality

Offline Robustness to Reward Shift

Emp

LCB

TS

DPT (ours)

(a)

0

100

200

300

400

500

Data

0.1

0.3

0.02

Suboptimality

Emp

LCB

TS

DPT (ours)

(b)

0

100

200

300

400

500

Data

0

5

10

15

20

25

Cumulative Regret

Emp

UCB

TS

DPT (ours)

(c)

Random

Expert

Dataset Composition

0

20

40

60

80

Return

Offline Dark Room (Train Tasks)

Dataset

AD

DPT (ours)

(a)

Random

Expert

Dataset Composition

0

20

40

60

80

Return

Offline Dark Room (Test Tasks)

Dataset

AD

DPT (ours)

(b)

0

5

10

15

20

25

30

35

40

Episodes

0

20

40

60

80

Return

DPT (ours)

AD

PPO

RL2

(c)

0

5

10

15

20

25

30

35

40

Episodes

0

20

40

60

80

Return

DPT (ours)

AD

PPO

RL2

(d)

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 Markov 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 of each method on the training tasks in Dark Room. Offline, DPT and AD demonstrate comparable performance as on the training tasks, indicating a minimal generalization gap to new goals. Online, DPT, AD, and RL 2 also achieve performance on the training tasks similar to that on the test tasks.

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

B.3 Sensitivity Analysis

We next seek to understand the sensitivity of DPT to different hyperparameters, including the model size and size of the pretraining dataset. These experiments are performed in the Dark Room environment. As shown in Fig. 9, the performance of DPT is robust to the model size; it is the same across different embedding sizes,

0

5

10

15

20

25

30

35

40

Episodes

0

10

20

30

40

50

60

70

80

Return

Online Dark Room Permuted

2 layers

4 layers

6 layers

8 layers

Layers

0

10

20

30

40

50

60

Return

Layers (Offline)

(a)

1 heads

2 heads

4 heads

8 heads

Heads

0

10

20

30

40

50

60

Return

Attention Heads (Offline)

(b)

16 dim

32 dim

64 dim

128 dim

Dimensions

0

10

20

30

40

50

60

Return

Embedding Dimension (Offline)

(c)

10k

50k

100k

200k

Samples

0

10

20

30

40

50

60

Return

Pretraining Samples (Offline)

(d)

number of layers, and number of attention heads. Notably, the performance is slightly worse with 8 attention heads, which may be attributed to slight overfitting. We do see that when the pretraining dataset is reduced to 10% of its original size ( 10000 samples) the performance degrades, but otherwise has similar performance with larger pretraining datasets.

C Additional Theory and Omitted Proofs

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

Proposition C.1 (Theorem 21 of [89]) . Let D and ˆ f be given as above under the aforementioned conditions. Then, with probability at least 1 -δ ,

The finiteness of F is done for simplicity, but we can see that this yields dependence on the log-cardinality, a common measure of complexity. Extensions to infinite F of bounded statistical complexity can be readily made to replace this. For our setting, the bound suggests that E P pre ∥ P pre ( ·| s query , D, ξ h ) -M θ ( ·| s query , D, ξ h ) ∥ 2 1 → 0 as N →∞ with high probability, provided the function class of M θ has bounded statistical complexity.

C.1 Posterior Sampling

Posterior sampling is most generally described with the following procedure [13]. Initialize a prior distribution T 1 = T pre and dataset D = {} . For k ∈ [ K ]

Sample τ k ∼ T k and compute ˆ π τ k

Execute π ⋆ τ k and add interactions to D

Update posterior distribution T k +1 ( τ ) = P ( τ | D ) .

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

C.2 Proof of Theorem 1

Theorem 1 (DPT ⇐⇒ PS) . Let the above assumptions hold. Then, P ps ( ξ H | D,τ c ) = P M θ ( ξ H | D,τ c ) for all trajectories ξ H .

Proof. Without loss of generality, for a task τ , we take π ⋆ τ ( ·| s ) to be deterministic and denote the optimal action in state s as π ⋆ τ ( s ) . Recall that we consider a fixed current task τ c and a fixed in-context dataset D . Define ξ h = ( s 1 , a 1 , . . . , s h , a h ) .

We now formally state the variant of the full joint distribution from which we sample during pretraining. Let τ and D ′ be an arbitrary task and dataset and let a ⋆ ∈ A , s query ∈ S , ξ H -1 ∈ ( S × A ) H -1 , and h ∈ [0 , H -1] be arbitrary.

The Unif [0 , H -1] is due to the fact that we sample h ∼ Unif [0 , H -1] and then truncate ξ h from ξ H -1 (or, equivalently, sample ξ h ∼ S h directly), marginalizing out the other variables. For h ′ ≤ h -1 , recall that we also use the notation S h ′ ( s 1: h ′ ) to denote the marginalization of the full joint S H . We will eventually work with the posterior of this distribution given the data D and history ξ h :

We define the following random sequences and subsequences:

where the variables are generated according to the following conditional process: τ ps ∼ P ( ·| D ) , S ps 1 ∼ ρ τ c , A ps h ∼ π ⋆ τ ps ( ·| S ps h ) , and S ps h +1 ∼ T τ c ( ·| S ps h , A ps h ) . We also define Ξ ps ( h ′ : h ; D ) to be the last h -h ′ elements of Ξ ps ( h ; D ) . Analogously, we define

where the variables are from the process: S pre 1 ∼ ρ τ c , A pre h ∼ P pre ( ·| S pre h , D, Ξ pre ( h -1; D )) , and S pre h +1 ∼ T τ c ( ·| S pre h , A pre h ) . Note that A pre h is sampled conditioned on the sequence Ξ pre ( h ; D ) so far.

We will show that Ξ ps ( h ; D ) and Ξ pre ( h ; D ) follow the same distribution for all h ∈ [ H ] . For convenience, we will drop notational dependence on D , except where it resolves ambiguity. Also, because of Assumption 1, we have that P pre ( ·| S pre h , D, Ξ pre ( h -1)) = M θ ( ·| S pre h , D, Ξ pre ( h -1)) , so we will just work with P pre for the remainder of the proof. We will also make use of the following lemma.

Proof. From the definition of posterior sampling (using the same prior, T pre), we have that

where the second line crucially uses the fact that posterior sampling chooses actions based only on the prior and history so far. Similarly, the third line uses the fact that D pre is compliant. Since the two sides are proportional in τ , they are equivalent.

We will prove Theorem 1 via induction for each h ∈ [ H ] . First, consider the base case for a sequence of length h = 1 . Recall that ρ τ c denotes the initial state distribution of τ c . We have that the densities can be written as

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

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

By the hypothesis, we have that P (Ξ ps ( h -1) = ξ h -1 ) = P (Ξ pre ( h -1) = ξ h -1 ) . For the second factor,

As before, we can further rewrite the last factor as

where

(38)

where ∝ denotes that the two sides are equal up to multiplicative factors independent of τ . In the first line, we used Bayes rule. In the second line, given that τ ps = τ (i.e. posterior sampling selected τ to

deploy), we decompose the probability of observing that sequence of states of actions. We also used Lemma C.2. The denominator does not depend on τ . Similarly, for the third and fourth lines, T τ c and S do not depend on τ . The final line follows from the definition of the joint pretraining distribution in this regime.

Therefore, we conclude that the posterior over the value of τ ps is the same as the posterior over the task in the pretraining distribution, given s h , D, ξ h -1 . Substituting back through all the previous equations, we have

This concludes the proof.

C.3 Proof of Corollary 6.2

Corollary C.3 (Finite MDPs) . Suppose that sup τ T test ( τ ) / T pre ( τ ) ≤ C for some C > 0 . For the above MDP setting, the pretrained model M θ satisfies E T test [ Reg τ ( M θ )] ≤ ˜ O ( C H 3 / 2 S √ AK ) .

Proof. Note that D pre is clearly compliant since it is generated by random sampling. We use the equivalence between M θ and posterior sampling established in Theorem 1. The proof then follows immediately from Theorem 1 of [13] to guarantee that

where the notation ˜ O omits polylogarithmic dependence. The bound on the test task distribution follows from the assumed bound on the likelihood ratio under the priors:

C.4 Proof of Corollary 6.3

Corollary C.4 (Latent representation learning in linear bandits) . For T test = T pre in the above linear bandit setting, M θ satisfies E T test [ Reg τ ( M θ )] ≤ ˜ O ( d √ K ) .

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

C.5 Proof of Proposition 6.4

Proposition C.5. Let P 1 pre and P 2 pre be pretraining distributions that differ only by their in-context dataset distributions, denoted by D 1 pre and D 2 pre . If D 1 pre and D 2 pre are compliant with the same support, then P 1 pre ( a ⋆ | s query , D, ξ h ) = P 2 pre ( a ⋆ | s query , D, ξ h ) for all a ⋆ , s query , D, ξ h .

Proof. The proof follows by direct inspection of the pretraining distributions. For P 1 pre , we have

The posterior distribution over tasks is simply

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

where the second equality holds because D 1 pre ( a j | s j , D j ; τ ) is assumed to be invariant to τ by compliance, and the fifth equality holds because D 2 pre ( a j | s j , D j ; τ ) is assumed to be invariant to τ . Therefore, we conclude that, for any s, D, ξ ,

Since also ∫ τ P 1 pre ( τ | s, D, ξ ) = 1 = ∫ τ P 2 pre ( τ | s, D, ξ ) , then

Substituting this back into Equation45 yields P 1 pre ( a ⋆ | s, D, ξ ) = P 1 pre ( a ⋆ | s, D, ξ ) .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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)?\n\nLIMITATIONS:\nThe authors acknowledge certain limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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.

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.

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

3. 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?

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "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.\n\nSTRENGTHS:\nMethod:\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\nWEAKNESSES:\nNovelty: 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.\n\nQUESTIONS:\nIt 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}