# Pretraining Decision Transformers with Reward Prediction for In-Context
Multi-task Structured Bandit Learning

Subhojyoti Mukherjee
ECE Department
UW-Madison
Wisconsin, Madison
<EMAIL_ADDRESS>Josiah P. Hanna
CS Department
UW-Madison
Wisconsin, Madison  Qiaomin Xie
ISyE Department
UW-Madison
Wisconsin, Madison
Robert Nowak
ECE Department
UW-Madison
Wisconsin, Madison

###### Abstract

In this paper, we study multi-task structured bandit problem where the goal is
to learn a near-optimal algorithm that minimizes cumulative regret. The tasks
share a common structure and the algorithm exploits the shared structure to
minimize the cumulative regret for an unseen but related test task. We use a
transformer as a decision-making algorithm to learn this shared structure so
as to generalize to the test task. The prior work of pretrained decision
transformers like DPT requires access to the optimal action during training
which may be hard in several scenarios. Diverging from these works, our
learning algorithm does not need the knowledge of optimal action per task
during training but predicts a reward vector for each of the actions using
only the observed offline data from the diverse training tasks. Finally,
during inference time, it selects action using the reward predictions
employing various exploration strategies in-context for an unseen test task.
We show that our model outperforms other SOTA methods like DPT, and
Algorithmic Distillation (AD) over a series of experiments on several
structured bandit problems (linear, bilinear, latent, non-linear).
Interestingly, we show that our algorithm, without the knowledge of the
underlying problem structure, can learn a near-optimal policy in-context by
leveraging the shared structure across diverse tasks. We further extend the
field of pre-trained decision transformers by showing that they can leverage
unseen tasks with new actions and still learn the underlying latent structure
to derive a near-optimal policy. We validate this over several experiments to
show that our proposed solution is very general and has wide applications to
potentially emergent online and offline strategies at test time. Finally, we
theoretically analyze the performance of our algorithm and obtain
generalization bounds in the in-context multi-task learning setting.

##  1 Introduction

In this paper, we study multi-task bandit learning with the goal of learning
an algorithm that discovers and exploits structure in a family of related
tasks. In multi-task bandit learning, we have multiple distinct bandit tasks
for which we want to learn a policy. Though distinct, the tasks share some
structure, which we hope to leverage to speed up learning on new instances in
this task family. Traditionally, the study of such structured bandit problems
has relied on knowledge of the problem structure like linear bandits (Li et
al., 2010; Abbasi-Yadkori et al., 2011; Degenne et al., 2020), bilinear
bandits (Jun et al., 2019), hierarchical bandits (Hong et al., 2022a, b),
Lipschitz bandits (Bubeck et al., 2008, 2011; Magureanu et al., 2014), other
structured bandits settings (Riquelme et al., 2018; Lattimore and Szepesvári,
2019; Dong et al., 2021) and even linear and bilinear multi-task bandit
settings (Yang et al., 2022a; Du et al., 2023; Mukherjee et al., 2023). When
structure is unknown an alternative is to adopt sophisticated model classes,
such as kernel machines or neural networks, exemplified by kernel or neural
bandits (Valko et al., 2013; Chowdhury and Gopalan, 2017; Zhou et al., 2020;
Dai et al., 2022). However, these approaches are also costly as they learn
complex, nonlinear models from the ground up without any prior data (Justus et
al., 2018; Zambaldi et al., 2018).

In this paper, we consider an alternative approach of synthesizing a bandit
algorithm from historical data where the data comes from recorded bandit
interactions with past instances of our target task family. Concretely, we are
given a set of state-action-reward tuples obtained by running some bandit
algorithm in various instances from the task family. We then aim to learn a
transformer (Vaswani et al., 2017) from this data that can in-context learn on
new task instances. Laskin et al. (2022) consider a similar goal and introduce
the Algorithm Distillation (AD) method, however, AD aims to copy the algorithm
used in the historical data and thus is limited by the ability of the data
collection algorithm. Lee et al. (2023) develop an approach, DPT, that enables
learning a transformer that obtains lower regret in-context bandit learning
compared to the algorithm used to produce the historical data. However, this
approach requires knowledge of the optimal action at each stage of the
decision process. In real problems, this assumption is hard to satisfy and we
will show that DPT performs poorly when the optimal action is only
approximately known. With this past work in mind, the goal of this paper is to
answer the question:

Can we learn an in-context bandit learning algorithm that obtains lower regret
than the algorithm used to produce training data without knowledge of optimal
actions in each training task?

To answer this question, we introduce a new transformer methodology that
obviates the need for knowledge of the optimal action in the in-context data
–- a piece of information that is often inaccessible. Our key observation is
that while the mean rewards of each action change from task to task, certain
probabilistic dependencies are persistent across all tasks with a given
structure (Yang et al., 2020, 2022a; Mukherjee et al., 2023). These
probabilistic dependencies can be learned from the pretraining data and
exploited to better estimate mean rewards and improve performance in a new
unknown test task. The nature of the probabilistic dependencies depends on the
specific structure of the bandit and can be complex (i.e., higher-order
dependencies beyond simple correlations). We propose to use transformer models
as a general-purpose architecture to capture the unknown dependencies by
training transformers to predict the mean rewards in each of the given
trajectories (Mirchandani et al., 2023; Zhao et al., 2023). The key idea is
that transformers have the capacity to discover and exploit complex
dependencies in order to predict the rewards of all possible actions in each
task from a _small_ history of action-reward pairs in a new task. This paper
demonstrates how such an approach can achieve state-of-the-art performance,
relying solely on historical data, without the need for any supplementary
information like the action features or knowledge of the complex reward
models. Our new decision transformer learns to adapt, in-context, to novel
actions as well as new tasks.

#### Contributions

  1. 1.

We introduce a new training procedure where predicting the next reward for all
arms circumvents the issue of requiring access to the optimal (or
approximately optimal) action during training time.

  2. 2.

We demonstrate empirically that this training procedure results in lower
regret in a wide series of tasks (such as linear, nonlinear, bilinear, and
latent bandits) compared to DPT, AD (and other in-context decision-making
algorithms).

  3. 3.

We also show that our training procedure leverages the latent structure even
when new actions are introduced both during training and testing time.

  4. 4.

Finally, we theoretically analyze the generalization ability of PreDeToR
through the lens of algorithmic stability and new results for the transformer
setting.

##  2 Background

In this section, we first introduce our notation and the multi-task,
structured bandit setting. We then formalize the in-context learning model
studied in Laskin et al. (2022); Lee et al. (2023); Sinii et al. (2023); Lin
et al. (2023); Ma et al. (2023); Liu et al. (2023c, a). We choose DPT (Lee et
al., 2023) as a representative example of this learning model and explain the
DPT model for the online setting.

###  2.1 Preliminaries

In this paper, we consider the multi-task linear bandit setting (Du et al.,
2023; Yang et al., 2020, 2022a). In the multi-task setting, we denote each
task as $m$. Each task $m$ constitutes of the action set $|\mathcal{A}|=A$ and
the feature space $\mathcal{X}$. The actions in $\mathcal{A}$ are indexed by
$a=1,2,\ldots,A$. The feature of each action is denoted by
$\mathbf{x}(m,a)\in\mathbb{R}^{d}$ and $d\ll A$. Denote
$\triangle(\mathcal{A})$ as the probability simplex over the action space
$\mathcal{A}$. We denote a policy as $\pi:\mathcal{A}\rightarrow[0,1]$, as a
probability distribution over actions such that
$\pi\in\triangle(\mathcal{A})$. This implies that $\sum_{a=1}^{A}\pi(a)=1$.

Define $[n]=\\{1,2,\ldots,n\\}$. In a multi-task structured bandit setting the
expected reward for each action in each task is assumed to be an unknown
function of the hidden parameter and action features (Lattimore and
Szepesvári, 2020; Gupta et al., 2020). The interaction proceeds iteratively
over $n$ rounds for each task $m\in[M]$. At each round $t\in[n]$ for each task
$m\in[M]$, the learner selects an action $I_{m,t}\in\mathcal{A}$ and observes
the reward $r_{m,t}=f(\mathbf{x}(I_{m,t}),\bm{\theta}_{m,*})+\eta_{m,t}$,
where $\bm{\theta}_{m,*}\in\mathbb{R}^{d}$ is the hidden parameter specific to
the task $m$ to be learned by the learner. The function $f(\cdot,\cdot)$ is
the unknown reward structure. This can be
$f(\mathbf{x}(I_{m,t}),\bm{\theta}_{m,*})=\mathbf{x}(I_{m,t})^{\top}\bm{\theta}_{m,*}$
for the linear setting or even more complex correlation between features and
$\bm{\theta}_{m,*}$ (Filippi et al., 2010; Abbasi-Yadkori et al., 2011;
Riquelme et al., 2018; Lattimore and Szepesvári, 2019; Dong et al., 2021). We
further assume that the noise $\eta_{m,t}$ is $\sigma^{2}$ sub-Gaussian.

In our paper, we assume that there exist weak demonstrators denoted by
$\pi^{w}$. These weak demonstrators are stochastic $A$-armed bandit algorithms
like Upper Confidence Bound (UCB) (Auer et al., 2002; Auer and Ortner, 2010)
or Thompson Sampling (Thompson, 1933; Agrawal and Goyal, 2012; Russo et al.,
2018; Zhu and Tan, 2020). The $\pi^{w}$ at every round $t$ for a task
$m\in[M]$ samples an action $I_{m,t}$ and observes the reward
$r_{m,t}=f(\mathbf{x}(I_{m,t}),\bm{\theta}_{m,*})+\eta_{m,t}$. Crucially note
that $\pi^{w}$ does not observe the feature $\mathbf{x}(I_{m,t})$. Moreover,
we assume that the $\pi^{w}$ does not have access to the feature space
$\mathcal{X}$ for any task $m$, and hence cannot plan its sampling policy
based on the feature vectors. The $\pi^{w}$ is used to collect data on which
we train our model. So we call these algorithms weak demonstrators. Note that
in many real-world world settings, like robotics (Fong et al., 2003; Lauri et
al., 2022), medical treatment (Murphy, 2003) the task is only partially
observed or have a complex correlated structure and therefore a learning agent
may not have access to the full set of action features, but only observe the
reward and action history. This is characterized by $\pi^{w}$. In the next
section, we propose the learning model for the training procedure.

###  2.2 In-Context Learning Model

Similar to Lee et al. (2023); Sinii et al. (2023); Lin et al. (2023); Ma et
al. (2023); Liu et al. (2023c, a) we assume the in-context learning model. We
first discuss the pretraining procedure.

Pretraining: Let $\mathcal{T}_{\text{pre}}$ denote the distribution over tasks
$m$ at the time of pretraining. Note that by our definition of the task $m$,
the distribution $\mathcal{T}_{\text{pre }}$ can span different reward and
action spaces. Let $\mathcal{D}_{\text{pre}}$ be the distribution over all
possible interactions that the $\pi^{w}$ can generate. We first sample a task
$m\sim\mathcal{T}_{\text{pre }}$ and then a context (or a prompt)
$\mathcal{H}_{m}$ which is a sequence of interactions for $n$ rounds
conditioned on the task $m$ such that
$\mathcal{H}_{m}\sim\mathcal{D}_{\text{pre}}(\cdot|m)$. So
$\mathcal{H}_{m}=\\{I_{m,t},r_{m,t}\\}_{t=1}^{n}$. We call this dataset
$\mathcal{H}_{m}$ an in-context dataset as it contains the contextual
information about the task $m$. We denote the samples in $\mathcal{H}_{m}$
till round $t$ as $\mathcal{H}_{m}^{t}=\\{I_{m,s},r_{m,s}\\}_{s=1}^{t-1}$.
This dataset $\mathcal{H}_{m}$ can be collected in several ways: (1) random
interactions within $m$, (2) demonstrations from an expert, and (3) rollouts
of an algorithm. Finally, we train a causal GPT-2 transformer model
$\mathrm{TF}$ parameterized by $\mathbf{\Theta}$ on this dataset
$\mathcal{D}_{\text{pre}}$. Specifically, we define
$\mathrm{TF}_{\mathbf{\Theta}}\left(\cdot\mid\mathcal{H}^{t}_{m}\right)$ as
the transformer model that observes the dataset $\mathcal{H}^{t}_{m}$ till
round $t$ and then produces a distribution over the actions. Our primary
novelty lies in our training procedure which we explain in detail in Section
3.1.

Testing: We now discuss the testing procedure for our setting. Let
$\mathcal{T}_{\text{test}}$ denote the distribution over test tasks
$m\in[M_{\text{test }}]$ at the time of testing. Let
$\mathcal{D}_{\text{test}}$ denote a distribution over all possible
interactions that can be generated by $\pi^{w}$ during test time. For online
training, the dataset $\mathcal{H}_{m}=\\{\emptyset\\}$ is initialized empty.
At each round $t$,
$\mathrm{TF}_{\mathbf{\Theta}}(\cdot\mid\mathcal{H}^{t}_{m})$ is deployed
where the model samples
$I_{t}\sim\mathrm{TF}_{\mathbf{\Theta}}(\cdot\mid\mathcal{H}^{t}_{m})$. The
model observes the reward $r_{t}$ which is added to $\mathcal{H}^{t}_{m}$. So
the $\mathcal{H}_{m}$ during online testing consists of
$\\{I_{t},r_{t}\\}_{t=1}^{n}$ collected during testing. This interaction
procedure is conducted for each test task $m\in[M_{\text{test }}]$. Finally,
note that in this testing phase, the model parameter $\mathbf{\Theta}$ is not
updated. Finally, the goal of the learner is to minimize cumulative regret for
all task $m\in[M_{\text{test}}]$ defined as follows:
$\mathbb{E}[R_{n}]=\frac{1}{M_{\text{test}}}\sum_{m=1}^{M_{\text{test}}}\sum_{t=1}^{n}\max_{a\in\mathcal{A}}f\left(\mathbf{x}(m,a)-\mathbf{x}(m,I_{t}),\bm{\theta}_{m,*}\right)$.

###  2.3 Related In-context Learning Algorithms

In this section, we discuss related algorithms for in-context decision-making.
For completeness, we describe the DPT and AD training procedure and algorithm
now. During training, DPT first samples $m\sim\mathcal{T}_{\text{pre}}$ and
then an in-context dataset
$\mathcal{H}_{m}\sim\mathcal{D}_{\text{pre}}(\cdot|,m)$. It adds this
$\mathcal{H}_{m}$ to the training dataset $\mathcal{H}_{\text{train}}$, and
repeats to collect $M_{\text{pre}}$ such training tasks. For each task $m$,
DPT requires the optimal action $a_{m,*}=\mathop{\rm
arg\,max}_{a}f(\mathbf{x}(m,a),\bm{\theta}_{m,*})$ where
$f(\mathbf{x}(m,a),\bm{\theta}_{m,*})$ is the expected reward for the action
$a$ in task $m$. Since the optimal action is usually not known in advance, in
Section 4 we introduce a practical variant of DPT that approximates the
optimal action with the best action identified during task interaction. During
training DPT minimizes the cross-entropy loss:

$\displaystyle\mathcal{L}^{\mathrm{{\color[rgb]{0,0.6,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.6,0}DPT}}}_{t}=\operatorname{cross-
entropy}(\mathrm{TF}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m}^{t}),p({a}_{m,*}))$
(1)

where $p({a}_{m,*})\\!\in\\!\triangle^{\mathcal{A}}$ is a one-hot vector such
that $p(j)\\!=\\!1$ when $j\\!=\\!{a}_{m,*}$ and $0$ otherwise. This loss is
then back-propagated and used to update the model parameter $\mathbf{\Theta}$.

