PaperlayerView exact revision on arXiv

Source-linked AI summary

Supervised Pretraining Can Learn In-Context Reinforcement Learning

Jonathan N. Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, Emma Brunskill

arXiv:2306.14892v1cs.LGcs.AI

TL;DR

The paper asks whether transformers can learn in-context reinforcement learning for bandits and Markov decision processes. It introduces DPT, which predicts optimal actions from interaction datasets, and finds that it supports exploration, generalization, and improved decision strategies while admitting a posterior-sampling interpretation. The main scope boundaries are the need for optimal-action supervision and mismatches between pretraining and deployment distributions.

  • Problem

    The paper studies how transformers can perform in-context decision-making for reinforcement learning from interaction data, including the online exploration and offline conservatism such decision-making requires.

  • Method

    DPT is pretrained with supervised learning to predict an optimal action from a query state and an in-context dataset of interactions across diverse tasks.

  • Results

    DPT supports online exploration, offline decision-making, generalization to unseen tasks, and improved strategies beyond its pretraining data; under some conditions, it is equivalent to posterior sampling.

  • Takeaways & Limitations

    Supervised pretraining can equip transformers with in-context reinforcement-learning abilities without explicitly hand-designing exploration or conservatism.

  • Takeaways & Limitations

    DPT requires optimal actions during pretraining, and its practical MDP implementation differs from true posterior sampling; deployment also faces distribution mismatch.

Abstract

from arXiv · show

Large transformer models trained on diverse datasets have shown a remarkable ability to learn in-context, achieving high few-shot performance on tasks they were not explicitly trained to solve. In this paper, we study the in-context learning capabilities of transformers in decision-making problems, i.e., reinforcement learning (RL) for bandits and Markov decision processes. To do so, we introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action given a query state and an in-context dataset of interactions, across a diverse set of tasks. This procedure, while simple, produces a model with several surprising capabilities. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline, despite not being explicitly trained to do so. The model also generalizes beyond the pretraining distribution to new tasks and automatically adapts its decision-making strategies to unknown structure. Theoretically, we show DPT can be viewed as an efficient implementation of Bayesian posterior sampling, a provably sample-efficient RL algorithm. We further leverage this connection to provide guarantees on the regret of the in-context algorithm yielded by DPT, and prove that it can learn faster than algorithms used to generate the pretraining data. These results suggest a promising yet simple path towards instilling strong in-context decision-making abilities in transformers.

1 Introduction

The paper studies whether transformers can perform in-context reinforcement learning from interaction data, including the exploration required online and conservatism required offline. It introduces DPT, a supervised objective whose predictions support generalization, improved strategies, and a posterior-sampling interpretation.

  • DPT: DPT trains a transformer to predict an optimal action from a query state and in-context interactions across diverse tasks.The same model can be deployed online or offline on new decision-making tasks.
  • Capabilities: Despite training only on optimal-action prediction, DPT discovers online exploration strategies comparable to hand-designed algorithms.Its predictions adapt to uncertainty in an unknown task and help discover optimal actions.
  • Capabilities: DPT generalizes to unseen reward distributions, goals, dynamics, and datasets without parameter updates at test time.These results indicate that its learned in-context strategies are robust across both bandit and MDP settings.
  • Capabilities: DPT can exploit latent structure and improve beyond the algorithms that generated its pretraining data.On linear bandits with unknown representations, it matches an efficient linear-bandit strategy even when source data comes from a suboptimal algorithm.
  • Theory: Under some conditions, DPT is theoretically equivalent in context to posterior sampling, a sample-efficient Bayesian RL algorithm.The model learns a posterior distribution over optimal actions, avoiding explicit posterior updates and sampling.

2 Related Work

