mlsanityData to 7 Oct 2026

Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning

arXiv 2406.05064 ↗First posted 7 Jun 2024cs.LGcited by 10

On the map: Offline RL and imitation learning from data → Decision transformers and sequence modeling

Abstract (arXiv)

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.

Key facts. Cited by 10 papers in the map (data through 7 Oct 2026). Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data. Builds on 87 of its 123 references that are in the map. First posted on arXiv 7 Jun 2024.

Builds on

87 of 123 references.

Closest ancestors · 8

Same direction, newest first

  1. 2023Rethinking Decision Transformer via Hierarchical Reinforcement Learning1 Nov
  2. 2023Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining12 Oct
  3. 2022Dichotomy of Control: Separating What You Can Control from What You Cannot24 Oct
  4. 2022Multi-Game Decision Transformers30 May
  5. 2021Offline Meta-Reinforcement Learning with Online Self-Supervision8 Jul
  6. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problem3 Jun
  7. 2021Decision Transformer: Reinforcement Learning via Sequence Modeling2 Jun
  8. 2020Offline Meta-Reinforcement Learning with Advantage Weighting13 Aug

Other references · 79

  1. 2017Attention Is All You Needcited by 59,932
  2. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  3. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
  4. 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
  5. 2022RT-1: Robotics Transformer for Real-World Control at Scalecited by 2,007
  6. 2020Conservative Q-Learning for Offline Reinforcement Learningcited by 1,777
  7. 2022Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?cited by 1,413
  8. 2018Off-Policy Deep Reinforcement Learning without Explorationcited by 1,227
  9. 2021An Explanation of In-context Learning as Implicit Bayesian Inferencecited by 915
  10. 2019Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reductioncited by 833
  11. 2023ExpeL: LLM Agents Are Experiential Learnerscited by 814
  12. 2022A Generalist Agentcited by 802
38 more
  1. 2019Behavior Regularized Offline Reinforcement Learningcited by 654
  2. 2017The Implicit Bias of Gradient Descent on Separable Datacited by 649
  3. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
  4. 2017A Tutorial on Thompson Samplingcited by 420
  5. 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 372
  6. 2016Learning to reinforcement learncited by 369
  7. 2021COMBO: Conservative Offline Model-Based Policy Optimizationcited by 365
  8. 2017On Kernelized Multi-armed Banditscited by 348
  9. 2021Transformers Can Do Bayesian Inferencecited by 338
  10. 2022Behavior Transformers: Cloning k modes with one stonecited by 299
  11. 2018Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learningcited by 270
  12. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 220
  13. 2020Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learningcited by 213
  14. 2019VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learningcited by 210
  15. 2018An Introduction to Deep Reinforcement Learningcited by 201
  16. 2019Reinforcement Learning in Feature Space: Matrix Bandit, Kernels, and Regret Boundcited by 199
  17. 2023Large Language Models as General Pattern Machinescited by 196
  18. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 186
  19. 2013Finite-Time Analysis of Kernelised Contextual Banditscited by 177
  20. 2022A Review of Safe Reinforcement Learning: Methods, Theory and Applicationscited by 169
  21. 2018Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Samplingcited by 166
  22. 2020Exploration-Exploitation in Constrained MDPscited by 146
  23. 2022On The Computational Complexity of Self-Attentioncited by 137
  24. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 137
  25. 2018Meta-Reinforcement Learning of Structured Exploration Strategiescited by 128
  26. 2023Large Language Models are Few-Shot Health Learnerscited by 115
  27. 2023Bayesian Low-rank Adaptation for Large Language Modelscited by 112
  28. 2023Structured State Space Models for In-Context Reinforcement Learningcited by 109
  29. 2018ProMP: Proximal Meta-Policy Searchcited by 91
  30. 2016Yelp Dataset Challenge: Review Rating Predictioncited by 84
  31. 2022Partially Observable Markov Decision Processes in Robotics: A Surveycited by 81
  32. 2022Why So Pessimistic? Estimating Uncertainties for Offline RL through Ensembles, and Why Their Independence Matterscited by 74
  33. 2022When does return-conditioned supervised learning work for offline reinforcement learning?cited by 73
  34. 2018Relational Deep Reinforcement Learningcited by 72
  35. 2023WikiChat: Stopping the Hallucination of Large Language Model Chatbots by Few-Shot Grounding on Wikipediacited by 67
  36. 2021Few-Shot Bot: Prompt-Based Learning for Dialogue Systemscited by 64
  37. 2020Gamification of Pure Exploration for Linear Banditscited by 61
  38. 2014Lipschitz Bandits: Regret Lower Bounds and Optimal Algorithmscited by 59

50 most cited of 79 All 79 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 1; 2025: 6; 2026: 3.

10 papers.

  1. 2025Should You Use Your Large Language Model to Explore or Exploit?cited by 8
  2. 2025Meta-Prompt Optimization for LLM-Based Sequential Decision Makingcited by 7
  3. 2025Large Language Model-Enhanced Multi-Armed Banditscited by 7
  4. 2025Contextual Thompson Sampling via Generation of Missing Datacited by 5
  5. 2024HVAC-DPT: A Decision Pretrained Transformer for HVAC Controlcited by 1
  6. 2026A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Banditscited by 1
  7. 2025In-Context Compositional Q-Learning for Offline Reinforcement Learningcited by 0
  8. 2025In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Taskscited by 0
  9. 2026When Do We Need LLMs? A Diagnostic for Language-Driven Banditscited by 0
  10. 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillationcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Contextual Bandit Optimization with Pre-Trained Neural Networks
  2. 2025Prompt Tuning Decision Transformers with Structured and Scalable Bandits
  3. 2019Meta-Learning for Contextual Bandit Exploration
  4. 2026From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning
  5. 2024Neural Dueling Bandits: Preference-Based Optimization with Human Feedback
  6. 2026Rethinking Reinforcement fine-tuning of LLMs: A Multi-armed Bandit Learning Perspective
  7. 2025Prompt-Tuning Bandits: Enabling Few-Shot Generalization for Efficient Multi-Task Offline RL
  8. 2024Diffusion Models Meet Contextual Bandits