During test time evaluation for online setting the DPT selects
$I_{t}\sim\mathrm{softmax}^{\tau}_{a}(\mathrm{TF}_{\mathbf{\Theta}}(\cdot|\mathcal{H}^{t}_{m}))$
where we define the $\mathrm{softmax}^{\tau}_{a}(\mathbf{v})$ over a $A$
dimensional vector $\mathbf{v}\in\mathbb{R}^{A}$ as
$\mathrm{softmax}^{\tau}_{a}(\mathbf{v}(a))=\exp(\mathbf{v}(a)/\tau)/\sum_{a^{\prime}=1}^{A}\exp(\mathbf{v}(a^{\prime})/\tau)$
which produces a distribution over actions weighted by the temperature
parameter $\tau>0$. Therefore this sampling procedure has a high probability
of choosing the predicted optimal action as well as induce sufficient
exploration. In the online setting, the DPT observes the reward $r_{t}(I_{t})$
which is added to $\mathcal{H}^{t}_{m}$. So the $\mathcal{H}_{m}$ during
online testing consists of $\\{I_{t},r_{t}\\}_{t=1}^{n}$ collected during
testing. This interaction procedure is conducted for each test task
$m\in[M_{\text{test }}]$. In the testing phase, the model parameter
$\mathbf{\Theta}$ is not updated.

An alternative to DPT that does not require knowledge of the optimal action is
the AD approach (Laskin et al., 2022; Lu et al., 2023). In AD, the learner
aims to predict the next action of the demonstrator. So it minimizes the
cross-entropy loss as follows:

$\displaystyle\mathcal{L}^{\mathrm{AD}}_{t}=\operatorname{cross-
entropy}(\mathrm{TF}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m}^{t}),p({I}_{m,t}))$
(2)

where $p({I}_{m,t})$ is a one-hot vector such that $p(j)=1$ when $j={I}_{m,t}$
(the true action taken by the demonstrator) and $0$ otherwise. During online
test time evaluation the AD selects
$I_{t}\sim\mathrm{softmax}^{\tau}_{a}(\mathrm{TF}_{\mathbf{\Theta}}(\cdot|\mathcal{H}^{t}_{m}))$.
Unfortunately, AD will, at best, match the regret of the demonstrator and
cannot obtain lower regret. In the next section, we introduce a new method
that can improve upon the demonstrator without knowledge of the optimal
action.

##  3 Proposed Algorithm PreDeToR

We now introduce our main algorithmic contribution, PreDeToR, which does not
require access to the optimal action during training in order to improve upon
the demonstrator.

###  3.1 Follow Demonstrator and Predict Next Reward

The key idea behind DPT is to leverage the in-context learning ability of
transformers to infer the reward of each arm in a given test task. By training
this in-context ability on a set of training tasks, the transformer can
implicitly learn structure in the task family and exploit this structure to
infer rewards without trying every single arm. Thus, in contrast to DPT and AD
that output actions direction, PreDeToR outputs a scalar value reward
prediction for each arm. To this effect, we append a linear layer of dimension
$A$ on top of a causal GPT2 model, denoted by
$\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m})$, and use
an MSE loss to train the transformer to predict the reward for each action
with these outputs. Note that we use
$\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m})$ to denote
a reward prediction transformer and
$\mathrm{TF}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m})$ as the transformer that
predicts a distribution over actions (as in DPT and AD ). Recall that the
dataset $\mathcal{H}_{m}\sim\mathcal{D}_{\text{pre}}(\cdot|m)$ where
$m\sim\mathcal{T}_{\text{pre}}$. Note that dataset consists of a history of
interactions $\mathcal{H}_{m}=\\{I_{m,t},r_{m,t}\\}_{t=1}^{n}$. Then we pass
this dataset to the transformer
$\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}(\cdot|\mathcal{H}_{m})$. At every
round $t$ the transformer predicts the _next reward_ for each of the actions
$a\in\mathcal{A}$ for the task $m$ based on
$\mathcal{H}_{m}^{t}=\\{I_{m,s},r_{m,s}\\}_{s=1}^{t-1}$. This predicted reward
is denoted by $\widehat{r}_{m,t+1}(a)$ for each $a\in\mathcal{A}$.

Loss calculation: For each training task, $m$, we calculate the loss at each
round, $t$, using the transformer’s prediction $\hat{r}_{m,t}(I_{m,t})$ and
the actual observed reward $r_{m,t}$ that followed action $I_{m,t}$. We use a
least-squares loss function:

$\displaystyle\mathcal{L}_{t}=\left(\widehat{r}_{m,t}(I_{m,t})-r_{m,t}\right)^{2}$
(3)

and hence minimizing this loss will minimize the mean squared-error of the
transformer’s predictions. The loss is calculated using (3) and is
backpropagated to update the model parameter $\mathbf{\Theta}$.

Exploratory Demonstrator: Observe from the loss definition in (3) that it is
calculated from the observed true reward and action from the dataset
$\mathcal{H}_{m}$. In order for the transformer to learn accurate reward
predictions during training, we require that the weak demonstrator is
sufficiently exploratory such that it collects $\mathcal{H}_{m}$ such that
$\mathcal{H}_{m}$ contains some reward $r_{m,t}$ for each action $a$. We
discuss in detail the impact of the demonstrator on PreDeToR (-$\tau$)
training in Section A.13.

###  3.2 Deploying PreDeToR for the online setting

In this section, we present the entire pseudocode of our algorithm and discuss
how it is deployed for online setting. In the online setting, PreDeToR, first
chooses the task $m\sim\mathcal{T}_{\text{test}}$, and the training dataset
$\mathcal{H}_{m}$ is initialized empty. Then at every round $t$, PreDeToR
chooses $I_{t}=\mathop{\rm
arg\,max}_{a\in\mathcal{A}}\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}\left(\widehat{r}_{m,t}(a)\mid\mathcal{H}^{t}_{m}\right)$
which is the action with the highest predicted reward and
$\widehat{r}_{m,t}(a)$ is the predicted reward of action $a$. Note that
PreDeToR is a greedy policy and may not conduct sufficient exploration.
Therefore we also introduce a soft variant, PreDeToR-$\tau$ that chooses
$I_{t}\sim\mathrm{softmax}^{\tau}_{a}\left(\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}\left(\mathbf{\widehat{r}}_{m,t}(a)\mid\mathcal{H}^{t}_{m}\right)\right)$.
Note that PreDeToR now adds the observed reward $r_{t}(I_{t})$ at round $t$ to
the dataset $\mathcal{H}_{m}$ and then uses that to act greedily for the next
round $t+1$. The full pseudocode of PreDeToR in the online setting is in
Algorithm 1. We discuss how PreDeToR (-$\tau$) is deployed in the offline
setting in Section A.15.

Algorithm 1 Pre-trained Decision Transformer with Reward Estimation (PreDeToR)

1: Collecting Pretraining Dataset

2: Initialize empty pretraining dataset $\mathcal{H}_{\text{train}}$

3: for $i$ in $[M_{\text{pre}}]$ do

4: Sample task $m\sim\mathcal{T}_{\text{pre}}$, in-context dataset
$\mathcal{H}_{m}\sim\mathcal{D}_{\text{pre}}(\cdot|m)$ and add this to
$\mathcal{H}_{\text{train}}$.

5: end for

6: Pretraining model on dataset

7: Initialize model $\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}$ with
parameters $\mathbf{\Theta}$

8: while not converged do

9: Sample $\mathcal{H}_{m}$ from $\mathcal{H}_{\text{train}}$ and predict
$\widehat{r}_{m,t}$ for action $(I_{m,t})$ for all $t\in[n]$

10: Compute loss in (3) with respect to $r_{m,t}$ and backpropagate to update
model parameter $\mathbf{\Theta}$.

11: end while

12: Online test-time deployment

13: Sample unknown task $m\sim\mathcal{T}_{\text{test}}$ and initialize empty
$\mathcal{H}_{m}^{0}=\\{\emptyset\\}$

14: for $t=1,2,\ldots,n$ do

15: Use $\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}$ on $m$ at round $t$ to
choose  $\displaystyle I_{t}\begin{cases}=\mathop{\rm
arg\,max}_{a\in\mathcal{A}}\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}\left(\widehat{r}_{m,t}(a)\mid\mathcal{H}^{t}_{m}\right),&\textbf{{\color[rgb]{0,0.6,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.6,0}PreDeToR}}\\\
\sim\textrm{softmax}^{\tau}_{a}\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}\left(\widehat{r}_{m,t}(a)\mid\mathcal{H}^{t}_{m}\right),&\textbf{{\color[rgb]{0,0.6,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.6,0}PreDeToR-$\tau$}}\end{cases}$

16: Add $\left\\{I_{t},r_{t}\right\\}$ to $\mathcal{H}_{m}^{t}$

17: end for

##  4 Empirical Study: Low Data Regime

In this section, we show that PreDeToR achieves lower regret than other in-
context algorithms for the low data regime. In a low data regime, every task
has a short horizon $n$ which makes this a challenging setup as there is less
data available per task. Therefore any good learner must learn to leverage the
structure across the tasks. This setting is quite common in recommender
settings where a user interaction with the system lasts only for a few rounds
(Kwon et al., 2022; Tomkins et al., 2020). We also analyze the exploration
strategy of PreDeToR (-$\tau$) in this setting. We study the performance of
PreDeToR in the large horizon setting in Section A.7.

Baselines: We first discuss the baselines used in this setting.

(1) PreDeToR: This is our proposed method shown in Algorithm 1.

(2) PreDeToR-$\tau$: This is the proposed exploratory method shown in
Algorithm 1 and we fix $\tau=0.05$.

(3) DPT-greedy: This baseline is the greedy approximation of the DPT algorithm
from Lee et al. (2023) which is discussed in Section 2.3. Note that we choose
DPT-greedy as a representative example of similar in-context decision-making
algorithms studied in Lee et al. (2023); Sinii et al. (2023); Lin et al.
(2023); Ma et al. (2023); Liu et al. (2023c, a) all of which require the
optimal action (or its greedy approximation). We choose DPT as a
representative example of this learning procedure. The DPT-greedy uses
training data $\mathcal{H}_{\text{train}}$ to estimate the optimal action for
each task $m$. For each task $m$, it DPT-greedy evaluates the optimal action
$\widehat{a}_{m,*}=\mathop{\rm
arg\,max}_{a}\sum_{t=1}^{n}\widehat{r}_{m,t}(a)$ where,
$\widehat{r}_{m,t}(a)=\sum_{t=1}^{n}\widehat{r}_{m,t}\mathbf{1}(I_{t}=a)$ is
the observed cumulative reward for the action $a$ in task $m$ and
$\mathbf{1}(\cdot)$ is the indicator function. Note that the optimal action
$\widehat{a}_{m,*}$ is evaluated using the entire dataset $\mathcal{H}_{m}$.
Finally, it implements the training procedure as discussed in Section 2.3.

(4) AD: This is the Algorithmic Distillation method (Laskin et al., 2022; Lu
et al., 2021) discussed in Section 2.3.

(5) Thomp: This baseline is the stochastic $A$-actioned bandit Thompson
Sampling algorithm from Thompson (1933); Agrawal and Goyal (2012); Russo et
al. (2018); Zhu and Tan (2020). We briefly describe the algorithm below: At
every round $t$ and each action $a$, Thomp samples
$\gamma_{m,t}(a)\sim\mathcal{N}(\widehat{\mu}_{m,t-1}(a),\sigma^{2}/N_{m,t-1}(a))$,
where $N_{m,t-1}(a)$ is the number of times the action $a$ has been selected
till $t-1$, and
$\widehat{\mu}_{m,t-1}(a)=\frac{\sum_{s=1}^{t-1}\widehat{r}_{m,s}\mathbf{1}(I_{s}=a)}{N_{m,t-1}(a)}$
is the empirical mean. Then the action selected at round $t$ is
$I_{t}=\mathop{\rm arg\,max}_{a}\gamma_{m,t}(a)$. Observe that Thomp is not a
deterministic algorithm like UCB (Auer et al., 2002). So we choose Thomp as
the weak demonstrator $\pi^{w}$ because it is more exploratory than UCB and
also chooses the optimal action $a_{m,*}$ sufficiently large number of times.
Thomp is a weak demonstrator as it does not have access to the feature set
$\mathcal{X}$ for any task $m$.

(6) LinUCB: (Linear Upper Confidence Bound): This baseline is the Upper
Confidence Bound algorithm for the linear bandit setting that selects the
action $I_{t}$ at round $t$ for task $m$ that is most optimistic and reduces
the uncertainty of the task unknown parameter $\bm{\theta}_{m,*}$. To balance
exploitation and exploration between choosing different items the LinUCB
computes an upper confidence value to the estimated mean of each action
$\mathbf{x}_{m,a}\in\mathcal{X}$. This is done as follows: At every round $t$
for task $m$, it calculates the ucb value $B_{m,a,t}$ for each action
$\mathbf{x}_{m,a}\in\mathcal{X}$ such that
$B_{m,a,t}=\mathbf{x}_{m,a}^{\top}\widehat{\bm{\theta}}_{m,t-1}+\alpha\|\mathbf{x}_{m,a}\|_{\mathbf{\Sigma}_{m,t-1}^{-1}}$
where $\alpha>0$ is a constant and $\widehat{\bm{\theta}}_{m,t}$ is the
estimate of the model parameter $\bm{\theta}_{m,*}$ at round $t$. Here,
$\mathbf{\Sigma}_{m,t-1}=\sum_{s=1}^{t-1}\mathbf{x}_{m,s}\mathbf{x}_{m,s}^{\top}+\lambda\mathbf{I}_{d}$
is the data covariance matrix or the arms already tried. Then it chooses
$I_{t}=\mathop{\rm arg\,max}_{a}B_{m,a,t}$. Note that LinUCB is a strong
demonstrator that we give oracle access to the features of each action; other
algorithms do not observe the features. Hence, in linear bandits, LinUCB
provides an approximate upper bound on the performance of all algorithms.

Outcomes: Before presenting the result we discuss the main outcomes from our
experimental results in this section:

(1) PreDeToR (-$\tau$) learns to exploit the latent structure of the
underlying tasks even when it is trained without the optimal action $a_{m,*}$
(or its approximation) and without the action features $\mathcal{X}$.

(2) PreDeToR (-$\tau$) learns to exploit the reward correlation between tasks
and is able to have an accurate reward estimate of the optimal action for each
task $m$.

(3) PreDeToR (-$\tau$) outperforms its weak demonstrator even when it is
trained without the optimal action $a_{m,*}$ and without the feature set
$\mathcal{X}$.

Experimental Result: These findings are reported in Figure 1. In Figure 1(a)
we show the linear bandit setting for horizon $n=25$, $M_{\text{pre}}=200000$,
$M_{\text{test}}=200$, $A=10$, and $d=2$. Further empirical setting details
are stated in Section A.2. The demonstrator $\pi^{w}$ is the Thomp algorithm.
We observe that PreDeToR (-$\tau$) has lower cumulative regret than DPT-
greedy, and AD. Note that for this low data regime the DPT-greedy does not
have a good estimation of $\widehat{a}_{m,*}$ which results in a poor
prediction of optimal action $\widehat{a}_{m,t,*}$. This results in higher
regret. Observe that PreDeToR (-$\tau$) performs quite similarly to LinUCB and
lowers regret compared to Thomp which also shows that PreDeToR is able to
exploit the latent linear structure and reward correlation of the underlying
tasks. AD is not able to outperform its demonstrator Thomp as it is trained to
predict the next action of the demonstrator. In Figure 1(b) we show the non-
linear bandit setting for horizon $n=50$, $M_{\text{pre}}=100000$,
$M_{\text{test}}=200$, $A=6$, and $d=2$. The demonstrator $\pi^{w}$ is the
Thomp algorithm. Again we observe that PreDeToR (-$\tau$) has lower cumulative
regret than DPT-greedy. Note that PreDeToR (-$\tau$) has lower regret than
LinUCB which fails to perform well in this non-linear setting due to its
algorithmic design. Finally, PreDeToR-$\tau$ performs slightly better than
PreDeToR in both settings as it conducts more exploration.

|  |  |
---|---|---|---

(a) Linear Bandit setting

(b) Non-linear bandit

(c) Test action distribution

(d) Test Prediction Error

Figure 1: Low-Data regime. The horizontal axis is the number of rounds.
Confidence bars show one standard error.

