Prompt Tuning Decision Transformers with Structured and Scalable Bandits

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.

W1: On the segment independence assumption

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

W2: On the computational overhead

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

W3: Cross-task knowledge transfer

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

Q1: How can we prove that global prompt dependencies don't increase $\varepsilon$

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

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

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

To further support this interpretation, we conducted an additional study in the 2D environment:

• For each task, we constructed prompts containing one highly informative segment and $J - 1$ less informative ones, then permuted their order across rollouts.
• PDT performance remained stable, with a return of $6.27 \pm 0.34$, indicating permutation invariance.
• Analyzing attention weights, we found that in 99% of rollouts, PDT focused on the informative segment, regardless of position.
• 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:

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

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

Lastly, 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.

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?

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

In 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:

$h(\mathbf{x}_1, \dots, \mathbf{x}_J) = \left| R(\mathbf{x}) - \frac{1}{J} \sum_j \phi_j(\mathbf{x}_j) \right|$

as suggested by the reviewer.

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

We will clarify this point in the revised version to make the connection to the regret bound more explicit.

Q3: Do prompt segment embeddings introduce a computational bottleneck?

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

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

We 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:

We believe this confirms that using PDT-derived segment embeddings offers a practical and scalable solution without introducing a computational bottleneck.

Q4: Does $O(\sqrt{K \log |\mathcal{P}|})$ deteriorate in high-dimensional states (and high $d$)?

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

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

We will clarify this interpretation in the revised manuscript.

Thanks for the author's response. I will increase my rating to 4.

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.

W1: ZORankSGD appears to underperform.

We appreciate pointing this out. Indeed, we observe that ZORankSGD performs similarly to or slightly below hill-climbing across several environments in Table 1.

Both methods apply black-box optimization over the prompt space using perturbed prompt variants, but differ in how they process feedback:

• Hill-climbing updates prompts based directly on absolute return values, preserving the full, informative feedback signal.
• ZORankSGD, by contrast, relies on rank-based comparisons among perturbations for gradient estimation, which discards magnitude information in favor of relative ordering.

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

W2: Hill-climbing over the (discrete) space of prompts?

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

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

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

That said, we appreciate the suggestion. We will clarify this distinction in the revision.

W3: The adoption of transformer-encoded trajectories is stated to be beneficial, but performance seems to decrease in Table 1.

We 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:

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

To 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:

We 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. This 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:

W4: Large standard deviations in Table 1.

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

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

We hope this helps clarify the source of variability and increases confidence in the robustness of our method.

Q1: Why assume simulator for policy evaluation but not improvement?

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

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

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

We believe this makes the simulator-for-evaluation assumption reasonable in settings where rapid, low-cost adaptation is required.

Q2: Equations 2 and 3 should be moved.

Thank you for spotting this! We've updated the manuscript to move them after the relevant notation is introduced.

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.

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

On the performance gain

We 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:

• 33/36 ID tasks and 35/36 OOD tasks (updated Table 2) show improvement over standard PDT with uniform prompt sampling.
• All 18 cases in the 2D environment (Table and Figure 3) show consistent improvements.
• Our method outperforms ZORankSGD, the strongest inference-time tuning baseline to date, across all environments.

To illustrate this, we include a summary of representative settings:

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

Regarding efficiency, our method:

• Requires no weight updates to the PDT backbone.
• Performs no backpropagation.
• Operates entirely in inference mode, with optional embedding caching.

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

Clarity of the presentation

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

On the practicality in real-world deployment

Indeed, 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$.

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

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

On the novelty

While grounded in existing components, our method introduces several distinct and novel contributions:

• We are, to our knowledge, the first to frame inference-time prompt selection in PDT as a structured contextual bandit problem.
• Our architecture uses independent reward models per prompt position, avoiding combinatorial explosion and enabling parallel, modular learning.
• It leverages PDT-derived embeddings to enable fast learning even in high-dimensional or pixel-based settings.

