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

Authors: Subhojyoti Mukherjee, Josiah P. Hanna, Qiaomin Xie, Robert D Nowak
Primary area: reinforcement learning
Keywords: Structured bandit, in-context learning, Decision Transformers

Abstract:
In this paper, we study the 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.

Decision: reject
Reviewer scores: 6, 3, 3

Reviews and discussion:
<thread>
creation date:
1730595102324

thread_type:
official_review

<thread_object>
creator: reviewer_5yut

<content>
summary:
This paper employs a Transformer model for decision-making in multi-armed bandit problems. Rather than generating actions directly, the authors propose training the Transformer to predict the reward for each arm. Actions are then selected greedily, with occasional random exploration, based on the reward estimates. The Transformer is trained on trajectories from a near-optimal expert, such as standard UCB or Thompson sampling algorithms. Results demonstrate that the trained Transformer outperforms traditional learning algorithms across a variety of bandit problems with shared structural features. The authors also present theoretical results quantifying generalization and transfer errors in relation to the training loss.

strengths:
The paper introduces an alternative pre-training method for simple decision-making environments, such as bandits. In my opinion, this design is well-suited for the bandit context, and the algorithm performs effectively in both synthetic and real environments. The authors also provide ample ablation studies to explore different setups and gain insight into the conditions under which this algorithm performs well. In general, I find no significant issues regarding the originality or clarity of this paper.

weaknesses:
I think the biggest weakness of the paper is that it fails to provide sufficient understanding of why this algorithm works well—specifically, it lacks insights or supporting evidence on why this design would outperform other algorithms and what key takeaways it offers for theoretical researchers and practitioners. For example, I find that the experimental section and the theory section feel somewhat disconnected. The theoretical results focus on generalization errors and transferability with respect to the pre-training loss. However, this pre-training loss does not directly correlate with the regret, which is the focus of the empirical study. This paper feels like a mixture of observations and results on different aspects of the problem, instead of a systematic effort to understand this pre-training pipeline.

questions:
1. Could you provide a systematic explanation of why this design benefits decision-making, particularly in the shared-structure bandit setting? After reading the paper, I am unclear about the authors' motivation for studying the bandit problem with a shared structure. In my opinion, there is a lack of insight into why a Transformer-based decision-maker would succeed even in a simple bandit environment.

2. While I understand that the proposed algorithm does not rely on the optimal action, it would be beneficial to include a comparison with the optimal action to assess the performance difference. This addition could help clarify how much of the performance improvement is attributable to the optimal decision in the training data itself. Furthermore, I believe the paper should cite the following related work, which is closely related.

Supervised Pretraining Can Learn In-Context Reinforcement Learning, Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, Emma Brunskill.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649980497

thread_type:
official_review

<thread_object>
creator: reviewer_hllk

<content>
summary:
This paper introduces a new training phase for a pre-trained Transformer model for structured bandit problems. Unlike existing approaches, such as Algorithm Distillation (AD) and Decision Pretrained Transformer (DPT), which predict the next (near-)optimal action using in-context data during training and testing, the proposed framework focuses on predicting the reward vector and selecting the arm based on this prediction.  This paper gives several experiments and a theoretical analysis of the proposed approach.

strengths:
The paper is well-written. The proposed training framework reduces the dependency on (near-)optimal solutions in training data/tasks, making it more flexible. Additionally, the authors provide a set of experiments and offer a theoretical analysis to understand the framework's effectiveness.

weaknesses:
[1]The primary advantage claimed for the proposed framework is its ability to function without requiring near-optimal solutions in the training data. However, for bandit problems specifically, this requirement is often not a major limitation. In simulated tasks, optimal arms are readily available, or near-optimal actions can be efficiently obtained through established algorithms like UCB or Thompson Sampling. To convincingly illustrate the advantage, experiments on more complex tasks, where identifying near-optimal actions is genuinely challenging, would strengthen the claim.


[2] The proposed framework appears to be limited to situations with a finite set of actions, which restricts its applicability in problems with a continuous action space. Is there a potential adaptation that could allow the framework to handle continuous action spaces? Additionally, the experiments are conducted with a relatively small action space (20 arms), which raises questions about performance in larger action spaces, such as those with over 100 arms.

[3]Although the framework outperforms other benchmarks in most experiments, the assumption that each arm’s features remain constant across training and testing tasks is atypical for structured bandit problems. In standard settings, arm features are generally assumed to be known and vary between tasks. It is unclear why the authors chose this fixed-feature assumption while not using these features as inputs in the transformer model, which seems to reduce the framework’s applicability. Can this assumption be relaxed to better reflect standard structured bandit settings?

[4]When tested with new actions, the framework’s performance declines significantly, as seen in Figure 3(c) and (d). This suggests that the pre-trained Transformer struggles to generalize in cases with out-of-distribution (OOD) actions because the pre-training data lacks these new actions. My question is: if the pre-training can relax the fixed-feature assumption—e.g., by allowing arm features to be randomly sampled in training— would this improve the model’s generalization capabilities?