We also show how the prediction error of the optimal action by PreDeToR is
small compared to LinUCB in the linear bandit setting. In Figure 1(c) we first
show how the $10$ actions are distributed in the $M_{\text{test}}=200$ test
tasks. In Figure 1(c) for each bar, the frequency indicates the number of
tasks where the action (shown in the x-axis) is the optimal action. Then, in
Figure 1(d), we show the prediction error of PreDeToR (-$\tau$) for each task
$m\in[M_{\text{test}}]$. The prediction error is calculated as
$(\widehat{\mu}_{m,n,*}(a)-\mu_{m,*}(a))^{2}$ where
$\widehat{\mu}_{m,n,*}(a)=\max_{a}\widehat{\bm{\theta}}_{m,n}^{\top}\mathbf{x}_{m}(a)$
is the empirical mean at the end of round $n$, and
$\mu_{*,m}(a)=\max_{a}\bm{\theta}_{m,*}^{\top}\mathbf{x}_{m}(a)$ is the true
mean of the optimal action in task $m$. Then we average the prediction error
for the action $a\in[A]$ by the number of times the action $a$ is the optimal
action in some task $m$. From the Figure 1(d), we see that for actions
$\\{2,3,5,6,7,10\\}$, the prediction error of PreDeToR is either close or
smaller than LinUCB. Note that LinUCB estimates the empirical mean directly
from the test task, whereas PreDeToR has a strong prior based on the training
data. So PreDeToR is able to estimate the reward of the optimal action quite
well from the training dataset $\mathcal{D}_{\text{pre}}$. This shows the
power of PreDeToR to go beyond the in-context decision-making setting studied
in Lee et al. (2023); Lin et al. (2023); Ma et al. (2023); Sinii et al.
(2023); Liu et al. (2023c) which require long horizons/trajectories and
optimal action during training to learn a near-optimal policy. We discuss how
exploration of PreDeToR (-$\tau$) results in low cumulative regret in Section
A.11. Importantly, we show that predicting the reward has significant
advantages in a low-data regime as it leads to the right notion of
exploration.

##  5 Empirical Study: New Actions

In this section, we discuss how new actions introduced both during test time
and data collection influence the performance of PreDeToR (-$\tau$) and other
baselines. These experiments go beyond what is studied in Lee et al. (2023);
Lin et al. (2023); Ma et al. (2023); Liu et al. (2023c) as they only study the
performance in setting when
$\mathcal{D}_{\text{test}}=\mathcal{D}_{\text{pre}}$. Moreover, these
experiments are particularly important as they show the extent to which
PreDeToR(-$\tau$) is leveraging the latent structure.

Invariant actions: We denote the set of actions fixed across the different
tasks in the pretraining in-context dataset as $\mathcal{A}^{\text{inv}}$.
Therefore these action features $\mathbf{x}(a)\in\mathbb{R}^{d}$ for
$a\in\mathcal{A}^{\text{inv}}$ are fixed across the different tasks $m$. Note
that these invariant actions help the transformer $\mathrm{TF}_{\mathbf{w}}$
to learn the latent structure and the reward correlation across the different
tasks.

New actions: However, we also want to test how robust is PreDeToR (-$\tau$) to
new actions not seen during training time. To this effect, for each task
$m\in[M_{\text{pre}}]$ and $m\in[M_{\text{test}}]$ we introduce
$A-|\mathcal{A}^{\text{inv}}|$ new actions. For each of these new actions
$a\in[A-|\mathcal{A}^{\text{inv}}|]$ we choose the features $\mathbf{x}(m,a)$
randomly from $\mathcal{X}\subseteq\mathbb{R}^{d}$. Note the transformer now
trains on a dataset
$\mathcal{H}_{m}\subseteq\mathcal{D}_{\text{pre}}\neq\mathcal{D}_{\text{test}}$.

|  |  |
---|---|---|---

(a) Linear $1$ new action

(b) Non-linear $1$ new action

(c) Linear $5$ new actions

(d) Non-linear $5$ new actions

Figure 2: New action experiments. The horizontal axis is the number of rounds.
Confidence bars show one standard error.

Baselines: We implement the same baselines discussed in Section 4.

Outcomes: Again before presenting the result we discuss the main outcomes from
our experimental results of introducing new actions during data collection and
evaluation:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent structure
and reward correlation between the different tasks even when
$\mathcal{D}_{\text{test}}\neq\mathcal{D}_{\text{pre}}$ and the number of
changing actions per task is small. So PreDeToR (-$\tau$) is more robust to
changing actions compared to AD and DPT-greedy.

(2) PreDeToR (-$\tau$) has lower regret than weak demonstrator even when it is
trained on $\mathcal{D}_{\text{pre}}\neq\mathcal{D}_{\text{test}}$.

(3) The PreDeToR (-$\tau$) cumulative regret increases as the number of
invariant actions decreases. This shows that invariant actions are important
for learning the underlying latent structure and reward correlation across
tasks.

Experimental Result: We observe these outcomes in Figure 2. We consider the
linear bandit setting of horizon $n=50$, $M_{\text{pre}}=200000$,
$M_{\text{test}}=200$, $A=20$, and $d=5$. Here during data collection and
during collecting the test data, we randomly select one new action from
$\mathbb{R}^{d}$ for each task $m$. So the number of invariant actions is
$|\mathcal{A}^{\text{inv}}|=19$. Again, the demonstrator $\pi^{w}$ is the
Thomp algorithm. We observe that PreDeToR (-$\tau$) has lower cumulative
regret than DPT-greedy, and AD. The weak performance of DPT-greedy can be
attributed to low data per task and inability to adapt when
$\mathcal{D}_{\text{pre}}\neq\mathcal{D}_{\text{test}}$. Observe that PreDeToR
(-$\tau$) almost matches LinUCB and has lower regret than Thomp. This shows
that PreDeToR (-$\tau$) is able to exploit the latent linear structure of the
underlying tasks. In Figure 2(b) we show the non-linear bandit setting for
horizon $n=50$, $M_{\text{pre}}=100000$, $M_{\text{test}}=200$,$A=6$, $d=2$,
and $|\mathcal{A}^{\text{inv}}|=5$. The demonstrator $\pi^{w}$ is the Thomp
algorithm. Again, we observe that PreDeToR has lower cumulative regret than
DPT-greedy and AD and furthermore improves upon LinUCB which fails to perform
well in this non-linear setting due to its assumption of linear rewards.
PreDeToR-$\tau$ performs slightly better than PreDeToR in $1$ new arm setting
as it conducts more exploration.

In the following two experiments, we test the limit of robustness of PreDeToR
(-$\tau$) by decreasing the number of invariant actions. We use
$|\mathcal{A}^{\text{inv}}|=15$ for linear, and $|\mathcal{A}^{\text{inv}}|=2$
for the non-linear setting. All the other task parameters are the same as in
the previous two experiments. The corresponding results are shown in Figure
2(c) and Figure 2(d). Observe that in the linear setting (where LinUCB is
effective), the performance of PreDeToR degrades compared to the case when
$|\mathcal{A}^{\text{inv}}|=15$. This shows that as the number of invariant
actions decreases, the ability of PreDeToR (-$\tau$) to estimate the
underlying latent structure and reward correlation also decreases. In Section
A.12, we analyze the exploration of PreDeToR (-$\tau$) in this setting, which
shows that the shared latent structure is essential for minimizing cumulative
regret.

We also empirically study the test performance of PreDeToR (-$\tau$) in other
bandit settings such as bilinear bandits (Section A.3), latent bandits
(Section A.4), draw a connection between PreDeToR and Bayesian estimators
(Section A.5), and perform sensitivity and ablation studies in Section A.6,
A.8, A.9, A.10. We discuss data collection algorithms in Section A.13 and the
offline setting in Section A.15. Due to space constraints, we refer the
interested reader to the relevant section in the appendices.

##  6 Theoretical Analysis

In this section, we present a theoretical analysis of how PreDeToR-$\tau$
generalizes to an unknown target task given a set of source tasks. We also
show that in low-data regimes the PreDeToR-$\tau$ has a low expected excess
risk for the unknown target task as the number of source tasks increases. We
proceed as follows: Suppose we have the training data set
$\mathcal{H}_{\mathrm{all}}=\\{\mathcal{H}_{m}\\}_{m=1}^{M_{\text{pre}}},$
where the task $m\sim\mathcal{T}$ with a distribution $\mathcal{T}$ and the
task data $\mathcal{H}_{m}$ is generated from a distribution
$\mathcal{D}_{\text{pre}}(\cdot|m).$ For illustration purposes, here we
consider the training data distribution $\mathcal{D}_{\text{pre}}(\cdot|m)$
where the actions are sampled following soft-LinUCB (Chu et al., 2011). Note
that soft-LinUCB produces a distribution over actions (see Section A.14 for a
detailed description). Given the loss function in (3), we can define the task
$m$ training loss of PreDeToR-$\tau$ as
$\widehat{\mathcal{L}}_{m}(\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}})=\frac{1}{n}\sum_{t=1}^{n}\ell(r_{m,t},\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}(\widehat{r}_{m,t}(I_{m,t})|\mathcal{H}^{t}_{m}))=\frac{1}{n}\sum_{t=1}^{n}(\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}(\widehat{r}_{m,t}(I_{m,t})|\mathcal{H}^{t}_{m})-r_{m,t})^{2}$.
We drop the notation $\mathbf{\Theta},\mathbf{r}$ from
$\mathrm{TF^{\mathbf{r}}}_{\mathbf{\Theta}}$ for simplicity and let
$M=M_{\text{pre}}$. We define

$\displaystyle\widehat{\mathrm{TF}}=$
$\displaystyle\underset{\mathrm{TF}\in\mathrm{Alg}}{\arg\min}\widehat{\mathcal{L}}_{\mathcal{H}_{\mathrm{all}}}(\mathrm{TF}):=\frac{1}{M}\sum_{m=1}^{M}\widehat{\mathcal{L}}_{m}(\mathrm{TF}),\quad\text{(ERM)}$
(4)

where $\mathrm{Alg}$ denotes the space of algorithms induced by the
$\mathrm{TF}$. Let
$\mathcal{L}_{m}(\mathrm{TF})=\mathbb{E}_{\mathcal{H}_{m}}\big{[}\widehat{\mathcal{L}}_{m}(\mathrm{TF})\big{]}$
and
$\mathcal{L}_{\mathrm{MTL}}(\mathrm{TF})=\mathbb{E}\big{[}\widehat{\mathcal{L}}_{\mathcal{H}_{\mathrm{all}}}(\mathrm{TF})\big{]}=$
$\frac{1}{M}\sum_{m=1}^{M}\mathcal{L}_{m}(\mathrm{TF})$ be the corresponding
population risks. For the ERM in (4), we want to bound the following excess
Multi-Task Learning (MTL) risk of PreDeToR-$\tau$

$\displaystyle\mathcal{R}_{\mathrm{MTL}}(\widehat{\mathrm{TF}})=\mathcal{L}_{\mathrm{MTL}}(\widehat{\mathrm{TF}})-\min_{\mathrm{TF}\in\mathrm{Alg}}\mathcal{L}_{\mathrm{MTL}}(\mathrm{TF}).$
(5)

Note that in in-context learning, a training sample $\left(I_{t},r_{t}\right)$
impacts all future decisions of the algorithm from time step $t+1$ to $n$.
Therefore, we need to control the stability of the input perturbation of the
learning algorithm learned by the transformer. We introduce the following
stability condition.

######  Assumption 6.1.

(Error stability (Bousquet and Elisseeff, 2002; Li et al., 2023)). Let
$\mathcal{H}=\left(I_{t},r_{t}\right)_{t=1}^{n}$ be a sequence in
$[A]\times[0,1]$ with $n\geq 1$ and $\mathcal{H}^{\prime}$ be the sequence
where the $t^{\prime}$th sample of $\mathcal{H}$ is replaced by
$\left(I_{t}^{\prime},r_{t}^{\prime}\right)$. Error stability holds for a
distribution $(I,r)\sim\mathcal{D}$ if there exists a $K>0$ such that for any
$\mathcal{H},\left(I_{t}^{\prime},r_{t}^{\prime}\right)\in([A]\times[0,1]),t\leq
n$, and $\mathrm{TF}\in\mathrm{Alg}$, we have

$\displaystyle\big{\lvert}\mathbb{E}_{(I,r)}\left[\ell(r,\mathrm{TF}(\widehat{r}(I)|\mathcal{H}))-\ell\left(r,\mathrm{TF}(\widehat{r}(I)|\mathcal{H}^{\prime})\right)\right]\big{\rvert}\,\leq\tfrac{K}{n}$

Let $\rho$ be a distance metric on $\mathrm{Alg}$. Pairwise error stability
holds if for all $\mathrm{TF},\mathrm{TF}^{\prime}\in\mathrm{Alg}$ we have

$\displaystyle\begin{aligned}
&\big{\lvert}\mathbb{E}_{(x,y)}\left[\ell(r,\mathrm{TF}(\widehat{r}(I)|\mathcal{H}))-\ell\left(r,\mathrm{TF}^{\prime}(\widehat{r}(I)|\mathcal{H})\right)-\ell(r,\mathrm{TF}(\widehat{r}(I)|\mathcal{H}^{\prime}))+\ell\left(r,\mathrm{TF}^{\prime}(\widehat{r}(I)|\mathcal{H}^{\prime})\right)\right]\big{\rvert}\leq\tfrac{K\rho\left(\mathrm{TF},\mathrm{TF}^{\prime}\right)}{n}\end{aligned}$

Now we present the Multi-task learning (MTL) risk of PreDeToR-$\tau$.

######  Theorem 6.2.

(PreDeToR risk) Suppose error stability 6.1 holds and assume loss function
$\ell(\cdot,\cdot)$ is $C$-Lipschitz for all $r_{t}\in[0,B]$ and horizon
$n\geq 1$. Let $\widehat{\mathrm{TF}}$ be the empirical solution of (ERM) and
$\mathcal{N}(\mathcal{A},\rho,\epsilon)$ be the covering number of the
algorithm space $\mathrm{Alg}$ following Definition C.2 and C.3. Then with
probability at least $1-2\delta$, the excess MTL risk of PreDeToR-$\tau$ is
bounded by $\mathcal{R}_{\mathrm{MTL}}(\widehat{\mathrm{TF}})\leq
4\tfrac{C}{\sqrt{nM}}+2(B+K\log
n)\sqrt{\tfrac{\log(\mathcal{N}(\mathrm{Alg},\rho,\varepsilon)/\delta)}{cnM}},$
where $\mathcal{N}(\mathrm{Alg},\rho,\varepsilon)$ is the covering number of
transformer $\widehat{\mathrm{TF}}$ and $\epsilon=1/\sqrt{nM}$.

The proof of this theorem is provided in Section C.1. From Theorem 6.2 we see
that in low-data regime with a small horizon $n$, as the number of tasks $M$
increases the MTL risk decreases. We further discuss the stability factor $K$
and covering number $\mathcal{N}(\mathrm{Alg},\rho,\varepsilon)$ in Remark
C.4, and C.5.

We now present the transfer learning risk of PreDeToR-$\tau$ for an unknown
target task $g\sim\mathcal{T}$ with the test dataset
$\mathcal{H}_{g}\sim\mathcal{D}_{\text{test}}(\cdot|g)$. Note that the test
data distribution $\mathcal{D}_{\text{test}}(\cdot|g)$ is such that the
actions are sampled following soft-LinUCB.

######  Theorem 6.3.

(Transfer risk) Consider the setting of Theorem 6.2 and assume the training
source tasks are independently drawn from task distribution $\mathcal{T}$. Let
$\widehat{\text{ TF }}$ be the empirical solution of (ERM) and
$g\\!\sim\\!\mathcal{T}$. Define the expected excess transfer learning risk
$\mathbb{E}_{g}[\mathcal{R}_{g}]=\mathbb{E}_{{g}}\big{[}\mathcal{L}_{g}(\widehat{\mathrm{TF}})\big{]}\\!-\\!\arg\min_{\mathrm{TF}\in\mathrm{Alg}}\mathbb{E}_{{g}}\left[\mathcal{L}_{g}(\mathrm{TF})\right]$.
Then with probability at least $1\\!-\\!2\delta$, the
$\mathbb{E}_{g}\big{[}\mathcal{R}_{g}\big{]}\\!\leq\\!4\tfrac{C}{\sqrt{M}}\\!+\\!2B\sqrt{\tfrac{\log(\mathcal{N}(\mathrm{Alg},\rho,\varepsilon)/\delta)}{M}},$
where $\mathcal{N}(\mathrm{Alg},\rho,\varepsilon)$ is the covering number of
$\widehat{\mathrm{TF}}$ and $\epsilon=\tfrac{1}{\sqrt{M}}$.

