Research map / Reinforcement learning / Visual and representation RL
Prompt Tuning Decision Transformers with Structured and Scalable Bandits
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a structured contextual bandit that tunes trajectory prompts of a frozen Prompting Decision Transformer at inference time, using one reward model per prompt slot over PDT-encoded segments, with a regret bound under an additive interaction assumption.
Problem
Uniform random sampling of trajectory prompts in Prompting Decision Transformers (PDT) ignores prompt informativeness, and selecting the best prompt over combinatorial segment choices scales poorly for offline multi-task RL.
Approach
Factorize prompt selection into J slot-specific contextual bandits (one reward model per prompt position) that score demonstration segments, and use the frozen pre-trained PDT's token embeddings as fixed-size segment features, reducing search from combinatorial to linear in J.
Key results
- On MuJoCo Ant (J=1,H=5), TS bandit reaches 835.38 return vs 694.43 for uniform-sampled PDT and 535-551 on Meta-World, in-distribution.
- On OOD MuJoCo Half Cheetah (J=1,H=5), epsilon-greedy with encoded features reaches -32.61 vs -64.78 for PDT without tuning.
- In the synthetic bandit study, the structured method outperforms flat epsilon-greedy, UCB, and TS as problem size H^J grows, with small interaction term h.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Tune prompt selection with Thompson Sampling | PDT · MuJoCo Half Cheetah, J = 1, H = 5 (in-distribution) | mean return: -44.75 $\pm$ 0.91 -> -33.56 $\pm$ 13.48 | Helped |
| Tune prompt selection with Thompson Sampling | PDT · MuJoCo Ant, J = 1, H = 5 (in-distribution) | mean return: 694.43 $\pm$ 227.04 -> 835.38 $\pm$ 171.25 | Helped |
| Tune prompt selection with TS and PDT features | PDT · Meta-world Pick-place, J = 1, H = 5 (in-distribution) | mean return: 551.58 $\pm$ 26.09 -> 556.87 $\pm$ 24.11 | No clear effect |
| Tune prompt selection with Thompson Sampling | PDT · MuJoCo Half Cheetah, J = 1, H = 5 (OOD) | mean return: -64.78 $\pm$ 36.91 -> -29.26 $\pm$ 21.25 | Helped |
| Fine-tune PDT on target-task data for 250 epochs | PDT · MuJoCo Half Cheetah, J = 1, H = 5 (OOD) | mean return: -64.78 $\pm$ 36.91 -> -129.05 $\pm$ 65.91 | Hurt |
| Use Transformer-encoded segments instead of unencoded segments (ablation) | Bandit prompt tuning · MuJoCo Ant, J = 2, H = 20 | wall-clock inference time: $\ \ \sim 172$m -> $\sim 34$m | Helped |
| Increase prompt size from J = 1, H = 5 to J = 2, H = 20 (ablation) | PDT, no tuning · MuJoCo Half Cheetah (in-distribution) | mean return: -44.75 $\pm$ 0.91 -> -42.68 $\pm$ 1.3 | Helped |
| Increase expert-data share from 0% to 30% (ablation) | PDT, no tuning · Sparse 2D Point | mean return: -39.4 $\pm$ 16.8 -> -50.1 $\pm$ 27.4 | Mixed |
| Tune prompts with UCB on the mixed dataset | PDT · Sparse 2D Point, 0% expert data | mean return: -39.4 $\pm$ 16.8 -> -3.8 $\pm$ 10.0 | Helped |
| Tune prompts with UCB on the mixed dataset | PDT · Sparse 2D Point, 30% expert data | mean return: -50.1 $\pm$ 27.4 -> 9.8 $\pm$ 0.3 | Helped |
| Mask all but the informative prompt segment (ablation) | PDT · Sparse 2D Point | mean return: $6.27 \pm 0.34$ -> $5.88 \pm 0.23$ | No clear effect |
| Mask the informative prompt segment, leaving only uninformative segments (ablation) | PDT · Sparse 2D Point | mean return: $6.27 \pm 0.34$ -> $-34.06 \pm 4.49$ | Hurt |
Surprising findings
- Fine-tuning the PDT backbone on target task data for 250 epochs did not reliably improve and sometimes degraded performance.
- A single-task DT is outperformed by prompt tuning on the 2D environment, while uniform PDT without tuning falls well short of optimal when expert demonstrations are mixed with novice ones.
Limitations
- Combinatorial growth of the search space as the number of demonstrations increases makes online exploration impractical.
- The additive segment-independence assumption is only approximately supported by an attention analysis; interaction term may be non-negligible.
- Pure offline RL backbone struggles with extrapolation on OOD tasks, limiting gains.
Experimental setup
Datasets: MuJoCo Half Cheetah MuJoCo Ant Meta-World Pick-Place Sparse 2D Point (custom, PPO-generated)
Benchmarks: MuJoCo Half Cheetah MuJoCo Ant Meta-World Pick-Place Sparse 2D Point
Metrics: mean return cumulative regret
Baselines: CQL oracle PDT (no tuning) hill-climbing ZORankSGD PDT fine-tuned
Builds on: Prompting Decision Transformer Decision Transformer Conservative Q-Learning ZORankSGD
Closest papers by meaning
- Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization (NeurIPS 2024 · similarity 0.90)
- Pre-Trained Multi-Goal Transformers with Prompt Optimization for Efficient Online Adaptation (NeurIPS 2024 · similarity 0.87)
- Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining (ICLR 2024 · similarity 0.87)
- Critic-Guided Decision Transformer for Offline Reinforcement Learning (AAAI 2024 · similarity 0.87)
- In-Dataset Trajectory Return Regularization for Offline Preference-based Reinforcement Learning (AAAI 2025 · similarity 0.87)
- Rethinking Decision Transformer via Hierarchical Reinforcement Learning (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.