Prompt Tuning Decision Transformers with Structured and Scalable Bandits
Authors: Finn Rietz
Oleg Smirnov
Sara Karimi
Lele Cao
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Decision Transformer
Prompt tuning
Bandits
Multi-task transfer
Primary area: reinforcement_learning
TL;DR: We provide a scalable bandit architecture for prompt tuning of decision transformers for increased downstream performance.

Abstract: Prompt tuning has emerged as a key technique for adapting large pre-trained Decision Transformers (DTs) in offline Reinforcement Learning (RL), particularly in multi-task and few-shot settings. The Prompting Decision Transformer (PDT) enables task generalization via trajectory prompts sampled uniformly from expert demonstrations -- without accounting for prompt informativeness. In this work, we propose a bandit-based prompt-tuning method that learns to construct optimal trajectory prompts from demonstration data at inference time. We devise a structured bandit architecture operating in the trajectory prompt space, achieving linear rather than combinatorial scaling with prompt size.  Additionally, we show that the pre-trained PDT itself can serve as a powerful feature extractor for the bandit, enabling efficient reward modeling across various environments. We theoretically establish regret bounds and demonstrate empirically that our method consistently enhances performance across a wide range of tasks, high-dimensional environments, and out-of-distribution scenarios, outperforming existing baselines in prompt tuning.