The proof is given in Section C.2. This shows that for the transfer learning
risk of PreDeToR-$\tau$ (once trained on the $M$ source tasks) scales with
$\widetilde{O}(1/\sqrt{M})$. This is because the unseen target task
$g\sim\mathcal{T}$ induces a distribution shift, which, typically, cannot be
mitigated with more samples $n$ per task. A similar observation is provided in
Lin et al. (2023). We further discuss this in Remark C.7. We also observe a
similar phenomenon empirically; see the discussion in Section A.14.

##  7 Conclusions, Limitations and Future Works

In this paper, we studied the supervised pretraining of decision transformers
in the multi-task structured bandit setting when the knowledge of the optimal
action is unavailable. Moreover, our proposed methods PreDeToR (-$\tau$) do
not need to know the action representations or the reward structure and learn
these in-context with the help of offline data. The PreDeToR (-$\tau$) predict
the reward for the next action of each action during pretraining and can
generalize well in-context in several regimes spanning low-data, new actions,
and structured bandit settings like linear, non-linear, bilinear, latent
bandits. The PreDeToR (-$\tau$) outperforms other in-context algorithms like
AD, DPT-greedy in most of the experiments. Finally, we theoretically analyze
PreDeToR-$\tau$ and show that pretraining it in $M$ source tasks leads to a
low expected excess error on a target task drawn from the same task
distribution $\mathcal{T}$. In future, we want to extend our PreDeToR
(-$\tau$) to MDP setting (Sutton and Barto, 2018; Agarwal et al., 2019), and
constraint MDP setting (Efroni et al., 2020; Gu et al., 2022).

