mlsanityData to 7 Oct 2026

In-Context Reinforcement Learning for Variable Action Spaces

arXiv 2312.13327 ↗First posted 20 Dec 2023cs.LG cs.AIcited by 23

Not on the map. Why?

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.

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

Builds on

38 of 55 references.

References · 38

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  5. 2018Representation Learning with Contrastive Predictive Codingcited by 7,662
  6. 2015FaceNet: A Unified Embedding for Face Recognition and Clusteringcited by 2,833
  7. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,424
  8. 2021Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processingcited by 2,222
  9. 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,598
  10. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  11. 2022Toy Models of Superpositioncited by 928
  12. 2013Exact solutions to the nonlinear dynamics of learning in deep linear neural networkscited by 748
26 more
  1. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 636
  2. 2020A Survey on Contrastive Self-supervised Learningcited by 553
  3. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
  4. 2016Learning to reinforcement learncited by 369
  5. 2022Images Speak in Images: A Generalist Painter for In-Context Visual Learningcited by 317
  6. 2022Multi-Game Decision Transformerscited by 212
  7. 2023Large Language Models as General Pattern Machinescited by 196
  8. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 186
  9. 2021Open-Ended Learning Leads to Generally Capable Agentscited by 177
  10. 2015Deep Reinforcement Learning in Large Discrete Action Spacescited by 158
  11. 2022Prompting Decision Transformer for Few-Shot Policy Generalizationcited by 143
  12. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 137
  13. 2023Structured State Space Models for In-Context Reinforcement Learningcited by 109
  14. 2018A Dissection of Overfitting and Generalization in Continuous Reinforcement Learningcited by 86
  15. 2023Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretrainingcited by 82
  16. 2017Towards Generalization and Simplicity in Continuous Controlcited by 65
  17. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 54
  18. 2023A Survey on Transformers in Reinforcement Learningcited by 50
  19. 2023Generalization to New Sequential Decision Making Tasks with In-Context Learningcited by 39
  20. 2023Emergence of In-Context Reinforcement Learning from Noise Distillationcited by 26
  21. 2020Generalization to New Actions in Reinforcement Learningcited by 25
  22. 2021Introducing Symmetries to Black Box Meta Reinforcement Learningcited by 25
  23. 2019Lifelong Learning with a Changing Action Setcited by 17
  24. 2022Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learningcited by 11
  25. 2023Headless Language Models: Learning without Predicting with Contrastive Weight Tyingcited by 5
  26. 2023Action Pick-up in Dynamic Action Space Reinforcement Learningcited by 3

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 9; 2025: 11; 2026: 3.

23 papers.

Sort:

By area:

  1. 2024Retrieval-Augmented Decision Transformer: External Memory for In-context RLcited by 20
  2. 2024A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Taskscited by 14
  3. 2024XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learningcited by 13
  4. 2025Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planningcited by 12
  5. 2024Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learningcited by 10
  6. 2025Yes, Q-learning Helps Offline In-Context RLcited by 8
  7. 2025Mixture-of-Experts Meets In-Context Reinforcement Learningcited by 7
  8. 2024Benchmarking General-Purpose In-Context Learningcited by 6
  9. 2024N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needscited by 6
  10. 2025Scalable In-Context Q-Learningcited by 5
  11. 2025Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worldscited by 4
  12. 2025Filtering Learning Histories Enhances In-Context Reinforcement Learningcited by 3
11 more
  1. 2025Safe In-Context Reinforcement Learningcited by 2
  2. 2024Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Modelscited by 1
  3. 2024Random Policy Enables In-Context Reinforcement Learning within Trust Horizonscited by 1
  4. 2024HVAC-DPT: A Decision Pretrained Transformer for HVAC Controlcited by 1
  5. 2025Attacks and Defenses Against LLM Fingerprintingcited by 1
  6. 2025Beyond Distributions: Geometric Action Control for Continuous Reinforcement Learningcited by 1
  7. 2025Zero-Shot Action Generalization with Limited Observationscited by 0
  8. 2025Free Random Projection for In-Context Reinforcement Learningcited by 0
  9. 2026Learning-Guided Sparsification of Dynamic Graphs in Robotic Explorationcited by 0
  10. 2026Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learningcited by 0
  11. 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillationcited by 0