In-Context Compositional Q-Learning for Offline Reinforcement Learning

2025/09/28 by Q. N. Xu, Yuhao Huang, Xu, Qiushui +10
Neuroscience · #Artificial Intelligence (cs.AI) #EEG and Brain-Computer Interfaces #FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · pdf · doi:10.48550/arxiv.2509.24067

openalex publication_date 2025/09/28 · openalex created_date 2025/10/19 · openalex updated_date 2026/07/28

Abstract

Accurate estimation of the Q-function is a central challenge in offline reinforcement learning. However, existing approaches often rely on a shared global Q-function, which is inadequate for capturing the compositional structure of tasks that consist of diverse subtasks. We propose In-context Compositional Q-Learning (ICQL), an offline RL framework that formulates Q-learning as a contextual inference problem and uses linear Transformers to adaptively infer local Q-functions from retrieved transitions without explicit subtask labels. Theoretically, we show that, under two assumptions -- linear approximability of the local Q-function and accurate inference of weights from retrieved context -- ICQL achieves a bounded approximation error for the Q-function and enables near-optimal policy extraction. Empirically, ICQL substantially improves performance in offline settings, achieving gains of up to 16.4% on kitchen tasks and up to 8.8% and 6.3% on MuJoCo and Adroit tasks, respectively. These results highlight the underexplored potential of in-context learning for robust and compositional value estimation and establish ICQL as a principled and effective framework for offline RL.

Citations

• Reinforcement Learning with Action Chunking
• Flow Q-Learning
• PIKE-RAG: sPecIalized KnowledgE and Rationale Augmented Generation
• Doubly Mild Generalization for Offline Reinforcement Learning
• Hierarchical Universal Value Function Approximators
• Retrieval-Augmented Decision Transformer: External Memory for In-context RL
• Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
• Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning
• Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning
• Linear attention is (maybe) all you need (to understand transformer optimization)
• Supervised Pretraining Can Learn In-Context Reinforcement Learning
• Revisiting the Minimalist Approach to Offline Reinforcement Learning
• Transformers learn in-context by gradient descent
• In-context Reinforcement Learning with Algorithm Distillation
• Transformers are Sample-Efficient World Models
• TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second
• A Generalist Agent
• Offline Reinforcement Learning with Implicit Q-Learning
• Efficient Local Planning with Linear Function Approximation
• Bellman-consistent Pessimism for Offline Reinforcement Learning
• A Minimalist Approach to Offline Reinforcement Learning
• Decision Transformer: Reinforcement Learning via Sequence Modeling
• SCC: an efficient deep reinforcement learning agent mastering the game of StarCraft II
• OR-Gym: A Reinforcement Learning Library for Operations Research Problems
• Conservative Q-Learning for Offline Reinforcement Learning
• Finite-sample Analysis of Greedy-GQ with Linear Function Approximation under Markovian Noise
• Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
• D4RL: Datasets for Deep Data-Driven Reinforcement Learning
• Reinforcement Learning for Combinatorial Optimization: A Survey
• Is a Good Representation Sufficient for Sample Efficient Reinforcement Learning?
• Behavior Regularized Offline Reinforcement Learning
• Way Off-Policy Batch Deep Reinforcement Learning of Implicit Human Preferences in Dialog
• Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction
• Information-Theoretic Considerations in Batch Reinforcement Learning
• Off-Policy Deep Reinforcement Learning without Exploration
• QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation
• A Finite Time Analysis of Temporal Difference Learning With Linear Function Approximation
• Data-Efficient Hierarchical Reinforcement Learning
• Diversity is All You Need: Learning Skills without a Reward Function
• Hindsight Experience Replay
• The Option-Critic Architecture
• k-NN Regression Adapts to Local Intrinsic Dimension
• Hierarchical Reinforcement Learning with the MAXQ Value Function Decomposition
• Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers
• Reinforcement Learning: An Introduction
• Nearest neighbor pattern classification

Related

• Yes, Q-learning Helps Offline In-Context RL
• Scalable In-Context Q-Learning
• Adaptive Q-Aid for Conditional Supervised Learning in Offline Reinforcement Learning
• Generalisation in Multitask Fitted Q-Iteration and Offline Q-learning
• Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL