# Prompt-Tuning Bandits: Enabling Few-Shot Generalization for Efficient Multi-Task Offline RL

arXiv: https://arxiv.org/abs/2502.06358 · HTML: https://mlsanity.com/paper/2502.06358/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 13 of its 21 references that are in the map.
- First posted on arXiv 10 Feb 2025.

## Placed in

- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

## Builds on (most cited first, 13 of 13)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [InstructZero: Efficient Instruction Optimization for Black-Box Large Language Models](https://mlsanity.com/paper/2306.03082/) (2023)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [Zeroth-Order Optimization Meets Human Feedback: Provable Learning via Ranking Oracles](https://mlsanity.com/paper/2303.03751/) (2023)
- [Prompt-Tuning Decision Transformer with Preference Ranking](https://mlsanity.com/paper/2305.09648/) (2023)
- [Prompt Tuning with Diffusion for Few-Shot Pre-trained Policy Generalization](https://mlsanity.com/paper/2411.01168/) (2024)

## Abstract (arXiv)

Prompting has emerged as the dominant paradigm for adapting large, pre-trained transformer-based models to downstream tasks. The Prompting Decision Transformer (PDT) enables large-scale, multi-task offline Reinforcement Learning (RL) pre-training by leveraging stochastic trajectory prompts to identify the target task. However, these prompts are sampled uniformly from expert demonstrations, overlooking a critical limitation: not all prompts are equally informative for differentiating between tasks. This limits generalization and adaptation, especially in low-data or open-world settings where sample efficiency is crucial. To address this issue, we propose a lightweight, inference-time, bandit-based prompt-tuning framework. The bandit explores and optimizes trajectory prompt selection to enhance task performance, while avoiding costly fine-tuning of the transformer backbone. Our experiments indicate not only clear performance gains due to bandit-based prompt-tuning, but also better sample complexity, scalability, and prompt space exploration compared to prompt-tuning baselines. These results highlights the importance of adaptive prompt selection mechanisms for efficient generalization in offline multi-task RL.