Prompt Tuning Decision Transformers with Structured and Scalable Bandits

• URL: http://arxiv.org/abs/2502.04979v3
• Date: Wed, 01 Oct 2025 17:56:28 GMT
• Title: Prompt Tuning Decision Transformers with Structured and Scalable Bandits
• Authors: Finn Rietz, Oleg Smirnov, Sara Karimi, Lele Cao,
• Abstract summary: We propose a bandit-based prompt-tuning method that learns to construct optimal trajectory prompts from demonstration data at inference time.<n>We show that our method consistently enhances performance across a wide range of tasks, high-dimensional environments, and out-of-distribution scenarios.
• Score: 4.460057058209513
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• 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.

Related papers

• Bigger, Regularized, Categorical: High-Capacity Value Functions are Efficient Multi-Task Learners [60.75160178669076]
  We show that the use of high-capacity value models trained via cross-entropy and conditioned on learnable task embeddings addresses the problem of task interference in online reinforcement learning.<n>We test our approach on 7 multi-task benchmarks with over 280 unique tasks, spanning high degree-of-freedom humanoid control and discrete vision-based RL.
  arXiv  Detail & Related papers  (2025-05-29T06:41:45Z)
• Towards bandit-based prompt-tuning for in-the-wild foundation agents [2.6731152954002924]
  We propose an inference time bandit-based prompt-tuning framework to enhance task performance. Our experiments indicate not only clear performance gains due to bandit-based prompt-tuning, but also better sample complexity, scalability, and prompt space exploration.
  arXiv  Detail & Related papers  (2025-02-10T11:20:10Z)
• Prompt Tuning with Diffusion for Few-Shot Pre-trained Policy Generalization [55.14484317645865]
  We develop a conditional diffusion model to produce exceptional quality prompts for offline reinforcement learning tasks. We show that the Prompt diffuser is a robust and effective tool for the prompt-tuning process, demonstrating strong performance in the meta-RL tasks.
  arXiv  Detail & Related papers  (2024-11-02T07:38:02Z)
• Task-agnostic Pre-training and Task-guided Fine-tuning for Versatile Diffusion Planner [12.360598915420255]
  Diffusion models have demonstrated their capabilities in modeling trajectories of multi-tasks.<n>Existing multi-task planners or policies typically rely on task-specific demonstrations via multi-task imitation, or require task-specific reward labels.<n>We propose a versatile diffusion planner capable of leveraging large-scale inferior data that contains task-agnostic sub-optimal trajectories.
  arXiv  Detail & Related papers  (2024-09-30T05:05:37Z)
• Pre-trained Language Models Improve the Few-shot Prompt Ability of Decision Transformer [10.338170161831496]
  Decision Transformer (DT) has emerged as a promising class of algorithms in offline reinforcement learning (RL) tasks. We introduce the Language model-d Prompt Transformer (LPDT), which leverages pre-trained language models for meta-RL tasks and fine-tunes the model using Low-rank Adaptation (LoRA) Our approach integrates pre-trained language model and RL tasks seamlessly.
  arXiv  Detail & Related papers  (2024-08-02T17:25:34Z)
• Offline Reinforcement Learning from Datasets with Structured Non-Stationarity [50.35634234137108]
  Current Reinforcement Learning (RL) is often limited by the large amount of data needed to learn a successful policy. We address a novel Offline RL problem setting in which, while collecting the dataset, the transition and reward functions gradually change between episodes but stay constant within each episode. We propose a method based on Contrastive Predictive Coding that identifies this non-stationarity in the offline dataset, accounts for it when training a policy, and predicts it during evaluation.
  arXiv  Detail & Related papers  (2024-05-23T02:41:36Z)
• Task-Distributionally Robust Data-Free Meta-Learning [99.56612787882334]
  Data-Free Meta-Learning (DFML) aims to efficiently learn new tasks by leveraging multiple pre-trained models without requiring their original training data. For the first time, we reveal two major challenges hindering their practical deployments: Task-Distribution Shift ( TDS) and Task-Distribution Corruption (TDC)
  arXiv  Detail & Related papers  (2023-11-23T15:46:54Z)
• Self-regulating Prompts: Foundational Model Adaptation without Forgetting [112.66832145320434]
  We introduce a self-regularization framework for prompting called PromptSRC. PromptSRC guides the prompts to optimize for both task-specific and task-agnostic general representations.
  arXiv  Detail & Related papers  (2023-07-13T17:59:35Z)
• Approximated Prompt Tuning for Vision-Language Pre-trained Models [54.326232586461614]
  In vision-language pre-trained models, prompt tuning often requires a large number of learnable tokens to bridge the gap between the pre-training and downstream tasks. We propose a novel Approximated Prompt Tuning (APT) approach towards efficient VL transfer learning.
  arXiv  Detail & Related papers  (2023-06-27T05:43:47Z)
• Supervised Pretraining Can Learn In-Context Reinforcement Learning [96.62869749926415]
  In this paper, we study the in-context learning capabilities of transformers in decision-making problems. We introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline.
  arXiv  Detail & Related papers  (2023-06-26T17:58:50Z)
• Model-Based Reinforcement Learning with Multi-Task Offline Pretraining [59.82457030180094]
  We present a model-based RL method that learns to transfer potentially useful dynamics and action demonstrations from offline data to a novel task. The main idea is to use the world models not only as simulators for behavior learning but also as tools to measure the task relevance. We demonstrate the advantages of our approach compared with the state-of-the-art methods in Meta-World and DeepMind Control Suite.
  arXiv  Detail & Related papers  (2023-06-06T02:24:41Z)
• Prompt-Tuning Decision Transformer with Preference Ranking [83.76329715043205]
  We propose the Prompt-Tuning DT algorithm to address challenges by using trajectory segments as prompts to guide RL agents in acquiring environmental information. Our approach involves randomly sampling a Gaussian distribution to fine-tune the elements of the prompt trajectory and using preference ranking function to find the optimization direction. Our work contributes to the advancement of prompt-tuning approaches in RL, providing a promising direction for optimizing large RL agents for specific preference tasks.
  arXiv  Detail & Related papers  (2023-05-16T17:49:04Z)
• SMART: Self-supervised Multi-task pretrAining with contRol Transformers [34.604339091596884]
  Self-supervised pretraining has been extensively studied in language and vision domains. It is difficult to properly design such a pretraining approach for sequential decision-making tasks. We propose a generic pretraining framework for sequential decision making.
  arXiv  Detail & Related papers  (2023-01-24T05:01:23Z)
• Effective Adaptation in Multi-Task Co-Training for Unified Autonomous Driving [103.745551954983]
  In this paper, we investigate the transfer performance of various types of self-supervised methods, including MoCo and SimCLR, on three downstream tasks. We find that their performances are sub-optimal or even lag far behind the single-task baseline. We propose a simple yet effective pretrain-adapt-finetune paradigm for general multi-task training.
  arXiv  Detail & Related papers  (2022-09-19T12:15:31Z)
• Prompting Decision Transformer for Few-Shot Policy Generalization [98.0914217850999]
  We propose a Prompt-based Decision Transformer (Prompt-DT) to achieve few-shot adaptation in offline RL. Prompt-DT is a strong few-shot learner without any extra finetuning on unseen target tasks.
  arXiv  Detail & Related papers  (2022-06-27T17:59:17Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.