Title: Prompt Tuning Decision Transformers with Structured and Scalable Bandits
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Decision Transformer

Prompt tuning

Bandits

Multi-task transfer
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:  \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.\n\nQUESTIONS:\nIn 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?\n\nLIMITATIONS:\nThe limitations are discussed in the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n### 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

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

STRENGTHS AND WEAKNESSES:
Strengths:

1. The paper is well-organized, and the illustrative example and figures are easy to understand 
2. The paper addresses an important research problem in the field of offline RL via DT. 
3. The use of the bandit algorithm in the prompt-tuning procedure is an interesting idea. 
4. The authors conduct a comprehensive set of experiments to evaluate the performance. 

Weaknesses: 

1. The motivation for using DT in offline RL problems and the introduction to the underlying DT framework need to be further explained. 

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

3.  In general, does the independence of the segments negatively impact the prompt reward? 

4. In the case of large $J$, how to address the scalability and weak signal in the bandit selection procedure for arms. 

5. The regret analysis is only for the bandit algorithm, but it is not for the optimal generalized policy

6. As the authors discussed OOD, what are the technical insights of the proposed model that perform well in this case?

QUESTIONS:
see strengths and weaknesses

LIMITATIONS:
see strengths and weaknesses
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03