Prompt Tuning Decision Transformers with Structured and Scalable Bandits
Finn Rietz"
Oleg Smirnov King AI Labs, Microsoft Gaming oleg. swi.novmmi www.oft...roooooooommmmmmoooooo. com
Örebro University f finn.rt rietz2oru..ee
Sara Karimi
Lele Cao
King Al Labs, Microsoft Gaming sara. sara.karimg karimiriiieii wwwiiiiiiii..mmmm
King Al Labs, Microsoft Gaming delecao@microsoft.. com
Abstract
Prompt tuning has emerged as a key technique for adapting large pre-trained De- cision Transformers (DTss in offline Reinforcement Learning (RL), particularly in multi-task and few-shot settings. The Prompting Decision Transformer (PDT) enables task generalization via trajectory prompts sampled uniformly from expert demonstrations without accounting for prompt informativeness. In this work, we propose a bandit-based prompt-tuning method that learns to construct optimal tra- jectory prompts from demonstration data at inference time. We devise structured bandit architecture operating in the trajectory prompt space, achieving linear rather than combinatorial scaling with prompt size. Additionally, we show that the pre- trained PDT itself can serve as a powerful feature extractor for the bandit, enabling efficient reward modeling across various environments. We theoretically establish regret bounds and demonstrate empirically that our method consistently enhances performance across wide range of tasks, high- -ighenmnniio environments, and out-of- -isttrriinttonnn scenarios, outperforming existing baselines in prompt tuning.
Introduction
The ability exploit large amounts offline data is crucial for training foundation models capable of generalizing across diverse tasks (Radford) 2018 Reed et al. 2022; Brohan et al. 2022). In Reinforcement Learning (RL), a seminal contribution the Decision Transformer (DT) (Chen et al.) 2021), which reframes offline RI. (Levine et al. 2020) as a sequence modeling problem, thereby unlocking powerful Transformer architectures for offline RI. DT particularly well-suited for offline RI. because it sidesteps the well-knows instabilities of temporal-differees function approximation (Sutton Barto (018), which are exacerbated under distribution shift the offline setting, by replacing discounted, moving-target value estimates with stationary return-to- go conditioning, autoregressive sequence modeling, and causal attention mechanisms (Chen et al. 202 The Prompting Decision Transformer (PDT) (Xu et al. (022) further extends DT from single-tast task to muulti-task settings, enabling large-scale models and generalized pre-training in offline multi-task and few-shot RL (Xu al. 2022 Mitchell et al. (021). Analogous prompting in Large Language Models (LLMs). PDT differentiates tasks through a stochastic trajectory prompt prepended to the context, allowing it to identify and model optimal action marginals for each task in the offline dataset. This makes PDT particularly appealing for offline multi-task RI,, as avoids additional learning
learning with
Work performed while the author was intern at King Labs (part Microsoft Gaming). Code available at: https://githut. com/ Baangg//tttt bandits
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

on the downstream task while enabling efficient and robust adaptation of the model's behavior at inference time, solely through adjustments the prompt.
However, PIDT samples prompts uniformly, overlooking that not all prompts are equally informative, in fully observable MDPs. We hypothesize that the sampling of non-informative prompts from expert demonstrations can diminish PDT's ability to differentiate between tasks, thereby leading to performance degradation.
While improving on PDT's uninformed coompt-sampling strategy, prior works on PDT prompt-tuning suffer from key limitations. The approach by Yuan et al. (2024) replaces trajectory prompts with less-expressive goal- conditinni and relies on hindsight relabeling. The generative approaches by Hu et al. (2033 2024) are not applicable in discrete settings and don't 't adhere causal relationships between prompt tokens. Furthermore, all of these works treat prompts as flat, unstructured inputs and operate directly on MDP modalities, which leads to poor scaling with prompt size and state- and action-space sizes.
To address these shortcomings, we introduce a scalable, robust, and computationally efficient bandit- based prompt-tuning method for PDT. Our method exploits the inherent structure of the prompt space to reduce the complexity of prompt selection, transforming it from combinatorial problem into one that scales linearly with prompt size. Moreover, we show how to leverage the pre-traines PDT as a feature extractor to obtain compact prompt representations that enable efficient deployment even in high-dimg nigh-dimensional (e.g., pixel-based) settings. By optimizing prompt selection, our approach boosts downstream task performance without costly weight updates to the underlying Transformer backbone. Our experiments reveal clear performance gains with the proposed method, effectively bridging the gap between pre-training and adaptation.
2
Preliminaries
In this section, we introduce key concepts and terminologies that form the foundation of this work.
2.1 Problem definition: Offline multi-task RL
The offline multi-task RL problem is formalized similarly to prior works (Xu et al. 2022 Mitchell et al. 2021). The objective to solve a set of training tasks Tirain with the option to evaluate task generalization capabilities
holdout test set T
For each task T E T a dataset D; is provided, consisting of trajectories sampled from the corresponding MDP M; = (S;, Ai, rı, di, Yi, /.)). Here, S; is the state space, A; is the action space, rs : A; - Th. represents the reward function, d; S, A; S; -> [0. 1] defines the discrete-time transition dynamics, %/2 (0, 1] is the discount factor, and μe is the initial state distribution of MDP i.
The goal is to learn generalized policy, (s, 1b - a, capable of solving all tasks in T Here, eb serves an auxiliary input to the policy that sufficiently describes task from T. In the case of PDT, & is the stochastic trajectory prompt sampled from the demonstration set for the target task. For each task, the optimal generalized policy is expected to maximize the corresponding expected discounted reward objective specific to task
x ( (n,i)=Bi =E Σ
(r at)
The expectation is taken over trajectories where actions are sampled from the policy π, and states are sampled using the transition dynamics d; of the MDP M, In the offline setting, only the pre-collected datasets D (Do, ... D are available for learning, with no additional data collection allowed. A simulator may be used for policy evaluation, but not for gathering new training data due to the offline nature of the problem.
2.2 Prompting Decision Transformer
Chen al. (202T) reframed RL sequence modeling problem, leveraging the capabilities of Transformer architectures to model trajectories. A trajectory in DT is represented as sequence of triplets (fz, st, az), where fr ET Te is the return-to-go, S, is the state, and a; is the action at

---

time step t. DT captures the temporal and causal dependencies between states, actions, and rewards, enabling the model to predict optimal actions directly, without relying on explicit value functions or policies.
Xu al. (022)) extended DT multi-task offline RL setting by introducing stochastic trajectory prompts as task-specific context, enabling the model to identify underlying MDPs. Such a prompi composed of / trajectory segments, each of length H,, resulting in total of И prompt tokens, as per Equation where the superscript denotes tokens associated with the prompt.
segment
segment
0
fg. aj
H aj+H H.........
Sk a
S H: a
whee L most recent transitions
(p) o PL-L, BI-L a
a a
P2, Sc,
For a particular training task T E T PDT learns to model the sequence in Equation 3 by autoregressively predicting the action tokens, where @ denotes concatenation. While PDT randomly samples the J segments that constitute the prompt p from a set expert demonstrations P, we propose optimize segment composition to enhance the downstream task performance with Multi-Armed Bandit approach.
2.3 Multi-Armed Bandits
Multi-Armed Bandits (MABs) provide a framework for optimizing stochastic reward functions, making them effective tool for tasks like prompt tuning. In the standard MAB setting, each time step k the agent selects an action (or "arm") ag E A. where A is the set of available arms. Here,
denote bandit time steps, to avoid confusion with the MDPs time t. The agent then receives a stochastic reward тı R(ax). with the goal of maximizing the cumulative reward Σ Tx over time horizon K (Auer et al. 2002). This requires balancing the exploration of arms to gather information about their reward distributions A(a) and exploitation of arms with known high rewards while minimizing cumulative regret, defined as:
K
Regret (K)
Σ
max E[R(a)] E[rx] aEA
(4)
Contextual MABs (CMAB) extend this framework by incorporating side information (or "context") The reward distribution then conditioned both
C₁ which is observed before selecting the arm and the context, ги ~ R(as Ca) )- The agent's 's objective is to learn a policy : C - A that
maximizes the expected reward E
Σ E R(T(c)) cx) By leveraging shared features across arms
through the context Cj2 CMABs enable efficient learning and better generalization, particularly in settings where arms share intrinsic characteristics (Li et al. 2010).
3 Related Work
Recently, there has been surge in methods across various domains aimed at enhancing the perfor- mance and generalization of Transformer- based approaches through automatic prompt tuning.
Prompting in LLMs. For LLMs, Chen et al. (2024) introduced InstructZero, which uses Bayesian LLMs,(Chen optimization to explore low-dimensional soft prompt vectors. These vectors are then passed to an open-source LLM to generate instructions for the black-box L.M. Building on this, Lin et al. (0223 proposed INSTINCT, which replaces the Gaussian process in Bayesian optimization with neural network surrogate, leveraging neural bandit algorithm to enhance expressivity. Additionally, (Shi et al. (2024a) demonstrated that the budget MAB framework enables learning the optimal prompt within limited number of LLM evaluations. These methods rely on optimization in continuous spaces, followed by prompt reconstruction, whereas our approach directly operates within the PDT prompt space.
3

---

Multi-task DT. For Hu et al. (2023) proposed generating prompt candidates for PDT by per- turbing initial trajectories with Gaussian noise and applying online or offline feedback guide ranking-based optimization. Prompt Diffuser (Hu et al. 2024) similarly frames instruction optimiza- tion as conditional generative modeling task, synthesizing prompts from random noise. In contrast, our method constructs prompts directly from expert demonstrations, which is particularly beneficial in settings like discrete state and action spaces, where noise-based approaches are less effective.
Wang et al. 2024a) introduced hierarchical prompting, using two levels of context: one for encoding task-specific information and another for guiding rollouts via demonstration segments. Hyper- Decision Transformer (Xu et al. (2023) augments DT with hyper-network for adaptation to novel tasks, but both methods rely on large quantities of demonstration data. Our method, by contrast, selects high-quality prompts from only a few expert trajectories. Other approaches focus on representation learning or adaptation, e.g.. Wang et al. (2024b) focus on disentangled world models, while Xie et al.. (2023) propose latent-variable conditioned RI. and fine-tuning. Our approach avoids updating the transformer backbone and relies solely on test-time prompt optimization.
Finally, MGPO (Yuan et al. 2024) proposes general framework for online prompt tuning. While sharing similar goal, we provide a method to identify the best prompt from a demonstration dataset, while MGPO iteratively refines the prompt with data collected during online rollouts.
LLM for RL. At the intersection of LLM and RL domains, Yang & Xu (2024) and the closely related LaMo method (Shi et al., 2024b) proposed leveraging a pre-trained LLM as an initializer for PDT, harnessing rich linguistic knowledge to boost performance
unseen tasks. In another work,
Zheng et al. (2024) introduced an approach to decompose the prompt into cross-task and task-specific components, ensuring more robust test-time adaptation to unseen tasks. Furthermore, the model is initialized by incorporating parameters from a pre-trained language model to provide it with prior knowledge. Compared to those, our approach avoids reliance on dedicated LLM, sidestepping the fine-tuning and scalability challenges inherent in such methods.
4 Method
We now present our inference time prompt-tuning bandit method for offline multi-task RL. Given
a dataset D (Du..... .... .,,)) of trajectories for n training tasks Tttain we utilize the original PDT method to learn the generalized policy described in Equation [T Details of the PDT algorithm and training process are provided by Xu et al. (2022).
An optimal PDT model x²(x; ()) is assumed to be trained until convergence on D. After training, the goal is to evaluate and enhance the model's performance on a specific task Te, either from the set of holdout test tasks 7 or from the training set T We assume access to a small set of demonstrations P. for the target task, which serve as a source for sampling prompts, as well as simulator of the corresponding M; to perform online evaluations. Our primary objective is to identify the stochastic trajectory prompt constructible from P; that maximizes performance when applied to the pre-trained generalized policy.
4.1 Bandit-based Prompt Tuning
We propose leveraging a bandit- based approach to efficiently identify the optimal trajectory prompt. A naive implementation of bandit-based prompt tuning would involve deploying bandit with one arm for each possible prompt P constructible from P., i.e., considering all possible combinations of segments among demonstration trajectories. This approach, however, scales extremely poorly, as the number of prompts (and consequently, the size of the bandit problem) grows linearly with the total number of transitions in P, and combinatorially with J, the number of segments in each prompt. Additionally, treating each prompt as separate and independent arm disregards prompt similarity, leading to inefficient sample complexity. Thus, we propose to optimize the selection of segments that form the stochastic trajectory prompt (7i F.) with a special contextual MAB algorithm.
As illustrated in Figure T each round kt the bandit selects a prompt Pk: consisting of J segments, balancing exploration and exploitation in the prompt space. The process of constructing the prompt Ps. through the bandit mechanism is discussed in detail in the next section. We then perform the k-th rollout of the PDT π "xx; 0) in M,, using the selected prompt Pk and the most recent transitions, as per Equation
x" (xxx 0) serves as 3 The resulting performance G Ti a:)
Te(Ss,

---

Few-shot Demonstrations
Go Eis
Rollout
Feature Extractor (Optiona))
Contextual Multi-armed Bandit
Prompt- based Causel Transformer
Prompps
Figure 1:
Overview of our Dandit-based prompt-tuning method for multi-task learning with PDT. Each 2; represents triplet (f,; a,), each represents prompt segment, each represents demonstration trajectory. The bandit explores the demonstration dataset Pi for the current task i to find the best prompt p (ff..... ff). The online return G& achieved by the underlying PDT model at round k and using prompt Pk serves as a reward for the bandit.
reward from the bandit's perspective and is stored together with the prompt pk for training the bandit. The algorithmic pscudocode for prompt-tuning with our bandit is provided in the Algorithm D in Appendix(A
4.2 Scalable and Sample-Efficien Bandit Architecture
To address the aforementioned scalability issues in naïve bandit, design structured bandit architecture that factorizes the problem across prompt segments. Our structured architecture features J arms, for each segment in the stochastic trajectory prompt. Each arm j {1, (1, J) maintains an independent reward model éj. which predicts the performance of the underlying PDT when given segment is placed at position j in the prompt. This decomposition reduces the complexity of
the search space from combinatorial to linear in J. Exploration can be performed using strategies such e-greedy, Upper Confidence Bounds (UCB) (Li et al. 2010: Zhou al. 2020). Thompson Sampling (TS) (Thompson, 1933). or other mechanisms.
To select the prompt Pk. each model Φ predicts the reward for each segment (Fo...., 7n) P its position, j. in the prompt, resulting prediction matrix Y I The J segments with the highest predicted performance are then selected by computing arg max over the segment dimension of Y, subject the chosen exploration strategy. This yields the prompt pk (7 (Ff.... ..., =f). After the rollout & of the PDT with prompt in M,, all reward models are independently updated on their accumulated (7f G pairs. The complete algorithmic pseudocode for prompt selection and update with our structured bandit architecture is provided in Algorithms and in Appendix A
4.3 Regret Analysis
We present regret bound for our bandit architecture, assuming that the reward for a prompt (TI,..., T.) is estimated as the average of J independent reward models (F). This estimate incurs an approximation error due to unmodeled correlations between segments, which is bounded by small constant 8.
The prompt-segment independence assumption is motivated by the modular structure of prompts in PDT. Prompt segments correspond to localized behaviors, and many MDPs can be effectively identified by small number of informative state-action pairs. In such cases, inter-segment interactions are minimal, since the presence of few discriminative (f, s, a) transitions is sufficient for task identification. Moreover, during PDT pretraining, prompt segments are sampled independently
Note that if P. contains M expert trajectories, each of length L. and each prompt segment has a length H. the total number of segments is given by [P]] M (L H +
5

---

from the demonstration pool without constraints on order or co-ocuurrence. As a result, the model not encouraged to rely on global interactions between segments, but instead learns attend to individually informative transitions within each segment. This training setup implicitly promotes invariance to inter-segment dependencies and limits the influence of global prompt structure on downstream behavior. See Section E for an attention weight analysis that further supports this interpretation.
Theorem 4.1. Assume that the reward function G: pJ + I. for a prompt p = (fi...., f.) decom- poses as the mean of J independent reward models (5(F;)
G(p) Σ o5( h(Fi....,
(5)
and that the interaction term h is uniformly bounded by Ih( dr)l S E, VF, E Let
p" = -Fi..... denote the optimal prompt, and suppose that for each slot j. bandit algorithm guarantees a slot-specific regret
Regret, (K) Σ
-
(6)
over rounds. Then the cumulative regret after K rounds is bounded as:
Regret (K) A E -(C("")GG(p)))) - Regret (Kgret,(K)=r + 2Ke.
suppose each slot-specific reward model Ø5 learned using standard contextual bandit algorithm (e.g., UCB, Thompson Sampling, e-greedy) over the set of segments P. with [P] segments. Without loss of generality, assume that each 6j bounded in [0, either via normalization of returns bounded regression targets. If the regret for each oy satisfies Regret, O (v (vK log then the total regret is bounded by
Corollary 4.2. Under the same assumptions as Theorem 4.1
Regret(K) O (VK log PI 1P
(8)
Corollary 4.2 shows that the proposed bandit architecture preserves the sublincar regret bound of standard algorithms, while introducing an additional error term that grows linearly with the number of rounds K.
4.4 Arm Features
Learning reward models dy 7 R for raw trajectory segments becomes impractical for MDPs with large state spaces, such as those involving pixel-space observations and high-dimensional actions. The key issue lies in the input size of these reward models, which scales as H ([S) + |A| + 1), growing linearly with S, A, and segment length H. This limitation can be addressed by integrating a feature extractor V + Rd that encodes raw trajectory segments from the MDP modalities into latent feature space.
We propose to leverage the pre-trained PDT model as feature extractor by taking the latent representation of prompt tokens as the embedding for the prompt. This approach not only mitigates the scaling issue but also aligns with the inductive biases of the pre-trained PDT, which is expected to encode meaningful, task-relevant information. The reward models then operate on these fixed-size embeddings rather than raw, flattened segments, significantly improving scalability.
In summary, we introduce a specialized bandit architecture for PDT prompt tuning. By leveraging the structure of the prompt space, we optimize J independent reward models, each corresponding to a segment position in the trajectory prompt, which reduces the search complexity from combinatorial to linear. Additionally, we utilize the pre-trained Transformer backbone to embed prompt segments, allowing the reward model input size to remain fixed even in high-dimensional state and action spaces.

---

Experiments
We now present comprehensive evaluation of our method. We first assess performance across multiple MuJoCo (odorov et al. 2012) tasks from standard multi-task benchmarks (Yu et al. 2020 Finn et al. 2017). We then analyze prompt-space exploration custom 2D environment. Finally, we evaluate the regret behavior of our bandit architecture and compare its scalability to standard algorithms.
5.1 Environments, Datasets, Baselines
Environments.
We consider the following environments: (i) MuJoCo Half Cheetah: Joint control of a half-cheetah agent. Tasks involve matching varying target velocities. (ii) MuJoCo Ant: Joint control of an ant agent. Tasks involve movement along different target directions. (iii) Meta-world Pick-place: End-effector con- trol of simulated Sawyer arm. Tasks require picking objects from varying locations and placing them at distinct target positions. (iv) Sparse 2D Point: planar point agent navigates to a goal position. Tasks are distinguished by varying goal coordinates, dis- tributed on circles with different radii and angles, as depicted Figure 2a The agent must issue stop action to end the episode and receive sparse reward, defined as the negative distance to the goal. A bonus of 100 is awarded for stopping near the target. Crucially, for all settings, the task parameter (target velocity, direction, goal position) not part of the state it must be inferred from the trajectory prompt. A detailed description of the environments is provided in Appendix B
Ixpect
Cood lave
Goeti hape
Cools task
Exámole tasks
Start:
(a) Multiple tasks.
(b) Trajectories.
Figure 2: Sparse 2D point environment: Multiple tasks via parameters (radius) and a (angle); (b) Expert vs. novice trajecto- ries.
Offline RL dataset. For each training task PDT requires dataset of trajectories Tú and a dataset of (expert) demonstrations P.. For the MuJoCo tasks, we use datasets and corresponding training and testing tasks from Xu al. (2022). To generate datasets for the Sparse employ PPO (Schulman et al. 2017) on 60 tasks with three discrete radii and 20 discrete angles. The first 48 tasks are used for training, and the remaining
point environment, we
are reserved for testing generalization, PPO is run for 1M steps on each training task, and the resulting trajectories are stored as D, The highest- return trajectories are selected form P: To reduce the computational cost of computing arg max over all possible prompt segments, P. is limited to 10 randomly sampled high-return trajectories.
Baselines. We evaluate our method against the following RL baselines: (i) optimal policy oracle, either scripted or learned using Conservative Q-Learning (CQL) (Kumar et al,, 2020); (ii) Decision Transformer (DT) trained across multiple tasks without further enhancements; (ii) a vanilla Prompt Decision Transformer (PDT) that samples prompts uniformly at random; (iv) Gaussian perturbation-
based method that performs hill climbing in the prompt space; and (v) the ranking- based prompt tuner from Hu et al. (2023). referred to as "ZORankSGD".
5.2 Prompt Tuning Results and Analysis
Does the proposed method reliably improve the performance a frozen PDT backbone? We investigate this question empirically on three MuJoCo environments, where we compare the inference-time performance of a standard PDT, that samples prompts randomly from P;, to a PDT enhanced by our bandit-based prompt-tuning method. Each method is evaluated over 250 online rollouts on in-distribution training tasks, and we report the mean return from the final 10 rollouts, averaged over three seeds and all tasks. Results are presented in Table []
Our bandit-based based prompt-tuning method consistently and substantially improves the performance of the frozen PDT backbone, even surpassing the single-task CQL oracle on the MuJoCo Ant environment. While PDT performance generally increases with larger prompt sizes (i.e., higher J and H), our method yields even greater gains, demonstrating the clear advantage of optimized prompt
7

---

MaJoCo Half Cheetah
MoJoCoo Ant
Meta-vorld
Pick-place
Method
J=I.H=
J=2,H H
J=1.8=
J=2.
CQL oracle
15.8211233
-25.82 12.53
760.54 288.20
760.54 1 288.20
535.84 + 31.02
535,84 31.02
PDT, no uning
.44.5 + 0.91
-42.6 + 1.3
694,43 = 227.04
754.22 = 175.1
551.58 26.09
535.52 + 24.86
Hill-chibing
40.00 26.39
29,93 22.84
738.56 181.56
740.17
189.38
555.79
22.72
540.15 23.49
ZORankSGD
-43.22 21.14
34.77
735.47 = 180.20
731.22 + 172.48
554.26
23.00
537.20 - 25,
r-greddy"
-2.60 3.77
+31212
815.17 = 182.28
819.52 +: 191.39
555.35
+
TS
-38.52 = 11.21
-27.62
12.90
800.95 = 184.75
791.06 + 192.75
556.87 24.11
540.82
+ 22.83 22,83
UCB"
-36.55 - 11.93
-26.87
14.10
751.35 = 192.25
744.55 + 168.50
552.68 24.77
538.84
23.14
e-greedy
-76.50 6.68
-58.76
+ 8:21
812.14 176.88
816.77 + 268.42
556.22 25.16
541.80 23.76
TS
-33.56 +: 13.48
-26.28
+ 10.14
835_18 171.25
729.65 + 175.41
556. + 24.56
$41.33 + 22.79
UCB
-35.72 + 14.96
-26.57
+ 10.79
732.22 + 180.25
777.29 180.77
554.92 26.06
538.50
23.64
Table I: Inference time performance (mean - std) across environments with varying trajectory
segments (J) and lengths (H). Bold values indicate best performance.
Method
MuJoCo Half Cheetah
MuJoCo Ant
Meta-vorld
Pick-Place
J= =3
J= H 20
J 1, =
J 2. = 20
J=
J, =2
Oracle
-23.8 + 10.39
23.81 + 10.39
508.09 + 231.90
508.09 + 231.90
525.07 + 60.15
525.07 60.15
PDT, no tuming
-64.78 36.91
-40.55
43.19
363.90
105.42
360.07
72.36
502.8 + 63.98
524.37 39.56
PDT,, finetuned
-129.05 +: 65.91
39.30
14.66
306.29
63,91
360,96
138.58
495.37 57.87
488.17
50.15
Hill-elimbing
-53.84 + 23.56
-34.39
26.05
355.80
+ 135.17
344.46 +: 57.21
560.92
+: 27.04
544.19-
28.83
ZOR
-59.85
32.37
-6.66
19.45
383.57
193.35
340.68
44.60
503.56
66.16
$38.05
31,46
e-greddy*
-32.61 19.85
-23,93
14.14
477.24 + 84.64
431.52
531.49 49.86
552.11
23.83
TS
-38.56 + 21.61
-31.80
+: 14.20
468.76 + 79.50
441.44 = 80.25
549.38 36.12
553.12
20.81
UCB"
-42.76 1: 15.77
-33.27 +: 17.27
392.39
+ 51.57
346.59 + 49.50
506.43 1: 65.38
539.42
31.12
r-greddy
-2.977 21.10
25.18
22.17
480.85
84.76
530.26
550.06
22.38
TS
-29.26 21.25
-34.23
+ 14.51
466.11
+ 79.50
438.82 + 125.40
$51.33 +: 34.52
553,34 18.82
UCB
.42.99
15.67
33.78
16.83
376.68
82.25
413.31 46.31
512.25 69.62
534.58
35.35
Table 2: OOD performance (mean + std) across environments with varying trajectory segments (J)
and lengths (lf). Bold values indicate best performance.
selection over uniform sampling from P,. When using Transformer-encooded trajectory segments (denoted by 1,), our method achieves comparable performance in terms of return to the unencoded variants, demonstrating that the PDT provides a robust and compact representation for bandit-based prompt tuming. The main benefit of this encoding lies in its fixed-dimensional representation, which makes the approach scalable to high-dimensional settings, as further evidenced by the inference-time results in Appendix g
Although hill-climbing and molking-based prompt-tuning (Hu et al. 2023) can also improve perfor- mance, their results are less consistent and typically weaker. We attribute this instability to their reliance on iterative, noise-driven perturbations of initial prompts.
In summary, bandit-based prompt-tuning reliably enhances frozen PDT performance, outperforming existing baselines, particularly when leveraging Transformer-encodeed prompt segments.
Does the bandit-based method improve inference-time performance on OOD tasks? To assess the benefits of prompt tuning in out-of-distribution (OOD) scenarios, we evaluate the frozen PDT model on OOD tasks from the MuJoCo suite using the same procedure as before. Results are shown in Table 2 We find that our bandit-based method significantly improves PDT's generalization in all tasks, often reaching performance levels comparable to in-distribution tasks (MuJoCo Half Cheetah, Meta-world Pick-Place). though not always achieving optimal returns. This suggests that the effectiveness of our method extends to OOD tasks, provided the underlying model retains sufficient generalization capability.
In contrast, hill-climbing and ranking- sanking-based prompt-tuning (Hu et al. 2024) exhibit reduced robust- ness in the OOD setting, falling considerably short of the performance gains of our method in MuJoCo Half Cheetah and MuJoCo Ant environments. We hypothesize that this degradation arises from their reliance on the initial sampled prompt and the increased optimization difficulty faced by the PDT in unfamiliar environments.
Finally, we observe that fine-tuning the PDT backbone solely on target task data for 250 epochs does not reliably improve performance and even degrades performance in some settings. This degradation is hypothesized to stem from the large size of the Transformer backbone, the absence of regularization provided by diverse multi-task training data, and less efficient full-model updates. similar phenomenon was previously reported by Yuan et al. (2024).
8

---

Method
==
==
J
ST" oracle
10
0.0
10:00.0
PDT. no tuning
0.0
6.3
8.3
Hill-climiig
58
7.9
6.2-40
ZORankSGD
-0.6 + 30,9
4.4 16.7
311
e-greedy
9.0 0.6
9,4
9.6 0.2
Santau 1ozt
UCB
9.2 0.5
8.8 +.9
8.8 = 0.9
TS
9.7 0.1
9.7
9.4 0.5
e-greedy
9.5
8.9 0.7
UCB
9.4 0.5
9.5
9.3 0.4
TS
9.9 + 0.0
9.9 0.0
9.8 0.1
(b) Inference-time performance (mean + std) in the Sparse 2D Point environment. A standard single- task Decision DT achieves -64.3 1 24.1 average return, underscoring
(a) Top: First 20 prompts. Bottom: Last 20 prompts. Each dot represents the prompt's mean state co- ordinate, colored according to the achieved return
need for multi-task models like PDT. Results are averaged over training tasks, three seeds, and the final 50 rollouts.
Stars indicate task goals, the diamond marks initial state, and the shaded re- gion indicates OOD test tasks.
Figure 3: (a) Visualization of prompt selection across tasks in the Sparse 2D Point environment. (b) Inference-time performance showing the benefit of prompt tuning over single-task baseline.
How do different methods explore the prompt space? We use the Sparse 2D Point environment to qualitatively analyze how different methods explore the prompt space. For interpretability, we consider PDT with single prompt segment (J = 1), allowing prompts to be visualized via the mean spatial coordinate their states. Each method run for 250 rollouts per task, and we visualize the first and last 20 selected prompts in Figure 3a
The standard PDT samples the prompt space uniformly without refining its selection strategy over time. Hill-climbing gradually discovers high-performing prompts in the local neighborhood of the initial sample by iteratively applying Gaussian noise. ZORankSGD moves prompts toward the task goal, sometimes overshooting it and generating out-of-distribution prompts due to the unconstrained generative approach Hu et al. (2023).
Our bandit method (with e-greedy exploration) initially samples prompts uniformly, capturing both examples. Over time, it leverages accumulated reward estimates to
low- and avoid low-reward regions and increasingly selects segments near the task goal, which provide more informative cues for task identification.
We also report the performance all methods on the 2D environment tasks with the largest radius in Table 3b] The standard PDT without prompt tuning falls drastically short of the optimal return, due to the overlap between prompt datasets. The performance of the PDT increases with larger prompts, which implies that excessive random sampling suffices to find informative prompts in this environment. Best prompt-tuning results are again achieved by our bandits compared to the hill-climbing and ranking-based prompt-tuning baselines, with Thompson Sampling (TS) performing slightly better than UCB or e-greedy exploration. We find no considerable difference between using Transformer encoded prompts (marked by Vf or not, which can be attributed the simplicity of the 2D environment.
5.3 Bandit Regret Analysis
We empirically evaluate the proposed bandit-based method with respect to the interaction bound s, and compare its performance to standard MAB algorithms that learn directly over the full combinatorial space. This evaluation is conducted on synthetic task carefully designed reflect key characteristies
of PDT prompt tuning. Specifically, the task involves identifying an optimal vector of length J, where each dimension takes values in f1..... H}. The reward function R(x) is defined only for complete vectors and is bounded within [0. 1]. making the contribution of individual components unobservable, mirroring structure of real-world setting. While the synthetic setup is necessarily simplified, it enables controlled investigation of core assumptions underlying our theoretical analysis, particularly the independence of segment- wise reward estimates. We report cumulative regret after K 2000 rounds for varying values of J and H. comparing structured method to standard

---

greedy
Thompse Sampping
UCB
Interaction termi
1400
Standard
1400
structured Stardard
1400
Structured
Standard
0.175
Thampan Sumplng
greedy
140C
2400
Cumuletion mumueetn aagret
1440
01550
UCB
Cumulciio
80C
820
800
0.100
0.075
430
000
600
0.00
200
230
200
0.025
13
20"
10t
10+
10
-01
10
-0P
10
10P
10-
10
10'
10'
125
10-
10"
10"
101
10"
Problem Stss
Problom Stzo
em Size
Probfom Stzo
Figure 4: Performance of the structured (ours) and standard MAB methods on synthetic prompt tuning task. Problem instances are generated by sweeping over {2, 4, 4,5) segments and H = (3, 5, 7, 10] choices, with problem size reported as H Shaded regions indicate one standard deviation around the mean; results are averaged over three random seeds.
e-greedy, UCB, and TS strategies. For our method, we also estimate the neglected interaction term l(x1...... x.) [(x) Σ, (x;))) where R(x) denotes the true reward.
As shown in Figure our method matches the performance of standard, "flat" bandit baselines on small problems and consistently outperforms them as problem size grows. This demonstrates more efficient exploration of the combinatorial space by leveraging its structure. Additionally, the interaction term (x,,,...,, X,) remains low and increases only marginally with problem size, empirically supporting the bounded interaction assumption Theorem 4.1 These findings reinforce the validity of our method even in settings where interactions between prompt segments exist while
global reward must inferred from segment-leve estimates.
6 Conclusion
This work introduces bandit-based prompt-tuning method extending PDT and addressing the lim- itations of uniform prompt sampling. By leveraging a contextual MAB framework, our approach optimizes the selection of trajectory segments maximize task performance and enhance adapta- tion to unseen tasks. Experimental results highlight consistent improvements across diverse tasks, demonstrating the efficacy and robustness of the proposed method in multiple environments. This method not only advances the state of prompt optimization in PDT but also contributes to the broader integration of offline RI. and sequence modeling paradigms.
Limitation and broader impact. A key limitation is the combinatorial expansion of the search space the number of demonstrations increases, which renders online exploration impractical. A promising direction to learn a sampler that pre-selects high-potential segments or clusters of similar prompts reduce redundancy. While bandit-based tuning enhances OOD performance, pure offline RL struggles with extrapolation. Incorporating meta-learning techniques, such as In-Context RL (Laskin et al. 2022), into PDT pretraining offers compelling extension of this work.
References
Peter Auer, Nicolò Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Mach. Learn., 47(2-3):2235 256, May 2002. ISSN 0885-6125. doi: 10.102//A: 1013689704352. URL.https://ddo /02
Nikhil Barhate. Minimal implementation of decision transformer. https://shooturrl.c attYelpppla.//tllllll 2022.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishhan, Karol Hausman, Alex Herzog, Jasmine Hsu, et Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.0688 2022.
10

---

Lichang Chen, Jiuhai Chen, Tom Goldstein, Heng Huang, and Tianyi Zhou. InstructZero: Effi- cient instruction optimization for black-box large language models. In Proceedings of the 41st International Conference Machine Learning, volume 235, pp. 6503-6518, 2024.
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,
2021.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International conference on machine learning, PMLR, 2016.
pp.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International conference on machine learning, pp. 1126-1135. PMLR, 2017.
Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arXiv preprinl arXiv:2305. 09648,, 2023.
Shengchao Hu, Wanru Zhao, Weixiong Lin, Li Shen, Ya Zhang, and Dacheng Tao. Prompt tuning with diffusion for few-shot pre-trained policy generalization. arXiv preprint 2rX 2024.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning, Advances in neural information processing systems, 33:1 33:1179-1191, 2020.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint 2XXi:2220..121155 2022.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint wwXiv:2005.016433 2020.
Lihong Li. Wei Chu, John Langford, and Robert Schapire. contextual-bandit approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pp. 661 661-670, 2010.
Xiaoqiang Lin, Zhaoxuan Wu, Zhongxiang Dai, Wenyang Hu, Yao Shu, See-Kiong Ng, Patrick Jaillet, and Bryan Kian Hsiang Low. Use your instinct: Instruction optimization using neural bandits coupled with transformers. In NeurlPS 2023 Workshop on Instruction Tuning and Instruction Following, 2023.
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, pp. 7780-7791. PMLR, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Alec Radford. Improving language understanding by generative pre-training. 2018.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. Transactions Machine Learning Research, 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06344 2017.
Chengshuai Shi, Kun Yang, Jing Yang, and Cong Shen. Best arm identification for prompt learning under limited budget. In ICLR 2024 Workshop on Understanding of Foundation Model, 2024a.
Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon Shaolei Du, and Huazhe Xu. Unleashing the power of pre-trained language models for offline reinforcement learning. In The Twelfth International Conference Learning Representations, 2024b.

---

Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. A Bradford Book, Cambridge, MA, USA, 2018 ISBN 0262039249.
William R Thompson. On the likelihood that one unknown probability exceeds another view of the evidence of two samples. Biometrika, 25(3/4):285-29) 1933.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033.
IEEE, 2012.
Zhe Wang, Haozhu Wang, and Yanjun Qi. Hierarchical prompt decision transformer: Improving few-shot policy generalization with global and adaptive. arXiv preprint arXiv:2412.0097,, 2024a.
Zhi Wang, Li Zhang, Wenhao Wu, Yuanheng Zhu, Dongbin Zhao, and Chunlin Chen. Meta-dt: Offline meta-rl conditional sequence modeling with world model disentanglement. Advances in Neural Information Processing Systems, 37:44845 44870, 2024b.
Zhihui Xie, Zichuan Lin, Deheng Ye, Qiang Fu, Yang Wei, and Shuai Li. Future-conditioned unsupervised pretraining for decision transformer. In International Conference on Machine Learning, pp. 38187-38203. PMLR, 2023.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learning, pp. 24631-24645 PMLR, 2022.
Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyper-decision transformer for efficient online policy adaptation. 2023.
Yu Yang and Pan Xu. Pre-trained language models improve the few-shot prompt ability of decision transformer. In Workshop on Training Agents with Foundation Models at RLC 2024, 2024.
Tianhe Yu, Deirdre Quillen, Zhanpeng Hc, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, PMLR, 2020.
Haoqi Yuan, Yuhui Fu, Feiyang Xie, and Zongqing Lu. Pre-trained multi-goal transformers with prompt optimization for efficient online adaptation. Advances in Neural Information Processing Systems, 37:55086-55114 2024.
Hongling Zheng, Li Shen, Yong Luo, Tongliang Liu, Jialie Shen, and Dacheng Tao. Decomposed prompt decision transformer for efficient unseen task generalization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Dongruo Zhou, Lihong Li, and Quanquan Gu. Neural contextual bandits with ucb-based exploration. In International Conference on Machine Learning, pp. 11492- 11502. PMLR, 2020.
12

---

Supplementary Material:
A Prompt-Tuning Algorithm
We provide the full pseudocode of our method. The overall prompt-tuning procedure outlined in Algorithm _ prompt selection via the structured bandit is shown in Algorithm 2 and Algorithm details the update of the bandit's internal reward models.
Algorithm Serompt-Tuning Bandit
Input: Pre-trained PDT parameter Simulator M,, expert demonstrations Initialize: Bandit parameter Φ", dataset B +-
for bandit steps & K do 2 Select prompt Pl: for current rollout & using parameters o² 3: Performance metric G =0
2: Algorithm
4: for MDP steps E T do
5: Make PDT imput (Equation B): X = Pk.
O whtt
TT" (xx:0) 6: Sample action from PDT: Mi(s,,
7: Step environment: rt, S
8: Log reward: G.
G +TI
9: end for
Store data B BU (pk. G
Algorithm 3 Update of using B, yielding d++1
12: end for
Return: Final bandit parameter 6K
Algorithm 2 Prompt Selection
Algorithm 3 Bandit Update
Input: Expert demonstrations P., bandit parameter é (011... ..)
Input: Bandit dataset B X. y, bandit pa-
rameter @ (Ø1, ... ",))), learning rate
Initialize: Prediction matrix Y
1: for segment 7 E P, do
1: for reward model E J do Get segments at j-th index X, X()
2: Predict reward (6) (7)..... .... .,. 2
prediction matrix Y
Predict reward y, c (X,)
3: Append
4: Define L(o,) = MSE y y.
4: end for Return: P PPllarg max( Y[j.:))] Vj [}
for gradient steps of ((
... do
end for
8: end for
Return: New parameter (0f Φ5)

---

B Environment Details
This environment involves mixed control of planar point agent, where the state representation consists of the agent's s current 2D coordinates. The agent always begins at (0, 0) with the objective of reaching (hidden) goal coordinate. has two continuous actions for movement across the plane and a binary stop action. The task requires the agent to navigate to the goal coordinate and appropriately select the stop action. sparse reward is provided upon selecting the stop action, proportional to the agent's distance from the goal. When stop is selected within close proximity of the goal, the environment provides a reward bonus of 10, which discounted based on the number of wasted steps. While the goal is not explicitly part of the state, it is implicitly encoded through the reward function, ensuring that each individual task remains a fully observable MDP.
Sparse 2D point:
The environment offers a continuous task space, parameterized by the angle and radius, for arbi- trary goal locations on the 2D plane. We discretize the task space by using three discrete radii, (0.9, 1.9, 2.9), and 20 discrete angles (0.0 0.1 ...... 1.9 x). instead of sampling continuous task parameters. This primarily done to separate datasets for different tasks, since PDT requires expert demonstration for each training task. To separate these 3 20 60 tasks into the training set T and testing set T tasks with an angles greater than 1.5 a (independently of the radius) are treated as testing task and are not part of the training set. This split yields 48 training tasks and 12 testing tasks. Spatially, the test set indicated by the shaded area in Figure R
This task involves continuous joint control of a planar half-cheetah agent, The state and action spaces have dimensions 20 and respectively. Tasks vary by target velocity, and the reward is proportional to the deviation from this target. Although the target velocity is not part of the observable state, it implicitly encoded via the reward, ensuring that each task defines fully observable MDP. There are 40 tasks in total, 35 are used for training, and 5 are used for testing. For more environment details, see (Duan al. 2016 Xu al., (022).
MuJoCo Half Cheetah:
MuJoCo Ant: As in Xu et al. (0022). this task requires continuous control of an ant agent via eight joint actuators. The 27-dimensional state space includes positions and velocities of the agent's body. Tasks involve moving in different directions, with reward based on velocity along the target heading, There are 50 tasks total, 45 are used for training, and are used for testing. For more environment details, see Duan et al. 2016); Xu et al. (202)).
Meta-World Pick-Place: Following Yu et al. 2020), this task involves Cartesian control of simulated Sawyer robot's end-effector. The d-dimensional action space includes 3D position deltas and gripper torque. The 39-dimensional state space encodes gripper and object positions, and quaternions, but excludes the goal location. The benchmark consists of 50 task variations with different placement goals; are held out during training for evaluation.

---

C Training Details
All experiments were conducted on an instance equipped with an NVIDIA T4 GPU with 8GB of memory, utilizing the PyTTrcch library (Paszke et al. 2019). Pre-traing Pre-training the multi-task PDT backbone took approximately 5 hours. Performing 250 online rollouts and prompt tuning took approximately 30 to 120 minutes per task, depending environment and algorithm.
Details of the hyperparameters for DT and PDT models are provided in Table while those for the bandit model are listed in Table 4 We used standard values for each hyperparameter instead of performing an expensive hyperparameter optimization. invironmentt-peccfic hyperparameters are reported separately in Table The implementations of DT and PDT were adopted from [Barhate (2022) with minimal modifications to integrate seamlessly with the CMAB prompt-tuning framework. The MuJoCo experiments were performed by integrating PDT repository Xu et al. (2022). For the hill-climbing and ZORankSGD (Hu et al. (023) baselines, sample initial prompt randomly from the expert demonstration dataset for the target task, then apply Gaussian noise and perform hill climbing. To encourage convergence to a final prompt, we linearly anneal the scale of the exploration noise from 1.1 to 0.1 over the course of the 250 online rollouts.
prompt-tuning bandits into the official
Hyperparameter
Value
Hyperparameter
Value
Number of transformer blocks
Batch size
All data
Number of attention heads
1
Learning rate
le-3
Embedding dimension
128
Evaluation trials
250
Transformer activation function
GELU
in e-greedy
0.1
MLP activation function
ReLU
(exploration parameter) in UCB
3
MLP hidden layers
2
Reward model type
MLP
MLP width
128
MLP hidden layers
2
Batch size (per task)
8
MIP width
16
Learning rate Learning rate decay weight Optimizer
le-4
MI.P activation function Optimizer
ReLU
le-4
Adam
AdamW
Table 4: Hyperparameters of Contextual Multi-
Table 3: Hyperparameter values for PDT and DT models.
armed Bandit
Environment
Target Return
Prompt Length (J H)
Context Length
DD-Point-Sparse
10
(1 x 3), (2 x 3). (4 x 3)
Cheetah-vel
x 5), (2 20)
20
Table 5: Environment-speccfic hyperparameters of DT and PDT.
3

---

D
Additional Results
D.1
Online Rollouts in Sparse point
Total reward
=1
Total reward
=
Total reward
20
20
20
-2
-20
-20
100
200
100
200
100
200
Online rollouts
Online rollouts
Online rollouts
(a) Online performance of standard PDT, without prompt-tuning
Total reward
Total reward
==2
Total reward
4
20
20
20
0
20
-2
-20
100
200
100
100
200
Online rollouts
Online rollouts
Online rollouts
(b) Online performance with prompt-tuning, using e-greedy exploration and transformer features & segment representation for the bandit's reward models.
Total reward
J=1
Total reward
==2
Total reward
==
20
0
-0
20
-2
0
100
200
100
200
100
200
Online rollouts
Online rollouts
Online rollouts
(c) Online performance with prompt-tuning, using UCB exploration and transformer features & as segment representation for the bandit's reward models.
Total reward
Total reward
==2
Total reward
4
20
20
20
0
-2
-20
-20
200
Online rollouts
Online rollouts
Online rollouts
(d) Online performance with prompt-tuning, using e-greedy exploration and raw trajectory segments instead
of transformer features for the bandit's reward models.
Total reward
==1
Total reward
==2
Total reward
= 4
20
20
20
0
-20
-20
-20
0
100
200
100
200
100
200
Online rollouts
Online rollouts
Online rollouts
(e) Online performance with prompt-tuning, using UCB exploration and raw trajectory segments instead of transformer features for the bandit's reward models.
Total reward
-=1
Total reward
=
Total reward
=
20
20
20
mw
-20
-0
-20
0
100
200
100
200
100
200
Online rollouts
Online rollouts
Online rollouts
(f) Online performance with prompt-tuning, using Gaussian noise and hill climbing to optimize prompt selection.
The dashed line marks the optimal return +10, the shaded area corresponds to standard deviation around the mean.
Figure 5: Plots of data used Table 3b

---

==2
20
20
20
T0Ca reward
TOCK reward
100
150
200
250
50
100
150
200
250
100
150
200
250
Online rolouts
Online rolluuss
Online rolouts
(a) Online performance after 250 episode with prompt-tuning, using ZORankSGD for prompt tuning.
20
20
20
Tobal roward
Total reward
Tota reward
-20
-20
-20
200
600
800
1000
1200
200
400
900
800
1000
1200
200
400
600
800
1000
1200
Online rolouts
Onlnn rolluus
Online rolluuts
(b) Online performance after 1500 episodes with prompt-tuning, using ZORankSGD for prompt tuning.
Figure 6: Plots of ZORankSGD prompt-tuning as reported in Table 3b The dashed line marks the optimal return +10. The shaded corresponds one standard deviation around the mean. ZORankSGD requires the online evaluation of m = 5 prompt perturbations to compute one prompt update. Results using budget of 250 online rollouts (and 50 prompt improvement steps) shown in subfig (a). Results using 1250 online rollouts (and 250 prompt improvements) are shown subfig (b).

---

D.2 Prompt Quality Experiment for the Sparse 2D point Environment
Can prompt-tuning exploit non-expert datasets? While the PDT model with random prompt selection strategy relies on access to expert demonstrations, our method demonstrates greater robust- ness to the quality of demonstration data. To examine this, we generate additional prompt datasets (P E {0, 10, 100). which combine 1% of expert data along with novice demonstrations, selected from trajectories in the bottom 5th percentile of task returns (see Figure 2b).
When sampling prompts from these mixed datasets, the bandit-based prompt optimization significantly enhances the PDT model's robustness and improves its performance, shown in Table 6 This approach reduces reliance on pure expert demonstrations by learning to identify the optimal prompt from any arbitrary mixture dataset.
Expert percentage 5%
No tuning
e-greedy
UCB
0%
-3944 16.8
-12.1 = 19.4
-3.8 10.0
10%
-40.8 14.1
4.4 3.4
7.3
20%
-49.6 32.8
5.5 3.4
9.4
30%
-50.1 27.4
6.5
9.8+0.3
40%
41.4 26.4
2.1 6.4
8.6
50%
4.0
5.1
8.7
60%
5.0
1.5
8.2
70%
-28.6 17.0 17.0
0.4
1.3
80%
-12.3 11.1 11.1
8.6
9.8 = 0.3
90%
0.9 = 0.5
8.6 0.9
8.6 1.0
100%
0.7 1.7
9.7 0.4
9.7 0.4
Table 6: Without prompt-tuning, PDTTs performance deteriorates with the percentage of expert trajectories in P,. Our prompt-tuning method is robust with respect to the percentage of expert data and achieves near-optimal performance with as little 10% expert demonstrations, Results are averaged over three seeds for single training task.
Optimal
40%
60%
100%
0-
Total reward
reward
-25
-25
-25
-50
-50
Total
-50
-75
-75
-75
-100
20
40
60
80
20
60
80
20
60
80
Online rollouts
Online rallouts
Online rollouts
(a) Baseline PDT, no tuning
(b) e-grecdy prompt tuning
(c) UCB prompt tuning
Figure Prompt quality experiment. Color indicates percentage of expert demonstrations P; Without prompt-tuning, PDTTs performance degrades and is roughly proportional to the percentage of expert demonstrations the dataset. Our prompt-tuning method quickly recovers and converges to near-optimal performance by finding the high-performance prompts in the mixture dataset. The shaded region denotes standard deviation around the mean, averaged over three random seeds for single training task.

---

E Segment Independence Attention Analysis
The key assumption based on which we design our structured bandit architecture is that prompt segments contribute independently and additively to prompt informativeness. This assumption implies that PDT can identify the target tasks by attending key (f,s, s. a)-pairs the prompt, as opposed to deriving task identity from the interaction effects between prompt segments, or by attending to the prompt globally. This assumption crucial, because it allows us to decompose the prompt tuning problem into bandit that maintains J reward models, one for each prompt segment, which results in reduction of the size of the exploration problem from combinatorial basis for the assumption is that a) many MDPs can be characterized by their optimal state-action marginal (i.c., by key (s, a) -pairs) and b) that PDT's pre-training paradigm discourages reliance segment interaction or arrangement, since prompt segments sampled uniformly at random during pre-training, not by enforcing order or other constraints between segments. This training setup implicitly encourages invariance permutation and inter-segment dependencies, and thus, believe that this limits the influence of global prompt structure on the model's behavior.
to linear in The
To provide additional support for this interpretation, we conduct an analysis of PDT's attention weight on our 2D environment. Here, for each task, we construct prompts that contain exactly one segment that is highly informative for task identification, based our domain knowledge, while segments carry little information for discriminating between tasks. We then performed rollouts using different permutations of these segments prompts. The performance of PDT remained robust under those permutations with 6.27 + 0.34 return on average over tasks, prompts, and model configurations (J 2, H = 3 and J 4, H = 3), which implies prompt segment permutation invariance PDT. We visualize a representative PDT attention mask for different prompt permutations in Figure 8 to confirm qualitatively that PDT indeed identifies tasks by attending mostly to the individual, informative prompt segment.
remaining
Furthermore, we quantify attention to each prompt segment by computing the mean of the token attention score of tokens each prompt segment. We then sum the per-timestep mean segment attention scores, and find that in 99% of rollouts conducted during this analysis, PDT assigned the highest attention to the informative segment. Additionally, we manually set attention weights for specific prompt segments to zero, and find that PDT's performance remains stable when uninformative segments are masked out, while it degrades drastically when the informative segment is removed, as per Table 7
Full prompt Only informative segment Only uninformative segments -34.0 + 4.49
6.27 + 0.34 Table 7: PDT performance with different attention scopes.
5.88 + 0.23
These results suggest that PDT behavior is largely driven by the most informative segment, and that inter-segment interactions are negligible in practice, which supports and justifies our assumption and bandit architecture. However, we also observed instances where PDT attended to all prompt segments rather uniformly, despite J - segments not being useful for task discrimination per our domain knowledge. We hypothesize that these instances were due to slight overfitting to the prompt dataset, and see the design of regularization that prevents PDT from exploiting such nuances in the prompt dataset as important future work.
7

---

Prompt Segment < 2 +>
Prompt Segment x
Segment <
History
Querry token index
0.3
Attention weiaht
35
Online state idxs
0.2
40
0.1
0.0
10
20
30
50
Causal token attention
(a) Prompt segment is informative.
0.4
30
Prompt Segment
Prompt Segment
Prompt Segment
History
Querry token index
0.3
Attention weight
Online state idxs
0.2
40
45
0.1
0.0
10
20
30
40
50
Causal token attention
(b) Prompt segment informative.
0.4
30
Prompt Segment
Promet Segment
Prompt Segment
History
Querry token index
0.3
35
0.2
Attention
0.1
50
0.0
Causal token attention
(c) Prompt segment is informative.
Segment
gingent
Segment
Querry token index
2 >
0.3
Attention weight
35
Online state idxs
40
0.2
0.1
50
0.0
30
Causal token attention
(d) Prompt segment 4 is informative.
Figure 8: Attention weights from the first transformer block (after softmax), darker values indicate higher attention scores. The first 20 rows have been cropped to improve readability. The blue horizontal lines indicate the indices of state tokens at times H, H + 1,... from which action the corresponds is predicted. It can be that PDT attends mostly to the informative prompt token, independently of it's position or neighboring segments.

---

Method
MuJoCo Ant (J 1, = 5)
MuJoCo Ant (J 2, H=20)
PDT (full-model fine-tuning)
28m
38m
Hill climbing
16m
24m
ZORankSGD
52m
120m
TS
28m
34m
TS
33m
172m
Table 8: Wall-clock inference times with different methods.
Wall-clock inference time
Table reports representative wall-clock inference times for different methods on the MuJoCo Ant environment. Comparing bandit architecture with PDT-encoded prompt representations (denoted by to to its raw unencoded counterpart highlights the efficiency benefits of using the PDT as a feature extractor. When operating PDT-encoded prompts, the Thompson Sampling bandit maintains low inference times even for large prompts (J = 2, H 20), since the PDT'-encoded representation has a fixed dimension Idd determined by the PDT's token embedding space, independent of the prompt length. In contrast, raw unencoded prompts lead to input sizes that scale H x ([S) + |A| + 1) for each segment, substantially increasing inference time due to the higher cost updating the reward models.
G Proof of Theorem 4.1
Theorem. Assume that the reward function G: pJ_ - R for prompt (fi...., F.) decomposes as the mean of J independent reward models (5(f)):
G(p) - - Σ (5 (.))
and that the interaction term is uniformly bounded by (h(Fi....,,.) S VF, E P. Let p" -,] denote the optimal prompt, and suppose that for each slot j. bundit algo- rithm guarantees a slot-specity regret
K Regret (Kgret,(KK)) Σ -
over rounds. Then the cumulative regret after K rounds is bounded as:
K E - - - Regret (K) + 2Ks.
Regret (K) 4
Proof. We begin with the decomposition of the reward function. The instantaneous regret at round
Moi......)) .3)
Σ
(,( + M(Ti.
ト
Σ j=1
(Te.)) + N(TT.1...., (T.)))
J
j=1
Rearranging terms yields:
( + | - [..)).
G(p') - G(p:) - 1=1
By the bounded interaction assumption, we have:
((Ti..... ,,)) (T......,,7.))))))) 22.

---

Thus, taking expectations and summing over = 1 to K, the cumulative expected regret satisfies
K
Regret(K)
E B G(p))
t=1
K
E ( - ()(()))) + 2s
[=1
J=1
E [og(T]) o5(T,s)] + 2KE.
By definition, the slot-specific cumulative regret is given by
K
Regret, (K) =
[T}( - ()(),
so we obtain the final bound:
Regret(K) S J j=1
Regret, (K) + 2Ke.
10

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The contributions outlined in the abstract and introduction consistent with
the prompt-tuning bandit method Section the evaluations in Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the main limitations of our method, and potential remedies, in
Section 6 Potential limitations in the evaluation procedure are discussed Section 5
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?

---

Answer: [Yes]
Justification:
Our theoretical result the regret bound of our bandit architecture in Theo-
rem 4.1
We provide proof for this regret bound in Section G in the Appendix and state all assumptions.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide the complete pseudocode of our method in Section A in the Appendix. We provide a detailed description of our custom Sparse 2D point environment in Section Training details and hyperparameters are stated in Section Clin in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
12

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide codebase with instructions for reproducing each experiment https://githin cow/ hatpg//pt-bandidi bandits
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ httppp.ccppubcccguuuceeeee Codeeubmii ssi criionPooiccy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We state the split between in-distribution training tasks and -isttiiiittiittinn testing tasks for each environment in Section B in the Appendix. Hyperparameters stated in Section[[]
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the
the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report all results in terms terms of the mean + one standard deviation, either in tabular form as confidence intervals in plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
13

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the hardware specification and approximate execution time for our experiments in Section C
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurrpps wwtpp///bwwccccceccciiieeeeeeeee
Answer: [Yes]
Justification: We have reviewed and adhere the NeurIPS Code of Ethics. We experiment mainly on publically available meta RI. datasets (Yu et al., 2020 Duan et al. (0166). additional experiments were conducted on simulated 2D environment and on a synthetic bandit benchmark.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
14

---

We are contributing foundational method to improve the performance of PDT models and can therefore not estimate the true societal impact. However, we don't see any application deployment with negative societal impact resulting directly from this work. Guidelines:
Justification:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
This work does not release data or models with high risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All used assets (public datasets, algorithm implementations) are cited appro- priately.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
15

---

The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: The source code released along with our paper is properly documented and contains the license terms.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work did not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This work did not involve human or animal subjects.
Guidelines:
16

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
This work does not involve any LLM components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy httpps//haarrpps.c///ofeeeeeeee for what should or should not be described.
17


Reviews:
[{"review_id": "eNcfFacbrD", "reviewer": "Reviewer_1Mcn", "summary": "In this paper, the main goal is to improve prompt tuning for Decision Transformers (DTs) in offline reinforcement learning. The authors propose a bandit-based prompt-tuning method. This new approach can select optimal trajectory prompts at inference time using a structured bandit framework and scales linearly with prompt size. The empirical evaluations are conducted over benchmarks and demonstrate the applicability of the model in real-world scenarios.", "questions": "see strengths and weaknesses", "limitations": "see strengths and weaknesses", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper is well-organized, and the illustrative example and figures are easy to understand \n2. The paper addresses an important research problem in the field of offline RL via DT. \n3. The use of the bandit algorithm in the prompt-tuning procedure is an interesting idea. \n4. The authors conduct a comprehensive set of experiments to evaluate the performance. \n\nWeaknesses: \n\n1. The motivation for using DT in offline RL problems and the introduction to the underlying DT framework need to be further explained. \n\n2. The selection of the prompt is heavily dependent on the online return. How to ensure the quality of the estimated reward and causal correlation with the bandit selection. \n\n3.  In general, does the independence of the segments negatively impact the prompt reward? \n\n4. In the case of large $J$, how to address the scalability and weak signal in the bandit selection procedure for arms. \n\n5. The regret analysis is only for the bandit algorithm, but it is not for the optimal generalized policy\n\n6. As the authors discussed OOD, what are the technical insights of the proposed model that perform well in this case?", "was_revised": true, "final_justification": "I am grateful for the author's detailed responses.", "rebuttal": "We thank the reviewer for their thoughtful comments and constructive feedback. We are particularly grateful for their recognition of the significance of our method and the value of the empirical insights presented. In the following, we address each of the reviewer’s concerns and questions point by point:\n\n## **W1.** Clarify motivation.\n\nWe appreciate the feedback and will revise the Introduction and Background sections to improve clarity and motivation. We will highlight why PDT is effective for offline, multi-task pre-training without relying on policy iteration or value functions. We will also emphasize how prompting supports fast, reliable adaptation to downstream tasks compared to full-model fine-tuning, which is more sensitive to distribution shift.\n\n## **W2.** Causal relation between bandit and PDT.\n\nWe thank for the comment and would like to clarify that the reward feedback for the bandit is the actual return obtained by the pre-trained PDT when rolled out with a given prompt. Since this is the performance metric we aim to maximize, it provides a direct and causally grounded reward signal for the bandit.\n\nThe goal of the bandit is to identify the prompt that elicits the highest return as efficiently as possible, minimizing regret over rollout budget. The feedback signal (i.e., return) is thus fully aligned with the optimization objective, and does not rely on proxy metrics or estimated values. As such, we believe the return is both reliable and causally valid for guiding prompt selection.\n\nWe hope this addresses the concern, and we welcome further clarification if we’ve misunderstood the intent of the question.\n\n## **W3.** Does the assumption of independent segments negatively impact the prompt reward?\n\nThank you for raising this important point. Our method assumes that prompt segments contribute independently to the return -- an assumption that could, in theory, lead to performance degradation if inter-segment interactions were critical. However, both the design of PDT and our empirical analysis support the validity of this assumption in practice.\n\nFirst, as described in the Preliminaries, PDT is pre-trained using independently sampled prompt segments drawn from expert trajectories. No ordering or structural constraints are enforced across segments. As a result, the model is unlikely to rely on cross-segment dependencies and instead learns to attend to individually informative $(\\mathbf{s}, \\mathbf{a})$ pairs.\n\nTo test this empirically, we conducted an additional study in the 2D environment:\n- For each task, we constructed prompts containing one highly informative segment and $J - 1$ less informative ones, then permuted their order across rollouts.\n- **PDT performance remained stable**, with a return of $6.27 \\pm 0.34$, **indicating permutation invariance**.\n- Analyzing attention weights, we found that **in 99% of rollouts, PDT focused on the informative segment**, regardless of position.\n- When we masked attention to prompt segments, performance dropped sharply if the informative segment was excluded, but remained near-optimal when less informative segments were masked out:\n\n| Attention Scope | Return |\n| ---- | ---- |\n| Full prompt | $6.27 \\pm 0.34$ |\n| Only informative segment | $5.88 \\pm 0.23$ |\n| Only uninformative segments | $-34.06 \\pm 4.49$ |\n\nThese results suggest that PDT behavior is largely driven by the most informative segment, and that inter-segment interactions are negligible in practice.\n\nFinally, in Section 5.3, we explicitly measure the interaction term $h(\\cdot)$, which quantifies the deviation from full segment independence, and find that it remains small across problem sizes. Combined with strong empirical performance across tasks, we believe this offers strong support for the independence assumption underlying our bandit design.\n\n## **W4.** How to address the bandit scalability with large $J$?\n   \nScalability with respect to the number of prompt segments $J$ is a central consideration in our design, since the total prompt space grows combinatorially with $J$. Our method addresses this through two key architectural choices:\n1. Structured decomposition: We maintain $J$ independent reward models, one per prompt slot. This reduces the effective search space from combinatorial in $J$ to linear, making the exploration problem significantly more tractable.\n2. Contextual modeling: Our use of a contextual multi-armed bandit allows the model to exploit feature-based similarities between prompt segments. This shared representation improves generalization across arms, which is particularly important in settings with large $J$ or sparse downstream rewards.\n\nWe further demonstrate in Section 5.3 that our architecture remains effective at scale. In synthetic experiments, our method handles prompt spaces with up to $|P|^J = 10^5$ combinations and maintains sublinear regret, validating both the scalability and sample efficiency of the approach.\n\nIn summary, while larger $J$ naturally increases the search space, our structured and contextual bandit design is specifically built to address this challenge.\n\n## **W5.** No regret analysis for the optimal policy.\n   \nIndeed, the formal regret bound we present is defined with respect to the bandit's ability to identify high-return prompts. However, because the bandit selects prompts that are passed directly to the frozen PDT, the regret incurred by the bandit translates directly into regret in the downstream return of the policy.\n\nThus, our regret analysis characterizes how efficiently the combined system of bandit + fixed PDT converges to optimal behavior in the prompt-tuning setting, which is the focus of our contribution.\n\n## **W6.** Why does the proposed method perform well OOD?\n      \nOur main insight is that bandit-based prompt-tuning identifies optimal trajectory prompts reliably and efficiently, without requiring prior knowledge of task structure, reward function, or dynamics. This is especially useful OOD, where knowledge about the task is often limited.\n\nEmpirically, we find that our bandit-based prompt-tuning consistently and considerably outperforms standard PDT and ZORankSGD prompt-tuning in OOD settings, as shown in the following table:\n\n|             | Half Cheetah (J=1, H=5)      | Half Cheetah (J=2, H=20)     | Ant (J=1, H=5)               | Ant (J=2, H=20)              | Pick-Place (J=1, H=5) | Pick-Place (J=2, H=2) |\n|:------------|:-----------------------------|:-----------------------------|:-----------------------------|:-----------------------------|:----------------------------|:----------------------------|\n| PDT         | -64.78 $\\pm$ 36.91           | -40.95 $\\pm$ 43.19           | 363.90 $\\pm$ 105.42          | 360.07 $\\pm$ 72.36           | 502.80 $\\pm$ 63.98          | 524.37 $\\pm$ 39.56          |\n| ZORankSGD   | -59.85 $\\pm$ 32.37 (+7.6\\%)  | -36.60 $\\pm$ 19.45 (+10.6\\%) | 383.57 $\\pm$ 193.35 (+5.4\\%) | 340.68 $\\pm$ 44.60 (-5.4\\%) | 503.56 $\\pm$ 66.16 (+0.2\\%) | 538.05 $\\pm$ 31.46 (+2.6\\%) |\n| Bandit (best) | -29.26 $\\pm$ 21.25 (+54.8\\%) | -23.93 $\\pm$ 14.14 (+41.6\\%) | 480.85 $\\pm$ 84.76 (+32.1\\%) | 441.44 $\\pm$ 80.25 (+22.6\\%) | 551.33 $\\pm$ 34.52 (+9.7\\%) | 553.34 $\\pm$ 18.82 (+5.5\\%) |\n\nNote that we have updated the results for the ML1 Pick-Place environment after identifying an issue in the official PDT codebase that affected PDT pretraining on this benchmark. The corrected results are now included in the revised tables and are fully consistent with the trends observed on other environments:\n\n| Method   / Env                   | ID Pick-Place (J=1, H=5)   | ID Pick-Place (J=2, H=2   | OOD Pick-Place (J=1, H=5)   | OOD Pick-Place (J=2, H=2)   |\n|:---------------------------|:---------------------------|:--------------------------|:----------------------------|:----------------------------|\n| CQL oracle                 | 535.84 $\\pm$ 31.02         | 535.84 $\\pm$ 31.02        | 525.07 $\\pm$ 60.15          | 525.07 $\\pm$ 60.15          |\n| PDT, no tuning             | 551.58 $\\pm$ 26.09         | 535.52 $\\pm$ 24.86        | 502.8 $\\pm$ 63.98           | 524.37 $\\pm$ 39.56          |\n| PDT, finetuned             | N/A                        | N/A                       | 495.37 $\\pm$ 57.87          | 488.17 $\\pm$ 50.15          |\n| Hill-climbing              | 555.79 $\\pm$ 22.72         | 540.15 $\\pm$ 23.49        | **560.92 $\\pm$ 27.04**      | 544.19 $\\pm$ 28.83          |\n| ZORankSGD                  | 554.26 $\\pm$ 23.00         | 537.20 $\\pm$ 25.5         | 503.56 $\\pm$ 66.16          | 538.05 $\\pm$ 31.46          |\n| $\\epsilon$-greedy$^{\\Psi}$ | 555.35 $\\pm$ 24.15         | 541.32 $\\pm$ 22.91        | 531.49 $\\pm$ 49.86          | 552.11 $\\pm$ 23.83          |\n| TS$^{\\Psi}$                | **556.87 $\\pm$ 24.11**     | 540.82 $\\pm$ 22.83        | 549.38 $\\pm$ 36.12          | 553.12 $\\pm$ 20.81          |\n| UCB$^{\\Psi}$               | 552.68 $\\pm$ 24.77         | 538.84 $\\pm$ 23.14        | 506.43 $\\pm$ 65.38          | 539.42 $\\pm$ 31.12          |\n| $\\epsilon$-greedy          | 556.22 $\\pm$ 25.16         | **541.80 $\\pm$ 23.76**    | 530.26 $\\pm$ 51.61          | 550.06 $\\pm$ 22.38          |\n| TS                         | 556.11 $\\pm$ 24.56         | 541.33 $\\pm$ 22.79        | 551.33 $\\pm$ 34.52          | **553.34 $\\pm$ 18.82**          |\n| UCB                        | 554.92 $\\pm$ 26.06         | 538.50 $\\pm$ 23.64        | 512.25 $\\pm$ 69.62          | 534.58 $\\pm$ 35.35          |\n\nWhile all prompt-tuning methods show only modest performance gains on the in-distribution tasks for ML1 Pick-Place, we attribute this to the nature of the dataset: the publicly available prompt trajectories are short, homogeneous, and densely clustered, leaving limited headroom for prompt optimization. \n\nWe believe that the updated results strengthen the overall conclusions and confirm that our method remains effective and competitive across diverse settings, both on ID and OOD tasks."}, {"review_id": "A1rAEqOc3F", "reviewer": "Reviewer_8UyU", "summary": "This paper introduces a prompt tuning method for Prompt Decision Transformers to improve performance in solving multitask and meta-reinforcement learning (meta-RL) problems. Compared with previous methods, it proposes leveraging a Multi-Armed Bandit algorithm to select the best prompt for Prompt Decision Transformers. The paper provides a theoretical analysis of regret and presents empirical results on the HalfCheetah, Ant, Meta-World, and Sparse 2D Point environments.", "questions": "1. Can you elaborate more on the Scalable and Sample-Efficient Bandit Architecture section? How is the contextual multi-armed bandit modeled, and how does the bandit algorithm work in the prompt tuning?\n2. Can you elaborate more on the assumption of section 4.3? Why do you suppose the reward approximation error is bounded by $\\epsilon$?\n3. Can you explain why you use so many bandit algorithms in the experiments? The empirical results show that not all the bandit algorithms can improve the performance.\n4. Can you provide the inference time of the proposed method compared with the baseline methods?\n5. Why use the multi-arm bandit algorithms in the prompt tuning? I think the best arm identification is more appropriate for the prompt tuning. Can you explain more about this?\n6. Can you provide more baseline results, such as the Meta-DT [1]?\n\n\n[1] Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement. 2024", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n\n**Quality:** This paper proposes to model the prompt selection as a contextual bandit problem and uses bandit algorithms to select the best prompt for the decision transformers, achieving good performance. Claims are supported by the experiments compared with baseline methods.\n\n**Clarity:** The paper is well-organized in the methodology and experiments.\n\n**Significance:** The paper provides a novel prompt tuning method for the prompt decision transformer in solving the RL tasks.\n\n**Originality:** This paper demonstrates that prompt tuning based on the bandit algorithm can also improve the performance of prompt decision transformers.\n\n**Weaknesses:**\n\n**Quality:** Although the paper provides the theoretical analysis of the regret for prompt tuning, it assumes that the reward approximation error is bounded by $\\epsilon$ which seems to be a strong assumption. The empirical results show that the proposed method outperforms the baselines. However, it demonstrates that not all the bandit tuning methods can improve the performance. The best results appear across all the bandit algorithms in Table 1 and 2. This observation leads to the question of whether the bandit algorithm can truly improve the performance. I doubt it’s worth spending significant inference time for only a marginal performance gain.\n\n**Clarity:** This paper spends a lot of space introducing the previous methods. It should elaborate more on the Scalable and Sample-Efficient Bandit Architecture section.\n\n**Significance:** The paper admits that the limitation of the combinatorial explosion is on the increasing prompt dataset size. It impedes the practical deployment of the method.\n\n**Originality:** This paper leverages the bandit-based prompt tuning method in the context of prompt decision transformers. The method is a combination of the existing techniques and does not provide any novel perspective or insights on the prompt tuning or decision transformers.", "was_revised": true, "final_justification": "I read the rebuttal and discussion. While some clarifications were helpful, two major concerns remain: (1) missing baselines and (2) an error in the ML1 pick-place environments that affects evaluation validity. Given their impact on significance and reproducibility, I maintain my score and confidence.", "rebuttal": "We thank the reviewer for their thoughtful review and for raising important points about the scope, assumptions, and empirical evaluation of our method. We address each concern below and outline planned additions and clarifications to improve the manuscript accordingly.\n\nPlease note that we have updated results for the **ML1-pick-place** environment due to an issue in the reference implementation. The updated results can be seen in our replies to Reviewers 1Mcn and KYsS, and Tables 1 & 2 have been revised accordingly.\n\n## On the performance gain\n\nWe appreciate the opportunity to clarify both the magnitude and practical significance of our contribution. We emphasize that the performance gains from our bandit-based method are both consistent and substantial across a diverse and challenging set of environments:\n- **33/36** ID tasks and **35/36** OOD tasks (updated Table 2) show improvement over standard PDT with uniform prompt sampling.\n- **All 18 cases** in the 2D environment (Table and Figure 3) show consistent improvements.\n- Our method outperforms ZORankSGD, the strongest inference-time tuning baseline to date, across all environments.\n\nTo illustrate this, we include a summary of representative settings:\n\n|             | HC (J=1, H=5)      | HC (J=2, H=20)     | Ant (J=1, H=5)     | Ant (J=2, H=20)    | Pick-Place (J=1)   | Pick-Place (J=2)   | 2D (J=1)          | 2D (J=2)          | 2D (J=5)          |\n|-------------|--------------------|--------------------|--------------------|--------------------|--------------------|--------------------|-------------------|-------------------|-------------------|\n| PDT         | -64.8 ± 36.9       | -40.9 ± 43.2       | 363.9 ± 105.4      | 360.1 ± 72.4       | 502.8 ± 64.0       | 524.4 ± 39.6       | 0.0 ± 2.1         | 6.3 ± 0.8         | 8.3 ± 0.6         |\n| ZORankSGD   | -59.9 (+7.6%)      | -36.6 (+10.6%)     | 383.6 (+5.4%)      | 340.7 (−5.4%)      | 503.6 (+0.2%)      | 538.1 (+2.6%)      | −0.6              | 4.4 (−30.2%)      | 3.1 (−62.7%)      |\n| Bandit (best)     | -29.3 (+54.8%)     | -23.9 (+41.6%)     | 480.8 (+32.1%)     | 441.4 (+22.6%)     | 551.3 (+9.7%)      | 553.3 (+5.5%)      | 9.9               | 9.9 (+57.1%)      | 9.8 (+18.1%)      |\n\nIn many cases, our method lifts performance from suboptimal to near-oracle levels, especially in high-variance tasks. Furthermore, Figure 3 demonstrates the qualitative gap: while uniform sampling fails to identify informative demonstrations, our method rapidly concentrates on high-quality prompts, significantly improving downstream return.\n\nRegarding efficiency, our method:\n- Requires no weight updates to the PDT backbone.\n- Performs no backpropagation.\n- Operates entirely in inference mode, with optional embedding caching.\n\nThis makes the runtime burden low. Full rollout time comparisons are included in Q4. In summary, our approach offers strong empirical gains at minimal additional cost, making it a practical solution for high-stakes or resource-constrained deployment.\n\n## Clarity of the presentation\n\nWe appreciate the suggestion. In the revision, we will enhance Section 4.2 with additional commentary and examples to make the bandit architecture and prompt selection process more transparent. In particular, we will emphasize the full algorithm pseudocode in Appendix A and reallocate space from the Related Work section if necessary to maintain clarity and making the algorithm more accessible.\n\n## On the practicality in real-world deployment\n\nIndeed, the combinatorial growth of the prompt space poses a challenge for naïve prompt selection strategies. However, our method is explicitly designed to address this issue. By decomposing prompt selection into slot-wise decisions and modeling the reward of each segment independently, we reduce the problem from combinatorial to linear in the number of prompt slots $J$.\n\nWe validate this scalability across high-dimensional environments such as MuJoCo Ant (27 states, 8 actions), HalfCheetah (20 states, 7 actions), and Meta-World ML1 (39 states, 3 actions), where our method consistently improves performance over PDT and prompt-tuning baselines. We further demonstrate that our bandit scales efficiently to prompt spaces with up to $10^5$ candidates, far exceeding the size of typical PDT prompt datasets.\n\nWe believe these results show that our method not only addresses the combinatorial issue in theory but also scales in practice, enabling deployment in realistic offline RL settings.\n\n## On the novelty\n\nWhile grounded in existing components, our method introduces several distinct and novel contributions:\n- We are, to our knowledge, the first to frame inference-time prompt selection in PDT as a structured contextual bandit problem.\n- Our architecture uses independent reward models per prompt position, avoiding combinatorial explosion and enabling parallel, modular learning.\n- It leverages PDT-derived embeddings to enable fast learning even in high-dimensional or pixel-based settings.\n\nThese components form a new class of inference-time adaptive prompt tuning, distinct from prior work in finetuning or gradient-based selection.\n\nWe will revise both the Introduction and Related Work sections to highlight these conceptual distinctions and novel contributions more clearly.\n\n## Q1: Method elaboration\n\nOur method maintains $J$ independent reward models $\\phi_1, …, \\phi_J$, where each model scores prompt segments for a specific slot position. This is motivated by the observation that prompts in PDT are structured concatenations of $J$ segments.\n\nAs shown in Figure 1 and described in Section 4.2:\n- Each $\\phi_j$ is trained to predict returns conditioned on the segment occupying slot $j$.\n- This decomposition ensures modularity and linear scaling, reducing the complexity of learning over full prompts.\n- During inference, a candidate prompt $\\rho_k$ is selected via exploration or exploitation and passed to the frozen PDT for rollout.\n- The rollout return $G(\\rho_k)$ is then used as feedback, and each $\\phi_j$ is updated with the segment it contributed.\n\nFull procedural details are available in Appendix A, Algorithms 1-3.\n\n## Q2: On the bounded $\\varepsilon$ assumption\n\nThe interaction term $h(\\tilde{\\tau}_1, …, \\tilde{\\tau}_J)$, introduced in Equation 5, captures residual dependencies between prompt segments that are not modeled by the additive architecture.\n\nWhile exact theoretical bounds are intractable for real PDTs, we offer empirical support in Section 5.3 using a synthetic benchmark. There we find:\n- The residual term is small and stable across a range of segment pool sizes and prompt lengths.\n- The regret bound in our theorem holds empirically.\n\nFurthermore, we argue that prompt segment informativeness is additive in practice. This is due to:\n- The ability to infer task identity from a few key $(\\mathbf{s,a})$ pairs\n- The random and unordered segment sampling during PDT pre-training, which precludes reliance on inter-segment dependencies.\n\nThese points supports the assumption that $|h(\\cdot)| \\le \\varepsilon$ holds in both theory and practice. We will add a discussion of this reasoning in the revision of Section 4.3.\n\n## Q3: On the bandit exploration strategies\n\nWe included multiple exploration strategies (UCB, Thompson Sampling, and $\\epsilon$-greedy) to demonstrate that the architecture is robust across exploration mechanisms and that the gains are not due to a specific heuristic.\n\nResults:\n- Across 90 total experiments, our method improves over PDT in 86 cases.\n- Even in the few settings where a strategy underperforms, the others compensate, which shows architectural robustness.\n- Thompson Sampling delivers the most consistent gains and is our default choice.\n\nThis broad-based evaluation allows users to choose strategies that best match environment properties (e.g. reward sparsity, task diversity).\n\n## Q4: On the inference time\n\nOur method is computationally lightweight by design:\n- It operates entirely at inference time.\n- Requires no gradient updates or fine-tuning of PDT.\n- Trains only small MLP-based reward models.\n- When using PDT embeddings, segment encodings can be cached and reused.\n\nBelow is the wall-clock time for representative Ant configurations (equal rollout iteration counts):\n\n|             | Ant (J=1, H=5) | Ant (J=2, H=20) |\n|-------------|----------------|-----------------|\n| PDT         | ~28m           | ~38m            |\n| HillClimb   | ~16m           | ~24m            |\n| ZORankSGD   | ~52m           | ~120m           |\n| TS$^\\Psi$   | ~28m           | ~34m            |\n| TS          | ~33m           | ~172m           |\n\nThese numbers confirm that our approach is either competitive or significantly faster than prior baselines while achieving stronger performance. We will include these comparisons in the updated paper to help contextualize our method’s cost-effectiveness.\n\n## Q5: On CMAB vs. best-arm identification\n\nBest-arm identification (BAI) assumes a fixed reward distribution per arm, which is unsuitable for prompt selection:\n- Prompt quality varies dynamically with the rollout context.\n- Rewards are stochastic and contextual, influenced by both task identity and trajectory history.\n- CMABs enable conditioning on rich segment embeddings and incorporate contextual similarity, allowing effective generalization across prompts.\n- Our structured CMAB architecture ensures scalability via linear slot-based decomposition.\n\nWe will elaborate to discuss these differences explicitly and clarify our methodological choice.\n\n## Q6: On the Meta-DT baseline\n\nWe appreciate the reviewer highlighting Meta-DT as a relevant baseline. Due to its architectural and training complexity, we could not incorporate it within the rebuttal timeframe. However, we see Meta-DT as highly complementary: it augments PDT with a contextualized world model, and we believe our bandit could be applied on top of it for additional gains. We plan to explore this integration in future work."}, {"review_id": "bo0SABFcGB", "reviewer": "Reviewer_KYsS", "summary": "This work proposes a method for tuning prompts for Decision Transformers. The idea is to outsource the selection of sub-trajectories to be used as prompts to a multi-arm contextual bandit. Crucially, the method maintains an independent reward model for each hard, thus alleviating the combinatorial growth of the optimization space. Furthermore, the authors suggest that arm fathers can be extracted from the embedding function already existing within the decision transformer. An experimental evaluation is provided, spanning from a simple 2D integrator, to standard mujoco tasks (ant, half-cheetah, meta-world). The method is compared to perturbation-based methods, and demonstrates better performance. Several exploration algorithms for the bandit are compared. Finally, an empirical regret evaluation is carried out to motivate the structured nature of the bandit.", "questions": "### Questions:\nOn top of the ones listed above:\n- Authors assume that a simulator of the environment is available for online evaluation of the policy, but not for data collection. Why is this a reasonable assumption? Trajectories observed during policy evaluation could easily be used for policy improvement.\n\n### Others:\n- equations 2 and 3 should be moved below (between lines 72 and 73), as the notation is not yet introduced.", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths:\n- The submission is very well written, and motivated.\n- The application of bandits for optimizing prompts of decision transformers is, to the best of my knowledge novel.\n- The experiments are well designed and motivated. In particular, Figure 3 nicely illustrated the method at work.\n\n### Weaknesses:\n- The experimental section has contradictory or confusing results:\n    - Baselines appear to be rather weak: ZORankSGD appears to underperform with respect to a naive Hill-climbing strategy in Table 1. Why is this the case? Furthermore, baseline (ii) (line 246) seems to be missing from the table, but I would expect it to represent a bottom line.\n    - If I understand correctly the design of the current hill-climbing baseline, an interesting baseline would be hill-climbing over the (discrete) space of prompts ($\\mathcal{P}_i$). Can the authors comment on whether this would be a reasonable inclusion? Essentially, it would ablate whether the bandit is better at optimizing than a simple hill-climbing routine over a similar space.\n    - The adoption of transformer-encoded trajectories is stated to be beneficial (262). However, average performance seems to decrease in Table 1. Can the authors comment on this?\n    - Results in Table 1 have very large standard deviations, to the point that it’s hard to claim significant improvement. How are the standard deviations computed?\n- The experimental selection of 4 relatively simple tasks could be expanded.", "was_revised": true, "final_justification": "My issues were partially resolved during the rebuttal.\nThe authors explained baselines more in details, and promised to correct an inconsistency between what is claimed in the submission (line 262) and their experimental results. I am not entirely convinced by the way in which performance is aggregated over environments, as standard deviation capture cross-task variability, rather than actual confidence intervals on, say, mean performance. A more informative aggregation, and an evaluation on more environments, would make this submission more convincing.\n\nNevertheless, this paper remains an interesting submission and I would still vote for acceptance.", "rebuttal": "We thank the reviewer for the detailed and constructive feedback. We appreciate the recognition of our paper’s motivation, clarity, and the novelty of applying structured bandit-based prompt tuning to PDT. Below, we respond to the specific concerns raised, and outline how we plan to clarify or extend the paper in the final version.\n\n## W1: ZORankSGD appears to underperform.\n\nWe appreciate pointing this out. Indeed, we observe that ZORankSGD performs similarly to or slightly below hill-climbing across several environments in Table 1.\n\nBoth methods apply black-box optimization over the prompt space using perturbed prompt variants, but differ in how they process feedback:\n- Hill-climbing updates prompts based directly on absolute return values, preserving the full, informative feedback signal.\n- ZORankSGD, by contrast, relies on rank-based comparisons among perturbations for gradient estimation, which discards magnitude information in favor of relative ordering.\n\nThis design may make ZORankSGD more robust in noisy settings, but also less sample-efficient, and explains why it often performs similarly to, or slightly worse than, the simpler hill-climbing baseline in our experiments.\n\n## W2: Hill-climbing over the (discrete) space of prompts?\n\nHill-climbing over the discrete prompt space is an interesting idea, but there are practical limitations that make it challenging and potentially less effective than our proposed approach.\n\nThe space of possible prompts is discrete and combinatorial, growing rapidly with the number of prompt segments $J$ and the size of the prompt dataset $|P|$. Standard hill-climbing requires defining a notion of local neighborhood (e.g., single-segment substitutions), but this is difficult to implement in a high-dimensional, unordered prompt space.\n\nOur method avoids the combinatorial growth by structuring the exploration space: we model each prompt slot independently, and leverage **contextual** bandits to generalize across similar segments. This enables scalable and sample-efficient search, while exploiting structure in the prompt space.\n\nThat said, we appreciate the suggestion. We will clarify this distinction in the revision.\n\n## W3: The adoption of transformer-encoded trajectories is stated to be beneficial, but performance seems to decrease in Table 1.\n  \nWe appreciate the opportunity to clarify. First, we would like to note that we have updated Tables 1 and 2 after identifying an issue in the PDT reference codebase that affected pretraining on the **ML1 Pick-Place** environment. All relevant experiments have been rerun, and the updated results are included in the revision:\n \n| Method / Env| ID Pick-Place (J=1, H=5)   | ID Pick-Place (J=2, H=2   | OOD Pick-Place (J=1, H=5)   | OOD Pick-Place (J=2, H=2)   |\n|:---------------------------|:---------------------------|:--------------------------|:----------------------------|:----------------------------|\n| CQL oracle                 | 535.84 $\\pm$ 31.02         | 535.84 $\\pm$ 31.02        | 525.07 $\\pm$ 60.15          | 525.07 $\\pm$ 60.15          |\n| PDT, no tuning             | 551.58 $\\pm$ 26.09         | 535.52 $\\pm$ 24.86        | 502.8 $\\pm$ 63.98           | 524.37 $\\pm$ 39.56          |\n| PDT, finetuned             | N/A                        | N/A                       | 495.37 $\\pm$ 57.87          | 488.17 $\\pm$ 50.15          |\n| Hill-climbing              | 555.79 $\\pm$ 22.72         | 540.15 $\\pm$ 23.49        | **560.92 $\\pm$ 27.04**      | 544.19 $\\pm$ 28.83          |\n| ZORankSGD                  | 554.26 $\\pm$ 23.00         | 537.20 $\\pm$ 25.5         | 503.56 $\\pm$ 66.16          | 538.05 $\\pm$ 31.46          |\n| $\\epsilon$-greedy$^{\\Psi}$ | 555.35 $\\pm$ 24.15         | 541.32 $\\pm$ 22.91        | 531.49 $\\pm$ 49.86          | 552.11 $\\pm$ 23.83          |\n| TS$^{\\Psi}$                | **556.87 $\\pm$ 24.11**     | 540.82 $\\pm$ 22.83        | 549.38 $\\pm$ 36.12          | 553.12 $\\pm$ 20.81          |\n| UCB$^{\\Psi}$               | 552.68 $\\pm$ 24.77         | 538.84 $\\pm$ 23.14        | 506.43 $\\pm$ 65.38          | 539.42 $\\pm$ 31.12          |\n| $\\epsilon$-greedy          | 556.22 $\\pm$ 25.16         | **541.80 $\\pm$ 23.76**    | 530.26 $\\pm$ 51.61          | 550.06 $\\pm$ 22.38          |\n| TS                         | 556.11 $\\pm$ 24.56         | 541.33 $\\pm$ 22.79        | 551.33 $\\pm$ 34.52          | **553.34 $\\pm$ 18.82**          |\n| UCB                        | 554.92 $\\pm$ 26.06         | 538.50 $\\pm$ 23.64        | 512.25 $\\pm$ 69.62          | 534.58 $\\pm$ 35.35          |\n\nThe results are consistent with the trends observed in other environments. The relatively modest gains on ML1 Pick-Place in the ID setting are due to characteristics of the prompt dataset: prompt trajectories are short, homogeneous, and already near-optimal. This limits the available headroom for further improvement through prompt tuning.\n\nTo support our claim that the pre-trained PDT provides useful embeddings for bandit-based prompt-tuning, we report relative performance differences when using Transformer-encoded prompt segments (denoted as $\\Psi$) and unencoded prompts:\n\n|   ID tasks                            | Half Cheetah (J = 1, H = 5)   | Half Cheetah (J = 2, H = 20)   | Ant (J = 1, H = 5)           | Ant (J = 2, H = 20)          | Pick-place (J = 1, H = 5)   | Pick-place (J = 2, H = 2)   |\n|:------------------------------|:------------------------------|:-------------------------------|:-----------------------------|:-----------------------------|:----------------------------|:----------------------------|\n| Best bandit (no encoding)     | -33.56 $\\pm$ 13.48            | -26.28 $\\pm$ 10.14             | 835.38 $\\pm$ 171.25          | 816.77 $\\pm$ 268.42          | 556.22 $\\pm$ 25.16          | 541.80 $\\pm$ 23.76          |\n| Best bandit ($\\Psi$ encoding) | -36.55 $\\pm$ 11.93 (-8.9\\%)   | -26.87 $\\pm$ 14.10 (-2.2\\%)    | 815.17 $\\pm$ 182.28 (-2.4\\%) | 819.52 $\\pm$ 191.39 (+0.3\\%) | 556.87 $\\pm$ 24.11 (+0.1\\%) | 541.32 $\\pm$ 22.91 (-0.1\\%) |\n\n\n|     OOD tasks                          | Half Cheetah (J=1, H=5)      | Half Cheetah (J=2, H=20)    | Ant (J=1, H=5)              | Ant (J=2, H=20)             | Pick-Place (J=1, H=5)       | Pick-Place (J=2, H=2)       |\n|:------------------------------|:-----------------------------|:----------------------------|:----------------------------|:----------------------------|:----------------------------|:----------------------------|\n| Best bandit (no encoding)     | -29.26 $\\pm$ 21.25           | -25.18 $\\pm$ 22.17          | 480.85 $\\pm$ 84.76          | 438.82 $\\pm$ 125.40         | 551.33 $\\pm$ 34.52          | 553.34 $\\pm$ 18.82          |\n| Best bandit ($\\Psi$ encoding) | -32.61 $\\pm$ 19.85 (-11.4\\%) | -23.93 $\\pm$ 14.14 (+5.0\\%) | 477.24 $\\pm$ 84.64 (-0.8\\%) | 441.44 $\\pm$ 80.25 (+0.6\\%) | 549.38 $\\pm$ 36.12 (-0.4\\%) | 553.12 $\\pm$ 20.81 (-0.0\\%) |\n\n\nWe see this as a clear positive result: it shows that PDT’s latent representations provide a compact and task-relevant embedding of trajectory segments, enabling the reward models to scale more effectively in high-dimensional domains. This is especially valuable because, without using $\\Psi$, the input dimensionality of each reward model would scale as $(|\\mathcal{A}| + |\\mathcal{S}| + 1) \\times H$. Using Transformer-derived embeddings allows us to decouple model input size from the raw environment dimensionality while preserving informative structure. \nThis results in near-constant inference time independently of prompt size, as can be seen in the following table that shows average rollout times for representative configurations in the MuJoCo Ant environment:\n\n\n|             | Ant (J=1, H=5)               | Ant (J=2, H=20)              |\n|:------------|:-----------------------------|:-----------------------------|\n| TS$^\\Psi$   |    ~28m       |  ~34m         | \n| TS  |     ~33m     |      ~172m     | \n\n## W4: Large standard deviations in Table 1.\n\nThe standard deviations in Table 1 reflect variation across many downstream tasks and random seeds, following the evaluation setup. Specifically, we report the mean and standard deviation of the final 10 rollout returns, averaged over three seeds and all tasks.\n\nWe acknowledge that the earlier results for ML1 Pick-Place had particularly high variance. As noted above, we identified and fixed an issue in the reference PDT codebase affecting that environment, and reran the relevant experiments. The updated results exhibit more stable behavior and continue to support the performance trends reported.\n\nWe hope this helps clarify the source of variability and increases confidence in the robustness of our method.\n\n## Q1: Why assume simulator for policy evaluation but not improvement?\n\nThank you for this very insightful question. We agree -- in principle, online evaluation rollouts could be repurposed for policy improvement. In fact, this corresponds to freezing the prompt and fine-tuning the full PDT using the observed trajectories, which we include as a baseline in Table 2.\n\nCompared to fine-tuning, bandit-based prompt selection is significantly more sample-efficient: in practice, optimal prompts are often found within the first 10–20 rollouts (see Appendix D), whereas full-model fine-tuning on the same small dataset leads to degraded performance, as shown in our results.\n\nThis is consistent with observations from the offline-to-online RL literature, where limited online data, distributional shift, and noisy gradients often hinder fine-tuning. Prompt-tuning, in contrast, leverages a strong pre-trained policy and requires only selection, rather than learning, at inference time.\n\nWe believe this makes the simulator-for-evaluation assumption reasonable in settings where rapid, low-cost adaptation is required.\n\n## Q2: Equations 2 and 3 should be moved.  \n\nThank you for spotting this! We've updated the manuscript to move them after the relevant notation is introduced."}, {"review_id": "ZDT5B7c0tw", "reviewer": "Reviewer_kSsg", "summary": "This paper presents a bandit-based prompt-tuning method for the Prompting Decision Transformer (PDT) to enhance task generalisation in offline Reinforcement Learning, particularly in multi-task and few-shot settings. PDT, which utilises trajectory prompts from expert demonstrations, has limitations due to its uniform sampling of prompts, thereby ignoring their varying informativeness. The proposed method optimises prompt selection using a structured bandit architecture, enabling efficient reward modelling and linear scaling concerning prompt size. The method outperforms existing baselines across tasks, high-dimensional environments, and out-of-distribution scenarios without requiring costly weight updates to the underlying Transformer backbone.", "questions": "In Theorem 4.1, the assumption is made that the error of the interaction term between segments, $h(\\tilde{\\tau}_1, ..., \\tilde{\\tau}_J)$, is bounded ($|\\cdot| \\leq \\varepsilon$), and it claims that this assumption holds due to PDT's \"modular structure.\" However:\n\n- The core of the Transformer architecture is the self-attention mechanism, which inherently models cross-segment dependencies. How can we prove that these global dependencies do not significantly increase $\\varepsilon$? \n- If $\\varepsilon$ is indeed large, the $K\\varepsilon$ term in equation (8) would cause regret to grow linearly (especially when $K$ is large). Does this fundamentally undermine the theoretical advantages of the method? \n- Are there any experimental quantifications of the actual $\\varepsilon$ value? For example, by comparing the independent segment reward predictions with the actual return of the complete prompt?\n\nSection 4.4 suggests using the pre-trained PDT model as a feature extractor $\\Psi$, which maps original segments to embeddings:\n\n-  The core objective of this method is to \"quickly tune at inference time,\" but obtaining embeddings for each candidate segment $\\tilde{\\tau}$ requires a forward pass through a computationally intensive Transformer. When $|\\mathcal{P}_i| = M \\times (L - H + 1)$ becomes large (e.g., $M = 50, L = 100, H = 5 \\rightarrow 4800$ segments), does this embedding computation become a new bottleneck? \n- On the other hand, in high-dimensional state spaces (e.g., pixel-based observations), would the $d$-dimensional embeddings need to increase to preserve useful information significantly? Could this cause the $O(\\sqrt{K \\log |\\mathcal{P}|})$ term of the bandit method to deteriorate?", "limitations": "The limitations are discussed in the conclusion.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:  \nThis work introduces a scalable and efficient bandit-based prompt-tuning method for PDT, optimizing prompt selection to improve task performance without requiring costly weight updates. It shows consistent performance gains across various tasks and environments.\n\nWeaknesses:  \nThe core limitations stem from its oversimplified assumption of inter-segment dependency and an insufficient credit assignment mechanism. While it demonstrates better scalability than naive bandit approaches, its application in large-scale multi-task environments remains constrained by sample efficiency issues, computational overhead, and lack of cross-task knowledge transfer.", "was_revised": true, "final_justification": "I appreciate the explanations regarding the segment independence assumption, computational overhead, and the empirical quantification of the interaction term. The additional experiments further strengthen the justification for the independence assumption and demonstrate the robustness of the proposed method. Therefore, I have raised my rating to 4.", "rebuttal": "We thank the reviewer for their thorough and thoughtful review, as well as for raising important points about the core aspects of our method. Below, we address the concerns and suggestions in detail, and outline the modifications we plan to incorporate in the revised manuscript.\n\n## W1: On the segment independence assumption\n\nWe elaborate further in our detailed response below, but in brief: we do not believe the segment independence assumption is oversimplified. Many MDPs can be effectively identified by their optimal state-action marginals, meaning that the inclusion of a few informative $(\\hat{r}, \\mathbf{s}, \\mathbf{a})$ pairs in the prompt ought to be sufficient for downstream adaptation. In such settings, inter-segment interactions are minimal. Empirically, our method consistently outperforms strong baselines across environments (Tables 1-3), indicating that this assumption holds well in practice. In Section 5.3, we further quantify the interaction term $h(\\cdot)$ and find that it remains small, supporting the theoretical justification of our bandit decomposition.\n\n## W2: On the computational overhead\n\nWhen using Transformer-derived prompt embeddings, segment features are computed once per segment and cached, removing the need for repeated forward passes during tuning. The bandit updates themselves involve only small MLPs and incur negligible overhead compared to full-model updates. As shown in the results, our method delivers better performance than full-model fine-tuning, despite being significantly more efficient in wall-clock time.\n\n## W3: Cross-task knowledge transfer\n\nOur setting follows the standard PDT prompt-tuning framework, where a multi-task pre-trained model is adapted to a single target task via prompt selection. In this formulation, knowledge transfer occurs from pretraining to online inference, and our method directly builds on this assumption. While extending the framework to continual or multi-task online adaptation is an exciting direction, it goes beyond the scope of our current study and the literature we aim to advance.\n\n## Q1: How can we prove that global prompt dependencies don't increase $\\varepsilon$\n\nThank you for raising this insightful question. While PDT does attend to the full input sequence, including all prompt segments and the current history (as described in Equation 3), the pretraining procedure samples each prompt segment independently from multiple trajectories in the demonstration pool $\\mathcal{P}$. That is, prompts are constructed without enforcing order or co-occurrence constraints between segments.\n\nAs a result, the model is not trained to rely on segment interactions or their specific arrangement. Instead, it learns to attend to individual informative $(\\hat{r}, \\mathbf{s}, \\mathbf{a})$ pairs, which are sufficient for identifying the task in many offline RL settings. This training setup implicitly encourages invariance to inter-segment dependencies, and thus limits the influence of global prompt structure on the model’s behavior.\n\nUnder this assumption, the interaction term $\\varepsilon$, which captures residual dependencies between segments, is expected to remain small. This is supported by our empirical analysis in Section 5.3 and Figure 4, where we quantify the deviation and show that the independence-based approximation holds well in practice.\n\nTo further support this interpretation, we conducted an additional study in the 2D environment:\n\n- For each task, we constructed prompts containing one highly informative segment and $J - 1$ less informative ones, then permuted their order across rollouts.\n- **PDT performance remained stable**, with a return of $6.27 \\pm 0.34$, **indicating permutation invariance**.\n- Analyzing attention weights, we found that **in 99% of rollouts, PDT focused on the informative segment**, regardless of position.\n- When we masked attention to prompt segments, performance dropped sharply if the informative segment was excluded, but remained near-optimal when only it was retained:\n\n| Attention Scope | Return |\n| ---- | ---- |\n| Full prompt | $6.27 \\pm 0.34$ |\n| Only informative segment | $5.88 \\pm 0.23$ |\n| Only uninformative segments | $-34.06 \\pm 4.49$ |\n\nThese results suggest that PDT behavior is largely driven by the most informative segment, and that inter-segment interactions are negligible in practice. We will also include a visualization of attention weights in the appendix of the final paper, but we are not allowed to link additional image material during the rebuttal.\n\nWith that being said, it is true that if an MDP is characterized by cross-segment interactions (e.g., temporal ordering across segments), the independence assumption may not hold, and $\\varepsilon$ could be larger. This would weaken the regret bound. However, such tasks are not aligned with PDT's training procedure, which uses uniform, position-independent sampling of prompt segments. As a result, PDT is not trained to model or exploit inter-segment dependencies, and its performance is inherently limited in those settings. We will clarify this limitation in the revised manuscript.\n\nLastly, when using PDT to encode prompt segments, we ensure that each segment is encoded independently: we pad the segment into a fixed-length sequence and use the representation of the last token, which only attends to tokens within the same segment. This prevents leakage of information from other segments and preserves the independence assumption in the segment embeddings used by the bandit.\n\n## Q2 & Q3: If epsilon is indeed large, the term in equation (8) would cause regret to grow linearly (especially when K is large). Are there any experimental quantifications of the actual value epsilon? \n\nWe agree that a large $\\varepsilon$ would lead to linear regret growth and undermine the benefit of our structured bandit formulation. To assess this directly, we quantify the interaction term $\\varepsilon$ empirically in Section 5.3 using a synthetic prompt-tuning task designed to introduce segment dependencies.\n\nIn this setup, the bandit must identify a correct integer sequence of length $J$ that maximizes reward, but it receives feedback only for the full sequence, while modeling each position with an independent reward model -- exactly mirroring the structure of our method. We compute the interaction term as:\n\n$h(\\mathbf{x}_1, \\dots, \\mathbf{x}_J) = \\left| R(\\mathbf{x}) - \\frac{1}{J} \\sum_j \\phi_j(\\mathbf{x}_j) \\right|$\n\nas suggested by the reviewer.\n\nEven in this adversarial setting, the measured interaction term remains small and increases slowly with $J$, as shown in the right-most plot of Figure 4. This suggests that the independence assumption introduces only a minor approximation error in practice.\n\nWe will clarify this point in the revised version to make the connection to the regret bound more explicit.\n\n## Q3: Do prompt segment embeddings introduce a computational bottleneck?\n\nWhile it is true that encoding each prompt segment with the PDT requires a forward pass through the Transformer, this cost is incurred only once per segment. All segment embeddings are pre-computed and cached prior to bandit tuning, so no further Transformer forward passes are required during prompt selection or reward modeling.\n\nIn our experiments, the number of candidate segments remains modest, typically on the order of a few hundred (with $M \\approx 10$ expert trajectories). As a result, the one-time embedding cost does not constitute a bottleneck in practice.\n\nWe also note that while PDT-encoded segments (denoted $\\Psi$) do not always outperform raw segments in terms of return (see updated results for KYsS), they offer a practical advantage: they decouple the input dimensionality of the reward models from the environment’s state and action spaces. Without embeddings, the input size would scale with $(|\\mathcal{S}| + |\\mathcal{A}| + 1) \\times H$; with embeddings, it remains constant. This leads to significant efficiency benefits in high-dimensional domains. As shown below, average rollout times in the MuJoCo Ant environment are notably lower when using cached embeddings:\n\n|             | Ant (J=1, H=5)               | Ant (J=2, H=20)              |\n|:------------|:-----------------------------|:-----------------------------|\n| TS$^\\Psi$   |    ~28m       |  ~34m         | \n| TS  |     ~33m     |      ~172m     |\n\nWe believe this confirms that using PDT-derived segment embeddings offers a practical and scalable solution without introducing a computational bottleneck.\n\n## Q4: Does $O(\\sqrt{K \\log |\\mathcal{P}|})$ deteriorate in high-dimensional states (and high $d$)?\n\nThank you for the question. While high-dimensional inputs (e.g., raw pixels) can be challenging, we rely on pre-trained PDT embeddings to compress each segment into a fixed-length representation. This keeps the input to the reward models constant, regardless of the original state/action dimensions. Our experiments already involve high-dimensional settings, e.g., Ant ($(27 + 8 + 1) * H$), Meta-World ($(39 + 3 + 1) * H$), Cheetah ($(20 + 7 + 1) * H$), and show that PDT embeddings perform comparably to raw inputs for prompt tuning.\n\nRegarding the regret term $O(\\sqrt{K \\log |\\mathcal{P}|})$, we clarify that the term $|\\mathcal{P}|$ refers to the size of the prompt dataset, not the dimensionality of the prompt or segment embeddings. Therefore, the regret is neither directly influenced by the embedding dimension $d$ nor by the dimensionality of unencoded modalities. There is, however, an implicit dependence on the PDT's capacity: if the PDT cannot solve a downstream task well regardless of the prompt, then the maximum achievable reward by the bandit is capped. In such cases, observed regret may be high, but this reflects a limitation of the underlying model, not in the bandit optimization or its theoretical guarantees.\n\nWe will clarify this interpretation in the revised manuscript."}]

Author rebuttal:

Meta-review: This paper introduces a bandit-based prompt-tuning framework for Decision Transformers (DTs) in offline Reinforcement Learning, addressing the limitations of uniform prompt sampling in the Prompting Decision Transformer (PDT). The method learns to construct informative trajectory prompts at inference time, using a structured bandit architecture that scales linearly with prompt size rather than combinatorially. Leveraging the pre-trained PDT as a feature extractor, the approach enables efficient reward modeling and decision-making across diverse environments. The authors provide theoretical regret guarantees and empirically demonstrate significant performance gains over existing prompt-tuning baselines across multi-task, high-dimensional, and out-of-distribution RL settings. 

During the rebuttal phase, the authors successfully addressed most of the reviewers’ concerns. While one reviewer noted the absence of a comparison with the Meta-DT baseline, the authors explained that this baseline could not be reproduced within the rebuttal period. Since the paper already compares against other important baselines and demonstrates effectiveness, we strongly encourage the authors to include Meta-DT in the final version. In addition, one reviewer pointed out that aggregating results across environments using standard deviation captures cross-task variability but does not reflect confidence intervals on mean performance. We therefore encourage the authors to adopt more informative aggregation methods and evaluate on additional environments in the final version.

Despite these limitations, the work demonstrates novelty and makes valuable contributions to the community. Taken together, we have decided to accept this paper.
Decision: Accept (poster)