# Scalable In-Context Q-Learning

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 45 of its 59 references that are in the map.
- First posted on arXiv 2 Jun 2025.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Learning to Reason under Off-Policy Guidance](https://mlsanity.com/paper/2504.14945/) (2025)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Offline RL for Natural Language Generation with Implicit Language Q Learning](https://mlsanity.com/paper/2206.11871/) (2022)
- [What to Hide from Your Students: Attention-Guided Masked Image Modeling](https://mlsanity.com/paper/2203.12719/) (2022)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [A Tutorial on Meta-Reinforcement Learning](https://mlsanity.com/paper/2301.08028/) (2023)
- [Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL](https://mlsanity.com/paper/2209.03993/) (2022)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [Offline Meta-Reinforcement Learning with Advantage Weighting](https://mlsanity.com/paper/2008.06043/) (2020)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL](https://mlsanity.com/paper/2305.19923/) (2023)
- [Visual In-Context Prompting](https://mlsanity.com/paper/2311.13601/) (2023)
- [AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents](https://mlsanity.com/paper/2310.09971/) (2023)
- [A Survey of In-Context Reinforcement Learning](https://mlsanity.com/paper/2502.07978/) (2025)
- [Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2206.10442/) (2022)

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

- [Mixture-of-Experts Meets In-Context Reinforcement Learning](https://mlsanity.com/paper/2506.05426/) (2025)
- [DADP: Domain Adaptive Diffusion Policy](https://mlsanity.com/paper/2602.04037/) (2026)
- [Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.25267/) (2026)
- [Context Value Informed In-Context Reinforcement Learning](https://mlsanity.com/paper/2509.23209/) (2025)
- [Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation](https://mlsanity.com/paper/2609.35333/) (2026)

## Abstract (arXiv)

Recent advancements in language models have demonstrated remarkable in-context learning abilities, prompting the exploration of in-context reinforcement learning (ICRL) to extend the promise to decision domains. Due to involving more complex dynamics and temporal correlations, existing ICRL approaches may face challenges in learning from suboptimal trajectories and achieving precise in-context inference. In the paper, we propose Scalable In-Context Q-Learning (S-ICQL), an innovative framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization, while retaining the scalability and stability of supervised pretraining. We design a prompt-based multi-head transformer architecture that simultaneously predicts optimal policies and in-context value functions using separate heads. We pretrain a generalized world model to capture task-relevant information, enabling the construction of a compact prompt that facilitates fast and precise in-context inference. During training, we perform iterative policy improvement by fitting a state value function to an upper-expectile of the Q-function, and distill the in-context value functions into policy extraction using advantage-weighted regression. Extensive experiments across a range of discrete and continuous environments show consistent performance gains over various types of baselines, especially when learning from suboptimal data. Our code is available at textcolormagentahrefhttps://github.com/NJU-RL/SICQLhttps://github.com/NJU-RL/SICQL.