# In-Context Reinforcement Learning for Variable Action Spaces

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

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 55 references that are in the map.
- First posted on arXiv 20 Dec 2023.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [FaceNet: A Unified Embedding for Face Recognition and Clustering](https://mlsanity.com/paper/1503.03832/) (2015)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [A Survey on Contrastive Self-supervised Learning](https://mlsanity.com/paper/2011.00362/) (2020)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Images Speak in Images: A Generalist Painter for In-Context Visual Learning](https://mlsanity.com/paper/2212.02499/) (2022)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [Large Language Models as General Pattern Machines](https://mlsanity.com/paper/2307.04721/) (2023)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Open-Ended Learning Leads to Generally Capable Agents](https://mlsanity.com/paper/2107.12808/) (2021)
- [Deep Reinforcement Learning in Large Discrete Action Spaces](https://mlsanity.com/paper/1512.07679/) (2015)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [A Dissection of Overfitting and Generalization in Continuous Reinforcement Learning](https://mlsanity.com/paper/1806.07937/) (2018)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [Towards Generalization and Simplicity in Continuous Control](https://mlsanity.com/paper/1703.02660/) (2017)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)

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

- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning](https://mlsanity.com/paper/2406.08973/) (2024)
- [A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks](https://mlsanity.com/paper/2410.22391/) (2024)
- [Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning](https://mlsanity.com/paper/2502.19009/) (2025)
- [Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning](https://mlsanity.com/paper/2406.05064/) (2024)
- [Yes, Q-learning Helps Offline In-Context RL](https://mlsanity.com/paper/2502.17666/) (2025)
- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs](https://mlsanity.com/paper/2411.01958/) (2024)
- [Benchmarking General-Purpose In-Context Learning](https://mlsanity.com/paper/2405.17234/) (2024)
- [Filtering Learning Histories Enhances In-Context Reinforcement Learning](https://mlsanity.com/paper/2505.15143/) (2025)
- [Scalable In-Context Q-Learning](https://mlsanity.com/paper/2506.01299/) (2025)
- [Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds](https://mlsanity.com/paper/2502.02869/) (2025)
- [Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.25582/) (2025)
- [Random Policy Enables In-Context Reinforcement Learning within Trust Horizons](https://mlsanity.com/paper/2410.19982/) (2024)
- [Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models](https://mlsanity.com/paper/2410.09701/) (2024)
- [HVAC-DPT: A Decision Pretrained Transformer for HVAC Control](https://mlsanity.com/paper/2411.19746/) (2024)
- [Attacks and Defenses Against LLM Fingerprinting](https://mlsanity.com/paper/2508.09021/) (2025)
- [Beyond Distributions: Geometric Action Control for Continuous Reinforcement Learning](https://mlsanity.com/paper/2511.08234/) (2025)
- [Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner](https://mlsanity.com/paper/2604.05112/) (2026)
- [Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.25267/) (2026)
- [Zero-Shot Action Generalization with Limited Observations](https://mlsanity.com/paper/2503.08867/) (2025)
- [Free Random Projection for In-Context Reinforcement Learning](https://mlsanity.com/paper/2504.06983/) (2025)
- [Learning-Guided Sparsification of Dynamic Graphs in Robotic Exploration](https://mlsanity.com/paper/2604.16509/) (2026)
- [Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation](https://mlsanity.com/paper/2609.35333/) (2026)

## Abstract (arXiv)

Recently, it has been shown that transformers pre-trained on diverse datasets with multi-episode contexts can generalize to new reinforcement learning tasks in-context. A key limitation of previously proposed models is their reliance on a predefined action space size and structure. The introduction of a new action space often requires data re-collection and model re-training, which can be costly for some applications. In our work, we show that it is possible to mitigate this issue by proposing the Headless-AD model that, despite being trained only once, is capable of generalizing to discrete action spaces of variable size, semantic content and order. By experimenting with Bernoulli and contextual bandits, as well as a gridworld environment, we show that Headless-AD exhibits significant capability to generalize to action spaces it has never encountered, even outperforming specialized models trained for a specific set of actions on several environment configurations. Implementation is available at: https://github.com/corl-team/headless-ad.