Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning
This addresses the challenge of efficient exploration in bandit learning for tasks with shared structure, offering a general solution that improves over prior methods.
The paper tackles the multi-task structured bandit problem by pretraining a decision transformer to learn a near-optimal policy from demonstrator data, enabling it to outperform the demonstrator on unseen test tasks without requiring access to optimal actions.
We study learning to learn for the multi-task structured bandit problem where the goal is to learn a near-optimal algorithm that minimizes cumulative regret. The tasks share a common structure and an algorithm should exploit the shared structure to minimize the cumulative regret for an unseen but related test task. We use a transformer as a decision-making algorithm to learn this shared structure from data collected by a demonstrator on a set of training task instances. Our objective is to devise a training procedure such that the transformer will learn to outperform the demonstrator's learning algorithm on unseen test task instances. Prior work on pretraining decision transformers either requires privileged information like access to optimal arms or cannot outperform the demonstrator. Going beyond these approaches, we introduce a pre-training approach that trains a transformer network to learn a near-optimal policy in-context. This approach leverages the shared structure across tasks, does not require access to optimal actions, and can outperform the demonstrator. We validate these claims over a wide variety of structured bandit problems to show that our proposed solution is general and can quickly identify expected rewards on unseen test tasks to support effective exploration.