These components form a new class of inference-time adaptive prompt tuning, distinct from prior work in finetuning or gradient-based selection.

We will revise both the Introduction and Related Work sections to highlight these conceptual distinctions and novel contributions more clearly.

Q1: Method elaboration

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

As shown in Figure 1 and described in Section 4.2:

• Each $\phi_j$ is trained to predict returns conditioned on the segment occupying slot $j$.
• This decomposition ensures modularity and linear scaling, reducing the complexity of learning over full prompts.
• During inference, a candidate prompt $\rho_k$ is selected via exploration or exploitation and passed to the frozen PDT for rollout.
• The rollout return $G(\rho_k)$ is then used as feedback, and each $\phi_j$ is updated with the segment it contributed.

Full procedural details are available in Appendix A, Algorithms 1-3.

Q2: On the bounded $\varepsilon$ assumption

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

While exact theoretical bounds are intractable for real PDTs, we offer empirical support in Section 5.3 using a synthetic benchmark. There we find:

• The residual term is small and stable across a range of segment pool sizes and prompt lengths.
• The regret bound in our theorem holds empirically.

Furthermore, we argue that prompt segment informativeness is additive in practice. This is due to:

• The ability to infer task identity from a few key $(\mathbf{s,a})$ pairs
• The random and unordered segment sampling during PDT pre-training, which precludes reliance on inter-segment dependencies.

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

Q3: On the bandit exploration strategies

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

Results:

• Across 90 total experiments, our method improves over PDT in 86 cases.
• Even in the few settings where a strategy underperforms, the others compensate, which shows architectural robustness.
• Thompson Sampling delivers the most consistent gains and is our default choice.

This broad-based evaluation allows users to choose strategies that best match environment properties (e.g. reward sparsity, task diversity).

Q4: On the inference time

Our method is computationally lightweight by design:

• It operates entirely at inference time.
• Requires no gradient updates or fine-tuning of PDT.
• Trains only small MLP-based reward models.
• When using PDT embeddings, segment encodings can be cached and reused.

Below is the wall-clock time for representative Ant configurations (equal rollout iteration counts):

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

Q5: On CMAB vs. best-arm identification

Best-arm identification (BAI) assumes a fixed reward distribution per arm, which is unsuitable for prompt selection:

• Prompt quality varies dynamically with the rollout context.
• Rewards are stochastic and contextual, influenced by both task identity and trajectory history.
• CMABs enable conditioning on rich segment embeddings and incorporate contextual similarity, allowing effective generalization across prompts.
• Our structured CMAB architecture ensures scalability via linear slot-based decomposition.

We will elaborate to discuss these differences explicitly and clarify our methodological choice.

Q6: On the Meta-DT baseline

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

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:

W1. Clarify motivation.

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

W2. Causal relation between bandit and PDT.

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

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

We hope this addresses the concern, and we welcome further clarification if we’ve misunderstood the intent of the question.

W3. Does the assumption of independent segments negatively impact the prompt reward?

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

First, 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.

To test this empirically, we conducted an additional study in the 2D environment:

• For each task, we constructed prompts containing one highly informative segment and $J - 1$ less informative ones, then permuted their order across rollouts.
• PDT performance remained stable, with a return of $6.27 \pm 0.34$, indicating permutation invariance.
• Analyzing attention weights, we found that in 99% of rollouts, PDT focused on the informative segment, regardless of position.
• 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:

These results suggest that PDT behavior is largely driven by the most informative segment, and that inter-segment interactions are negligible in practice.

Finally, 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.

W4. How to address the bandit scalability with large $J$?

Scalability 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:

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

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

In summary, while larger $J$ naturally increases the search space, our structured and contextual bandit design is specifically built to address this challenge.

W5. No regret analysis for the optimal policy.

Indeed, 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.

Thus, 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.

W6. Why does the proposed method perform well OOD?

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

Empirically, 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:

Note 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:

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

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