## References

  * Abbasi-Yadkori et al. [2011] Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári.  Improved algorithms for linear stochastic bandits.  _Advances in neural information processing systems_ , 24, 2011.
  * Agarwal et al. [2019] Alekh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun.  Reinforcement learning: Theory and algorithms.  _CS Dept., UW Seattle, Seattle, WA, USA, Tech. Rep_ , 32, 2019.
  * Agrawal and Goyal [2012] Shipra Agrawal and Navin Goyal.  Analysis of thompson sampling for the multi-armed bandit problem.  In _Conference on learning theory_ , pages 39–1. JMLR Workshop and Conference Proceedings, 2012.
  * Auer and Ortner [2010] Peter Auer and Ronald Ortner.  Ucb revisited: Improved regret bounds for the stochastic multi-armed bandit problem.  _Periodica Mathematica Hungarica_ , 61(1-2):55–65, 2010.
  * Auer et al. [2002] Peter Auer, Nicolò Cesa-Bianchi, and Paul Fischer.  Finite-time Analysis of the Multiarmed Bandit Problem.  _Machine Learning_ , 47(2):235–256, May 2002.  ISSN 1573-0565.  doi: 10.1023/A:1013689704352.  URL https://doi.org/10.1023/A:1013689704352.
  * Bengio et al. [1990] Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier.  _Learning a synaptic learning rule_.  Université de Montréal, Département d’informatique et de recherche …, 1990.
  * Bishop [2006] C Bishop.  Pattern recognition and machine learning.  _Springer google schola_ , 2:531–537, 2006.
  * Bousquet and Elisseeff [2002] Olivier Bousquet and André Elisseeff.  Stability and generalization.  _The Journal of Machine Learning Research_ , 2:499–526, 2002.
  * Box and Tiao [2011] George EP Box and George C Tiao.  _Bayesian inference in statistical analysis_.  John Wiley & Sons, 2011.
  * Brandfonbrener et al. [2022] David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna.  When does return-conditioned supervised learning work for offline reinforcement learning?  _Advances in Neural Information Processing Systems_ , 35:1542–1553, 2022.
  * Brohan et al. [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.
  * Bubeck et al. [2008] Sébastien Bubeck, Gilles Stoltz, Csaba Szepesvári, and Rémi Munos.  Online optimization in x-armed bandits.  _Advances in Neural Information Processing Systems_ , 21, 2008.
  * Bubeck et al. [2011] Sébastien Bubeck, Gilles Stoltz, and Jia Yuan Yu.  Lipschitz bandits without the lipschitz constant.  In _Algorithmic Learning Theory: 22nd International Conference, ALT 2011, Espoo, Finland, October 5-7, 2011. Proceedings 22_ , pages 144–158. Springer, 2011.
  * Carlin and Louis [2008] Bradley P Carlin and Thomas A Louis.  _Bayesian methods for data analysis_.  CRC press, 2008.
  * Chaudhuri et al. [2017] Kamalika Chaudhuri, Prateek Jain, and Nagarajan Natarajan.  Active heteroscedastic regression.  In _International Conference on Machine Learning_ , pages 694–702. PMLR, 2017.
  * Chen et al. [2021] 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.
  * Chowdhury and Gopalan [2017] Sayak Ray Chowdhury and Aditya Gopalan.  On kernelized multi-armed bandits.  In _International Conference on Machine Learning_ , pages 844–853. PMLR, 2017.
  * Chu et al. [2011] Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire.  Contextual bandits with linear payoff functions.  In _Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics_ , pages 208–214. JMLR Workshop and Conference Proceedings, 2011.
  * Dai et al. [2022] Zhongxiang Dai, Yao Shu, Arun Verma, Flint Xiaofeng Fan, Bryan Kian Hsiang Low, and Patrick Jaillet.  Federated neural bandit.  _arXiv preprint arXiv:2205.14309_ , 2022.
  * Degenne et al. [2020] Rémy Degenne, Pierre Ménard, Xuedong Shang, and Michal Valko.  Gamification of pure exploration for linear bandits.  In _International Conference on Machine Learning_ , pages 2432–2442. PMLR, 2020.
  * Dong et al. [2021] Kefan Dong, Jiaqi Yang, and Tengyu Ma.  Provable model-based nonlinear bandit and reinforcement learning: Shelve optimism, embrace virtual curvature.  _Advances in neural information processing systems_ , 34:26168–26182, 2021.
  * Du et al. [2023] Yihan Du, Longbo Huang, and Wen Sun.  Multi-task representation learning for pure exploration in linear bandits.  _arXiv preprint arXiv:2302.04441_ , 2023.
  * Duan et al. [2016] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel.  Rl $\\\ 2$: Fast reinforcement learning via slow reinforcement learning.  _arXiv preprint arXiv:1611.02779_ , 2016.
  * Efroni et al. [2020] Yonathan Efroni, Shie Mannor, and Matteo Pirotta.  Exploration-exploitation in constrained mdps.  _arXiv preprint arXiv:2003.02189_ , 2020.
  * Fedorov [2013] Valerii Vadimovich Fedorov.  _Theory of optimal experiments_.  Elsevier, 2013.
  * Filippi et al. [2010] Sarah Filippi, Olivier Cappe, Aurélien Garivier, and Csaba Szepesvári.  Parametric bandits: The generalized linear case.  _Advances in neural information processing systems_ , 23, 2010.
  * Finn et al. [2017] 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.
  * Fong et al. [2003] Terrence Fong, Illah Nourbakhsh, and Kerstin Dautenhahn.  A survey of socially interactive robots.  _Robotics and autonomous systems_ , 42(3-4):143–166, 2003.
  * François-Lavet et al. [2018] Vincent François-Lavet, Peter Henderson, Riashat Islam, Marc G Bellemare, Joelle Pineau, et al.  An introduction to deep reinforcement learning.  _Foundations and Trends® in Machine Learning_ , 11(3-4):219–354, 2018.
  * Fu et al. [2016] 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.
  * Fujimoto et al. [2019] 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.
  * Ge et al. [2022] Yao Ge, Yuting Guo, Yuan-Chi Yang, Mohammed Ali Al-Garadi, and Abeed Sarker.  Few-shot learning for medical text: A systematic review.  _arXiv preprint arXiv:2204.14081_ , 2022.
  * Ghasemipour et al. [2022] 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.
  * Gu et al. [2022] Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, Yaodong Yang, and Alois Knoll.  A review of safe reinforcement learning: Methods, theory and applications.  _arXiv preprint arXiv:2205.10330_ , 2022.
  * Gupta et al. [2018] Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine.  Meta-reinforcement learning of structured exploration strategies.  _Advances in neural information processing systems_ , 31, 2018.
  * Gupta et al. [2020] Samarth Gupta, Shreyas Chaudhari, Subhojyoti Mukherjee, Gauri Joshi, and Osman Yağan.  A unified approach to translate classical bandit algorithms to the structured bandit setting.  _IEEE Journal on Selected Areas in Information Theory_ , 1(3):840–853, 2020.
  * Hong et al. [2020] Joey Hong, Branislav Kveton, Manzil Zaheer, Yinlam Chow, Amr Ahmed, and Craig Boutilier.  Latent bandits revisited.  _Advances in Neural Information Processing Systems_ , 33:13423–13433, 2020.
  * Hong et al. [2022a] Joey Hong, Branislav Kveton, Sumeet Katariya, Manzil Zaheer, and Mohammad Ghavamzadeh.  Deep hierarchy in bandits.  In _International Conference on Machine Learning_ , pages 8833–8851. PMLR, 2022a.
  * Hong et al. [2022b] Joey Hong, Branislav Kveton, Manzil Zaheer, and Mohammad Ghavamzadeh.  Hierarchical bayesian bandits.  In _International Conference on Artificial Intelligence and Statistics_ , pages 7724–7741. PMLR, 2022b.
  * Janner et al. [2021] Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  _Advances in neural information processing systems_ , 34:1273–1286, 2021.
  * Jiang et al. [2022] 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.
  * Johnson et al. [2002] Richard Arnold Johnson, Dean W Wichern, et al.  Applied multivariate statistical analysis.  2002\.
  * Jun et al. [2019] Kwang-Sung Jun, Rebecca Willett, Stephen Wright, and Robert Nowak.  Bilinear bandits with low-rank structure.  In _International Conference on Machine Learning_ , pages 3163–3172. PMLR, 2019.
  * Justus et al. [2018] Daniel Justus, John Brennan, Stephen Bonner, and Andrew Stephen McGough.  Predicting the computational cost of deep learning models.  In _2018 IEEE international conference on big data (Big Data)_ , pages 3873–3882. IEEE, 2018.
  * Kang et al. [2022] Yue Kang, Cho-Jui Hsieh, and Thomas Chun Man Lee.  Efficient frameworks for generalized low-rank matrix bandit problems.  _Advances in Neural Information Processing Systems_ , 35:19971–19983, 2022.
  * Kumar et al. [2019] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Kumar et al. [2020] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Kveton et al. [2017] Branislav Kveton, Csaba Szepesvári, Anup Rao, Zheng Wen, Yasin Abbasi-Yadkori, and S Muthukrishnan.  Stochastic low-rank bandits.  _arXiv preprint arXiv:1712.04644_ , 2017.
  * Kwon et al. [2022] Jeongyeol Kwon, Yonathan Efroni, Constantine Caramanis, and Shie Mannor.  Tractable optimality in episodic latent mabs.  _Advances in Neural Information Processing Systems_ , 35:23634–23645, 2022.
  * Landolfi et al. [2019] Nicholas C Landolfi, Garrett Thomas, and Tengyu Ma.  A model-based approach for sample-efficient multi-task reinforcement learning.  _arXiv preprint arXiv:1907.04964_ , 2019.
  * Laskin et al. [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.
  * Lattimore and Szepesvári [2019] Tor Lattimore and Csaba Szepesvári.  An information-theoretic approach to minimax regret in partial monitoring.  In _Conference on Learning Theory_ , pages 2111–2139. PMLR, 2019.
  * Lattimore and Szepesvári [2020] Tor Lattimore and Csaba Szepesvári.  _Bandit algorithms_.  Cambridge University Press, 2020.
  * Lauri et al. [2022] Mikko Lauri, David Hsu, and Joni Pajarinen.  Partially observable markov decision processes in robotics: A survey.  _IEEE Transactions on Robotics_ , 39(1):21–40, 2022.
  * Lee et al. [2023] Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill.  Supervised pretraining can learn in-context reinforcement learning.  _arXiv preprint arXiv:2306.14892_ , 2023.
  * Lee et al. [2022] 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.
  * Li et al. [2020] 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.
  * Li et al. [2010] Lihong Li, Wei Chu, John Langford, and Robert E Schapire.  A contextual-bandit approach to personalized news article recommendation.  In _Proceedings of the 19th international conference on World wide web_ , pages 661–670, 2010.
  * Li et al. [2017] Lihong Li, Yu Lu, and Dengyong Zhou.  Provably optimal algorithms for generalized linear contextual bandits.  In _International Conference on Machine Learning_ , pages 2071–2080. PMLR, 2017.
  * Li et al. [2023] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.  Transformers as algorithms: Generalization and stability in in-context learning.  In _International Conference on Machine Learning_ , pages 19565–19594. PMLR, 2023.
  * Lin et al. [2023] Licong Lin, Yu Bai, and Song Mei.  Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining.  _arXiv preprint arXiv:2310.08566_ , 2023.
  * Liu et al. [2021] 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.
  * Liu et al. [2023a] Xiaoqian Liu, Jianbin Jiao, and Junge Zhang.  Self-supervised pretraining for decision foundation model: Formulation, pipeline and challenges.  _arXiv preprint arXiv:2401.00031_ , 2023a.
  * Liu et al. [2023b] Xin Liu, Daniel McDuff, Geza Kovacs, Isaac Galatzer-Levy, Jacob Sunshine, Jiening Zhan, Ming-Zher Poh, Shun Liao, Paolo Di Achille, and Shwetak Patel.  Large language models are few-shot health learners.  _arXiv preprint arXiv:2305.15525_ , 2023b.
  * Liu et al. [2019] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill.  Off-policy policy gradient with state distribution correction.  _arXiv preprint arXiv:1904.08473_ , 2019.
  * Liu et al. [2020] 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.
  * Liu et al. [2023c] Zhihan Liu, Hao Hu, Shenao Zhang, Hongyi Guo, Shuqi Ke, Boyi Liu, and Zhaoran Wang.  Reason for future, act for now: A principled framework for autonomous llm agents with provable sample efficiency.  _arXiv preprint arXiv:2309.17382_ , 2023c.
  * Lu et al. [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.
  * Lu et al. [2021] Yangyi Lu, Amirhossein Meisami, and Ambuj Tewari.  Low-rank generalized linear bandit problems.  In _International Conference on Artificial Intelligence and Statistics_ , pages 460–468. PMLR, 2021.
  * Ma et al. [2023] Yi Ma, Chenjun Xiao, Hebin Liang, and Jianye Hao.  Rethinking decision transformer via hierarchical reinforcement learning.  _arXiv preprint arXiv:2311.00267_ , 2023.
  * Madotto et al. [2021] Andrea Madotto, Zhaojiang Lin, Genta Indra Winata, and Pascale Fung.  Few-shot bot: Prompt-based learning for dialogue systems.  _arXiv preprint arXiv:2110.08118_ , 2021.
  * Magureanu et al. [2014] Stefan Magureanu, Richard Combes, and Alexandre Proutiere.  Lipschitz bandits: Regret lower bound and optimal algorithms.  In _Conference on Learning Theory_ , pages 975–999. PMLR, 2014.
  * Maillard and Mannor [2014] Odalric-Ambrym Maillard and Shie Mannor.  Latent bandits.  In _International Conference on Machine Learning_ , pages 136–144. PMLR, 2014.
  * Min et al. [2022] 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.
  * Mirchandani et al. [2023] Suvir Mirchandani, Fei Xia, Pete Florence, Brian Ichter, Danny Driess, Montserrat Gonzalez Arenas, Kanishka Rao, Dorsa Sadigh, and Andy Zeng.  Large language models as general pattern machines.  _arXiv preprint arXiv:2307.04721_ , 2023.
  * Mishra et al. [2017] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel.  A simple neural attentive meta-learner.  _arXiv preprint arXiv:1707.03141_ , 2017.
  * Mitchell et al. [2021] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn.  Offline meta-reinforcement learning with advantage weighting.  In _International Conference on Machine Learning_ , pages 7780–7791. PMLR, 2021.
  * Mnih et al. [2013] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Mukherjee et al. [2023] Subhojyoti Mukherjee, Qiaomin Xie, Josiah P Hanna, and Robert Nowak.  Multi-task representation learning for pure exploration in bilinear bandits.  _arXiv preprint arXiv:2311.00327_ , 2023.
  * Müller et al. [2021] Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter.  Transformers can do bayesian inference.  _arXiv preprint arXiv:2112.10510_ , 2021.
  * Murphy [2003] Susan A Murphy.  Optimal dynamic treatment regimes.  _Journal of the Royal Statistical Society Series B: Statistical Methodology_ , 65(2):331–355, 2003.
  * Nagabandi et al. [2018] 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 meta-reinforcement learning.  _arXiv preprint arXiv:1803.11347_ , 2018.
  * Neyshabur et al. [2017] Behnam Neyshabur, Ryota Tomioka, Ruslan Salakhutdinov, and Nathan Srebro.  Geometry of optimization and implicit regularization in deep learning.  _arXiv preprint arXiv:1705.03071_ , 2017.
  * Pal et al. [2023] Soumyabrata Pal, Arun Sai Suggala, Karthikeyan Shanmugam, and Prateek Jain.  Optimal algorithms for latent bandits with cluster structure.  In _International Conference on Artificial Intelligence and Statistics_ , pages 7540–7577. PMLR, 2023.
  * Perkins and Precup [1999] Theodore J Perkins and Doina Precup.  Using options for knowledge transfer in reinforcement learning title2, 1999.
  * Pong et al. [2022] 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.
  * Pukelsheim [2006] Friedrich Pukelsheim.  _Optimal design of experiments_.  SIAM, 2006.
  * Rakelly et al. [2019] 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.
  * Reed et al. [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.
  * Riquelme et al. [2018] Carlos Riquelme, George Tucker, and Jasper Snoek.  Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling.  _arXiv preprint arXiv:1802.09127_ , 2018.
  * Roberts et al. [2019] Adam Roberts, Colin Raffel, Katherine Lee, Michael Matena, Noam Shazeer, Peter J Liu, Sharan Narang, Wei Li, and Yanqi Zhou.  Exploring the limits of transfer learning with a unified text-to-text transformer.  2019\.
  * Rothfuss et al. [2018] Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel.  Promp: Proximal meta-policy search.  _arXiv preprint arXiv:1810.06784_ , 2018.
  * Russo et al. [2018] 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.
  * Schaul and Schmidhuber [2010] Tom Schaul and Jürgen Schmidhuber.  Metalearning.  _Scholarpedia_ , 5(6):4650, 2010.
  * Semnani et al. [2023] Sina Semnani, Violet Yao, Heidi Zhang, and Monica Lam.  Wikichat: Stopping the hallucination of large language model chatbots by few-shot grounding on wikipedia.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pages 2387–2413, 2023.
  * Shafiullah et al. [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.
  * Siegel et al. [2020] 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.
  * Sinii et al. [2023] Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov.  In-context reinforcement learning for variable action spaces.  _arXiv preprint arXiv:2312.13327_ , 2023.
  * Soudry et al. [2018] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro.  The implicit bias of gradient descent on separable data.  _Journal of Machine Learning Research_ , 19(70):1–57, 2018.
  * Sutton and Barto [2018] Richard S Sutton and Andrew G Barto.  _Reinforcement learning: An introduction_.  MIT press, 2018.
  * Thompson [1933] 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.
  * Tomkins et al. [2020] Sabina Tomkins, Peng Liao, Predrag Klasnja, Serena Yeung, and Susan Murphy.  Rapidly personalizing mobile health treatment policies with limited data.  _arXiv preprint arXiv:2002.09971_ , 2020.
  * Valko et al. [2013] Michal Valko, Nathaniel Korda, Rémi Munos, Ilias Flaounas, and Nelo Cristianini.  Finite-time analysis of kernelised contextual bandits.  _arXiv preprint arXiv:1309.6869_ , 2013.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang et al. [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.
  * Wu et al. [2019] Yifan Wu, George Tucker, and Ofir Nachum.  Behavior regularized offline reinforcement learning.  _arXiv preprint arXiv:1911.11361_ , 2019.
  * Xie et al. [2021] 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.
  * Yang et al. [2023] Adam X Yang, Maxime Robeyns, Xi Wang, and Laurence Aitchison.  Bayesian low-rank adaptation for large language models.  _arXiv preprint arXiv:2308.13111_ , 2023.
  * Yang et al. [2020] Jiaqi Yang, Wei Hu, Jason D Lee, and Simon S Du.  Impact of representation learning in linear bandits.  _arXiv preprint arXiv:2010.06531_ , 2020.
  * Yang et al. [2022a] Jiaqi Yang, Qi Lei, Jason D Lee, and Simon S Du.  Nearly minimax algorithms for linear bandits with shared representation.  _arXiv preprint arXiv:2203.15664_ , 2022a.
  * Yang and Wang [2020] Lin Yang and Mengdi Wang.  Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound.  In _International Conference on Machine Learning_ , pages 10746–10756. PMLR, 2020.
  * Yang et al. [2022b] 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_ , 2022b.
  * Yu et al. [2021] 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.
  * Zambaldi et al. [2018] Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al.  Relational deep reinforcement learning.  _arXiv preprint arXiv:1806.01830_ , 2018.
  * Zhao et al. [2023] Andrew Zhao, Daniel Huang, Quentin Xu, Matthieu Lin, Yong-Jin Liu, and Gao Huang.  Expel: Llm agents are experiential learners.  _arXiv preprint arXiv:2308.10144_ , 2023.
  * Zhou et al. [2020] Dongruo Zhou, Lihong Li, and Quanquan Gu.  Neural contextual bandits with ucb-based exploration.  In _International Conference on Machine Learning_ , pages 11492–11502. PMLR, 2020.
  * Zhu and Tan [2020] Qiuyu Zhu and Vincent Tan.  Thompson sampling algorithms for mean-variance bandits.  In _International Conference on Machine Learning_ , pages 11599–11608. PMLR, 2020.
  * Zintgraf et al. [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 meta-learning.  _arXiv preprint arXiv:1910.08348_ , 2019.

##  Appendix A Appendix

###  A.1 Related Works

In this section, we briefly discuss related works.

In-context decision making [Laskin et al., 2022, Lee et al., 2023] has emerged
as an attractive alternative in Reinforcement Learning (RL) compared to
updating the model parameters after collection of new data [Mnih et al., 2013,
François-Lavet et al., 2018]. In RL the contextual data takes the form of
state-action-reward tuples representing a dataset of interactions with an
unknown environment (task). In this paper, we will refer to this as the in-
context data. Recall that in many real-world settings, the underlying task can
be structured with correlated features, and the reward can be highly non-
linear. So specialized bandit algorithms fail to learn in these tasks. To
circumvent this issue, a learner can first collect in-context data consisting
of just action indices $I_{t}$ and rewards $r_{t}$. Then it can leverage the
representation learning capability of deep neural networks to learn a pattern
across the in-context data and subsequently derive a near-optimal policy [Lee
et al., 2023, Mirchandani et al., 2023]. We refer to this learning framework
as an in-context decision-making setting.

The in-context decision-making setting of Sinii et al. [2023] also allows
changing the action space by learning an embedding over the action space yet
also requires the optimal action during training. In contrast we do not
require the optimal action as well as show that we can generalize to new
actions without learning an embedding over them. Similarly, Lin et al. [2023]
study the in-context decision-making setting of Laskin et al. [2022], Lee et
al. [2023], but they also require a greedy approximation of the optimal
action. The Ma et al. [2023] also studies a similar setting for hierarchical
RL where they stitch together sub-optimal trajectories and predict the next
action during test time. Similarly, Liu et al. [2023c] studies the in-context
decision-making setting to predict action instead of learning a reward
correlation from a short horizon setting. In contrast we do not require a
greedy approximation of the optimal action, deal with short horizon setting
and changing action sets during training and testing, and predict the
estimated means of the actions instead of predicting the optimal action. A
survey of the in-context decision-making approaches can be found in Liu et al.
[2023a].

In the in-context decision-making setting, the learning model is first trained
on supervised input-output examples with the in-context data during training.
Then during test time, the model is asked to complete a new input (related to
the context provided) without any update to the model parameters [Xie et al.,
2021, Min et al., 2022]. Motivated by this, Lee et al. [2023] recently
proposed the Decision Pretrained Transformers (DPT) that exhibit the following
properties: (1) During supervised pretraining of DPT, predicting optimal
actions alone gives rise to near-optimal decision-making algorithms for
unforeseen task during test time. Note that DPT does not update model
parameters during test time and, therefore, conducts in-context learning on
the unforeseen task. (2) DPT improves over the in-context data used to
pretrain it by exploiting latent structure. However, DPT either requires the
optimal action during training or if it needs to approximate the optimal
action. For approximating the optimal action, it requires a large amount of
data from the underlying task.

At the same time, learning the underlying data pattern from a few examples
during training is becoming more relevant in many domains like chatbot
interaction [Madotto et al., 2021, Semnani et al., 2023], recommendation
systems, healthcare [Ge et al., 2022, Liu et al., 2023b], etc. This is
referred to as few-shot learning. However, most current RL decision-making
systems (including in-context learners like DPT) require an enormous amount of
data to learn a good policy.

The in-context learning framework is related to the meta-learning framework
[Bengio et al., 1990, Schaul and Schmidhuber, 2010]. Broadly, these techniques
aim to learn the underlying latent shared structure within the training
distribution of tasks, facilitating faster learning of novel tasks during test
time. In the context of decision-making and reinforcement learning (RL), there
exists a frequent choice regarding the specific ’structure’ to be learned, be
it the task dynamics [Fu et al., 2016, Nagabandi et al., 2018, Landolfi et
al., 2019], a task context identifier [Rakelly et al., 2019, Zintgraf et al.,
2019, Liu et al., 2021], or temporally extended skills and options [Perkins
and Precup, 1999, Gupta et al., 2018, Jiang et al., 2022].

However, as we noted in the Section 1, one can do a greedy approximation of
the optimal action from the historical data using a weak demonstrator and a
neural network policy [Finn et al., 2017, Rothfuss et al., 2018]. Moreover,
the in-context framework generally is more agnostic where it learns the policy
of the demonstrator [Duan et al., 2016, Wang et al., 2016, Mishra et al.,
2017]. Note that both DPT-greedy and PreDeToR are different than algorithmic
distillation [Laskin et al., 2022, Lu et al., 2023] as they do not distill an
existing RL algorithm. moreover, in contrast to DPT-greedy which is trained to
predict the optimal action, the PreDeToR is trained to predict the reward for
each of the actions. This enables the PreDeToR (similar to DPT-greedy) to show
to potentially emergent online and offline strategies at test time that
automatically align with the task structure, resembling posterior sampling.

As we discussed in the Section 1, in decision-making, RL, and imitation
learning the transformer models are trained using autoregressive action
prediction [Yang et al., 2023]. Similar methods have also been used in Large
language models [Vaswani et al., 2017, Roberts et al., 2019]. One of the more
notable examples is the Decision Transformers (abbreviated as DT) which
utilizes a transformer to autoregressively model sequences of actions from
offline experience data, conditioned on the achieved return [Chen et al.,
2021, Janner et al., 2021]. This approach has also been shown to be effective
for multi-task settings [Lee et al., 2022], and multi-task imitation learning
with transformers [Reed et al., 2022, Brohan et al., 2022, Shafiullah et al.,
2022]. However, the DT methods are not known to improve upon their in-context
data, which is the main thrust of this paper [Brandfonbrener et al., 2022,
Yang et al., 2022b].

Our work is also closely related to the offline RL setting. In offline RL, the
algorithms can formulate a policy from existing data sets of state, action,
reward, and next-state interactions. Recently, the idea of pessimism has also
been introduced in an offline setting to address the challenge of distribution
shift [Kumar et al., 2020, Yu et al., 2021, Liu et al., 2020, Ghasemipour et
al., 2022]. Another approach to solve this issue is policy regularization
[Fujimoto et al., 2019, Kumar et al., 2019, Wu et al., 2019, Siegel et al.,
2020, Liu et al., 2019], or reuse data for related task [Li et al., 2020,
Mitchell et al., 2021], or additional collection of data along with offline
data [Pong et al., 2022]. However, all of these approaches still have to take
into account the issue of distributional shifts. In contrast PreDeToR and DPT-
greedy leverages the decision transformers to avoid these issues. Both of
these methods can also be linked to posterior sampling. Such connections
between sequence modeling with transformers and posterior sampling have also
been made in Chen et al. [2021], Müller et al. [2021], Lee et al. [2023], Yang
et al. [2023].

###  A.2 Experimental Setting Information

Linear Bandit: We consider the setting when
$f(\mathbf{x},\bm{\theta}_{*})=\mathbf{x}^{\top}\bm{\theta}_{*}$. Here
$\mathbf{x}\in\mathbb{R}^{d}$ is the action feature and
$\bm{\theta}_{*}\in\mathbb{R}^{d}$ is the hidden parameter. For every
experiment, we first generate tasks from $\mathcal{T}_{\text{pre}}$. Then we
sample a fixed set of actions from
$\mathcal{N}\left(\mathbf{0},\mathbf{I}_{d}/d\right)$ in $\mathbb{R}^{d}$ and
this constitutes the features. Then for each task $m\in[M]$ we sample
$\bm{\theta}_{m,*}\sim\mathcal{N}\left(\mathbf{0},\mathbf{I}_{d}/d\right)$ to
produce the means
$\mu(m,a)=\left\langle\bm{\theta}_{m,*},\mathbf{x}(m,a)\right\rangle$ for
$a\in\mathcal{A}$ and $m\in[M]$. Finally, note that we do not shuffle the data
as the order matters. Also in this setting $\mathbf{x}(m,a)$ for each
$a\in\mathcal{A}$ is fixed for all tasks $m$.

Non-Linear Bandit: We now consider the setting when
$f(\mathbf{x},\bm{\theta}_{*})=1/(1+0.5\cdot\exp(2\cdot\exp(-\mathbf{x}^{\top}\bm{\theta}_{*})))$.
Again, here $\mathbf{x}\in\mathbb{R}^{d}$ is the action feature, and
$\bm{\theta}_{*}\in\mathbb{R}^{d}$ is the hidden parameter. Note that this is
different than the generalized linear bandit setting [Filippi et al., 2010, Li
et al., 2017]. Again for every experiment, we first generate tasks from
$\mathcal{T}_{\text{pre}}$. Then we sample a fixed set of actions from
$\mathcal{N}\left(\mathbf{0},\mathbf{I}_{d}/d\right)$ in $\mathbb{R}^{d}$ and
this constitutes the features. Then for each task $m\in[M]$ we sample
$\bm{\theta}_{m,*}\sim\mathcal{N}\left(\mathbf{0},\mathbf{I}_{d}/d\right)$ to
produce the means
$\mu(m,a)=1/(1+0.5\cdot\exp(2\cdot\exp(-\mathbf{x}(m,a)^{\top}\bm{\theta}_{m,*})))$
for $a\in\mathcal{A}$ and $m\in[M]$. Again note that in this setting
$\mathbf{x}(m,a)$ for each $a\in\mathcal{A}$ is fixed for all tasks $m$.

We use NVIDIA GeForce RTX 3090 GPU with 24GB RAM to load the GPT 2 Large
Language Model. This requires less than 2GB RAM without data, and with large
context may require as much as 20GB RAM.

###  A.3 Empirical Study: Bilinear Bandits

In this section, we discuss the performance of PreDeToR against the other
baselines in the bilinear setting. Again note that the number of tasks
$M_{\text{pre}}\gg A\geq n$. Through this experiment, we want to evaluate the
performance of PreDeToR to exploit the underlying latent structure and reward
correlation when the horizon is small, the number of tasks is large, and
understand its performance in the bilinear bandit setting [Jun et al., 2019,
Lu et al., 2021, Kang et al., 2022, Mukherjee et al., 2023]. Note that this
setting also goes beyond the linear feedback model [Abbasi-Yadkori et al.,
2011, Lattimore and Szepesvári, 2020] and is related to matrix bandits [Yang
and Wang, 2020].

Bilinear bandit setting: In the bilinear bandits the learner is provided with
two sets of action sets, $\mathcal{X}\subseteq\mathbb{R}^{d_{1}}$ and
$\mathcal{Z}\subseteq\mathbb{R}^{d_{2}}$ which are referred to as the left and
right action sets. At every round $t$ the learner chooses a pair of actions
$\mathbf{x}_{t}\in\mathcal{X}$ and $\mathbf{z}_{t}\in\mathcal{Z}$ and observes
a reward

$\displaystyle
r_{t}=\mathbf{x}_{t}^{\top}\mathbf{\Theta}_{*}\mathbf{z}_{t}+\eta_{t}$

where $\mathbf{\Theta}_{*}\in\mathbb{R}^{d_{1}\times d_{2}}$ is the unknown
hidden matrix which is also low-rank. The $\eta_{t}$ is a $\sigma^{2}$ sub-
Gaussian noise. In the multi-task bilinear bandit setting we now have a set of
$M$ tasks where the reward for the $m$-th task at round $t$ is given by

$\displaystyle
r_{m,t}=\mathbf{x}_{m,t}^{\top}\mathbf{\Theta}_{m,*}\mathbf{z}_{m,t}+\eta_{m,t}.$

Here $\mathbf{\Theta}_{m,*}\in\mathbb{R}^{d_{1}\times d_{2}}$ is the unknown
hidden matrix for each task $m$, which is also low-rank. The $\eta_{m,t}$ is a
$\sigma^{2}$ sub-Gaussian noise. Let $\kappa$ be the rank of each of these
matrices $\mathbf{\Theta}_{m,*}$.

A special case is the rank $1$ structure where
$\mathbf{\Theta}_{m,*}=\bm{\theta}_{m,*}\bm{\theta}_{m,*}^{\top}$ where
$\mathbf{\Theta}_{m,*}\in\mathbb{R}^{d\times d}$ and
$\bm{\theta}_{m,*}\in\mathbb{R}^{d}$ for each task $m$. Let the left and right
action sets be also same such that
$\mathbf{x}_{m,t}\in\mathcal{X}\subseteq\mathbb{R}^{d}$. Observe then that the
reward for the $m$-th task at round $t$ is given by

$\displaystyle
r_{m,t}=\mathbf{x}_{m,t}^{\top}\mathbf{\Theta}_{m,*}\mathbf{x}_{m,t}+\eta_{m,t}=(\mathbf{x}_{m,t}^{\top}\bm{\theta}_{m,*})^{2}+\eta_{m,t}.$

This special case is studied in Chaudhuri et al. [2017].

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, and Thomp. However, we
now implement the LowOFUL [Jun et al., 2019] which is optimal in the bilinear
bandit setting.

LowOFUL: The LowOFUL algorithm first estimates the unknown parameter
$\mathbf{\Theta}_{m,*}$ for each task $m$ using E-optimal design [Pukelsheim,
2006, Fedorov, 2013, Jun et al., 2019] for $n_{1}$ rounds. Let
$\widehat{\mathbf{\Theta}}_{m,n_{1}}$ be the estimate of
$\mathbf{\Theta}_{m,*}$ at the end of $n_{1}$ rounds. Let the SVD of
$\widehat{\mathbf{\Theta}}_{m,n_{1}}$ be given by
$\operatorname{SVD}(\widehat{\mathbf{\Theta}}_{m,n_{1}})=\widehat{\mathbf{U}}_{m,n_{1}}\widehat{\mathbf{S}}_{m,n_{1}}\widehat{\mathbf{V}}_{m,n_{1}}$.
Then LowOFUL rotates the actions as follows:

$\displaystyle\mathcal{X}^{\prime}_{m}=\left\\{\left[\widehat{\mathbf{U}}_{m,n_{1}}\widehat{\mathbf{U}}^{\perp}_{m,n_{1}}\right]^{\top}\mathbf{x}_{m}:\mathbf{x}_{m}\in\mathcal{X}\right\\}\textbf{
and
}\mathcal{Z}^{\prime}=\left\\{\left[\widehat{\mathbf{V}}_{m,n_{1}}\widehat{\mathbf{V}}^{\perp}_{m,n_{1}}\right]^{\top}\mathbf{z}_{m}:\mathbf{z}_{m}\in\mathcal{Z}\right\\}.$

Then defines a vectorized action set for each task $m$ so that the last
$\left(d_{1}-\kappa\right)\cdot\left(d_{2}-\kappa\right)$ components are from
the complementary subspaces:

$\displaystyle\widetilde{\mathcal{A}}_{m}$
$\displaystyle=\left\\{\left[\operatorname{vec}\left(\mathbf{x}_{m,1:\kappa}\mathbf{z}_{m,1:\kappa}^{\top}\right);\operatorname{vec}\left(\mathbf{x}_{m,\kappa+1:d_{1}}\mathbf{z}_{m,1:\kappa}^{\top}\right);\operatorname{vec}\left(\mathbf{x}_{m,1:\kappa}\mathbf{z}_{m,\kappa+1:d_{2}}^{\top}\right);\right.\right.$
$\displaystyle\quad\quad\left.\left.\operatorname{vec}\left(\mathbf{x}_{m,\kappa+1:d_{1}}\mathbf{z}_{m,\kappa+1:d_{2}}^{\top}\right)\right]\in\mathbb{R}^{d_{1}d_{2}}:\mathbf{x}_{m}\in\mathcal{X}^{\prime}_{m},\mathbf{z}_{m}\in\mathcal{Z}^{\prime}_{m}\right\\}.$

Finally for $n_{2}=n-n_{1}$ rounds, LowOFUL invokes the specialized OFUL
algorithm [Abbasi-Yadkori et al., 2011] for the rotated action set
$\widetilde{\mathcal{A}}_{m}$ with the low dimension
$k=\left(d_{1}+d_{2}\right)\kappa-\kappa^{2}$. Note that the LowOFUL runs the
per-task low dimensional OFUL algorithm rather than learning the underlying
structure across the tasks [Mukherjee et al., 2023].

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent bilinear
structure and reward correlation between the different tasks even when $A\geq
n$ but $M_{\text{pre}}\gg A$.

(2) The PreDeToR (-$\tau$) has lower regret than its weak demonstrator Thomp
which does not know the underlying bilinear structure.

(3) The PreDeToR (-$\tau$) has lower regret compared to DPT-greedy and AD.
Also, observe that PreDeToR matches the performance of LowOFUL.

|
---|---

(a) Rank $1$ $\mathbf{\Theta}_{m,*}$

(b) Rank $2$ $\mathbf{\Theta}_{m,*}$

Figure 3: Experiment with bilinear bandits. The y-axis shows the cumulative
regret.

Experimental Result: We observe these outcomes in Figure 3. In Figure 3(a) we
experiment with rank $1$ hidden parameter $\mathbf{\Theta}_{m,*}$ and set
horizon $n=20$, $M_{\text{pre}}=200000$, $M_{\text{test}}=200$, $A=30$, and
$d=5$. In Figure 3(b) we experiment with rank $2$ hidden parameter
$\mathbf{\Theta}_{m,*}$ and set horizon $n=20$, $M_{\text{pre}}=250000$,
$M_{\text{test}}=200$, $A=25$, and $d=5$. Again, the demonstrator $\pi^{w}$ is
the Thomp algorithm. We observe that PreDeToR has lower cumulative regret than
DPT-greedy, AD and Thomp. Note that for any task $m$ for the horizon $20$ the
Thomp will be able to sample all the actions at most once. Note that for this
small horizon setting the DPT-greedy does not have a good estimation of
$\widehat{a}_{m,*}$ which results in a poor prediction of optimal action
$\widehat{a}_{m,t,*}$. In contrast PreDeToR learns the correlation of rewards
across tasks and can perform well. Observe from Figure 3(a), and 3(b) that
PreDeToR has lower regret than Thomp and matches LowOFUL. Also, in this low-
data regime it is not enough for LowOFUL to learn the underlying
$\mathbf{\Theta}_{m,*}$ with high precision. Hence, PreDeToR also has slightly
lower regret than LowOFUL. The training of AD ensures that it does not
outperform the demonstrator Thomp. Most importantly it shows that PreDeToR can
exploit the underlying latent structure and reward correlation better than
DPT-greedy, and AD.

###  A.4 Empirical Study: Latent Bandits

In this section, we discuss the performance of PreDeToR (-$\tau$) against the
other baselines in the latent bandit setting and create a generalized bilinear
bandit setting. Note that the number of tasks $M_{\text{pre}}\gg A\geq n$.
Using this experiment, we want to evaluate the ability of PreDeToR (-$\tau$)
to exploit the underlying reward correlation when the horizon is small, the
number of tasks is large, and understand its performance in the latent bandit
setting [Hong et al., 2020, Maillard and Mannor, 2014, Pal et al., 2023,
Kveton et al., 2017]. We create a latent bandit setting which generalizes the
bilinear bandit setting [Jun et al., 2019, Lu et al., 2021, Kang et al., 2022,
Mukherjee et al., 2023]. Again note that this setting also goes beyond the
linear feedback model [Abbasi-Yadkori et al., 2011, Lattimore and Szepesvári,
2020] and is related to matrix bandits [Yang and Wang, 2020].

Latent bandit setting: In this special multi-task latent bandits the learner
is again provided with two sets of action sets,
$\mathcal{X}\subseteq\mathbb{R}^{d_{1}}$ and
$\mathcal{Z}\subseteq\mathbb{R}^{d_{2}}$ which are referred to as the left and
right action sets. The reward for the $m$-th task at round $t$ is given by

$\displaystyle
r_{m,t}=\mathbf{x}_{m,t}^{\top}\underbrace{(\mathbf{\Theta}_{m,*}+\mathbf{U}\mathbf{V}^{\top})}_{\mathbf{Z}_{m_{*}}}\mathbf{z}_{m,t}+\eta_{m,t}.$

Here $\mathbf{\Theta}_{m,*}\in\mathbb{R}^{d_{1}\times d_{2}}$ is the unknown
hidden matrix for each task $m$, which is also low-rank. Additionally, all the
tasks share a _common latent parameter matrix_
$\mathbf{U}\mathbf{V}^{\top}\in\mathbb{R}^{d_{1}\times d_{2}}$ which is also
low rank. Hence the learner needs to learn the latent parameter across the
tasks hence the name latent bandits. Finally, the $\eta_{m,t}$ is a
$\sigma^{2}$ sub-Gaussian noise. Let $\kappa$ be the rank of each of these
matrices $\mathbf{\Theta}_{m,*}$ and $\mathbf{U}\mathbf{V}^{\top}$. Again
special case is the rank $1$ structure where the reward for the $m$-th task at
round $t$ is given by

$\displaystyle
r_{m,t}=\mathbf{x}_{m,t}^{\top}\underbrace{(\bm{\theta}_{m,*}\bm{\theta}_{m,*}^{\top}+\mathbf{u}\mathbf{v}^{\top})}_{\mathbf{Z}_{m,*}}\mathbf{x}_{m,t}+\eta_{m,t}.$

where $\bm{\theta}_{m,*}\in\mathbb{R}^{d}$ for each task $m$ and
$\mathbf{u},\mathbf{v}\in\mathbb{R}^{d}$. Note that the left and right action
sets are the same such that
$\mathbf{x}_{m,t}\in\mathcal{X}\subseteq\mathbb{R}^{d}$.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, AD, Thomp, and LowOFUL.
However, we now implement a special LowOFUL (stated in Section A.3) which has
knowledge of the shared latent parameters $\mathbf{U}$, and $\mathbf{V}$. We
call this the LowOFUL (oracle) algorithm. Therefore LowOFUL (oracle) has
knowledge of the problem parameters in the latent bandit setting and hence the
name.

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent bilinear
structure and reward correlation between the different tasks even when $A\geq
n$ but $M_{\text{pre}}\gg A$.

(2) The PreDeToR (-$\tau$) has lower regret than its weak demonstrator Thomp
which does not know the underlying latent structure. The PreDeToR (-$\tau$)
has lower regret compared to DPT-greedy.

(3) The PreDeToR (-$\tau$) has regret that is closer to LowOFUL (oracle).

|
---|---

(a) Rank $1$ $\mathbf{Z}_{m,*}$

(b) Rank $2$ $\mathbf{Z}_{m,*}$

(c) Rank $3$ $\mathbf{Z}_{m,*}$

Figure 4: Experiment with latent bandits. The y-axis shows the cumulative
regret.

Experimental Result: We observe these outcomes in Figure 4. In Figure 4(a) we
experiment with rank $1$ hidden parameter
$\bm{\theta}_{m,*}\bm{\theta}_{m,*}^{\top}$ and latent parameters
$\mathbf{u}\mathbf{v}^{\top}$ shared across the tasks and set horizon $n=20$,
$M_{\text{pre}}=200000$, $M_{\text{test}}=200$, $A=30$, and $d=5$. In Figure
4(b) we experiment with rank $2$ hidden parameter $\mathbf{\Theta}_{m,*}$, and
latent parameters $\mathbf{U}\mathbf{V}^{\top}$ and set horizon $n=20$,
$M_{\text{pre}}=250000$, $M_{\text{test}}=200$, $A=25$, and $d=5$. In Figure
4(c) we experiment with rank $3$ hidden parameter $\mathbf{\Theta}_{m,*}$, and
latent parameters $\mathbf{U}\mathbf{V}^{\top}$ and set horizon $n=20$,
$M_{\text{pre}}=300000$, $M_{\text{test}}=200$, $A=25$, and $d=5$. Again, the
demonstrator $\pi^{w}$ is the Thomp algorithm. We observe that PreDeToR
(-$\tau$) has lower cumulative regret than DPT-greedy, AD and Thomp. Note that
for any task $m$ for the horizon $20$ the Thomp will be able to sample all the
actions at most once. Note that for this small horizon setting the DPT-greedy
does not have a good estimation of $\widehat{a}_{m,*}$ which results in a poor
prediction of optimal action $\widehat{a}_{m,t,*}$. In contrast PreDeToR
(-$\tau$) learns the correlation of rewards across tasks and is able to
perform well. Observe from Figure 4(a), 4(b), and 4(c) that PreDeToR has lower
regret than Thomp and has regret closer to LowOFUL (oracle)which has access to
the problem-dependent parameters. Hence. LowOFUL (oracle) outperforms PreDeToR
(-$\tau$) in this setting. The training of AD ensures that it does not
outperform the demonstrator Thomp. This shows that PreDeToR is able to exploit
the underlying latent structure and reward correlation better than DPT-greedy,
and AD.

###  A.5 Understanding PreDeToR

In this section, we try to understand the behavior of PreDeToR and its ability
to exploit the reward correlation across tasks under a _linear multivariate
Gaussian model_. In this model, the hidden task parameter, $\bm{\theta}_{*}$,
is a random variable drawn from a multi-variate Gaussian distribution [Bishop,
2006] and the feedback follows a linear model. We study this setting since we
can estimate the Linear Minimum Mean Square Estimator (LMMSE) in this setting
[Carlin and Louis, 2008, Box and Tiao, 2011]. This yields a posterior
prediction for the mean of each action over all tasks on average, by
leveraging the linear structure when $\bm{\theta}_{*}$ is drawn from a multi-
variate Gaussian distribution. So we can compare the performance of PreDeToR
against such an LMMSE and evaluate whether it is exploiting the underlying
linear structure and the reward correlation across tasks.

Figure 5: BayesPred Performance

Consider the linear feedback setting consisting of $A$ actions and the hidden
task parameter
$\bm{\theta}_{*}\sim\mathcal{N}(0,\sigma^{2}_{\bm{\theta}}\mathbf{I}_{d})$.
The reward of the action $\mathbf{x}_{t}$ at round $t$ is given by
$r_{t}=\mathbf{x}_{t}^{\top}\bm{\theta}_{*}+\eta_{t}$, where $\eta_{t}$ is
$\sigma^{2}$ sub-Gaussian. Let $\pi^{w}$ collect $n$ rounds of pretraining in-
context data and observe $\\{I_{t},r_{t}\\}_{t=1}^{n}$. Let $N_{n}(a)$ denote
the total number of times the action $a$ is sampled for $n$ rounds. Note that
we drop the task index $m$ in these notations as the random variable
$\bm{\theta}_{*}$ corresponds to the task. Define the matrix
$\mathbf{H}_{n}\in\mathbb{R}^{n\times A}$ where the $t$-th row represents the
action $I_{t}$ for $t\in[n]$. The $t$-th row of $\mathbf{H}_{n}$ is a one-hot
vector with the $I_{t}$-th component being 1. We represent each action by one
hot vector because we assume that this LMMSE does not have access to the
feature vectors of the actions similar to the PreDeToR for fair comparison.
Then define the reward vector $\mathbf{Y}_{n}\in\mathbb{R}^{n}$ where the
$t$-th component is the reward $r_{t}$ observed for the action $I_{t}$ for
$t\in[n]$ in the pretraining data. Define the diagonal matrix
$\mathbf{D}_{A}\in\mathbb{R}^{A\times A}$ estimated from pretraining data as
follows

$\displaystyle\mathbf{D}_{A}(i,i)$
$\displaystyle=\begin{cases}\frac{\sigma^{2}}{N_{n}(a)},\text{ if
}N_{n}(a)>0\\\ =0,\text{ if }N_{n}(a)=0\end{cases}\vspace*{-1em}$ (6)

where the reward noise being $\sigma^{2}$ sub-Gaussian is known. Finally
define the estimated reward covariance matrix
$\mathbf{S}_{A}\in\mathbb{R}^{A\times A}$ as
$\mathbf{S}_{A}(a,a^{\prime})=\widehat{\mu}_{n}(a)\widehat{\mu}_{n}(a^{\prime})$,
where $\widehat{\mu}_{n}(a)$ is the empirical mean of action $a$ estimated
from the pretraining data. This matrix captures the reward correlation between
the pairs of actions $a,a^{\prime}\in[A]$. Then the posterior average mean
estimator $\widehat{\mu}\in\mathbb{R}^{A}$ over all tasks is given by the
following lemma. The proof is given in Section B.1.

######  Lemma 1.

Let $\mathbf{H}_{n}$ be the action matrix, $\mathbf{Y}_{n}$ be the reward
vector and $\mathbf{S}_{A}$ be the estimated reward covariance matrix. Then
the posterior prediction of the average mean reward vector $\widehat{\mu}$
over all tasks is given by

$\displaystyle\widehat{\mu}=\sigma^{2}_{\bm{\theta}}\mathbf{S}_{A}\mathbf{H}_{n}^{\top}\left(\sigma^{2}_{\bm{\theta}}\mathbf{H}_{n}(\mathbf{S}_{A}+\mathbf{D}_{A})\mathbf{H}_{n}^{\top}\right)^{-1}\mathbf{Y}_{n}.$
(7)

The $\widehat{\mu}$ in (7) represents the posterior mean vector averaged on
all tasks. So if some action $a\in[A]$ consistently yields high rewards in the
pretraining data then $\widehat{\mu}(a)$ has high value. Since the test
distribution is the same as pretraining, this action on average will yield a
high reward during test time.

We hypothesize that the PreDeToR is learning the reward correlation covariance
matrix from the training data $\mathcal{H}_{\text{train}}$ and acting greedily
on it. To test this hypothesis, we consider the greedy BayesPred algorithm
that first estimates $\mathbf{S}_{A}$ from the pretraining data. It then uses
the LMMSE estimator in 1 to calculate the posterior mean vector
$\widehat{\mu}$, and then selects $I_{t}=\mathop{\rm
arg\,max}_{a}\widehat{\mu}(a)$ at each round $t$. Note that BayesPred is a
greedy algorithm that always selects the most rewarding action (exploitation)
without any exploration of sub-optimal actions. Also the BayesPred is an LMMSE
estimator that leverages the linear reward structure and estimates the reward
covariance matrix, and therefore can be interpreted as a lower bound to the
regret of PreDeToR. The hypothesis that BayesPred is a lower bound to PreDeToR
is supported by Figure 5. In Figure 5 the reward covariance matrix for
BayesPred is estimated from the $\mathcal{H}_{\text{train}}$ by first running
the Thomp ($\pi^{w}$). Observe that the BayesPred has a lower cumulative
regret than PreDeToR and almost matches the regret of PreDeToR towards the end
of the horizon. Also note that LinUCB has lower cumulative regret towards the
end of horizon as it leverages the linear structure and the feature of the
actions in selecting the next action.

###  A.6 Empirical Study: Increasing number of Actions

In this section, we discuss the performance of PreDeToR when the number of
actions is very high so that the weak demonstrator $\pi^{w}$ does not have
sufficient samples for each action. However, the number of tasks
$M_{\text{pre}}\gg A>n$.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, AD, Thomp, and LinUCB.

Outcomes: We first discuss the main outcomes from our experimental results of
introducing more actions than the horizon (or more dimensions than actions)
during data collection and evaluation:

(1) The PreDeToR (-$\tau$) exploits the underlying latent structure and reward
correlation between the different tasks even when $A>n$ but $M_{\text{pre}}\gg
A$.

(2) The PreDeToR (-$\tau$) has lower regret than its weak demonstrator even
when it is trained on
$\mathcal{H}_{\text{train}}\subseteq\mathcal{D}_{\text{pre}}$.

(3) The PreDeToR (-$\tau$) performance drops for complex non-linear bandit
setting with increasing number of dimensions.

|
---|---

(a) Linear Bandit

(b) Non-linear Bandit

Figure 6: Testing the limit experiments. The horizontal axis is the number of
rounds. Confidence bars show one standard error.

Experimental Result: We observe these outcomes in Figure 6. In Figure 6(a) we
show the linear bandit setting for $M_{\text{pre}}=250000$,
$M_{\text{test}}=200$, $A=100$, $n=50$ and $d=5$. Again, the demonstrator
$\pi^{w}$ is the Thomp algorithm. We observe that PreDeToR (-$\tau$) has lower
cumulative regret than DPT-greedy and AD. Note that for any task $m$ the Thomp
will not be able to sample all the actions even once. The weak performance of
DPT-greedy can be attributed to both short horizons and the inability to
estimate the optimal action for such a short horizon $n<A$. The AD performs
similar to the demonstrator Thomp because of its training. Observe that
PreDeToR (-$\tau$) has similar regret to LinUCB and lower regret than Thomp
which also shows that PreDeToR is exploiting the latent linear structure of
the underlying tasks. In Figure 6(b) we show the non-linear bandit setting for
horizon $n=40$, $M_{\text{pre}}=200000$, $A=60$, $d=2$, and
$|\mathcal{A}^{\text{inv}}|=5$. The demonstrator $\pi^{w}$ is the Thomp
algorithm. The training of AD ensures that it does not outperform the
demonstrator Thomp. Again we observe that PreDeToR (-$\tau$) has lower
cumulative regret than DPT-greedy, AD and LinUCB which fails to perform well
in this non-linear setting due to its algorithmic design.

###  A.7 Empirical Study: Increasing Horizon

In this section, we discuss the performance of PreDeToR with respect to an
increasing horizon for each task $m\in[M]$. However, note that the number of
tasks $M_{\text{pre}}\geq n$. Note that Lee et al. [2023] studied linear
bandit setting for $n=200$. We study the setting up to a similar horizon
scale.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, AD, Thomp, and LinUCB.

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent structure
and reward correlation between the different tasks even when $A\geq n$,
$M_{\text{pre}}\gg A$ with increasing horizon.

(2) The PreDeToR (-$\tau$) outperforms its weak demonstrator even when trained
on $\mathcal{H}_{\text{train}}\subseteq\mathcal{D}_{\text{pre}}$.

|  |
---|---|---
|  |

(a) Horizon $20$

(b) Horizon $40$

(c) Horizon $60$

(d) Horizon $100$

(e) Horizon $120$

(f) Horizon $140$

|
---|---

(g) Horizon $200$

(h) Increasing Horizon

Figure 7: Experiment with increasing horizon. The y-axis shows the cumulative
regret.

Experimental Result: We observe these outcomes in Figure 7. In Figure 7 we
show the linear bandit setting for $M_{\text{pre}}=150000$,
$M_{\text{test}}=200$, $A=20$, $n=\\{20,40,60,100,120,140,200\\}$ and $d=5$.
Again, the demonstrator $\pi^{w}$ is the Thomp algorithm. We observe that
PreDeToR (-$\tau$) has lower cumulative regret than DPT-greedy, and AD. Note
that for any task $m$ for the horizon $20$ the Thomp will be able to sample
all the actions at most once. Observe from Figure 7(a), 7(b), 7(c), Figure
7(d), 7(e), 7(f) and 7(g) that PreDeToR (-$\tau$) is closer to LinUCB and
outperforms Thomp which also shows that PreDeToR (-$\tau$) is learning the
latent linear structure of the underlying tasks. In Figure 7(h) we plot the
regret of all the baselines with respect to the increasing horizon. Again we
see that PreDeToR (-$\tau$) is closer to LinUCB and outperforms DPT-greedy, AD
and Thomp. The training of AD ensures that it does not outperform the
demonstrator Thomp. This shows that PreDeToR (-$\tau$) is able to exploit the
latent structure and reward correlation across the tasks for varying horizon
length.

###  A.8 Empirical Study: Increasing Dimension

In this section, we discuss the performance of PreDeToR with respect to an
increasing dimension for each task $m\in[M]$. Again note that the number of
tasks $M_{\text{pre}}\gg A\geq n$. Through this experiment, we want to
evaluate the performance of PreDeToR and see how it exploits the underlying
reward correlation when the horizon is small as well as for increasing
dimensions.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$ DPT-greedy, AD, Thomp, and LinUCB.

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent structure
and reward correlation between the different tasks even when $A\geq n$ but
$M_{\text{pre}}\gg A$.

(2) The PreDeToR (-$\tau$) has lower regret than its weak demonstrator even
when trained on $\mathcal{D}_{\text{train}}\subseteq\mathcal{D}_{\text{pre}}$.

(3) The PreDeToR (-$\tau$) has lower regret than LinUCB for larger dimension.

|  |
---|---|---

(a) Dimension $10$

(b) Dimension $20$

(c) Dimension $30$

|
---|---

(d) Dimension $40$

(e) Increasing Dimension

Figure 8: Experiment with increasing dimension. The y-axis shows the
cumulative regret.

Experimental Result: We observe these outcomes in Figure 7. In Figure 7 we
show the linear bandit setting for horizon $n=20$, $M_{\text{pre}}=160000$,
$M_{\text{test}}=200$, $A=20$, and $d=\\{10,20,30,40\\}$. Again, the
demonstrator $\pi^{w}$ is the Thomp algorithm. We observe that PreDeToR
(-$\tau$) has lower cumulative regret than DPT-greedy, AD. Note that for any
task $m$ for the horizon $20$ the Thomp will be able to sample all the actions
at most once. Observe from Figure 8(a), 8(b), 8(c), and 8(d) that PreDeToR
(-$\tau$) is closer to LinUCB and has lower regret than Thomp which also shows
that PreDeToR (-$\tau$) is exploiting the latent linear structure of the
underlying tasks. In Figure 8(e) we plot the regret of all the baselines with
respect to the increasing dimension. Again we see that PreDeToR (-$\tau$) has
lower regret than DPT-greedy, AD and Thomp. Observe that with increasing
dimension PreDeToR is able to outperform LinUCB. This shows that the PreDeToR
(-$\tau$) is able to exploit reward correlation across tasks for varying
dimensions.

###  A.9 Empirical Study: Increasing Attention Heads

In this section, we discuss the performance of PreDeToR with respect to an
increasing attention heads for the transformer model for the non-linear
feedback model. Again note that the number of tasks $M_{\text{pre}}\gg A\geq
n$. Through this experiment, we want to evaluate the performance of PreDeToR
to exploit the underlying reward correlation when the horizon is small and
understand the representative power of the transformer by increasing the
attention heads. Note that we choose the non-linear feedback model and low
data regime to leverage the representative power of the transformer.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, AD, Thomp, and LinUCB.

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) learns to exploit the underlying latent structure
and reward correlation between the different tasks even when $A\geq n$ but
$M_{\text{pre}}\gg A$.

(2) The PreDeToR (-$\tau$) has lower regret than its weak demonstrator even
when trained on $\mathcal{D}_{\text{train}}\subseteq\mathcal{D}_{\text{pre}}$.

(3) The PreDeToR (-$\tau$) has lower regret compared to DPT-greedy, AD with
increasing attention heads.

|  |
---|---|---
|  |

(a) Attention Heads $2$

(b) Attention Heads $4$

(c) Attention Heads $6$

(d) Attention Heads $8$

(e) Attention Heads $12$

(f) Increasing Attention Heads

Figure 9: Experiment with increasing attention heads. The y-axis shows the
cumulative regret.

Experimental Result: We observe these outcomes in Figure 9. In Figure 9 we
show the non-linear bandit setting for horizon $n=20$,
$M_{\text{pre}}=160000$, $M_{\text{test}}=200$, $A=20$,
$\mathrm{heads}=\\{2,4,6,8\\}$ and $d=5$. Again, the demonstrator $\pi^{w}$ is
the Thomp algorithm. We observe that PreDeToR (-$\tau$) has lower cumulative
regret than DPT-greedy, AD. Note that for any task $m$ for the horizon $20$
the Thomp will be able to sample all the actions atmost once. Observe from
Figure 9(a), 9(b), 9(c), and 9(d) that PreDeToR (-$\tau$) has lower regret
than AD, Thomp and LinUCB which also shows that PreDeToR (-$\tau$) is
exploiting the latent linear structure of the underlying tasks for the non-
linear setting. The training of AD ensures that it does not outperform the
demonstrator Thomp. In Figure 9(f) we plot the regret of all the baselines
with respect to the increasing attention heads. Again we see that PreDeToR
(-$\tau$) regret decreases as we increase the attention heads.

###  A.10 Empirical Study: Increasing Number of Tasks

In this section, we discuss the performance of PreDeToR with respect to the
increasing number of tasks for the linear bandit setting. Again note that the
number of tasks $M_{\text{pre}}\gg A\geq n$. Through this experiment, we want
to evaluate the performance of PreDeToR to exploit the underlying reward
correlation when the horizon is small and the number of tasks is changing.
Finally, recall that when the horizon is small the weak demonstrator $\pi^{w}$
does not have sufficient samples for each action. This leads to a poor
approximation of the greedy action.

Baselines: We again implement the same baselines discussed in Section 4. The
baselines are PreDeToR, PreDeToR-$\tau$, DPT-greedy, AD, Thomp, and LinUCB.

Outcomes: We first discuss the main outcomes of our experimental results for
increasing the horizon:

(1) The PreDeToR (-$\tau$) fails to exploit the underlying latent structure
and reward correlation when the number of tasks is small.

(2) As the number of tasks increases the PreDeToR (-$\tau$) learns the latent
structure and the underlying reward correlation.

(3) The PreDeToR (-$\tau$) has lower regret compared to DPT-greedy, AD with
increasing number of tasks.

|  |
---|---|---
|  |

(a) Tasks $M_{\text{train}}=5000$

(b) Tasks $M_{\text{train}}=10000$

(c) Tasks $M_{\text{train}}=50000$

(d) Tasks $M_{\text{train}}=100000$

(e) Tasks $M_{\text{train}}=150000$

(f) Increasing tasks

Figure 10: Experiment with an increasing number of tasks. The y-axis shows the
cumulative regret.

Experimental Result: We observe these outcomes in Figure 10. In Figure 10 we
show the linear bandit setting for horizon $n=20$,
$M_{\text{pre}}\in\\{5000,10000,50000,100000,150000\\}$,
$M_{\text{test}}=200$, $A=20$, and $d=40$. Again, the demonstrator $\pi^{w}$
is the Thomp algorithm. We observe that PreDeToR (-$\tau$), AD and DPT-greedy
suffer more regret than the LinUCB when the number of tasks is small
($M_{\text{train}}\in\\{5000,10000\\}$ in Figure 10(a), and 10(b). However in
Figure 10(c), 10(d), 10(e), and 10(f) we show that PreDeToR has lower regret
than Thomp and matches LinUCB. This shows that PreDeToR (-$\tau$) is
exploiting the latent linear structure of the underlying tasks for the non-
linear setting. Moreover, observe that as $M_{\text{train}}$ increases the
PreDeToR has lower cumulative regret than DPT-greedy, AD. Note that for any
task $m$ for the horizon $20$ the Thomp will be able to sample all the actions
at most once. Therefore DPT-greedy does not perform as well as PreDeToR. The
training of AD ensures that it does not outperform the demonstrator Thomp.
Finally, note that the result shows that PreDeToR (-$\tau$) is able to exploit
the reward correlation across the tasks better as the number of tasks
increases.

###  A.11 Exploration of PreDeToR(-$\tau$)

In this section, we discuss the exploration of PreDeToR in the linear bandit
setting discussed in Section 4. Recall that the linear bandit setting consist
of horizon $n=25$, $M_{\text{pre}}=200000$, $M_{\text{test}}=200$, $A=10$, and
$d=2$. The demonstrator $\pi^{w}$ is the Thomp algorithm and we observe that
PreDeToR (-$\tau$) has lower cumulative regret than DPT-greedy, AD and matches
the performance of LinUCB. Therefore PreDeToR (-$\tau$) behaves almost
optimally in this setting and so we analyze how PreDeToR conducts exploration
for this setting.

Outcomes: We first discuss the main outcomes of our analysis of exploration in
the low-data regime:

(1) The PreDeToR (-$\tau$) has a two phase exploration. In the first phase it
explores with a strong prior over the training data. In the second phase, once
the task data has been observed for a few rounds it switches to task-based
exploration.

(2) PreDeToR (-$\tau$) samples the optimal action in each task more than DPT-
greedy. PreDeToR(-$\tau$) considers a larger set of possible optimal action
candidate set than DPT-greedy.

We first show in Figure 11(a) the training distribution of the optimal
actions. For each bar, the frequency indicates the number of tasks where the
action (shown in the x-axis) is the optimal action.

Then in Figure 11(b) we show how the sampling distribution of DPT-greedy,
PreDeToR and PreDeToR-$\tau$ change in the first $10$ and last $10$ rounds for
all the tasks where action $5$ is optimal. To plot this graph we first sum
over the individual pulls of the action taken by each algorithm over the first
$10$ and last $10$ rounds. Then we average these counts over all test tasks
where action $5$ is optimal. From the figure Figure 11(b) we see that
PreDeToR(-$\tau$) consistently pulls the action $5$ more than DPT-greedy. It
also explores other optimal actions like $\\{2,3,6,7,10\\}$ but discards them
quickly in favor of the optimal action $5$ in these tasks. This shows that
PreDeToR (-$\tau$) only considers the optimal actions seen from the training
data. Once sufficient observation have been observed for the task it switches
to task-based exploration and samples the optimal action more than DPT-greedy.

Finally, we plot the feasible action set considered by DPT-greedy, PreDeToR,
and PreDeToR-$\tau$ in Figure 11(c). To plot this graph again we consider the
test tasks where the optimal action is $5$. Then we count the number of
distinct actions that are taken from round $t$ up until horizon $n$. Finally
we average this over all the considered tasks where the optimal action is $5$.
We call this the candidate action set considered by the algorithm. From the
Figure 11(c) we see that DPT-greedy explores the least and gets stuck with few
actions quickly (by round $10$). Note that the actions DPT-greedy samples are
sub-optimal and so it suffers a high cumulative regret (see Figure 1(a)).
PreDeToR explore slightly more than DPT-greedy, but PreDeToR-$\tau$ explores
the most.

|  |
---|---|---

(a) Train Optimal Action Distribution

(b) Distribution of action sampling in all tasks where action $5$ is optimal

(c) Candidate Action Set in Time averaged over all tasks where action $5$ is
optimal

Figure 11: Exploration Analysis of PreDeToR(-$\tau$)

###  A.12 Exploration of PreDeToR(-$\tau$) in New Arms Setting

In this section, we discuss the exploration of PreDeToR (-$\tau$) in the
linear and non-linear new arms bandit setting discussed in Section 5. Recall
that we consider the linear bandit setting of horizon $n=50$,
$M_{\text{pre}}=200000$, $M_{\text{test}}=200$, $A=20$, and $d=5$. Here during
data collection and during collecting the test data, we randomly select one
new action from $\mathbb{R}^{d}$ for each task $m$. So the number of invariant
actions is $|\mathcal{A}^{\text{inv}}|=19$.

Outcomes: We first discuss the main outcomes of our analysis of exploration in
the low-data regime:

(1) The PreDeToR (-$\tau$) is robust to changes when the number of changing
arms is small. It samples the optimal action in each task more than DPT-
greedy.

(2) When the number of in-variant actions is small then all algorithms
PreDeToR, PreDeToR-$\tau$, and AD sample the optimal action very less. This
shows that the latent shared structure across the tasks is very important.

We first show in Figure 12(a) the training distribution of the optimal
actions. For each bar, the frequency indicates the number of tasks where the
action (shown in the x-axis) is the optimal action.

Then in Figure 12(b) we show how the sampling distribution of DPT-greedy,
PreDeToR and PreDeToR-$\tau$ change in the first $10$ and last $10$ rounds for
all the tasks where action $17$ is optimal. We plot this graph the same way as
discussed in Section A.11. From the figure Figure 12(b) we see that
PreDeToR(-$\tau$) consistently pulls the action $17$ more than DPT-greedy. It
also explores other optimal actions like $\\{1,2,3,8,9,15\\}$ but discards
them quickly in favor of the optimal action $17$ in these tasks.

Finally, we plot the feasible action set considered by DPT-greedy, PreDeToR,
and PreDeToR-$\tau$ in Figure 12(c). To plot this graph again we consider the
test tasks where the optimal action is $17$. Then we count the number of
distinct actions that are taken from round $t$ up until horizon $n$. Finally
we average this over all the considered tasks where the optimal action is
$17$. We call this the candidate action set considered by the algorithm. From
the Figure 12(c) we see that PreDeToR-$\tau$ explores more than PreDeToR in
this setting.

We also show how the prediction error of the optimal action by PreDeToR
compared to LinUCB in this $1$ new arm linear bandit setting. In Figure 13(a)
we first show how the $20$ actions are distributed in the
$M_{\text{test}}=200$ test tasks. In Figure 13(a) for each bar, the frequency
indicates the number of tasks where the action (shown in the x-axis) is the
optimal action. Then in Figure 13(b) we show the prediction error of PreDeToR
(-$\tau$) for each task $m\in[M_{\text{test}}]$. The prediction error is
calculated the same way as stated in Section 5 From the Figure 13(b) we see
that for most actions the prediction error of PreDeToR (-$\tau$) is closer to
LinUCB showing that the introduction of $1$ new action does not alter the
prediction error much. Note that LinUCB estimates the empirical mean directly
from the test task, whereas PreDeToR has a strong prior based on the training
data. Therefore we see that PreDeToR is able to estimate the reward of the
optimal action quite well from the training dataset
$\mathcal{D}_{\text{pre}}$.

|  |
---|---|---

(a) Train Optimal Action Distribution

(b) Distribution of action sampling in all tasks where action $17$ is optimal

(c) Candidate Action Set in Time averaged over all tasks where action $17$ is
optimal

Figure 12: Exploration Analysis of PreDeToR(-$\tau$) in linear 1 new arm
setting

|
---|---

(a) Test action distribution

(b) Test Prediction Error

Figure 13: Prediction error of PreDeToR(-$\tau$) in linear 1 new arm setting

We now consider the setting where the number of invariant actions is
$|\mathcal{A}^{\text{inv}}|=15$. We again show in Figure 14(a) the training
distribution of the optimal actions. For each bar, the frequency indicates the
number of tasks where the action (shown in the x-axis) is the optimal action.
Then in Figure 14(b) we show how the sampling distribution of DPT-greedy,
PreDeToR and PreDeToR-$\tau$ change in the first $10$ and last $10$ rounds for
all the tasks where action $17$ is optimal. We plot this graph the same way as
discussed in Section A.11. From the figure Figure 14(b) we see that none of
the algorithms PreDeToR, PreDeToR-$\tau$, DPT-greedy consistently pulls the
action $17$ more than other actions. This shows that the common underlying
actions across the tasks matter for learning the epxloration.

Finally, we plot the feasible action set considered by DPT-greedy, PreDeToR,
and PreDeToR-$\tau$ in Figure 14(c). To plot this graph again we consider the
test tasks where the optimal action is $17$. We build the candidate set the
same way as before. From the Figure 14(c) we see that none of the three
algorithms DPT-greedy, PreDeToR, PreDeToR-$\tau$, is able to sample the
optimal action $17$ sufficiently high number of times.

We also show how the prediction error of the optimal action by PreDeToR
compared to LinUCB in this $1$ new arm linear bandit setting. In Figure 15(a)
we first show how the $20$ actions are distributed in the
$M_{\text{test}}=200$ test tasks. In Figure 15(a) for each bar, the frequency
indicates the number of tasks where the action (shown in the x-axis) is the
optimal action. Then in Figure 15(b) we show the prediction error of PreDeToR
(-$\tau$) for each task $m\in[M_{\text{test}}]$. The prediction error is
calculated the same way as stated in Section 5. From the Figure 15(b) we see
that for most actions the prediction error is higher than LinUCB showing that
the introduction of $5$ new actions (and thereby decreasing the invariant
action set) significantly alters the prediction error.

|  |
---|---|---

(a) Train Optimal Action Distribution

(b) Distribution of action sampling in all tasks where action $17$ is optimal

(c) Candidate Action Set in Time averaged all tasks where action $17$ is
optimal

Figure 14: Exploration Analysis of PreDeToR(-$\tau$) in linear 5 new arm
setting

|
---|---

(a) Test action distribution

(b) Test Prediction Error

Figure 15: Prediction error of PreDeToR(-$\tau$) in linear 1 new arm setting

###  A.13 Data Collection Analysis

In this section, we analyze the performance of PreDeToR, PreDeToR-$\tau$, DPT-
greedy, AD, Thomp, and LinUCBwhen the weak demonstrator $\pi^{w}$ is Thomp,
LinUCB, or Uniform. We again consider the linear bandit setting discussed in
Section 4. Recall that the linear bandit setting consist of horizon $n=25$,
$M_{\text{pre}}=200000$, $M_{\text{test}}=200$, $A=10$, and $d=2$. Finally, we
show the cumulative regret by the above baselines in Figure 16(a), 16(b), and
16(b) when data is collected through Thomp, LinUCB, and Uniform respectively.

Outcomes: We first discuss the main outcomes of our experimental results for
different data collection:

(1) The PreDeToR (-$\tau$) excels in exploiting the underlying latent
structure and reward correlation when the data diversity is high.

(2) The PreDeToR (-$\tau$) fails to exploit the underlying latent structure
and reward correlation when the data diversity is less.

(3) The PreDeToR (-$\tau$) has lower regret compared to DPT-greedy in all
settings. AD outperforms PreDeToR (-$\tau$) when data is collected by LinUCB.

|  |
---|---|---

(a) Thomp data collection

(b) LinUCB data collection

(c) Uniform data collection

Figure 16: Data Collection with various algorithms and Performance analysis

Experimental Result: We observe these outcomes in Figure 16. In Figure 16(a)
we see that the $A$-actioned Thomp is explorative enough as it does not
explore with the knowledge of feature representation. So it pulls the sub-
optimal actions sufficiently high number of times before discarding them in
favor of the optimal action. Therefore the training data is diverse enough so
that PreDeToR (-$\tau$) can predict the reward vectors for actions
sufficiently well. Consequently, PreDeToR (-$\tau$) almost matches the LinUCB
algorithm. Both DPT-greedy and ADperform poorly in this setting.

In Figure 16(b) we see that the LinUCB algorithm is not explorative enough as
it explores with the knowledge of feature representation and quickly discards
the sub-optimal actions in favor of the optimal action. Therefore the training
data is not diverse enough so that PreDeToR (-$\tau$) is not able to correctly
predict the reward vectors for actions. Note that DPT-greedy also performs
poorly in this setting when it is not provided with the optimal action
information during training. The AD matches the performance of its
demonstrator LinUCB because of its training procedure of predicting the next
action of the demonstrator.

Finally, in Figure 16(c) we see that the $A$-armed Uniform is fully
explorative as it does not intend to minimize regret (as opposed to Thomp) and
does not explore with the knowledge of feature representation. Therefore the
training data is very diverse which results in PreDeToR (-$\tau$) being able
to predict the reward vectors for actions very well. Consequently, PreDeToR
(-$\tau$) perfectly matches the LinUCB algorithm. Note that AD performs the
worst as it matches the performance of its demonstrator whereas the
performance of DPT-greedy suffers due to the lack of information on the
optimal action during training.

###  A.14 Empirical Validation of Theoretical Result

In this section, we empirically validate the theoretical result proved in
Section 6. We again consider the linear bandit setting discussed in Section 4.
Recall that the linear bandit setting consist of horizon $n=25$,
$M_{\text{pre}}=\\{100000,200000\\}$, $M_{\text{test}}=200$, $A=10$, and
$d=2$. The demonstrator $\pi^{w}$ is the Thomp algorithm and we observe that
PreDeToR (-$\tau$) has lower cumulative regret than DPT-greedy, AD and matches
the performance of LinUCB.

Baseline (LinUCB-$\tau$): We define soft LinUCB (LinUCB-$\tau$) as follows: At
every round $t$ for task $m$, it calculates the ucb value $B_{m,a,t}$ for each
action $\mathbf{x}_{m,a}\in\mathcal{X}$ such that
$B_{m,a,t}=\mathbf{x}_{m,a}^{\top}\widehat{\bm{\theta}}_{m,t-1}+\alpha\|\mathbf{x}_{m,a}\|_{\mathbf{\Sigma}_{m,t-1}^{-1}}$
where $\alpha>0$ is a constant and $\widehat{\bm{\theta}}_{m,t}$ is the
estimate of the model parameter $\bm{\theta}_{m,*}$ at round $t$. Here,
$\mathbf{\Sigma}_{m,t-1}=\sum_{s=1}^{t-1}\mathbf{x}_{m,s}\mathbf{x}_{m,s}^{\top}+\lambda\mathbf{I}_{d}$
is the data covariance matrix or the arms already tried. Then it chooses
$I_{t}\sim\mathrm{softmax}_{a}^{\tau}(B_{m,a,t})$, where
$\mathrm{softmax}_{a}^{\tau}(\cdot)\in\triangle^{A}$ denotes a softmax
distribution over the actions and $\tau$ is a temperature parameter (See
Section 4 for definition of $\mathrm{softmax}_{a}^{\tau}(\cdot)$).

Outcomes: We first discuss the main outcomes of our experimental results:

(1) The PreDeToR (-$\tau$) excels in predicting the rewards for test tasks
when the number of training (source) tasks is large.

(2) The PreDeToR (-$\tau$) has a poor prediction of the rewards for test tasks
when the number of training (source) tasks is small.

|  |
---|---|---

(a) Prediction Error for $100000$ tasks

(b) Prediction Error for $200000$ tasks

(c) Cumulative Regret of PreDeToR (-$\tau$) compared against LinUCB-$\tau$

Figure 17: Empirical validation of theoretical analysis

Experimental Result: These findings are reported in Figure 17. In Figure 17(a)
we show the prediction error of PreDeToR (-$\tau$) for each task
$m\in[M_{\text{test}}]$. The prediction error is calculated as