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

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

## Key facts

- Cited by 10 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 87 of its 123 references that are in the map.
- First posted on arXiv 7 Jun 2024.

## 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, 30 of 87)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [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)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [ExpeL: LLM Agents Are Experiential Learners](https://mlsanity.com/paper/2308.10144/) (2023)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [A Tutorial on Thompson Sampling](https://mlsanity.com/paper/1707.02038/) (2017)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [On Kernelized Multi-armed Bandits](https://mlsanity.com/paper/1704.00445/) (2017)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learning](https://mlsanity.com/paper/1803.11347/) (2018)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning](https://mlsanity.com/paper/2002.08396/) (2020)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [An Introduction to Deep Reinforcement Learning](https://mlsanity.com/paper/1811.12560/) (2018)

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

- [Should You Use Your Large Language Model to Explore or Exploit?](https://mlsanity.com/paper/2502.00225/) (2025)
- [Meta-Prompt Optimization for LLM-Based Sequential Decision Making](https://mlsanity.com/paper/2502.00728/) (2025)
- [Large Language Model-Enhanced Multi-Armed Bandits](https://mlsanity.com/paper/2502.01118/) (2025)
- [Contextual Thompson Sampling via Generation of Missing Data](https://mlsanity.com/paper/2502.07064/) (2025)
- [HVAC-DPT: A Decision Pretrained Transformer for HVAC Control](https://mlsanity.com/paper/2411.19746/) (2024)
- [In-Context Compositional Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2509.24067/) (2025)
- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)
- [In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Tasks](https://mlsanity.com/paper/2510.00347/) (2025)
- [When Do We Need LLMs? A Diagnostic for Language-Driven Bandits](https://mlsanity.com/paper/2604.05859/) (2026)
- [Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation](https://mlsanity.com/paper/2609.35333/) (2026)

## 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.