Reproduce “Prompt Tuning Decision Transformers with Structured and Scalable Bandits” on OpenResearch

About this paper

Finn Rietz, Oleg Smirnov, Sara Karimi, Lele Cao

Researchers from King AI Labs, Microsoft Gaming, and Örebro University developed a structured bandit-based method to optimize trajectory prompt selection for the Prompting Decision Transformer (PDT), enhancing performance and generalization on both in-distribution and out-of-distribution tasks while maintaining computational efficiency. The approach leverages the pre-trained PDT as a feature extractor, provides theoretical regret bounds, and demonstrates robust improvements over uniform sampling and other baselines.

Key insights

• Factorizing the prompt selection problem into independent segment-level decisions with a structured bandit enables scalable and efficient optimization of discrete trajectory prompts.
• The pre-trained Transformer within PDT can serve as a powerful and compact feature extractor, providing robust latent representations of trajectory segments for bandit reward models.
• Optimizing context selection at inference time significantly boosts a pre-trained model's performance and generalization without requiring costly fine-tuning or weight updates to the Transformer backbone.

Summary by alphaXiv.