Research map / Reinforcement learning / Offline RL

Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining

ICLR 2024 · Reinforcement learning › Offline RL · approach: Decision Transformer sequence models · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Provides a theoretical framework for supervised pretraining of transformers for in-context RL (algorithm distillation and DPT), proving imitation and regret guarantees via ReLU-attention constructions of LinUCB, Thompson sampling, and UCB-VI.

Problem

Theoretical understanding of which RL algorithms supervised-pretrained transformers can implement in context, how well they perform online (regret), and how offline data distribution mismatch affects learned algorithms.

Approach

Frames algorithm distillation and decision-pretrained transformers as supervised imitation of a posterior-averaged expert; proves imitation error bounds scaling with covering number and a distribution ratio, and constructs ReLU-attention transformers that approximate soft LinUCB, Thompson sampling (via Pade matrix square roots), and soft UCB-VI with near-optimal regret.

Key results

• On linear bandits (d=5, A=10), the pretrained transformer outperforms Thompson sampling and empirical average and is comparable to LinUCB.
• On Bernoulli bandits, the transformer aligns with Thompson sampling, validating the DPT theory.
• Transformers with ReLU attention provably approximate soft LinUCB, Thompson sampling, and soft UCB-VI, yielding regret bounds like d*sqrt(T) log T for the distilled LinUCB learner.

Measured effects reported in the paper

Surprising findings

• Transformer imitated Thompson sampling even with a uniform context policy (alpha=0) in early stages, beyond theoretical predictions.
• Learned transformers do not generalize well out of distribution, e.g. shifted reward distributions or larger horizon T.

Limitations

• Regret bounds depend on distribution ratio, which can be exponentially large in T in the worst case.
• Learned transformers only guaranteed to match expert on training distribution; poor OOD generalization observed.
• Experiments are small-scale simulations with GPT-2 with 8 layers and limited capacity.

Experimental setup

Datasets: synthetic stochastic linear bandit trajectories synthetic Bernoulli bandit trajectories

Benchmarks: synthetic stochastic linear bandits (d=5, A=10) synthetic Bernoulli bandits (d=5)

Metrics: cumulative regret suboptimality

Baselines: LinUCB Thompson sampling empirical average UCB

Builds on: Algorithm Distillation Decision-Pretrained Transformers GPT-2 UCB-VI LinUCB

Closest papers by meaning

• Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement (NeurIPS 2026 · similarity 0.92)
• In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.92)
• Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.91)
• When should we prefer Decision Transformers for Offline Reinforcement Learning? (ICLR 2024 · similarity 0.90)
• From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning (NeurIPS 2026 · similarity 0.90)
• Towards Provable Emergence of In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.90)