Related work places in-context decision-making within meta-learning and contrasts DPT with methods that distill algorithms, model actions from offline data, or address distribution shift through pessimism or regularization.

  • Meta-learning: In-context learning is framed as meta-learning that learns shared task structure to accelerate learning on new tasks.Decision-making methods differ in whether they learn dynamics, task identifiers, skills, or the learning algorithm itself.
  • Meta-learning: Algorithm Distillation uses autoregressive supervision to distill traces from a single-task RL algorithm into a task-agnostic model.DPT also uses autoregressive supervised learning but is distinguished by its optimal-action prediction objective.
  • Autoregressive transformers: Decision Transformer models offline action sequences autoregressively while conditioning on achieved return.This represents a separate transformer-based decision-making approach from DPT’s interaction-conditioned optimal-action prediction.
  • Offline RL: Offline RL methods commonly address distribution shift with value pessimism or policy regularization.Offline meta-RL additionally reuses interactions across related tasks but must still handle distribution shift.

3 In-Context Learning Model

DPT represents tasks through interaction datasets and trains a transformer to predict optimal actions for query states. At test time, it uses the same fixed model either with static offline data or with online-collected interactions, while distribution mismatch remains a central scope boundary.

  • Decision model: An MDP is defined by state and action spaces, transition and reward functions, horizon, and an initial-state distribution.The framework also includes bandits as the single-state, one-step case.
  • Pretraining: During pretraining, DPT samples tasks, interaction datasets, query states, and optimal-action labels from a task distribution.Contexts may come from random interactions, expert demonstrations, or algorithm rollouts.
  • Pretraining: The transformer minimizes negative log-likelihood for the optimal action conditioned on the query state and partial interaction dataset.For discrete actions, the model uses a softmax distribution and treats prediction as classification.
  • Testing: DPT deploys offline by selecting the most probable action from a static dataset and online by sampling actions while appending new interactions.Its parameters are not updated during testing; the forward pass computes actions from the in-context data.
  • Limitations: Downstream deployment faces mismatch in induced state distributions, task distributions, and test-time datasets.These mismatches arise between pretraining and testing and can be especially pronounced when online data is collected by DPT itself.

4 Learning in Bandits

In bandits, DPT learns to reason about uncertainty from in-context interactions, supporting strong offline decisions and online exploration. It also exploits latent linear structure and adapts to expert-biased datasets.

  • Basic Bandits: DPT significantly exceeds Emp and LCB offline while matching TS on in-distribution bandits.This indicates that the transformer can hedge against noisy, undersampled actions rather than relying only on empirical means.
  • Basic Bandits: DPT matches UCB and TS in online cumulative regret when sampling actions, despite no explicit exploration training.Its online performance remains robust to reward-noise standard deviations unseen during pretraining.
  • Leveraging Structure from Suboptimal Data: DPT pretrained on TS data nearly matches LinUCB while significantly outperforming TS on linear bandits.The model exploits unknown linear structure, effectively learning a surrogate for the shared feature representation and enabling more informed exploration.
  • Adapting to Expert-Biased Datasets: DPT automatically changes its offline behavior with the pretraining distribution: DPT-Exp resembles LCB on expert-biased datasets, while DPT resembles TS.DPT-Exp was pretrained with mixtures containing varying fractions of expert data.

5 Learning in Markov Decision Processes

Across held-out MDP tasks, DPT generalizes from offline interactions to effective offline and online decision-making, including image-based navigation. It can also combine partial demonstrations into novel task-solving trajectories and learn from PPO-generated data.

  • Generalizing to New Offline Datasets and Tasks: On 20 held-out Dark Room goals, DPT reaches average return 61.5 from random datasets averaging 1.1 total reward.With expert data, DPT achieves near-optimal performance and can immediately exploit a transition to the goal.
  • Generalizing to New Offline Datasets and Tasks: Online, DPT solves Dark Room tasks faster than AD and reaches a higher final return than RL2 after 40 episodes.PPO makes no significant progress under the same sparse-interaction setting.
  • Learning from Image-Based Observations: DPT solves Miniworld offline from both random and expert datasets and learns online more efficiently than AD and RL2.Miniworld uses 25 × 25 RGB image observations, testing scalability to visual inputs.
  • Stitching Novel Trajectories from In-Context Subsequences: Given demonstrations of two Dark Room tasks, DPT infers a path for a third unseen task whose offline data contains only the original demonstrations.This demonstrates stitching of in-context subsequences into a new trajectory with higher task relevance.
  • Learning from Algorithm-Generated Policies and Rollouts: DPT can be pretrained with PPO-generated action labels and replay-buffer datasets, with random-context variants performing only marginally worse in some settings.This compares the DPT objective with algorithm-generated training data rather than requiring optimal-policy labels.