questions:
Please see the weaknesses. In addition, Line 272 states, “DPT-greedy estimates the optimal arm using the reward estimates for each arm during each task,” which appears to differ from the definition in Section 2.3 and Lee et al. (2023). Could you clarify whether DPT-greedy indeed uses the reward estimate or the best-arm estimate?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696323942

thread_type:
official_review

<thread_object>
creator: reviewer_6sdc

<content>
summary:
This paper introduces a new supervised pretraining procedure to learn the underlying reward structure in a multi-task structured bandit setting, which can be particularly useful when the optimal action is unknown in the training phase. This training method leads to improved performance across various bandit tasks, including linear, nonlinear, bilinear, and latent bandits. The authors also provide a theoretical generalization bound for the proposed algorithm.

strengths:
I think the idea of predicting reward function rather than the optimal action has its merit. The numerical experiment is thorough, and the intuitions are well-explained.

weaknesses:
I mainly have two concerns of the paper:

A distribution shift issue that doesn't arise under the in-context learning setting:
- The paper draws the analogy between the pretrained decision transformer with in-context learning multiple times throughout the paper, for both algorithm design and the theoretical bound. For in-context learning problems, we can think the task from the pretraining phase as being i.i.d. as the task during the test phase. But for the bandits problem, things are quite different. Specifically, the actions in $H_m$ generated during the pre-training phase is from weak demonstrators, whereas the actions in the test phase are generated from the transformer model itself. This will result in a distribution shift between the training and the test environment. And it matters both numerically and theoretically:
- Numerically, if the transformer's policy takes actions similar to the demonstrators such as UCB and TS, then why would it has a better regret performance than them? If the transformer's policy doesn't take actions similar to the demonstrators such as UCB and TS, how does the training through samples of $H_m$ from demonstrators extrapolate to the cases where the actions in $H_m$ is sampled from the transformer's policy? Given the large space of $H_m$, I don't believe the samples of $H_m$ following demonstrators can cover the whole space.
- Theoretically, I don't see any discussion of this distribution shift matter in the theoretical results in Section 7 (please correct me if I miss anything). To derive generalization bound in the face of this distribution shift, I believe notions such as the distribution ratio in the definition 5 of  https://arxiv.org/pdf/2310.08566 or others are necessary. 

Medium- to large-size action space or continuous action space: As I read from the algorithm and the experiments, the proposed method does not well handle tasks with continuous actions or with a large action space well, which are typical in linear and bilinear bandits. For tasks with a large number of actions, the proposed method becomes challenging to train because it requires H_m to include some reward for every possible action.

questions:
See above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732265310831

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Message to Area chair

comment:
Hi Area chair,

We have put significant effort into the paper and there are some serious misunderstandings on the part of the reviewers which we again point out below. It will be great if you urge some discussion on their part.

1. Reviewer 5yut: We have pointed out that we indeed discuss in detail how PreDeToR conducts a two-phase exploration in Appendix A.11 and Appendix A.12 which has not been shown by previous works. In fact, reviewer 5yut has not read the work properly and has asked us to cite the DPT paper which actually has been cited many times in the draft (and is our closest baseline).

2. Reviewer hLLK: The reviewer has a misunderstanding of how multi-task representation learning works for structured/linear bandit settings. We pointed out several works that have the same setting as ours. The reviewer has a misunderstanding on what we are showing in OOD generalization. We argue that all the current baselines (including ours) fail in the OOD setting, and shared structure is crucial in learning the optimal algorithm. This is an insight from our work. This has been counted as a weakness which is not a correct evaluation of the work.

We have made this comment public to everyone so it can lead to some discussion.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this work the authors study how to solve multi-task structured bandit problems via a meta learning procedure with similarity to the decision pretrained transformer. The main modification of this work is that in contrast with DPT, in this work the authors propose a method based on predicting the rewards of query arms. This means it this method does not require optimal arm supervision. It is also shown the method can learn the structure of the reward class such as when these are linear. 

I like the direction of this work, but it is a bit hard to argue for acceptance given the reviewer’s comments. I went over the author’s comments and the discussion with the reviewers and I also looked at the paper myself. One of the main issues of this work is presentation. It would help the authors to present their work in and slightly easier to read, by describing their algorithms as a meta-learning procedure, where the objective is to learn a good bandit history dependent reward predictor by pretraining on a distribution of source bandit tasks. These predictors can be used to construct something that looks like a Thompson sampling algorithm. I am somewhat worried that the proposed strategies, PreDeToR and PreDeToR-tau do not have probable exploration schemes. Both greedy and softmax strategies may not yield provable exploration guarantees. Another issue is the experimental evaluation. The reviewers were worried that the scenarios discussed in this work are limited. I understand the authors replicated many of the DPT experimental scenarios and I sympathize with that. Nonetheless having multiple readers of the work point out to the limited experimental results, does represent a signal that more experiments could be conducted.