# Prompt Tuning Decision Transformers with Structured and Scalable Bandits

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

## Key facts

- Cited by 2 papers in the map (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 22 of its 37 references that are in the map.
- First posted on arXiv 7 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, 22 of 22)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Benchmarking Deep Reinforcement Learning for Continuous Control](https://mlsanity.com/paper/1604.06778/) (2016)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Offline Meta-Reinforcement Learning with Advantage Weighting](https://mlsanity.com/paper/2008.06043/) (2020)
- [InstructZero: Efficient Instruction Optimization for Black-Box Large Language Models](https://mlsanity.com/paper/2306.03082/) (2023)
- [Hyper-Decision Transformer for Efficient Online Policy Adaptation](https://mlsanity.com/paper/2304.08487/) (2023)
- [Future-conditioned Unsupervised Pretraining for Decision Transformer](https://mlsanity.com/paper/2305.16683/) (2023)
- [Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2310.20587/) (2023)
- [Prompt-Tuning Decision Transformer with Preference Ranking](https://mlsanity.com/paper/2305.09648/) (2023)
- [Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement](https://mlsanity.com/paper/2410.11448/) (2024)
- [Pre-trained Language Models Improve the Few-shot Prompt Ability of Decision Transformer](https://mlsanity.com/paper/2408.01402/) (2024)
- [Prompt Tuning with Diffusion for Few-Shot Pre-trained Policy Generalization](https://mlsanity.com/paper/2411.01168/) (2024)
- [Hierarchical Prompt Decision Transformer: Improving Few-Shot Policy Generalization with Global and Adaptive Guidance](https://mlsanity.com/paper/2412.00979/) (2024)

## Built on by (most cited first, 2 of 2)

- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)
- [PYTHALAB-MERA: Validation-Grounded Memory, Retrieval, and Acceptance Control for Frozen-LLM Coding Agents](https://mlsanity.com/paper/2605.08468/) (2026)

## Abstract (arXiv)

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.