6 Theory

The theory analyzes DPT as an in-context posterior-sampling algorithm and derives regret guarantees under explicit modeling and data assumptions. It also shows that compliant pretraining data distributions can yield the same learned model, while latent structure can produce tighter bounds than the algorithms generating the data.

  • Assumptions: The analysis assumes exact consistency between the pretrained model and the pretraining conditional action distribution, with compliant interaction datasets.Compliance requires dataset actions to depend only on observed states and prior interaction history, not hidden task information.
  • Posterior-sampling equivalence: DPT can implement in-context posterior sampling, matching the trajectory distribution of a well-specified PS algorithm under the stated assumptions.PS maintains a posterior over tasks and executes optimal policies, but posterior computation is generally intractable.
  • Regret guarantees: Under a bounded test-to-pretraining task-density ratio, DPT receives a finite-MDP expected regret guarantee of e O(CH3/2SThe supplied passage truncates the bound after S, so the remaining factors are not specified here.
  • Latent structure: For latent linear bandits, DPT receives a regret bound scaling with d and |A|K, improving over Thompson Sampling’s e O(|A|K) bound when TS does not use linear structure.The result applies when test and pretraining task distributions coincide and the in-context data are generated by Thompson Sampling.
  • Pretraining-data invariance: All compliant pretraining dataset distributions with the same support generate the same DPT model, whereas expert-biased datasets can produce qualitatively different behavior.The invariance result applies when datasets are generated from observed task data; expert-biased data violate compliance by using privileged task knowledge.

7 Discussion

The discussion presents DPT as a simple supervised-pretraining route to in-context reinforcement learning. Across bandits and MDPs, the method supports online exploration and offline decision-making while motivating further work on data requirements, theory-practice gaps, and foundation-model deployment.

  • Main findings: DPT predicts optimal actions from interaction datasets yet develops both online exploration and offline decision-making capabilities.These behaviors arise in evaluations on classic bandit and MDP problems without hand-designing exploration or conservatism.
  • Implications: The method’s simplicity lets transformers derive strategies that leverage problem structure instead of relying entirely on manually designed RL algorithms.The discussion identifies this as a central strength of supervised pretraining.
  • Limitations and future work: DPT requires optimal actions during pretraining, although actions from another RL-trained agent cause only a slight empirical performance loss.How to use multi-task decision-making datasets effectively remains an open problem.
  • Limitations and future work: The practical MDP implementation differs from true posterior sampling, leaving an empirical-theoretical gap for future analysis.The discussion also proposes studying how broader task distributions affect generalization and how these findings transfer to foundation models.

Additional Related Work

The paper situates DPT within broader in-context learning and posterior-sampling research. It connects transformer-based context adaptation with prior work on meta-learning and Bayesian decision-making for bandits and reinforcement learning.

  • In-context learning: DPT is motivated by in-context learning, where large autoregressive models exhibit meta-learning capabilities from large-scale training.This literature primarily studies models that infer task behavior from context without parameter updates.
  • Posterior sampling: Posterior sampling maintains a posterior over possible reward and transition models and uses new interactions to update decisions.For bandits, this framework is commonly called Thompson Sampling and extends naturally to reinforcement learning.

A Implementation and Experiment Details

The appendix details DPT’s transformer architecture, training procedure, deployment modes, baselines, and evaluations under shifted bandit and MDP conditions.

  • Training procedure: DPT is trained by sampling tasks, interaction datasets, query states, and optimal action labels, then minimizing cross-entropy predictions over partial contexts.The model is trained on examples collected across a task distribution and interaction-dataset distribution.
  • Deployment: Offline deployment selects the most likely action from a fixed dataset, whereas online deployment samples predicted actions and appends new experience after each episode.The offline dataset is generated without learner control; online learning begins with an empty dataset.
  • Architecture: The architecture stacks transition tuples with a padded query-state vector, omits positional encoding, and uses GPT-2 attention to predict actions.For discrete actions, transformer outputs become logits for action distributions or one-hot predictions.
  • Bandit evaluations: DPT remains robust to shifted Gaussian reward noise and performs strongly on Bernoulli bandits despite Gaussian-only pretraining, both offline and online.The Bernoulli evaluation uses random in-context datasets, and the reported results are shown in Figure 6.
  • MDP evaluations: On Dark Room training tasks, DPT and AD show minimal offline generalization gaps, while DPT, AD, and RL2 retain similar online performance across training and test tasks.The appendix also evaluates generalization to new action permutations, where DPT identifies the optimal policy offline from 100 samples.

B.3 Sensitivity Analysis

DPT is generally robust to model-size changes, while substantially reducing the pretraining dataset degrades performance.

  • Model size: DPT performance is unchanged across embedding sizes, layer counts, and attention-head counts, except for a slight decline with 8 attention heads.The authors attribute the small decline with 8 heads to possible overfitting.
  • Pretraining data: Reducing the pretraining dataset to 10% of its original size, or 10,000 samples, degrades performance; larger datasets perform similarly.The analysis is conducted on the offline Dark Room task.

C Additional Theory and Omitted Proofs

The appendix formalizes the approximation underlying DPT and describes posterior sampling as the theoretical reference procedure.

  • Approximation guarantee: A concentration result implies the expected model-distribution error approaches zero as pretraining samples N grows, with high probability, for bounded-complexity transformer classes.The finite-function-class presentation can be extended to infinite classes with bounded statistical complexity.
  • Posterior sampling: Posterior sampling initializes its task distribution with the pretraining prior, samples a task, executes its optimal policy, and adds the resulting interactions.The prior and posterior may concern reward functions in bandits or transition dynamics in MDPs.

C.2 Proof of Theorem 1

The proof establishes that, under the stated assumptions, DPT and posterior sampling induce the same trajectory distribution by matching their conditional task posteriors at each step.

  • Theorem: Theorem 1 states that DPT and posterior sampling assign identical probabilities to every trajectory given the offline dataset and current task.This is expressed as equality between the trajectory distributions P_ps and P_Mθ.
  • Proof strategy: The proof defines sequences of states and actions for posterior sampling and pretraining, then proves equality of their distributions by induction over trajectory length.The base case uses the initial-state distribution, while the inductive step preserves equality after each transition and action.
  • Posterior matching: The key inductive step identifies the posterior over the sampled posterior-sampling task with the posterior over pretraining tasks conditioned on the current state, dataset, and history.The argument uses Bayes’ rule, the shared prior, posterior-sampling dependence on history, and compliance of the pretraining data.

C.3 Proof of Corollary 6.2

The appendix derives regret guarantees for pretrained models in finite MDPs and linear bandits by invoking their equivalence to posterior sampling. The arguments rely on compliance and support conditions for the pretraining distributions.

  • Finite MDPs: The finite-MDP corollary bounds the pretrained model's expected test regret under a bounded test-to-pretraining horizon ratio.The stated condition is supτ Ttest(τ)/Tpre(τ) ≤ C, with the regret bound expressed using e O notation.
  • Finite MDPs: The finite-MDP proof uses random-sampling compliance and the posterior-sampling equivalence established in Theorem 1.It then invokes Theorem 1 of [12] to obtain the guarantee.
  • Linear bandits: For linear bandits with Ttest = Tpre, the corollary gives an expected regret bound expressed as e O(d.The supplied passage truncates the remainder of the displayed bound.
  • Linear bandits: The linear-bandit proof uses compliance from adaptive Thompson Sampling and transfers a posterior-sampling regret result through Theorem 1.It also requires Gaussian Thompson Sampling to assign positive probability to every finite in-context dataset.
  • Distribution comparison: The later proposition compares pretraining distributions that differ only in their in-context dataset distributions while requiring common support and compliance-related invariance.The proof decomposes the posterior and dataset distributions before substituting into Equation45.
Loading 2306.14892v1…