In-Context Compositional Q-Learning for Offline  Reinforcement Learning
Authors: Qiushui Xu
Yu-Hao Huang
Yushu Jiang
Wenliang Zheng
Lei Song
Jinyu Wang
Jiang Bian
Venue: ICLR 2026 Poster (2026)
Keywords: In-context Learning
Reinforcement Learning
Primary area: reinforcement learning

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.

Published as a conference paper at ICLR 2026
IN- IN-CONTEXT COMPOSITIONAL Q-LEARNING FOR OF- FLINE REINFOR LEARNING
Qiushui Xu', Yuhao Huang?, Yushu Jiang', Lei Song', " .Jinu Wang", Wenliang Zheng', Jiang Bian
Penn State University, Nanjing University, 3 University of Toronto, Microsoft Research (qjx5019, 9, wmz5132 edsu.edu, huangyhe smannemmalll inmmaiinnj..... edw. cn, barryy jianggmail wttpeeiiiuooooottooooooooo utoronto.ca. (Lei. Wang Jinyu, Jiang Bian) Aaian//microoffiii @microsoft.com
ABSTRACT
Accurate estimation of the Q-function is central challenge in offline reinforcement learning, However, existing approaches often rely on shared global Q-function, which 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 contextual inference problem and uses linear Transformers to adaptively infer local Q O-functions from retrieved transitions without explicit subtask labels. Theoretically, we show that, under two assumptions-linese approximability of the local Q-function and accurate inference weights from retrieved context-ICOD 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 RI.
INTRODUCTION
Offline reinforcement learning (Offline RI. aims to learn effective policies from fixed datasets without further interaction with the environment (Fujimoto et al. 2019 Lange et (012). This
setting is especially important in real-world domains such as robotics Kalashnikov et al. 2018).
logistics (Wang al. 202T), and operations research (Hubbs et al. 2020 Mazyavkina et al. 2021
where environment access limited, data collection expensive or risky, and historical data is often the only available resouree. A central challenge is distributional shift: when learned policy queries state-action pairs outside the dataset support, value extrapolation can cause severe overestimation and degenerate performance. (Fu al. 2020 Kumar et al.) 2020)
Contemporary methods primarily employ policy constraints (Chen et al. 2021) or value regulariza- tion (Kumar et al. 2020 Kostrikov et al.) (2022) address this challenge. However, policy-constraint methods are largely limited by the behavior policies used to collect the offline data and exhibit a trade-off between generalization and safe adherence the constraint. Recent value-regularizztion methods aim to provide conservative estimates that impose softer penalties on cut-of-distributiin actions. Nevertheless, the optimality of the learned value function is not guaranteed when the static dataset is limited and potentially biased.
We observe that, in many RI. control tasks, the state space can often be naturally divided into multiple subtasks. Although an expressive action-value function may in principle capture state- action values accurately, the knowledge learned from one subtask may not be fully transferable to another. For example, in MuJoCo locomotion tasks, knowledge for increasing walking speed may not help with recovery from unexpected non-nome non-nominal states. A visualization of this phenomenon is provided in Figure I which shows the distribution of states after dimensionality reduction, with colors indicating their actual future returns the offline dataset. Moreover, although states in the dataset can often be grouped into coherent clusters, each of which typically corresponds to specific subtask,

---

Published as a conference paper at ICLR 2026
Query
Query
t-SNE State Distribution
40
20
250
200
Query A R2
Query A R3
Query B R2
Query B R3
-20-
Query A Query
Figure 1: Center: dimension-reduced states and SAC value
estimates
on
Walker2d-Medi uxpexerr.. Left and right: two groups of similar states.
two clusters that are geometrically similar may nevertheless correspond to semantically different behaviors and exhibit distinct long-horizon returns. When offline data are insufficient and exploration unavailable, this property is not naturally captured by an offline value-learning algorithm that fits shared global value function.
To address these challenges, we propose to cast value learning in offline reinforcement learning a contextual inference problem, thereby enabling local Q-function approximation through in- context learning. Specifically, we introduce In-context Compositional Q-Learning (ICQL), general framework for offline RL that leverages the in-context learning capability of linear Transformers to infer local Q-functions from small retrieved sets of transitions. Rather than fitting global approximator of the value function, ICQL leverages the compositional nature and local structure of the task to learn family of value functions, thereby enabling flexible local adaptation of value estimation within context windows. Our key contributions are summarized as follows:
the best our knowledge, we introduce the first offline RI. framework TCQL that formulates Q-learning as a contextual inference problem, leveraging in in-context learning with linear Transformers to adaptively infer local Q-functions without requiring explicit subtask labels or predefined subtask structure.
We provide a theoretical analysis showing that ICQL achieves bounded approximation error under two assumptions-lineet approximability of the local Q-function and accurate inference of the weights from retrieved context-and prove that the greedy policy with respect to the inferred Q-function is near-optimal.
ICQL improves performance in offline settings through in-context local approximation, and we demonstrate the effectiveness of ICQL in both offline Q-learning and offline actor-critic frameworks. On Gym and Adroit tasks, ICQL yields score improvements of 8.8% and 6.3 %, respectively. Notably, on Kitchen tasks, ICQL achieves a 16.4% performance improvement the second-best baseline. We further show that TCQL yields accurate value estimation than strong baselines. These results highlight the underexplored potential of linear attention for robust and compositional value estimation offline RL.
We conduct extensive ablation studies to isolate the contributions of in-context learning and localized value inference. In addition, we investigate the impact of different retrieval strategies, including similarity metrics and context selection criteria, overall performance and stability.
2 RELATED WORK
Offline Reinforcement Learning. Offline RL aims to learn effective policies from static datasets without further interaction with the environment. central challenge this setting is distributional shift, which can lead to severe value overestimation and degraded policy performance when the learned policy queries mntoof-distribution actions. Several influential approaches address this issue by modifying Q-learning objectives introducing conservative regularization. Representative examples include COL (Kumar et al., 2020). IOL (Kostrikov et al.) (022). and TD3+BC (Fujimoto Gu 2021). COL introduces conservative penalty on Q-values for out-of-distribution actions to mitigate
2

---

Published as a conference paper at ICLR 2026
overestimation. TD3+BC combines TD3 with a behavior cloning loss to bias policy updates toward actions in the dataset while retaining value-based learning. IQL removes explicit policy optimization and learns malue-weighted regression targets to implicitly extract high-value actions from offline data. More recent work has continued to improve offline RI. from several complementary directions. ReBRAC (Tarasov et al. 2023) revisits the minimalist actor-critic design in offinn RI. and shows that strong performance can be obtained through careful regularization and implementation choices. DMG (Ma et al. (024) studies how to obtain better generalization in offline RL under mild assumptions. FQL (Park et al. 2025) introduces flow-based perspective for Q-learning while QC (Li et al., 2025) explores temporally structured action representations through action chunking. In addition, Sarrivaal-augmentrr sequence-modeling approaches such as RA-DT (Schmied et al. 2024) enrich trisiiontransffomer-syyyy policies with external memory for in-context RL. Despite their differences, these methods still mainly rely on global value or policy models trained over the entire dataset. Such global modeling can be limiting in compositional environments, where local transition structure and subtask-spee sppccific value patterns may vary substantially across the state space. In contrast, our approach casts value learning as a collection of local estimation problems and uses in-context inference to adapt D-functions to retrieved local transition dynamics without requiring additional supervision.
In-context Learning in RL. Recent work has applied Transformers to offline RL, often through sequence modeling for return-conditooed policy learning (Zhao et al. (025)). For example, Decision Transformer (Chen et al. 2021) and Gato (Reed et al. 2022) treat trajectories as sequences, while replay-based in-context RL (Chen al. [021] Reed et al. (022) uses Transformers for behavior cloning and reward learning. These approaches leverage the ability of pre-trained Transformers to adapt through prompt conditioning or in-context learning. In-context learning has both strong theoretical foundation (Von Oswald et al.] [023]] Shen et al., [024 Wang et al. 202.5 and strong empirical performance across tasks (Hollmann et al. 20223 Micheli et al. 2023). and it is increasingly |
studied in supervised settings (Laskin et al. 2023 Lee et al. 2023 Mukherjee et al. 2025). (Laskin)
jet al. 2023) proposes Algorithm Distillation (AD) to mimic the data-collection policy, but this approach constrained by the quality of the original algorithm. DPT (Lee et al., 2023) improves regret in contextual bandits through in-context learning, but it assumes access to optimal actions,
which is often unrealistic in offline RL. PreDeToR (Mukherjee al.. (025) adds reward prediction to Decision Transformer models, yet it still focuses on action generation. While these approaches focus on directly generating actions or policies from trajectories, they do not explicitly target value estimation, which is outside the scope of this paper. Therefore, we do not include these methods as baselines. Although recent work has explored Transformers in offline RL primarily for trajectory modeling or return-conene conditiond generation (Chen et al. 2021; Laskin et al. 2023) Mukherjee et al., 2025), we instead study linear attention as a tool for in-context value learning. Our results suggest that linear attention, when used for local Q-function estimation, provides strong performance and generalization benefits. To our knowledge, this is the first work to demonstrate the potential of linear attention for compositional value-based offline RL.
METHODOLOOGY
In this section, we introduce the proposed ICQL framework, including its local value modeling, retrieval mechanism, learning procedure, and theoretical properties.
3.1 LOCAL Q-FUNCTIONS
In this section, we define local Q-functions for offline RL based on the local neighborhood associated with each state. We let D denote the dataset containing all offline transitions.
Definition 3.1. (Local Q-function Approximation) Given transition (s, a, r, s', a') E DD, for some
d, d > O, nearby transition (,, F. ) 'D is defined as
(3, ",F,5',aa) E (s) ai Tr, $/ d)) E D || - $1 < de and |||| s.1l2 s {P) e P(d.d) (I)
For any transition (3, "F.,,,,5'."" d') E d. there exists an optimal uniform local weight vector w such that the local fuuuction approximation defined as
2 (ujadd (5. a) A w=T (1,,a), ((5,0,F,8,d) e ndddd
3

---

Published as a conference paper at ICLR 2026
Top-k similar retrieved
Retrieval Dataset D
transitions
D Squrry
[(5), au sj. a)
(qquery 100
Feature Extractor +
(Squery, Aquery)
Linear Transformer
(Squery. Qquery In dx
Squaryy
Figure 2: An overview In In-Context Compositional Q-Learning (ICQL). Given a query state-action pair (Squery, @query), ). the model embeds it with the feature extractor 6, retrieves the top-k most similar transitions from the offline dataset D, and forms a local context set. local linear Q-function (ss, al02 Equery
is then estimated from the retrieved context and used to update the actor.
where the function S A + RR is the feature function of the state-action pair (3,a). The best approximation the local Q-function Ongar (8, a) is Q Paja. (8, a), that is, there exists some Fapprox such that
Pnga. (s. w o(3.a)
= appmxx
V(s,a,F,s', a') e Q(d.))
In the remainder of this paper, we omit d from the notation of (d.) in Equation 11. because the condition ||s - I d² for some d > be readily satisfied continuous domains. We use f $11
to denote 2 (d. instead. The local Q- Q-function defined in Equation 2 can be viewed localized formulation of general linear Q-function approximation, which has been widely used in prior work (Yin et al. 2022 Du et al.,, [2000] Pouparr et al., 2002 Parr et al.) 2008). We assume that each local domain 52% admits its own state-dependent local structure. This perspective has been studied both theoretically and empirically and has been shown yield improved Q-function approximation and strong performance on complex tasks; see Section B for additional discussion. In practice, the radius is not directly tunable: it depends on the underlying density and geometry of the dataset and is unknown to the algorithm. Therefore, we adopt retrieval mechanism with size parameter & to control locality in practice.
3.2 RETRIEVAL METHODS
In this section, we introduce the methods used to retrieve transitions from the offline dataset D. We focus on three strategies: state-similar retrieval, random retrieval, and state-similar-v with high-reward retrieval. Each retrieval strategy provides different level of coverage of the local neighborhood Qda Squory associated with the query state Squery- Both state-similar retrieval and waw-smmmmla-it-hhieeeeeee reward retrieval are expected to capture more accurate and comprehensive local information from the local neighborhood Q The difference is that state-similat retrieval can preserve greater diversity in the action space, whereas state-similan secillrr-itt--eee with-higgh-reward retrieval is intended to recover higher higher-quality transitions. We define state-similar retrieval in this section. See Section C for additional details and the definitions of the other two retrieval methods.
Definition 3.2 (State-Similar Retrieval). Given the query state Squery, ICQL retrieves k transitions with the smallest l2 -ddistance between the retrieved state s, and Squery, i.e.,
D Dequery
A (sL, ai (,, S ((,, E D/st E arg togtop-k/ ( 1-Ilsquery -s 111))
Let us define
Squary d
A
fllequery - 1: Then, we have 57
Squery
equurry
Daquery The quantity dequery depends on the query state Squery, but for case of presentation, dj to denote d d'quorry Because the main version of ICQL the fixed state-similar retrieval method, use S2d
Squery denote the retrieved context provided TCQL, for notational consistency. In the next section, we explain how the transitions in Sd Sqquer are used to learn the best local Q-function approximation Q Q (s, a) for all (s, r, 5', s',a)) E Sdd oorry
Squcrr through in-context learning.
4

---

Published as a conference paper at ICLR 2026
3.3 IN-CONTEXT COMPOSITIONAL Q-LEARNING
We now describe how compositional Q-functions are learned through contextual inference. We first define context-dependent weight function for estimating the optimal local weight vector w defined in Definition [.1 for each state 8.
Definition 3.3 (Context-dependen Weights). The local weight function w, P(2) + Rd is
context-dependent function inferred through in-context learning or retrieval-basee adaptation, where P(L) - A S? the power set of 52 and Q2 contains all possible transitions for given task.
We emphasize that the offline dataset S S2. Based on Definition 3.3 there exists some St S S2 such that w, (f2)) - w It is not necessary that 22 S D. Different retrieval methods can be used to cover Q as much as possible, thereby improving weight approximation. Then, for any query state Squery and action Qquery: suppose 22 Su is the set of size N containing the & retrieved transitions from D obtained by the state-similarity criterion defined in Section 3.2) Inspired by Waang al. (2025b), the input "prompt" matrix for linear Transformer is constructed as
d0 70%
NN-1 N
Qquery 0
Zo
(5)
To
MN-1
where ø is the state-action-pair feature extractor, o; - o(s; a,) and of & Φ(s', af) and Qquery - Φ (Squery aquery) for any aquery E A. Assume L-layer linear Transformer, for l = 0, L- each layer l has weight matrices Pt. and G2 defined as
C7 Oaxa O1xd
C7
0dx1
O2dx2d 01x2d
02dx1 1
Pe a
G2 a
Oaxa
0ax1
(6)
O1xd
0
where all the matrices {Ce) E Rdxd are trainable parameters and d' is the dimension of hid- den feature. By feeding Zo through the linear Transformer TF? with linear attention layers Lin P,G) & PZM (ZTG2) GZ) and take the bottom-right element on index [2d + 1, 1]. we obtain context- Q-function approximation denoted as
Q Squery; Equery [dd
wh
Sqquery yd 4 "S(quurry, (query)s
Squary
Squery
which approximates Q Rurt
(s, a) defined in Equation
Each linear attention layer updates
w/ Squary
(Ade
iteratively for each retrieved transition (s, a, s', a') E Dde
Squery
wh 1+1 Squurry
(dde (Quuurrry
wo Squary Squeery a r 12s' ( - Q1, alS2% (quarr)) yws wws (D 8guery
squeer V a squers) (0((''
=w'
query &quuryy a
(Ade
(Dile wt u &quurry P
(8)
where ar the learning rate, denotes the index of linear attention layer, the first equality follows from SARSA (Sutton & Barto| 2018), and the second equality follows from Equation (7). See Section D for details on the thcorem showing that the proposed ICQL can implement in-context TD learning.
For training ICQL, we follow IQL (Kostrikov et al. 2022) by performing value iteration via expectile regression and policy extraction via advantage-weighhted regression. Specifically, the critic loss is defined using our local Q-function approximation:
Lettic -(san [pr (Q(s, - -v)]
(9)
where y = y= ((31022)) Bx' QQs'' a al10) ) V is also context-dependent value estimator, and p+(+) denotes the expectile regression error. The policy is optimized via advinnage-weighted regression, using an advantage defined by local value estimation conditioned on the current state and its retrieved similar states:
Ea-r exp (3 (B.(((s, a( v(s(s(+))) log log<(al)]]
Cpolicy E.-D
(10)
After training, the extracted policy can be evaluated independently without any additional retrieval or contextual inference. The pseudocode of training is provided in Algorithm
5

---

Published as a conference paper at ICLR 2026
3.4 THEORETICAL ANALYSIS ON ICQL
In this section, we analyze the theoretical properties of our algorithm ICQL. ICQL captures the compositional and local structures of complex decision-making tasks by enabling the Q-function to vary flexibly across different state regions. However, the performance of such local approximators depends critically on two factors:
(i) the expressiveness of the feature representation ó(s,a).
(ii) the accuracy of the learned weight function w,( in approximating the optimal local weight w corresponding to the state & and the retrieved offline transition set Pd
To show that the greedy policy induced by ICQL is near-optimal, we first derive pointwise and expected bounds on the local Q-function upproximation error, highlighting how both approximation error and weight estimation error contribute to the total error. Building on these results, we further characterize how the approximation error propagates to policy sub-optimality through the performance difference lemma. These analyses provide theoretical justification for the importance of accurate local value estimation in achieving strong policy performance in offline RL settings. In this section, we present only the assumptions required for the analysis and the main near-optimality theorem for ICQL. See Section E for more detailed and comprehensive proofs.
Assumption 3.1. Let S A + Rd be a fixed feature map. We assume that for all (s,a) ESX.. the feature norm is bounded as (lo(s.a)) S Bg-
Remark 3.2. Assumption 3.1 is commonly adopted in prior work (Wang & Zou 2020 Bhandari et al., 2018, Shen et al. 2.23). In our experiments, we use a tanh activation function in the last layer of our feature extractor which implies that each component of the feature vector o(s, a) is bounded by 1. Hence, we can conclude that ( a) Assumption 3.1
& d, where d is the feature dimension. This remark supports
Assumption 3.3 (Set Coverage). For each query state Squery E S, let D
Equery denote the ideal local
transition set defined in Section 3.3 Suppose the retrieved set S2de satisfies Squary
Sd Squary
lisuurry
102 Signery
2
(11)
for some coverage ratio a E (0, Equivalently, at least m o 0l2 S transitions from Ω Equery are contained in D2d
Remark 3.4. Assumption
3.3 quantifies how many transitions from Ω Digunry are covered by the Squaryy This type of coverage condition is standard in nonparametric regression (Györfi
retrieved set Qd4
al. 2002 Devroye al. 1996 Cover & Hart 1967 Kpotufe, 2011) and has also been widely adopted in the analysis of offline RL through concentrability or coverage coefficients (Munos) 2003 2007 Antos et al.) 2008 Chen al. 2019: Xie et al.) 2021). The value of dę and the choice of retrieval method both affect Kg, We present an ablation study on the number of retrieved transitions and the retrieval method Section 4.3
We now present our main theorem, which shows that the performance of the greedy policy with respect to the learned local Q-function approximation ((s, af
Sawery is near-optimal.
Theorem 3.5 (Policy Performance Gap). Suppose Assumptions 3.7 and 3.3 hold, and the learned
policy π is greedy with respect to (ossaa[22) Then, with probability least the performance gap is bounded as
log(1/8)
- S - Eanan
Cuppar + Bの) + CB
(12)
|22
where C' > 0 depends on Bo, B, and the conditioning of the local Gram matrix.
Proof.
See Section E. for details.

---

Published as a conference paper at ICLR 2026
4 EXPERIMENTS
In this section, we evaluate ICQL standard offline RL benchmarks and analyze its empirical behavior. We first describe the benchmark environments and datasets used in our experiments. Then present the results and analyze the performance of ICQL across tasks and variants.
4.1 ENVIRONMENTS AND DATASETS
We evaluate our method on a diverse set of continuous control benchmarks from the D4RL suite (uu al. 2020), which includes three types of offline reinforcement learning environments:
MuJoCo tasks (c.g., Ha FCheet aaetaaih---iiii- um) are standard locomotion environments based on MuJoCo (Todorov et al. 2012). featuring smooth dynamics and dense rewards. These tasks commonly used to assess sample efficiency and stability.
Adroit tasks (e.g., PrenHHumnn involve nigh-dimensional dexterous manipulation using 24-DoF robotic hand. The action spaces are complex and the datasets are collected from human demonstration behavior imitation, making them challenging due limited action coverage.
Kitchen tasks (e.g., Ki Kittcher-Compleet are long-horizon goal-conditioned tasks that require solving compositional subtasks (e.g., turning on lights, opening cabinets). These tasks emphasize multi-stage behavior and compositional reasoning.
4.2
MAIN RESULTS
We compare our method against five widely adopted offline RL algorithms: BC, (Chen al. 2021), TD3+BC (Fujimoto & Gu 2021), COL (Kumar et al. 2020) and IQL (Kostrikov et al. (2022). These baselines represent two complementary paradigms: the first three represent policy constraints, and the last two represent value regularization. The experiment results are shown in Table [
Table Performance comparison across Mujoco, Adroit, and Kitchen tasks. Average and standard deviation of scores are reported over 5 random seeds.
Mujoco Tasks
BC
DT
TD3+BC
CQL
IQL
ICQL(
Gain( %
Walker2d- Expumm--2
107.5
70.7
109.2
98.7
109.8
113.3
3.1%
Walker2d
Maikeerdd-Mdiiun Medium-v2
75.3
70,2
77.0
79.2
71.5
80.3
1.4%
Malker2d-Medium- Replay-v2
26.0
54.8
41.5
77.2
61.0
81.9
6.1%
Hlopper-Medium -x
52.5
78.2
105.4
98.5
111.010.6
5.3%
Senpper-Medium--22
52.9
57.1
53.5
63.3
62.6 17.9
-1.5%
Beenee--MedummRRepphee
18.1
65.8
59.4
95.0
82.4
96.414.9
1.5%
HalfChhetabh-eede Expert-v2
55.2
70.8
62.8
62,4
83.4
89.1
6.8%
HalfCheetah-Mediue
42.6
42.8
43.1
44.4
42.5
45.9
3.5%
HallCheetah-MMedimm Medium-Replay--2
36.6
39.5
41.8
45.5
38.9
44.7 20.1
-1.8%
Average
51.9
58.8
62.9
74.0
72.4
80.6
8.8%
Adroit Tasks
BC
DT
TD3+BC
CQL
IQL
ICQL
Gain %
Pen-Human-vl
63.9
79.5
64.6
37.5
89.5
85.645.6
-4.3%
Pen-Cloned-vl
37.0
74.0
76.8
39.2
4.9
89.4+4.8
5,4%
Hammer-Human
1.2
1.7
1.5
4.4
7.2
3.74:3.2
-49.4%
Hammer-Cloned-vy
0.6
3.7
1.8
2.1
0.5
4.5
23.4%
2.0
5.5
0.2
9,9
9.8
17.1
73.1%
Door-Clln -lor-Cloned-vl
0.0
3.2
-0.1
0.T
7.6
11.7+4.4
53.6%
Average
17.45
27.9
24.2
15.5
33.2
35.3
6.3%
Kitchen Tasks
BC
DT
TD3+BC
CQL
IQL
ICQL
Gain( %
Sitchen-Compledd
65.0
52.5
57.5
43.8
59.2
79.3
22.0%
Kinchen-Mixed-vv)
51.5
60.0
53.5
51.0
53.3
59.5
40.8%
Sitchen-Partial-vl
38.0
55.0
46.7
49.8
45.8
61.5
11.8%
Average
51.5
55.8
52.6
48.2
52.8
66.8
16.4%

---

Published as a conference paper at ICLR 2026
Results demonstrate that, MuJoCo tasks, ICQL outperforms second best baseline COL by 8.8% on average. On Adroit tasks, ICQL improves IỘL by 6.3%. Notably, on Kitchen task, ICQL achieves 16.4% improvement over DT Kitchen tasks, highlighting the importance of compositional value estimation in environments with complex, multi-stage structure. However on Hammer- dataset, TCQL is inferior
baseline methods, which may relate to the dataset quality issue. In Hamner r -uumne, the size of the dataset is smaller and the distance between query states and retrieved similar states are larger than those of Hammer -Clooned, making harder for in-context learning. The detailed analysis on this issue is provided in Section H.4 Although ICQL improves performance through in-context learning, the additional computational overhead remains moderate; detailed analysis is provided in Section H.3 Overall, these results validate the general applicability of ICQL across both value-learning and actor-critic paradigms.
For investigating whether TCQL can produce more accurate value estimation than baseline methods, we conduct analysis on the learned Q function by comparing the Q prediction among ICQL, IQL and online RL method SAC. We plot their Q estimations of the same set of offline dataset entries, and leverage t-SNE for showing their respective Q-estimate distribution over the same state space. Figure shows the results Walker2d-Med um dataset, where ICQL shares similarity score of 0.69 with SAC on Q estimation, while IQL can only achieve a similarity score about 0.29. This indicates that the superior performance of ICQL IQL comes from better Q estimation, ensured by local Q function estimation, over the noisy dataset.
SAC
ICQL
IQL
0.75
0.75
0.75
0.25
0,25
025
Figure 3: Q-value distribution states after t-SNE dimension reduction. of Wa dataset. The partitioned value patterns support our hypothesis that Q-functions are inherently compositional, motivating localized value modeling.
4.3 ABLATION STUDIES
In this section, we study how different design choices affect the performance of ICQL, including the number of in-context learning layers, the context length, retrieval strategy and in-context learning module architecture.
4.3.1 NUMBER OF IN-CONTEXT LEARNING LAYERS
In this experiment, investigate the effect of in-context learning steps, which is controlled by the number of layers in the in-context critic network. The number of layers selected from {4, 8, 16, 20)- The experiments are conducted on MuJoCo tasks. Figure displays the experiment outcomes and detailed numerical results are provided in Appendix Table 8 From Figure the normalized scores generally increase the number of layers increase in most of the tasks, indicating that a larger number of layers may lead to more sufficient in-context value-learning. While the phenomenon not obvious in Hopper tasks, one possible reason the significant distribution shift in Hopper environment due to the high variance of transition dynamics.
4.3.2 INFLUENCE OF CONTEXT LENGTH
In this experiment, investigate the effect of context lengths in ICOL. The context lengths selected from {10, 20, 30, 40). As shown in Figure a context length of 20 yields the generally best performance for in-context TD-learnin in Gym tasks, where too long too short context lengths
8

---

Published as a conference paper at ICLR 2026
Soare
Nomalted Stotre
70
Nomalted
Wallor2d
Hopper
HalKChersh
Walkerdd
Hogper
HarCheriab
Walleed
Royper
HaltCheethh
(8) Medium -eeeunt-xpeezt
(b) Medium
(c) Medium- (e)duunn-Rell
Figure 4: Normalized scores of different number of in-context learning layers Mujoco tasks, Each color represents different number layers, and the y-axis represents the normalized score.
Ceatett Lenght
Cantert Lesgls
Crateat Lengh
1D
Nomelrad Scote
80
70
Nomaltzd
6.00
Nomallzd
60-
Hopper
Happer
Hopper
(b)
(c) Modium-Roplay
Figure 5: Normalized scores of context lengths on Mujoco tasks. Each color represents different context lengths, and the y-axis represents the normalized score.
lead to sub-optimal results. These results provide evidence that the "locality" of context crucial for in-context learning performance. As the context length increases, the distance between query and context transitions also gets larger, which may break the "local" definition and bring noise into the in-context learning process. Detailed numerical results are shown Table g
4.3.3 CONTEXT RETRIEVAL STRATEGIES
In this experiment, investigate the impact of retrieval quality by applying different context retrieval strategies to ICQL. In addition to the standard State-Similar Retrieval, we compare two additional retrieval strategies: (1) Random Retrieval, which selects transitions uniformly at random from the -ie---mmiiaa--t-----------e-
offline dataset; and (2)
Retrieval, which further filters similar- state candidates by selecting those with higher rewards. The definitions of these three retrieval methods are provided in Section 3.2 and Section
Table 2:
Ablation study on retrieval strategies of TCQL. We compare three variants: Random Retrieval, State-Similar Retrieval, and
Retrieval.
Dataset
Random
State-Similar
Mhate-Smilar--iit--High---- High-sewardss
Walker2d-Medium--2
78.1
80.3
83,9
Walker2d-Medium-RReeee Replay-v22
67.5
81.9
75.1
Hopper- Medium--
74.1
62.6
59.9
Hopper- Mcdiun-Reelay-y12
81.0
96.4
90.8
HalfCheetah-Meclle
45.5
45.9
46,4
43.4
44.7
43.2
1
75.1
85.6
84.8
3.7
4.4
Door-Human-vi
12.0
17.1
15.6
Sitchen-Compledd
70.0
79.3
71.3
Kitchen-Mixed-vo
53.8
59.5
60.0
Kitchen-Partiale
47.5
61.5
50.0
Our results show that State-Similar Retrieval provides the strongest overall performance and the pest-performing strategy most tasks, demonstrating the benefit of constructing context from locally

---

Published as a conference paper at ICLR 2026
similar states. Random Retrieval is generally less effective and often yields inferior performance, which highlights the importance of context relevance in local value estimation. Meanwhile, State- wigh------h---hhhhh--
Retrieval outperforms the other strategies several tasks, includ- ing wa 1 ker 2 d-medi um, ha Echeetahh a-eeddh i um, hammer-human and k tchen-mi xed. This suggests that incorporating reward information into retrieval can be beneficial when high-value transitions are especially informative, while pure state similarity remains the most robust choice overall.
4.3.4 COMPARISON ACROSS DIFFERENT IN-CONTEXT MODELING CHOICES
We compare different architectural choices for the in-context module in ICQL, including linear Transformer. a small linear MLP, and a standard self sulf-attention-basee based Transformer. The results are shown Table Overall, the linear Transformer achieves the best performance on the large majority of tasks and shows especially clear advantages on replay-style and long-horizon datasets, such as Wa walker2d-MMdd -m- Ram-Replay, Hopper-Ment um-Repla ay, and the Kitchen benchmarks. The standard Transformer is generally less stable and performs substantially worse several tasks, while the linear MLP is more competitive but still underperforms the linear Transformer in most cases. Although the linear MLP performs best Hammer Hammer-Human and Hammer Cloned, this advantage does not transfer consistently to other domains. These results suggest that the proposed linear in-context mechanism is not only theoretically convenient, but also empirically important for stable and effective local Q-function estimation across diverse offline RI. environments.
Table 3: Performance comparison across different local modeling choices: linear attention Transformer, linear MLP, and standard seff-attention-bseed Transformer.
Task
I Linear Transformer
Linear MLP
Standard Transformer
Walkerrdd-Medium-Exppe
113.3
109.5
108.8
Walker2d-Medium
80.3
76.7
77.4
Walker2d- Medium- Medium-Replay
81.9
60.2
42.9
111.0
109.9
70.3
Hopper-Mediue Medium
62.6
55.7
61.9
Moopee--Medium- Replay
96.4
89.9
42.1
HaliCheetah-Medlin Expettt
89.1
83.0
72.5
HalfCheetah Medium
45.9
43.3
42.0
HalfCheetah- -enirttaa--Miu-u------ Replay
44.7
39,2
36.1
Pen-Human
85.6
66.6
72.7
Pen-Cloned
89.4
80.7
83.8
Hammer-Human
3.7
6.1
4.2
Hlamee--Cloned
4.5
7.9
1.8
Door-Human
17.1
6.9
8,9
Door-Cloned
11.7
3.5
3.4
Kitchen-Complets
79.3
70.0
78.3
Kitchen-Mixed
59,5
57.5
55.8
Kitehen-Partin
61.5
48.3
55.8
5 CONCLUSION AND FUTURE WORK
We introduced ICQL, a novel offline RL framework that casts value estimation as an in-context inference problem using linear attention. By retrieving local transitions and fitting context-dependent local O-functions. ICQL enables compositional reasoning without requiring subtask supervision. We provide theoretical guarantees showing that greedy action extraction based on ICQL yields near-optimal policy. Experiments show that ICQL achieves strong performance gains and produces value estimates that are closer to those online RL algorithms. These results highlight the potential of in-context learning the methodology of ICQL is agnostic to the choice of distance metric, retrieval quality remains a practical concern in complex, high-dimensional state spaces. An important and promising direction for future work is to combine ICQL with more sophisticated retrieval methods, such as pre-trained state encoders value-aware learnable retrievers.
powerful inductive bias for offline reinforcement learning. Although
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, Advances in Neural Information Processing Systems, (6:45614-456577 2023.
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention (maybe) all you need (to understand transformer optimization). In The Twelfth International Conference on Learning Representations, 2024.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in Neural Information Processing Systems, 30:5048-5058 2017.
András Antos, Csaba Szepesvári, and Rémi Munos. Learning near-optimal policies with bellman- residual minimization based fitted policy iteration and a single sample path. Machine Learning. 71:89-1299 2008.
Rushiv Arora. Hierarchical universal value function approximators. arXiv preprint arXiv:2410.08997, 2024.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI conference artificial intelligence, volume 31, pp. 1726-1734,201
Glen Berseth, Daniel Geng, Coline Manon Devin, Nicholas Rhinehart, Chelsea Finn, Dinesh Ja- yaraman, and Sergey Levine. Smirl: Surprise minimizing reinforcement learning in unstable environments. In International Conference Leurning Representations, 20221.
Jalaj Bhandari, Daniel Russo, and Raghav Singal. A finite time analysis of temporal difference learning with linear function approximation. In Conference on learning theory, pp. 1691-1692, 2018.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems, 34:15084-15097, 2021.
Xi Chen, Nan Jiang, and Alckh Agarwal. informationntheorettt considerations in batch reinforcement learning. In International Conference on Machine Learning, pp. 1049-1058, 2019.
Thomas M Cover and Peter E Hart. Nearest neighbor pattern classification. IEEE Transactions on Information Theory, 13(1):21122) 1967.
Damai Dai, Yutao Sun, Li Dong, Yaru Hao. Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 4005-4019, 2023.
Luc Devroye, László Györfi, and Gábor Lugosi. A Probabilistic Theory of Pattern Recognition. Springer, 1996.
Thomas G Dietterich. Hierarchical reinforcement learning with the maxq value function decomposi- tion. Journal of artificial intelligence research, 13:227-303, 2000.
Simon S Du, Sham M Kakade, Ruosong Wang, and Lin F Yang. Is good representation sufficient for sample efficient reinforcement learnin?? In International Conference on Learning Representations, 2020.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations, 2019.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 34:20132-201455 2021.
11

---

Published as a conference paper at ICLR 2026
Scott Fujimoto, David Meger, and Doina Precup. Off- Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2062, 2019.
László Györf, Michael Kohler, Adam Krzyžak, and Harro Walk. A Distribution-Fre Free Theory of Nonparametric Regression. Springer, 2002.
Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. Tabpfn: transformer that solves small tabular classification problems in a second. In International Conference on Learning Representations, 2023.
Christian D Hubbs, Hector D Perez, Owais Sarwar, Nikolaos V Sahinidis, Ignacio E Grossmann, and John M Wassick. Or-gym: A reinforcement learning library for operations research problems. arXiv preprint ar arXiv:2008.06319, 2020.
Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind W. Picard. Way off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint arXiv:1907.0044: 2019.
Sham Kakade, Michael Kearns, and John Langford. Exploration in metric state spaces. In Interna- tional Conference on Machine Learning, 2003.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference robot learning, pp. 651-673, 2018.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022.
Samory Kpotufe. k-nn regression adapts to local intrinsic dimension. In Advances in Neural Information Processing Systems, volume 24, pp. 729-737, 2011.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative g-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33:1179-1191, 2020.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch Reinforcement Learning, pp. 45-73. Springer, 2012.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, maxime gazeau, Himanshu Sahni, Satinder Singh, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation. In The Eleventh International Conference Learning Representations, 2023.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36:43057-4933 430833, 2023.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.0 2020.
Qiyang Li, Zhiyuan Zhou, and Sergey Levine. Reinforcement learning with action chunking. The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Yixiu Mao, Qi Wang, Yun Qu, Yuhang Jiang, and Xiangyang Ji. Doubly mild generalization for offline reinforcement learning. Advances in Neural Information Processing Systems, 37:51436-5144 2024.
Nina Mazyavkina, Sergey Sviridov, Sergei Ivanov, and Evgeny Burnaev. Reinforcement learning for combinatorial optimization: A survey. Computers & Operations Research, 134:10540: 2021.
12

---

Published as a conference paper at ICLR 2026
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers sample-efficient world models. In nnterationaa Conference on Learning Representations, 2023.
Subhojyoti Mukherjee, Josiah Hanna, Qiaomin Xie, and Robert D Nowak. Pretraining decision transformers with reward prediction for in-context multi-task structured bandit learning. In Reinforcement Learning Conference, 2025.
Rémi Munos. Error bounds for approximate policy iteration. In International Conference on Machine Learning, pp. 560-567, 2003.
Rémi Munos. Performance bounds in nn-norm for approximate value iteration. SIAM Journal on Control and Optimization, 46(2):541-561, 2007.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances Neural Information Processing Systems, 31 31:3307-3311, 2018.
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. In International Conference on Machine Learning, 2025.
Ronald E. Parr, Lihong Li, Gavin Taylor, Christopher Painter- Wakefield, and Michael Littman.
An analysis of linear models, linear value funecifnn approximation, and feature selection for reinforcement learning. International Conference on Machine Learning, 2008.
Pascal Poupart, Craig Boutilier, Relu Patrascu, and Dale Schuurmans. Piecewise linear value function approximation for factored mdps. In AAAI/IAAI, 2002.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth-maron, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenherg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In International conference on machine learning, pp. 1312-1320, 2015.
Thomas Schmied, Fabian Paischer, Vihang Patil, Markus Hofmarcher, Razvan Pascanu, and Sepp Hochreiter. setrivval-augmentie decision transformer: External memory for in-context rl. arXiv preprint arXiv:2410.07071 2024.
Han Shen, Kaiqing Zhang, Mingyi Hong, and Tianyi Chen. Towards understanding asynchronous advantage actor-critic: Convergence and linear speedup. IEEE Transactions on Signal Processing, 71:2579-2594, 2023.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Position: Do pretrained transformers learn in- context by gradient descent? In International Conference on Machine Learning, pp. 44-400 2024.
Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction. MIT press, 2nd edition, 2018.
Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 36:1 1,52-11620.. 2023.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033, IEEE, 2012.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174, 2023.
Jinyu Wang. Jingjing Fu, Rui Wang, Lei Song, and Jiang Bian. Pike-rag: specialized knowledge and rationale augmented generation. arXiv preprint arXiv:2501.1155) 2025a.
13

---

Published as a conference paper at ICLR 2026
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers can learn temporal difference methods for in-context reinforcement learning. In International Conference on Learning Representations, 2025b.
Tianshi Wang, Qikai Yang, Ruijie Wang. Dachun Sun, Jinyang Li, Yizhuo Chen, Yigong Hu, Chaoqi Yang, Tomoyoshi Kimura, Denizhan Kara, et al. Fine-grained control of generative data augmentation in iot sensing. Advances Neural Information Processing Systems, 37:32787-32812, 2024.
Xiangjun Wang, Junxiao Song, Penghui Qi, Peng Peng, Zhenkun Tang, Wei Zhang, Weimin Li, Xiongjun Pi, Jujie He, Chao Gao, et al. Sec: an efficient deep reinforcement learning agent mastering the game starcraft ii In International conference on machine learning, pp. 10905- 10915, 2021.
Yue Wang and Shaofeng Zou. Finite- sample analysis of greedy-gg gq with linear function approximation under markovian noise. In Conference on Uncertainty Artificial Intelligence, pp. 11-20, 2020.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: /9/1. 1361, 2019.
Tengyang Xie, Yuzhe Ma, Zhuoran Yang, and Zhaoran Wang. Bellman-consisten pessimism for offline reinforcement learning. In Advances in Neural Information Processing Systems, volume 34, pp. 6683-6694, 2021.
Dong Yin, Botao Hao, Yasin Abbasi-Yadkori, Nevena Lazié, and Csaba Szepesvári. Efficient local planning with linear function approximation. International Conference on Algorithmic Learning Theory, pp. 1165-1192, 2022.
Wenhao Zhao, Qiushui Xu, Linjie Xu, Lei Song, Jinyu Wang, Chunlai Zhou, and Jiang Bian. Unveiling markov heads in pretrained language models for offline reinforcement learning. In International Conference on Machine Learning, pp. 77807-77821, 2025.
LLM USAGE STATEMENT
LLMs used aid the writing and polishing of the manuscript.
APPENDIX
A PRELIMINARY
A.1 REINFORCEMENT LEARNING
Consider an infinite-horizon Markov Decision Process (MDP) defined by the tuple (S,A, Po- PMDP, R), where S and A denote the state and action spaces, respectively. The re- ward function is R x A + R. and the transition dynamics are governed by PMue(s']s, a), which denotes the probability of transitioning
state s' from state after taking action a. The initial state distribution is Po S + [0, 1]. and E [0, 1) is the discount factor.
At each time step the agent observes state 8t, selects action a, π (-1ss)) according to stochastic
policy A x S + [0. 1]. receives a reward r: = R(st, a.), and transitions to the next state 8141 PMDP (- 1t, az). This interaction generates trajectories of the form (so. ag, ro, $1, 41 .).
Given a policy ग, the associated Q-function and value function quantify the expected cumulative discounted rewards starting from the state-action pair (st, 0z) and the state s,, respectively:
oc
("(88 A
>
(13a)
v (8() Ber !Q (st,
(13b)
14

---

Published as a conference paper at ICLR 2026
The Q-function satisfies the Bellman Expectation Equation:
(* = R(8, a) + Ee - (8))
(14)
Similarly, the value function satisfies:
(*(s) = (11s) 1Q (s)
(15)
The goal of reinforcement learning is to learn a policy To (als) )) that maximizes the expected cumulative
discounted reward. The optimal value functions satisfy the Bellman Optimality Equations:
Q' (s. =((s,a) E. -mme
max Q"('s
(16a)
V* max Q" (s)
(16b)
acA
In the offline setting, rather than interacting with the environment, the agent is given a fixed dataset }')} collected by a behavior policy Rg. Offline RL algorithms aim to learn effective
D = a,, policy entirely from this static dataset D, without any further interaction with the environment. A central challenge in offline RL is the distributional shift (Kumar et al., 2019 Jaques et al. 2019) Levine et al.) 2020] Wu et al.) 2019) between the learned policy æ and the behavior policy T3, which often leads to overestimation and poor generalization when estimating Q-values for -iitr--distribution state-action pairs.
a. r,
A.2 IN-CONTEXT LEARNING WITH LINEAR ATTENTIONS
Recently, there has been significant interest in understanding the theoretical capabilities of in-context learning with linear attention mechanisms (Dai et al.) 2023 Ahn et al., 2023] [024], particularly in random instances linear regression and simple classification tasks. formally introduce these problem settings in this section. Throughout this paper, all vectors are treated as column vectors. We denote the identity matrix in R" by In, and the m x n all-zero matrix by 0mxn- For any matrix Z, use Z to denote its transpose, and use both (I, y) and x y interchangeably to denote the inner product.
We define a prompt matrix Z E ((d+1)x(++1) as follows:
z a (0)
(1)) 11 2(1]]
(0)
(1))
(n-1) (n-11)
0 )
(17)
(0)
g())
where (x=(0, ( ()) nー are context examples, (n) is the query input with its corresponding response
value ynn) masked as zero, and each (1()) E Rd and ي( E R for all i 0, Following ( Von Oswald et al. 2023), we define linear self-attention over the same prompt as
LinAttn 2; P.G) & PZM (Z GZ),
(18)
where P, G E Rdd+1)x(d+1) are learnable parameter matrices, and M E R(++1)((++1) a fixed mask matrix defined as
In 01xn
Onx1 0
M
A
(19)
The goal of training linear Transformers in this setting is to recover the unknown response variable corresponding to (m), which is represented as zero in the prompt matrix Z. By appropriately constructing the parameter matrices and G. the linear attention model in Equation (18) can perform in-context learning for linear regression and simple classification tasks. However, the ability of such models to perform in-context learning for offline reinforcement learning remains poorly understood. Moreover, these analyses are purely theoretical and have not been empirically validated on practical tasks. Transformers can perform in-context supervised learning by mimicking gradient descent updates (Von Oswald et al. 2023), and they can perform in-context reinforcement learning through TD-like methods via appropriately constructed linear attention mechanisms (Wang et al., 20256). However, (Wang al. 2025b) considers only the simplified setting of Markov Reward Processes (MRPs), where transitions and rewards depend only on the current state, i.e., St+1 ~ p(-1s,) and r++1 r(s)), with time-dependent context representttinss More precisely, their formulation assumes
15

---

Published as a conference paper at ICLR 2026
that each trajectory consists only of temporally contiguous steps. These restrictive assumptions do not hold in real-world crrision-making problems, and the empirical results in that work are limited to synthetic MRPs, which makes it difficult to assess how well the method would perform on practical RL tasks. To bridge this gap, we extend the analysis from MRPs to the more general MDP setting by directly estimating the state-action value function Qss, a) and removing time dependence from the context representations.
B OTHER RELATED WORK
Goal-conditioned and Hierarchical RL. Goal-conditioned methods such as UVFA (Schaul et al. 2013) and HER (Andrychowicz et al. 2017) condition policies or value functions on explicit goal inputs to facilitate generalization across tasks. Extensions to compositional settings further decompose Q-functions into subgoal components (Arora 2024). However, these approaches assume access to goal specifications or subtask labels, which are typically unavailable in offline settings. ICQL addresses this limitation by learning Q-functions conditioned on retrieved transition contexts, thereby eliminating the need for task supervision and improving sample efficiency. Hierarchical reinforcement learning decomposes tasks into subgoals or options, thereby enabling temporal abstraction and (Bacon al. 2017), and HIRO (Nachum al. 2018) explicitly model subtask boundaries and learn separate value functions for each subtask. Although effective when task structure is known or can be discovered, these methods often rely subgoal specification auxiliary termination conditions. In contrast, TCQL operates without predefined subtask structure and efficiently leverages offline data to obtain a provably accurate local falue-function approximation. Unsupervised RL methods such DIAYN [Eysenbach et al., 2019) and SMiRL (Berseth et al. (021) aim to discover diverse behaviors or latent subpolicies without external rewards or supervision. Although these methods can implicitly uncover structure, they are typically designed for unsupervised exploration or pretraining rather than for accurate value estimation in offline settings. ICQL instead focuses on precise local Q Q-function inference conditioned on retrieved experiences, thereby improving compositional generalization and training stability offline RI.
subpolicy reuse. Classical methods such as MAXQ (Dietterich) 2000), Ootton-Critte
Linear Q-function Approximation. Linear Q-function approximation has been widely used in prior work Yin et al.) 2022, Du et al. 2020; Poupart et al. 2002, Parr et al. 2008). Metric MDPs (Kakade et al. 2003), which define the Q-function according
state-distance metric, a natural complement to more direct parametric assumptions on value functions and dynamics (Kakade et al., 2003). However, this line of work does not consider local linear Q- Q-function approximation based on a state-distance metric. In our work, we focus on learning improved approximations of local value functions, whereas Kakade al. (2003) studies accurate approximation of the local environment. We assume that each local domain 12 admits its own state-dependent local structure. This perspective has been studied both theoretically and empirically and has been shown yield improved Q-function approximation and strong performance on complex tasks.
C DETAILED DEFINITIONS OF RETRIEVAL STRATEGIES
Retrieval- -base methods have shown strong performance across a wide range of domains (Wang et al. 2024| 2025a). In Section 4.3] we compare the performance of ICQL under different context-retrieval strategies for approximating the localized Q-function. The retrieval strategies discussed in Section 4.3 are described as follows:
State-Similar Retrieval:
Given the current state s, search for 20 similar states s; from the offline dataset using cosine similarity, and retrieve the corresponding transitions
{ss, 0g, s/, df).
Random Retrieval: Given the current state 8, randomly select 20 transitions
{st, Us Fi, 8/, af} as context.
State-Simllar-with--iiiiee Righh- Given the current state s, search for 60 similar states 8; from the offline dataset using cosine similarity, and retrieve the corresponding transi- tions {s, a, rs 8 d.). Then sort the retrieved transitions by reward r; and select the 20 transitions with the highest rewards as context.
16

---

Published as a conference paper at ICLR 2026
We have provided the mathematical definition of state-similar retrieval in Definition 3.2 Here, further provide the definitions the other two retrieval methods random retrieval and state-similar- with-high-rewarc rewards retrieval.
Definition C.1 (Random Retrieval). Given the query state Squery, the randomly retrieved context for TCQL is defined as
n Frandom
A as D an Tn, ~ D
D] R-1 Sizo
(20)
Squeryy
high
Definition C.2 -aie--Smmmaaa----iiieeee Rrigh-Reeward Retrieval), Given the query state Squery S2 for ICQL is defined as the set of l transitions with the smallest 12 dg-distance between the retrieved state and Squery and the highest transition reward ri, i.e.
Aquary
Thigh Siuuery
A {(1 ai, Ti, S'c de) E D D (8), as, Ti, s a') E arg top-k (ri}
(21)
where 17 Eqqurry is defined in Equation
For the retrieval methods defined in Definitions C.I to we can relate them to Equation (I by letting - min (ss,,,,.,,(, Squaryy
||s Squery ll and d2 min a5
Squery
Squery lla Therefore, we can conclude that 5
aquery S S24 and &quuary C 52da "quarry" which implies
Thiggh
Squery
that both state-similar retrieval and state-similar- serhaa--wt--hhhh--eeee hiihr-rir- retrieval can bounded by some local neighborhood associated with the query state Squery.
D DETAILS OF LINEAR TRANSFORMERS IN ICQL
In this section, we show the validity that ICQL implements the weight update defined in Equation (). We first introduce the following lemma, which motivated by the work of (Wang al. 2025b) on MRPs.
Lemma D.1. Consider the input Zo and matrix weights Po and Go. where
(0)
(N-1)
(N)
Zo
(
6g N-1)
cNN)
Po =
O2ax2d
02dx1
Go =
Odxd
C
Oaxd
C
0ax1
0ax1
01x2d
-1)
Vo (N)
01xd
01xd:
Mo
(22)
and u(0), e(i) E RR, y E R. Let Z1 A LinAttn(Zu: Po. Go) = PoZo M (2) G Zo), and let y
(N)
(N) & Z1 [2d + 1, + 1]. Then it holds
be the bottom-right element of the next layer output, i.e., y) that y1 (N) - (WN,, w1), where
N-1 Co
W1 = wy
(yo
(
+
w
Te(
w 00 )00
(23)
N
1=
Using the lemma above, we are ready to prove Theorem
D.2
Theorem D.2. Consider the L-layer linear Transformer following Equation
(18 Let the matrices
(Pe, G L-1 the mask matrix M, and the inpul prompt matrix Zo be defined in Equations (5) (6)
and (19,, respectively. Then the bottom-right element of the l-th layer output, Ze|2d + 1, II.
satisfies Zz|2d + 1, 1] - (Qquerys : w (Ddd
Squery ( }), where {w/ saperry Squery }} is defined by
(S)da
w Squary Squarry
(Sd
and, for l 2 0.
w Sqquery (d (Qure
(24)
=w - Sauery (52da
Ce
Σ
(r)
w (2dh (Odh Tdd w )
17

---

Published as a conference paper at ICLR 2026
Algorithm In-context Q-Learning (ICQL)
1: Input: Offline dataset D, the number of retrieved transitions k, feature dimension d.
2: Initialize: Linear transformer F with parameters 8, feature extractor d.
3: Sample trajectory {(s), 04, r2) 1=0 D.
4: For each query state retrieve states sf.
using the state-similar retrieval method defined in Definition 3.2 and extract corresponding transitions {(5), ai r3, ol? , a?))
5: //In-contexT Q value estimation.
6: for 0, .,.... T
do
7: Construct the imput prompt matrix Z. by Equation (5)
8: Q: TT, (2,)[2d - k 1] by Equation (18
9: end for
10: Update the parameters using Equation and Equation
10
Proof. Let U/ Ø; ((ss, az), 50 70, 70(5%, a(), yo r, for i (0,.., 1} and
(N) = Qquery = (AAquery: Oquery) ) c 0ax 2o (N) = 0, we get
v10
(S (D
w Equaary
SS Equary
NN Co
(retrwe Squary (D )" B. w Squary Squary
(Dde )" $0)0s.
Equary
which is the update rule for pre-conditioned SARSA. We also have
(())
-(w! e (D ) Qquery).
equery
By induction on the number of layers f, the proof is complete.
E PROOFS
In this section, we first derive pointwise and expected bounds on the Q-function approximation error, highlighting how both approximation eтrоr and weight-estimation error contribute the total error. Building on these results, we then characterize how approximation error propagates to policy subopti- mality through the performance difference lemma. These analyses provide theoretical justification for the importance of accurate local value estimation in achieving strong policy performance, particularly in offline RL settings.
Theorem E.1 (Weight Error under Coverage). Suppose Assumption 3.3 holds, and that the feature vectors are bounded as [io(s, a)ll < B& and the rewards are bounded as Ir| S B,. Let w be the optimal local weight vector defined in Definition 3.1 and let w. (D)d) be the weight estimated from the retrieved set. Then, with probability least
8, the following holds:
[ - -will S c (V
d d+ (og(1))
Eapprox
(25)
a |S2
where C is a constant depending on Bs, B,, and the conditioning of the local Gram matrix, and c8 Eaproox is the local approximation error defined in Definition 3.1
Proof. Fix query state and its ideal local transition set S2.. By Definition 3.1 vector W such that
there exists weight
Q n (8 = " ō(s a) + as),a))
(e,(8,
6x Fapprox
(26)
for all (s, s', a') E Sdh. By Assumption 3.3 the retrieved set S2dx overlaps with the ideal set on
least m 2 transitions. Denote this intersection by D 22 n Thus, the estimation of w from Ωdd is guaranteed to include at least valid local transitions. Let X E Rmxd be the feature matrix of Do with columns d55,, a), and let y E I be the corresponding targets. Then
y w X +
(27)
18

---

Published as a conference paper at ICLR 2026
where & collects the local approximation error, with I1Eloo I E Spproxx retrieved set is
The estimator from the
w. ( arg min
Σ
(yi w ((s));))))
(28)
12
Define the population moments on 22 as
= Ea; TT
Es; 10 yl.
(29)
Let G, b be the corresponding empirical moments on 52 Since at least =|[2 samples in D come from the true local set, standard matrix concentration implies that, with probability at least 1 6,
| Gll S c1 Bo
d+oog(1/8)
(30)
o|52%*
||b - bll s c2 Br
d=1o0(1/0) a|525*1
(31)
for universal constants C1 02 > 0. The optimal weight satisfies w G b. The empirical solution satisfies w, ) T = (up to residuals). Subtracting these systems gives
||=,(0 w || - ||~ (166 bll + || []]l[=[[l)
+ E
(32)
Fapprox'
Since G' is well-conditioned, ||- - 1/1u for some μ 0. Substituting the concentration results yields
|w. [w,( will S C.
d+ d= x10%
+ approxx
(33)
where C' depends B2, By, || ||, and μa. This is exactly the desired bound in equation 25
Theorem E.2 (Pointwise Q-function Error). Suppose Assumption 3.1 and Assumption 3.3 hold. For any fixed s E S. with probability at least -- ô, the pointwise error of the estimated Q Q-function satisfies
@(. Ogga as)
< Spprox (1 B&) +a) CB
d+ log
V(s,a,n,'' 0
12
(34)
where C > 0 depends on B Br, and the conditioning of the local Gram matrix.
Proof.
Fix 8 E S and a E A. By definition,
a(s, = us ( TT(s,a),
(p = w ((s(s)) + E Euppoxx
(35)
Thus,
((s; - Up a)
= - ( ) ((s,a) - -w"" (((ss - Espprox
(36)
I [|w (1 - will 1|ll || a) || + c.e Spprrox
(37)
I BB w. ( will + Epproxx
(38)
By Theorem E.1 with probability at least - 8,
Iw (S2dl
will |l I CV
|oog oft
+ approx
(39)
Substituting this inequality into the bound above yields
[ - Ogg (s,a) < CBpy/
d+ d+g(111/8)
axpprox (1 + Ba),
(40)
which holds for all (s, a,, a') E Adee. This proves equation
19

---

Published as a conference paper at ICLR 2026
Corollary E.3 (Expected Q-function Error). Suppose Assumptions 3.1 and 3.3 hold. Let μu be reference distribution over (s, a) E S A, and let μss be its marginal over states. Then, with probability least 1 - 6, the expected Q-function approximation error restricted to the retrieved set satisfies
E (s,a)r [ 222 (s, (s, 22
(41)
d log(1/5)
< Esous Cappreel (1 BB) CBay
|52g4
Proof. From Theorem E.2 for any (s, E Sdde we have
[Q(s, al2 - agz (s, $ Saproxx (1 Ba) + CB
d+log(1/6)
(42)
o|
Taking expectation over (s, a) ~ μ. restricted to (s,a) E D2 and noting that the right- hand side depends only on obtain
E (s 10(s,a//22 up (s, (s, 2
(43)
Fapprox (1+ (1+Ba) + +CB d+toR(1/d)
S Esnus
o1
This proves the result.
E.1 PROOF THE THEOREM ON POLICY PERFORMANCE GAP
In this section, we provide the proof of Theorem 3.5
Lemma E.4 (Performance Difference Lemma). Let be a policy, and let d denote its discounted state distribution. Then the performance gap between π and the optimal policy x² satisfies
-(T) 1- Esndt anm (==
(44)
where a" arg max, (s,
Proof. From Equation for any s S.
('(s,=((s)) (s,
Q(s x" =>)))) (Q(s, - (ss)
+ (Q(5, (s)) (s,
(45)
Since ı is greedy with respect to Q the middle term is non non-positive Thus,
(", x (s)) (s, (s)) < (s - Qs, n (s)) + (Q" Q(s, T(s))
< 25(s),
(46)
where, by Theorem E.2
8(s) = E² Epprox (1 + B&) + CBay/
d+log(1/8)
(47)
Taking expectations in Equation 44 and applying Equation
46 yields
2 E E [(())]
J(x') J[T) <
(48)
which gives the desired bound in Equation 12
20

---

Published as a conference paper at ICLR 2026
F ICQL VARIANTS FOR TD3+BC
In this section, illustrate how to extend our method to TD3+BC (Fujimoto & Gu 2021 TD3+BB introduces a simple behavior cloning regularization into value-based learning. This algorithm is easy to integrate with our framework, remains stable across diverse tasks, and serves strong baseline in the literature. Its simplicity and effectiveness make am ideal testbed for evaluating the impact of localized Q-function estimation. Together, these properties provide sufficient coverage of common design choices in offline RI.. Other algorithms can be extended in similar manner, but we omit them here for clarity and focus.
Our proposed ICQL can be seamlessly integrated into existing offiine RL algorithms by replacing the global Q-function with local, context-dependent estimator defined in Definition We demonstrate this idea by instantiating ICQL with TD3+ BC; see Algorithm [ for additional details.
ICQL-TDD+BC. TD3+BC uses a standard Bellman backup for the critic and augments the actor objective with behavior cloning. We again use the locally estimated Q(8, a) in both components. The critic loss is:
=E (s.a.,a))) ((() -)) }
critic
(49)
where y = r + min Qunge(( (s', ((d)) The actor is trained to maximize the estimated Q-value while remaining close to the dataset policy:
2D3++BC actor
-END [Q(s,
(.s. ((5)))4))) + B(aD)DD [[](()) all
(50)
Experimental results are reported in Table
Table 4: Evaluation for TD3+BC based ICQL variant on Mujoco and Adroit tasks. Average normal- ized scores are reported over 5 random seeds.
Mujoco Tasks
TD3-BC
ICQL-TD3-BClours)
Gain( %
Walkerrd-MMediun-Exxeee Expert-v2
109.2
109.3
0.1%
Walker2d-Mecium--2
77.0
72.7
-5.7%
Walker2d-Medee Medium Medium-Replay--2
41.5
55.0
32.5%
Hopper- Hopper-Mediumt Medium-Expert-v2
78.2
87.2
11.5%
Ceopper-Medium--y
53.5
57.9
8.3%
Hopper-Medium- Mee--M--dM---e- Medium-RRplaa---2222
59.4
65.8
10.9%
HalfCheetah-
BerCChhetta-Mediimte Medium-
Expprrt-222
63.7
1.5%
43.1
42.7
-0.8%
CerCCheeah-Mediunee Redian-
41.8
45.9
9.8%
Average
62.9
66.7
6.0%
Adroit Tasks
TD3-BC
ICQL-TD3-BC(orr))
Gain( %
Pen-Human--v
64.6
68.3
5.7%
Pen-Cloned-vl
76.8
74.7
-2.8%
Hammer- Huitnerr-uua--lll
1.5
1.6
7.9%
Hammer-Cloned-vl
1.8
7.3
300.6%
Door-Humane
0.2
2.0
1253.3%
Door-Cloned-vl
-0.1
-0.1
-60.0%
Average
24.2
25.6
6.2%
G IMPLEMENTATION DETAILS
In this section, we present the detailed network architectures our in-context critic and actor. We also describe the hyperparameter settings used in this paper.
21

---

Published as a conference paper at ICLR 2026
G.1 IN-CONTEXT CRITIC NETWORK
The in-context critic consists feature extractor and linear Transformer. The feature extractor is a 3-layer MLP with 256 hidden units. A Tanh activation is applied in the last layer, while ReLU used in the other layers, followed by layer normalization. The output dimension of the feature extractor is 64. A dropout rate of 0.1 is applied during training. The linear Transformer is constructed as described Equation (88. where trainable parameters appear only in The definition of G is given in Equation 6). where Cj denotes the trainable parameters in the /-th layer. The shape of C₁ is 64 x 64. We use gradient normalization to stabilize training by scaling gradients to have a maximum L2 norm of 10. The number of linear Transformer layers is set to 20.
G.2 POLICY NETWORK
For ICQL-IQL, the policy network is implemented as MLP with 2 hidden layers and ReLU activations. The policy network also includes an additional learnable vector that represents the logarithmic standard deviation of the action distribution. A dropout rate of 0.1 is applied during training.
For ICQL-TDS3 TDL--DD++BB BC, the policy network implemented 3-layer MLP with ReLU activations.
G.3 HYPER-PARAMEN PARAMETER SETTINGS
For IQQL-IQ,, follow the original IQL paper and different values of the expectile parameter
T and the temperature parameter 3 for different offline datasets. We search over (0.5, 0.7, 0.9) for the expectile parameter and 2, for the temperature parameter. The detailed settings listed
in Table
Table 5: Expectile and temperature settings for ICQL experiments.
Tasks
Expectile
Temperature
Tasks
Expectile
Temperature
Walker2d- BeeeeddMMMeddu---Eeeeeee
Pen- Pen-Human--v
0.7
Saalkr2ddMMedumg
Pen Cloned-v
Walker2d-Medle Mediium Medium-Replay-v2
vl
0.5
Hopper-Medium-Fxyeeee Expert-v2
0.7
Hammer-Cloned-
0.9
Phopper-Medium-Ve
Door-Human-o
0.5
Hopper-Medium--e -
0.7
Clor-Cloned--ll
0.7
HalfCheetah- MarCheccth-Mediun -eee---eeuuuu--r----eeeee Expert-v2
0.5
DiichenCoompeete-v0
0.9
HalfCheetah- MeetChhhhett--iim---i
0.5
Metchen-MMixede
0.5
HalfCheetah- Mediaahh--e---------- Realuum---pll22222
0.7
Eitchen-Partiale
0.9
2
For ICQL-TD3+BC. we follow the settings in the original paper and use the same hyperparameter value 2.5 for all datasets.
Other common hyperparameters are listed Table
Table 6: Common hyperparameters for the main ICQL experiments.
Hyperparameter
Value
Hidden dimension
256
Batch size
256
Training steps
1,00,0000
Evaluation episodes
10
Discount factor
0.99
Policy learning rate
3.0e-4
Critic learning rate
3.0e-4
Context length
20
22

---

Published as a conference paper at ICLR 2026
H ADDITIONAL EXPERIMENT RESULTS
H.1 EXTENDED BASELINES
In this section, we extend our comparisons to additional methods, including RA-DT Schmied et al. (2024). ReBRAC et al. (025), following their official implementations. ICQL demonstrates competitive or superior performance on tasks. The results are shown in Table
al. 2023), DMG Mao et al. (2024), FQL| et al. (2025), ) and QC
Table Performance comparison across Mujoco, Adroit, and Kitchen tasks, Average and standard deviation of scores are reported over 5 random seeds.
Task
BC
TD3BC
CQL
IQL
DT
RADT
ReBRAC
DMG
FQL
QC
ICQL
Walker2d-ME
107.5
109.2
98.7
109.8
70.7
107.8
109.2
109.5
101.0
102.8
113.3
Walker2d-M
75.3
77.0
79.2
71.5
70.2
68.9
82.8
85.0
72.4
34.1
80.3
Walker2d-MR
26.0
41.5
77.2
61.0
54.8
67.2
39.4
81.9
60.9
46.6
81.9
Hopper-ME
52.5
78.2
105.4
98.5
57.5
109.4
98.7
109.8
60.1
44.0
111.0
Hopper-M
52.9
53.5
58.0
63.3
57.1
62.4
60.6
92.3
55.6
64.6
62.6
Hopper-MR
59.4
95.0
82.4
65.8
81.6
87.4
100.1
55.0
18.6
HallCheetahh-MB
90.9
84.6
93.6
92.9
94.2
HallCheetah-M
42.6
43.1
44.4
42.5
42.8
42,0
44.6
47,9
43.9
48.2
45.9
HaltCheetah-ME
36.6
41.8
45.5
38.9
39.5
38.9
36,9
44.6
40.0
40.5
44.7
Pen-Human
63.9
64.6
37.5
89.5
79.5
17.8
91.5
66.2
61.2
55.7
85.6
Pen-Cloned
37.0
76.8
39.2
84.9
74.0
32.4
68.9
67.5
23.5
54.8
Hammer-Human
7.2
18.4
3.7
Hammerr-Comme
0.6
1.8
2.1
0.5
3.7
1.3
0.2
13.4
1.7
2.2
4.5
Door-Human
2.0
0.2
9.9
9.8
5.5
13.2
-0.1
0.1
0.2
0.7
17.1
Donr-Cloned
0.0
-0.1
0.1
7.6
3.2
2.4
9.0
3.7
0.1
4,4
11.7
Combleeee
65.0
$7.5
43.8
59.2
52.5
32.5
60.0
22.5
16.3
27.5
79.3
Sitchen-Mixed
53.5
51.0
53.3
60.0
47.5
30.0
45.0
60.0
59.5
Kitchen-Partin
38.0
46.7
49.8
45.8
55.0
53.8
62.5
37.5
15.8
52.5
61.5
Overall Average
47.7
58.5
63.2
56.2
57.2
60.0
62.0
50.5
47.7
69.7
H.2 NUMERICAL RESULTS FOR ABLATION STUDIES ON THE NUMBER OF LAYERS AND CONTEXT LENGTHS
In this section, we provide numerical results corresponding to Section 4.3. and Section 4.3.2
Table 8: Normalized scores for Gym tasks with different lengths of contexts and different number of layers in ICQL-IQI..
Context Length
Number of Layers
Gym Tasks
10
20
30
40
4
8
16
20
Walker2d-Mediun-E-xae
111.1
113.3
111.7
110.2
102.3
103.3
104.1
113.3
79.6
80.3
70.9
80.7
78.0
78.4
74.9
80.3
Walker20-Mediim- Replay
77.5
69.4
74.4
77.0
75.8
Medium -xppert
103.7
111.0
106.0
103.4
104.8
111.8
107.0
111.0
Hopper-Medium Medium
73.8
62.6
60.2
59.4
65.7
67.6
67.3
62.6
Phopper-Mediug Replay
89.9
96.4
81.2
83.9
100.5
97.8
91.8
96.4
HalfChectah-
Saieettth-Meiium-. Medium- -xppert
89.2
89.1
88.8
83.5
71.3
63.3
74.8
89.1
HalfCheetah- Meddium
45.9
45.9
46.3
45.8
45.1
44.8
45.0
45.9
HalfChcetah-Medhe Medium-RRepaa
43.7
44.7
44.3
44.2
43.5
43.6
43.8
44.7
Average
79.4
80.6
75.4
76.2
76.4
76.4
76.0
80.6
23

---

Published as a conference paper at ICLR 2026
H.3 COMPUTATION OVERHEAD ANALYSIS
H.3.1 COMPARISON OF TRAINING TIME, INFERENCE TIME, GFLOPS, AND MEMORY CONSUMPTION
In this section, we compare training time, inference time, GFLOPs, and memory consumption across all baseline methods. The analysis is conducted on the Walker2d-Medium--B-ae the results are summarized in Table 9 This analysis shows that, although ICQL incurs moderate additional computational cost relative most advanced baselines, it remains more efficient than sequential models such as DT and RA-DT while achieving substantially stronger performance.
dataset, and
Table 9: Computation cost comparison across offline RL algorithms, including per-step train- ing/inference time, FLOPs, and peak memory consumption.
Algorithm
Train Time (ms)
Infer Time (ms)
Training GFLOPs
Peak Memory (MB)
TD3BC IQL
7.23
0.26
0.17
30
10.52
0.61
0.22
26
CQL
47.57
0.61
2.64
79
DT
68.42
2.89
151.40
1383
RA-DT
121.02
3.13
1103.79
1424
ReBRAC
13.91
0.26
0.18
38
DMG
32.33
0.42
0.55
FQL
19.63
0.37
4.53
126
QC
21.60
0,25
4.65
244
ICQL
70.73
0.51
1.03
375
H.3.2 ANALYSIS OF GFLOPS AND MEMORY CONSUMPTION SCALING OF ICQL
We further report training GFLOPs and memory consumption for context lengths in {10, 20, 30, 40} and for different numbers linear Transformer layers Table and and Table The required training time scales with both context length and the number of layers. Using context length of 20 and 20 linear Transformer layers remains comparatively efficient while providing competitive performance.
Table 10: Training FLOPs (in GFLOPs) for different numbers of layers and context lengths K.
Layers
K=10
K=20
K=30
K=40
10
0.25
0.51
0.81
1.14
20
0.50
1.03
1.62
2.28
30
0.75
1.54
2.43
3.42
40
1.00
2.06
3.24
4.56
Table 11: Peak memory consumption (in MB) for different numbers of layers and context lengths K.
Layers
K=10
K=20
K=30
K=40
10
171.28
306.56
445.57
590.39
20
209.71
375.38
549.51
738.00
30
248.39
443.29
655.26
879,30
40
288.94
511.58
758,94
1023.75
H.3.3 DETAILED COMPARISON OF RETRIEVAL AND TRAINING TIME OF ICQL ACROSS ALL DATASETS
To mitigate repeated computation, pre-compute all retrieval indices onee before training for three reasons: (1) the offline dataset fixed; (2) the retrieval rule deterministic; and (3) pre-computation does not affect the learning dynamics outcomes. This converts the per-step retrieval cost into
24

---

Published as a conference paper at ICLR 2026
amortized constant-time lookup during training. ICQL follows standard actor-critic training paradigm in which the critic uses retrieved context to estimate local Q-values and the policy learns from these Q-values. At evaluation time, only the policy is used, which is consistent with standard actor-critic practice.
We report the real-time retrieval cost, the lookup time with cached indices, and the training and inference speed for all datasets. The reported results are averaged over all datasets used in our experiments. As shown in Table [2 cached retrieval adds only approximately 0.03 ms per step, which is negligible relative to the overall training time. A detailed breakdown of retrieval time and training and inference time is provided Table 13 and Table 14
Table 12: Average TCQL runtime of retrieval, training with different context lengths, and inference,
across all datasets.
Time (ms)
Retrieval with Cached Index
0.03
Train with K=10
46.94
Train with K=20
72.15
Train with K=30
113.86
Train with K=40
171.95
Inference
0.54
Table 13: Detailed retrieval time (ms) analysis across tasks and context lengths. Cached index retrieval eliminates repeated searest-neighbor searches and greatly reduces overhead.
Task
Dataset Size
K=10
K=20
K=30
K=40
Cached
Walker2d-Meddum--Eppee
1998318
6.38
6.52
6.90
7.70
0.04
Walker2d-Meddium
999322
3.98
3.96
4.37
5.16
0.03
Walker2ddMeeium-RReee
301698
1.92
2.18
2.64
3.90
0.03
Hopper-Mediuee -xxpert
1998966
6.04
6.11
6,39
7.31
0.03
Hopper- Medium
999998
3.85
3.71
4.05
4.77
0.03
Hopper-Medium-Reyee
401598
2.10
2.16
2.56
3.11
0.03
Serthheetth-MMedinu Expert
1998000
6.27
6.41
6.75
7.37
0.03
HalfCheetah- Medium
999000
3.96
3,90
4.24
5.17
0.04
HallCheetah- MMedaun----- Replly
201798
1.58
1.61
1.81
2.53
0.03
Pen- Pen-Human Pen-Cloned
4975
0.89
0.81
0.99
1.15
0.03
496264
2.91
3.05
3.52
6.67
0.03
Tlammer-Hunnee
11285
0.88
0.89
1.05
0.03
Hammer-Clone
996394
4.56
4.54
4,93
5.82
0.03
Door-Human
6704
0.88
0.88
1.07
1.17
0.03
Door-Cloned
995642
4.39
4.53
4.92
5.94
0.03
Kitchen- Tichen-Coomlete
3679
0.89
0.82
0.95
1.12
0.03
Kitchen-Parie
136937
1.36
1.41
1.60
1.91
0.03
Kitchen-Mixed
136937
1.49
1.38
1.63
2.11
0.03
H.4 FAILURE ANALYSIS ON THE HAMMER DATASET
In this section, we provide a failure analysis the Hammer- Hammer-Human dataset. We find that Hammer- Human exhibits two properties that make it particularly challenging for ICQL.
First, the dataset size is small and the coverage is sparse. Hammer- Hummmn--Human contains only 24 trajectories (-11k transitions), which is substantially fewer than Hammer- Flammer-Cloned (-996k transitions). This leads to larger distances between the query state and its retrieved neighbors, which violates locality assummtions, and to poorer state-space coverage, which makes retrieval more likely to include semantically imelevant transitions.
25

---

Published as a conference paper at ICLR 2026
Table 14: Training and inference time (ms) for different context lengths across tasks. Training time grows approximately linearly with the context length, while inference time remains nearly constant.
Task
K=10
K=20
K=30
K=40
Inference
Walker2d- Medium- Expert
48.90 90
70.73
111.75
170.71
0.51
Walker2d-Medhe
46.63
71.77
113.82
170.85
0.50
Walker2d-Medium Medium-Repllay
48.68
74.75
115.43
171.97
0.52
Hopper-Mediun--Nve Expert
48.31
70.71
114.56
173.52
0.51
Hopper-Medium Medium
46.39
71.60
113.35
171.63
0.57
Hopper- Mediuu- Replay
46.08
72.32
112.89
170.58
0.56
48.33
73.27
115,90
171,46
0.58
HalfCheetah-Mediin
47.69
74.45
113.85
171.75
0.51
HalfCheetah-Meddie
47.30
71.81
114.32
172.86
0.57
Pen Pen-Human
45.65
72.41
114.23
171.73
0.56
Pen-Cloned
44,50
69.59
112.02
170.31
0.51
Hammer-Human
46.88
73.78
113.93
171.66
0.52
Hammer-Clon Cloned
47.31
72.55
114.13
171.94
0.57
Door-Human
46. 16
71.34
113.11
171.44
0.57
Door-Cloned
45.37
71.20
112.11
170.61
0.58
Gitchen-Complet
47.45
73.65
116.36
175.98
0.54
Kitchen-Parte Partial
48.11
72.35
116.50
175.42
0.52
Kitchen-No Mixxed
45.25
70.47
111.17
170.59
0.52
Second, the transitions of low quality and the rewards noisy. Most Hammer-Humee trajectories have very low returns. As a result, for each query state, the retrieved neighbors tend to provide weak reward signals, which makes it more difficult fit effective local Q-function.
We provide comparisons of dataset statistics in Table 15 and comparisons of the distributions of the mean distance between query states and retrieved states in Figure 6 Both support these observations.
Table 15: Dataset statistics for Hammer- Hammer--uman and Hammer- Cloneerr---ode...
Dataset
Hammer-Human
Hammer-Cloned
Number of trajectories
24
3605
Number transitions
11285
996394
Mean Trajectory Length (Min-Max)
455.2 (347-623)
276.4 (199-623)
Mean Trajectory Return (Min-Max)
2817.5 16022)
779.8 (-407-16022)
Clored
Huran
3000
300
2000
Count
250
4500
0.0
Figure 6: Distribution of mean distance between query states and retrieved states Hammer dataset.
26

---

Published as a conference paper at ICLR 2026
Although Hammer-Cloned Cloned also contains mostly low-return behavior, its much larger dataset size provides substantially denser coverage. As a result, ICQL can retrieve states that are much closer to the query state, which enables more reliable local linear approximation and yields slightly higher scores.
Moreover, for both Hlammer-Humal and Hammer-Cloned, the proportion of high-reward transitions is extremely low, which makes inherently difficult to retrieve a local neighborhood that provides strong positive supervision. As result, even the Q-network successfully fits a local linear approximation, it rarely observes transitions that reliably correspond to high-return behavior. Consequently, the learned Q-values cannot meaningfully distinguish truly rewarding actions, which leads uniformly low evaluation scores on both datasets.
H.5 ANALYSIS OF THE RELATIONSHIP BETWEEN THEORETICAL d AND k
In the theory, a local set Qd Syuery defined as the set of all transitions whose states fall within radius-d neighborhood around 8. This radius determines the intrinsic locality scale at which the Q-function is assumed to be approximately linear. However, in practice, the radius d is not directly tunable: it depends on the underlying density and geometry of the dataset and is unknown to the algorithm.
Instead, ICQL controls locality through the retrieval size k. Retrieving the top-k nearest neighbors is equivalent to selecting a data-adaptive radius, where dę )ax(s,,)) ||s: s|12 and dx (ax(s(,)) (a( Etop-k ||x/ 8; Il that the practical neighborhood is exactly the theoretical local set with radius (ds, ). The distribution the mean distance between query states and retrieved states for different values of k is visualized in Figure 7
Walkerdd Dataset State Distance Nean stribattion
Hopper: Dataset State distance Mean Ditr butien
0:55-
0.
Detardl
Delasd
(a) cum-xxpert
hallChestat: Mean Datr
Kicsen Dataset State stance Mean Distribution
06
Medium
Mediun
Pandil Dauseet
Delavol
Medium Expert
(d) Medium
Figure Distribution mean distance between query states and retrieved states of different
Thus, l determines the effective radius implicitly and monotonically: a larger & expands the radius (dg, dr)) and increases the size and heterogeneity of Ωd= whereas a smaller k leads to tighter neighborhoods with more consistent local value structure.
H.6 ADDITIONAL VISUALIZATION OF LEARNED Q-VALUE COMPARISON
We extend the visualization analysis of learned Q-values for ICQL and IQL by comparing them with Q-values learned by the online RI method SAC the Walker2d-MMeduu- Walker2d-
27

---

Published as a conference paper at ICLR 2026
Medium, and Walker2d- visualization. We also include additional scatter plots that compare the Q-values estimated by each method against the SAC oracle. The visualizations clearly show that the correlation between TCQL and SAC is stronger than that between TQL and SAC, which indicates that TCQL, produces more accurate value estimates than IQL.
datasets. We scale all Q estimates the range [0.1] before
shown in Figure S and Figure(y These plots
025
025
5.25
Med1
SAC
ICQL
IQL
05
05
455
(b) Medium
SAC
ICOL
IQL
075
0.5
5.75
(c) Modi um-Roplay
Figure 8: Q-value of -edium--xppeeeeeeee Walker2d-Medium, and Baakkerdd-eeduu-RReee dataset on t-SNE mapped state distribution.
28

---

Published as a conference paper at ICLR 2026
SAC vs ICQL (r=0.629)
SAC vs FQL 02-00
1030
830
4300
630
3300
2.200
50
120
250
200
202
32
150
250
150
SAG O-sals
SAC Qualie
un-Fxpe
SAC vs ECQL (r-0.593)
SAC vs KQL (r-.0.361)
1503
502
100 Ordaa
102
6-
-1003
102
-150
200
250
201
150
300
250
SAC (irul
560 Q mu
(b)
SAC ICOL (r-0.792)
SAC vs IOL (r=0.215)
503
1308 Qvide
202
-503
-000
320
200
SAC Divuue
SAC ualu
(c) un-playyy
Figure Q-value correlation of Walker2d-Medium-Exper Walker26-Medium, and Walker2d- Medium-Replay dataset. Red: Correlation between Q-values learned by IQL and SAC. Blue: Correlation between Q-values learned by ICQL and SAC.
29

---

Published as a conference paper at ICLR 2026
H.7 ANALYSIS OF IN-CONTEXT CRITICS
In this section, we conduct additional analysis of the functionality of our in-context Q estimator. By construction, the forward pass of our in-context Q estimator is equivalent to step-wise optimization of the TD error. We analyze the outputs and parameter distributions of each intermediate layer to validate its effectiveness. We randomly select 10 different states and their corresponding actions retrieve 20 relevant transitions using cosine
from the offline dataset of state similarity, and estimate the Q-values for these state-action pairs. We store the outputs of all intermediate layers, and the visualization results are shown in Figure TO From Figure I0] we observe that the Q estimates exhibit a converging trend as the layers become deeper, which validates the iterative refinement process.
200
150
Q Estimtts
10
12
15
19
Figure 10: Q-estimates at each intermediate layer.
30


Reviews:
[{"review_id": "dRMq0Dt0iQ", "reviewer": "Reviewer_YKZA", "summary": "This paper proposes ICQL, an offline RL method that reframes Q-learning as contextual inference. For each query state-action it retrieves a small, state-similar set of transitions from the offline dataset and uses a linear transformer to infer a local linear Q-function, then plugs this into an IQL-style training pipeline. The motivation is that offline datasets often have locally coherent regions of value, but adjacent regions can differ sharply. A single global critic tends to smooth these away, whereas ICQL tries to learn \"per-query\" value estimates from retrieved context. On D4RL benchmarks, ICQL often outperforms strong offline RL baselines and the ablations support the claim that retrieval/local context actually matters.", "questions": "- What happens in ICQL when the retrieved neighborhood is not well modeled by a single linear head (e.g. contact-rich or multimodal regions)?\n\n- Could the authors clarify what is meant by \"nearby clusters presenting as noise\" and how this maps to the formal local-set definition?\n\n- The authors introduce $(d,\\bar d)$ for local sets, but the experiments effectively tune $k$ in $k$-NN. How should we relate the theory parameters to the practical $k$?\n\n- Is the choice of a linear transformer/head essential for stability/theory, or would a slightly more expressive local model (small MLP over the retrieved set) also work?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Treats offline Q-learning as \"retrieve local neighbors, do in-context TD, act\", which matches the empirical observation that value structure is local and nonuniform across the dataset.\n- ICQL often improves over IQL/CQL-style baselines on MuJoCo.\n- Varying context length / $k$ and retrieval strategy shows performance is sensitive to the retrieved neighborhood, which supports the paper's core claim that which transitions you condition on matters.\n- The paper states assumptions (local linear approximability, coverage of the local set) and derives a performance bound.", "weaknesses": "- ICQL requires a retrieval step for every training query. For large offline buffers this $k$-NN step can dominate wall-clock, but the paper does not report complexity, index structure (exact vs. ANN), or training time vs. IQL/CQL. This is a real deployment concern the paper currently ignores.\n\n- The method’s notion of a \"nearby transition\" is Euclidean distance in the (low-dim) state space. This is reasonable for MuJoCo/Adroit/Kitchen but does not address high-dimensional observations (images) where $L_2$ is brittle and the coverage assumption becomes much harder to satisfy. A learned or task-aware retriever is not explored, only mentioned.\n\n- The theory leans on \"within a small neighborhood the Q-function is linearly approximable\". This is a strong assumption for tasks with multimodal returns, contact-rich transitions, or poorly aligned state features. The paper does not analyze what happens when the neighborhood is not well modeled by a single linear head.\n\n- The current wording about \"nearby clusters presenting as noise\" is hard to follow and mixes an empirical visualization (clusters in the dataset) with the formal local-set definition.\n\n- There are experiments varying context length / $k$, and they do show non-monotonic behavior, but the paper avoids discussing the $d,\\bar d$ definitions it introduces earlier and effectively hides the choice inside $k$-NN. A clearer statement of how performance scales with $k$ (and how they pick it) would help.\n\n- The paper instantiates the method with a linear transformer/head but does not compare against a slightly more expressive local model (e.g., a small MLP on the retrieved set). As a result, it is unclear whether \"in-context linear-TD\" is essential, or just convenient.\n\n- There are typos and phrasing that make already technical sections (the definitions and the \"in-context TD\" explanation) harder to read.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "djNcs96nYp", "reviewer": "Reviewer_qZiC", "summary": "The paper introduces In-Context Compositional Q-Learning (ICQL) for offline RL: a transformer retrieves top-k similar transitions and fits a local linear Q for each query, instead of learning a single global critic. Claims: (i) with bounded features and adequate retrieval coverage, the greedy policy is near-optimal (performance-gap bound); (ii) on D4RL (Mujoco/Adroit/Kitchen), ICQL yields consistent gains—largest on compositional tasks—supported by ablations on context length, transformer depth, and retrieval strategy. The work positions ICQL as retrieval-conditioned value learning (not return-conditioned action modeling), aiming to better exploit task locality than standard offline RL critics.", "questions": "\"We observe that, for each RL control task, the state space can be inherently divided into multiple sub-tasks\": what kind of sub-tasks, please add some background description? Different parts of the state space correspond to different “phases”\n\nFigure1: I don't see obvious difference between QueryB-R1,2,3. \n\nHave you tried learned value-aware similarity (contrastive encoders predicting Q-proximity) or state–action NN retrieval, and how does that affect the theory’s coverage term?\n\nCan you report FLOPs/memory vs. baselines and show compute–performance trade-offs for context length?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Clear idea & scope: reframes Q-learning as local, retrieval-conditioned regression—intuitive and timely.\nTheory with knobs: bound ties performance to local linearity and retrieval coverage, matching design levers.\nEmpirically persuasive where it should be: biggest wins on multi-stage “Kitchen,” aligned with the compositionality thesis.\nAblations with signal: context length / retrieval choices matter in the expected directions.", "weaknesses": "Baseline gap: no direct comparison to retrieval-augmented sequence models (e.g., Retrieval-Augmented Decision Transformer: External Memory for In-context RL). Both ICQL and retrieval-augmented Decision Transformer methods use external memory to condition decisions in offline RL\n\nCompute transparency: training-time cost/latency vs. IQL/CQL/TD3+BC not reported.\n\nRetrieval proxy: L2 proximity in state space may misalign with Q-similarity; robustness to metric choice is under-explored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oZovmxaLW3", "reviewer": "Reviewer_srbX", "summary": "This paper proposes an offline RL algorithm called ICQL which solves offline Q-learning as a contextual inference problem. It retrieves local transitions from the dataset and fits a local linear Q-function via a linear-attention critic. The authors theoretically prove the bounded Q-approximation error and near-optimality guarantee for the greedy policy. In the experiments, ICQL outperforms most of the offline tasks in D4RL against offline RL baselines.", "questions": "1. The Definition 3.2 is also not clear to me\n    1. In Equation 4, why is it $-||s_{query} - s_i||^2$? If I understand correctly, it should be $k$ nearest states to $s_{query}$, right?\n    1. $d_{min}^s$ is the minimal distance to $s_{query}$ in $\\overline{\\Omega}^k_{s_{query}}$. Shouldn't $\\Omega_{s_{query}}^{d_{min}}$ be empty since there is no closer states in this set?\n1. The algorithm is still a bit unclear to me. Should the local weight vector $w_s^*$ be trained for every $s_{query}$? If so, I recommend the authors make a comprehensive analysis in training cost as it will take enormous computational resource.\n1. Why does the algorithm fail on Hammer-Human-v1, could the authors make further justificaiton on this failure case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The authors provide a clear formulation of their method and theoretical justificaiton of Q-approximation error bound and performance gap bound. The theoretical contribution is solid. The idea is simple but quite interesting. Capturing the local structure is often easier than considering the global pattern. The authors gives a feasible approach to this direction.\n1. The experiments on the well-known D4RL benchmark are comprehensive. They demonstrate a clear comparison against popular offline RL algorithms.\n1. The supplementary material gives sufficient information for theory verificaiton and experiment verification. They are very helpful for follow-up research.", "weaknesses": "1. In offline RL, the coverage ratio $\\sigma$ is an important factor influencing algorithm performance. In section 4.3.3, the authors only provide an indirect ablation to analyze the effect of $\\sigma$. A direct analysis of the coverage of ideal local transition set would make the work more practical.\n1. I notice some failure cases. While I often have a good tolerance on failure cases in novel methods, the performance on Hammer-Human-v1 looks too bad. A deepr analysis of this environment would be helpful to clarify the method's scope of applicability.\n1. The Figure 3 does not make sense to me. \n    1. Does SAC work as an oracle Q value estimator in this figure? If so, since SAC has two Q networks, which one do the authors use? Additionally, do all the three methods share the same policy network to ensure fairness? I also suggest approximating Q values via Monte Carlo rollouts using the same policy as the oracle, which is less biased by the training algorithms.\n    1. The Q values should be normalized to the same scale across three plots so it can provide a uniformed color scheme to readers. The current plots are hard to analyze.\n    1. I suggest replacing the tSNE plots with scatter plots that directly show the correlation of the estimated Q values and oracle Q values.\nMinor issue\n1. In line 377, it should be \"Figure. 3\" instead of a stand alone \"3\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JedAhRtOFm", "reviewer": "Reviewer_MZEj", "summary": "This paper introduces ICQL , a novel offline RL framework that formulates Q-learning as a contextual reasoning problem. The core idea is to use a linear Transformer to learn local Q-functions from retrieved similar transition data, rather than fitting a single global Q-function. The authors argue that many tasks exhibit compositional structures and local value patterns, and they provide theoretical guarantees on bounded approximation error and near-optimal policy extraction. Experiments on the D4RL benchmark demonstrate performance improvements of 8.6% on Gym, 6.3% on Adroit, and 16.4% on Kitchen tasks.", "questions": "1.Has the paper conducted any comparison between linear Transformers and standard self-attention? Which specific properties of linear attention are deemed essential or indispensable for the proposed method’s effectiveness?\n\n2.Assumption 3.3 requires knowledge of the coverage parameter σ, but it remains unclear how this parameter is determined in practice. How is σ set or estimated empirically, and how well does the theoretical analysis based on this assumption correlate with actual performance?\n\n3.Why does ICQL perform catastrophically on the Hammer-Human task (−49.4%)? Is this performance drop solely due to dataset characteristics, or does it reveal a fundamental limitation of the proposed method? Clarifying this distinction would help assess the robustness and generality of ICQL.\n\n4.The paper states that, \"After training, the extracted policy can be evaluated on its own without extra retrieval process or\ncontextual inference.\" However, since the policy is trained using context-dependent Q-values, how does it function without access to retrieval during evaluation? Does this setup introduce a train–test mismatch or affect the consistency between training and inference behaviors?\n\n5.Why does the paper not include comparisons with more recent approaches (e.g., [1, 2, 3, 4]) or other modern offline RL algorithms? Such comparisons are essential to demonstrate the competitiveness and relevance of the proposed method in the current research landscape.\n\n[1]Tarasov, D., Kurenkov, V., Nikulin, A., & Kolesnikov, S. Revisiting the minimalist approach to offline reinforcement learning. NIPS， 2023.\n\n[2]Mao Y, Wang Q, Qu Y, et al. Doubly mild generalization for offline reinforcement learning. NIPS, 2024.\n\n[3]Park S, Li Q, Levine S. Flow q-learning. ICML, 2025.\n\n[4]Li Q, Zhou Z, Levine S. Reinforcement learning with action chunking. arXiv preprint arXiv:2507.07969, 2025.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Viewing value learning in offline RL as a In-contextual learning problem is highly innovative and well-motivated by the observation that the state space exhibits local compositional structure. The paper provides a formal analysis of bounded approximation error (Theorem 3.5) and connects retrieval quality to performance through a coverage assumption. Moreover, it successfully integrates with existing methods such as IQL and TD3+BC, demonstrating broad applicability of the proposed framework.", "weaknesses": "1.The local linear Q-function approximation (Definition 3.1) is a standard practice in the literature. The paper’s main theoretical contribution lies in connecting this formulation to in-context learning, but the underlying assumptions—particularly Assumption 3.3 on set coverage—are quite strong and may not hold in practical settings.\n\n2.Retrieving k-nearest neighbors for each query and processing them with a multi-layer Transformer introduces significant computational overhead compared to standard methods. However, the paper does not provide any runtime analysis or comparison to quantify this cost.\n\n3.Weak justification for using linear Transformers: Why specifically choose linear attention? Although the paper claims to reveal the “underexplored potential of linear attention,” it provides no comparison with standard self-attention, leaving the rationale insufficiently supported.\n\n4.Limitations of the experimental evaluation: The experiments are conducted only on the D4RL benchmark, which is not sufficiently comprehensive. More importantly, the paper does not compare against recent state-of-the-art offline reinforcement learning methods, limiting the strength of its empirical validation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)