Scalable In-Context Q-Learning
Authors: Jinmei Liu
Fuhong Liu
Zhenhong Sun
Jianye HAO
Huaxiong Li
Bo Wang
Daoyi Dong
Chunlin Chen
Zhi Wang
Venue: ICLR 2026 Poster (2026)
Keywords: In-context reinforcement learning
Q-learning
advantage-weighted regression
world model
Primary area: reinforcement learning
TL;DR: We propose an innovative framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization, while retaining the scalability and stability of the supervised pretraining paradigm.

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

Published as a conference paper at ICLR 2026
SCALABLE IN-CONTEXT Q-LEARNING
Jinmei Liu', Fuhong Liu', Zhenhong Jianye Hao", Huaxiong Li',
Bo Wang"*, Daoyi Dong'. Chunlin Chen', Zhi Wang
1 Nanjing University
Australian National University
3 Tianjin University
4 University of Technology Sydney
jm Smliussmail .ju. edu. on jianye. haoêt ju edu. en (huaxiongli, i, bowangsme, clcher, wewen,,wwwwae zhiwang
zhenhong suneanu, edu, au
daoyi dong@uts. www.eeutsseduud.a. eu
cn
ABSTRACT
Recent advancements in language models have demonstrated remarkable in- in-context learning abilities, prompting the exploration of in-context reinforcement learning (ICRL) to extend the promise decision domains. Due to involving more complex dynamics and temporal correlations, existing aCI.oaape may face challenges in learning from suboptimal trajectories and achieving precise in-context inference. In the paper, we propose Scalable In- I--Cotext Q Q-Learning (S-ICQL), an innovative framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization, while retaining the scalability and stability supervised pretraining. We design a prompt-base based multi- head transformer architecture that simultaneously predicts optimal policies and in-context value functions using separate heads. We pretrain a generalized world model to capture task-relevant information, enabling the construction of compact prompt that facilitates fast and precise in-context inference. During training, we perform iterative policy improvement by fitting state value function to upper- expectile of the Q-function, and distill the in-context value functions into policy extraction using advantage-weighted weighted regression. Extensive experiments across a range of discrete and continuous environments show consistent performance gains over various types of baselines, especially when learning from suboptimal data. Our code is available at
INTRODUCTION
RL is a pivotal mechanism for training autonomous agents to solve complex tasks in interactive envi- ronments (Mnih et al.] 2015), with expanding applications in frontier challenges such as fine-tuning
language models ( (Yan et al. 20225] Hu et al. 2026 Zhan et al. 2026) and diffusion models (Black al. 2024: Liu al. [2026). A longstanding goal of RL is learn from diverse experiences and generalize beyond its training environments, efficiently adapting to unseen situations, dynamics, or objectives (Ackley & Littman 1989 Kirk al.) 2023 Hu al. 2025). A promising avenue is in- context learning (Brown et al. 2020) that trains large-scale transformer models on massive datasets to achieve remarkable generalization capabilities, adapting to new tasks via prompt conditioning without any model updates (Wang et al., 2023a Li et al. 2024a). Accordingly, in-context RL (ICRL) seeks to extend this promise to decision domains and has seen rapid progress in recent years (Nikulin et al. 2025). Existing studies contain two typical branches: algorithm distillation (AD) (Laskin et al. (023) and decision-prer prettrind transformer (DPT) (Lee et al. (023). due to their simplicity and generality. They commonly employ cross-episode transitions as few-shot prompts and train transformer-basee policies under supervised pretraining (Lin et al. 2024), followed by various improvements from model-based planning (Son et al. 2025). hierarchical decomposition (Huang et al. (024)) importance
weighting (Dong et al. 2025), etc (Sinii et al. 2024 Tarasov et al. 2025 Dai et 2024).
Though, significant challenges may emerge when extending the promise of in-context learning from (self-) supervised learning to RL, since RL involves more complex dynamics and temporal
*Correspondence Zhi Wang <zhiwang@ www.unuunnu.euu..n...... and Bo Wang doowangsme@nju.edu.cco

---

Published as a conference paper at ICLR 2026
correlations (Silver al., 2021). First, previous studies usually adopt a supervised pretraining paradigm, failing to go beyond imitating collected data (Yamagata et al. 2023). AD can require long-horizon context and inherit suboptimal behaviors due to the gradual update rule (Son et al. 2025). DPT relies on an oracle for optimal action labeling that can often be infeasible in practice (Tarasov et al.. (025). These limitations may hinder the efficient learning of optimal policies, especially when only suboptimal datasets are available. Second, substantial scope remains to advance the design of efficient prompts that can precisely encode RL task information. In language communities, prompts are concise, precise text instructions rich in semantic information, as text naturally conveys high-level concepts like objects (nouns) and actions (verbs) (Kakogeorgiou et al. 2022). However, ICRL approaches generally take raw transitions as prompts that have many more tokens than sentence and can be highly redundant, Within transitions from offline datasets, the task information can be entangled with behavior policies, thus producing biased task inference at test time (Yuan & Lu 2022). Hence, this type of prompt may be insufficient to precisely capture relevant information about decision tasks. The aforementioned limitations raise a key question: Can design scalable ICRL framework using lightweight prompts that precisely capture task-relevant information, while unleashing the core potential of fundamental reward maximization to learn from suboptimal data?
To tackle these challenges. we draw upon two basic properties inherent in full RL. First, classical RL algorithms learn value function to backpropagate expected returns using dynamic programming updates, | showing appealing stitching property, i.e., the ability to combine parts of suboptimal trajectories for finding globally optimal behaviors (Sutton et al., 1998). Naturally, the stitching property offers a compelling avenue for unleashing the potential of ICRL architectures, attaining substantial improvement over suboptimal data. Second, as RL agents learn through active interactions with the outer environment, the decision task is fully characterized by the environment dynamics, i.e. the state transition and reward functions p(s', als..)) a). The world model (Hafner et al. (025) can learn an internal representation of the environment dynamics, and is intrinsically invariant to behavior policies or collected datasets (Wang et al, 2024). Hence, leveraging the world model holds promise for designing a lightweight prompt structure capable of precisely encoding task-relevant information.
Drawing inspiration from full RL, we propose Scalable In-Context Q-Learning (S-ICQL), an innovative framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization. First, we design a prompt- based multi-head transformer architecture to maintain scalability and parameter efficiency. The model simultaneously predicts optimal policies and in-context value functions using separate heads, given task prompt and corresponding query inputs. Second, we pretrain generalized world model to capture task-relevant information from the multi-task offline dataset, and use it to transform a small number of raw transitions into lightweight prompt for fast and precise in-context inference. Finally, we perform iterative policy improvement by fitting state value function to an upper-expectic of the Q-function, and distill the in-context value functions into policy extraction using advantage-weiihtted regression. This formulation allows for learning a policy to maximize the Q-values subject to an offline dataset constraint, while retaining the scalability and stability of the supervised pretraining paradigm. In summary, our main contributions are threefold:
We introduce dynamic programming to supervised ICRL architectures, unleashing its potential toward learning from suboptimal trajectories with efficient reward maximization.
We design lightweight prompt structure that leverages world modeling to accurately capture task-relevant information, enabling fast and precise in-context inference.
We propose a scalable and parameter-efficient ICRL framework that integrates the advantages of RL and supervised learning paradigms. Comprehensive experiments validate our superiority over a range of baselines, especially when learning from suboptimal data.
2 RELATED WORK
The concept of agents adapting their behaviors within the context without model updates builds on earlier work in meta-RL (Beck et al. 2025), such as the memory-based RL2 (Duan et al. 2016)
In this paper, we use dynamic programming to indicate the fundamental characteristic of any RL algorithm relying on the gallman-backun operation. It updates state (or state-action) values based on value estimates of successor states (or state-action pairs), i.c., bootstrapping (Sutton et al. 1998). We use dynamic programming and Q-learning interchangeably to refer to this fundamental property of RI..
2

---

Published as a conference paper at ICLR 2026
and LLIRL Waan et al. 2022; 2023b), the oppimization-based MAML (Finn al. 2017) and
MACAW (Mitchell et al. 2021 ). and the context-based VariBAD Zintgraf et al. 2021). Meta-
DT (Wang al. 2024), etc. (Li et al. 2024b; Zhang et al.] 2025). Recently, there has been a shift
towards employing transformers to implement ICRL (Moeini et al.
2025 Nikulin et al. 2025) driven
by their proven ability capture long-term dependencies and exhibit emergent in-context learning behaviors Brown et al. (2020) AMAGO (Grigsby et al. 2024a b) trains long-sequence transformers over entire rollouts with actor-critic learning to tackle in-context goal-conditioned problems, and (Lu et al. 2023) leverages the S4 model's 's ability to handle long- range sequences for ICRL tasks. In offline settings, two classical branches are AD (Laskin et al., 2023) and DPT (Lee et al. 2023). AD trains a causal transformer to autoregressively predict actions using preceding learning histories as context, while DPT predicts the optimal action based on a query state and a prompt of interaction transitions. Follow-up studies enhance in-context learning from different perspectives (Son et al. [2025] Huang
ct al. 2024| Dai et al., 2024] Zisman et al. 2024 Wu ct al. 2025). For example, IDT (Huang cf al. 2024 designs hierarchical decision structure, and DICP (Son et al. 2025 incorporates model-based
planning with a learned dynamics model. These approaches generally adopt a supervised paradigm with raw transitions as prompts, while we harness dynamic programming for reward maximization and leverage world modeling to construct more efficient prompts.
DIT (Dong et al. 2025) and IC-IQL (Tarasov al. 2025) are the most relevant to our work. DIT also uses a weighted maximum likelihood estimation loss to train the transformer policy, where the weights are directly calculated from observed rewards in collected datasets. In contrast, we learn separate in-context value functions for computing advantage weights, enabling more stable and sample-efficient policy extraction. IC-IQL integrates a Q-learning objective into AD. The policy and value networks are updated by optimizing only their heads, without propagating gradients to the transformer backbone. In contrast, we train the full multi-head transformer policy end-to-end releasing the transformer's scalability with simplified pipeline. We also design a precise, lightweight prompt structure to overcome the limitation of AD algorithms that require long training histories as context. Empirical results in Sec. demonstrate our superiority to these two baselines.
METHOD
In this section, we present S-ICQL (Scalable In- Context Q-Learning), an innovative ICRL framework that leverages dynamic programming and world modeling for efficient reward maximization and task generalization. We adopt a prompt-based multi-head transformer architecture that simultancously predicts optimal policy and in-context value functions using separate heads, ensuring scalability and parameter efficiency. Figure (illustraas the method overview. The algorithm pseudocodes are presented in Appendix B and detailed implementations are given as follows.
3.1 PROBLEM STATEMENT
We consider multi-task offline RL setting, where tasks follow distribution M'= P(M). Tasks share the same state-action spaces S,A A but differ in reward functions R or transition dynamics T. An offline dataset D = Es a), rj. sfy) is collected by arbitrary behavior policies for each task out of a total of N training ones. The agent can only access the offline datasets (D') to (af]= train in-context policy as Ty (af [ ), where 3 is prompt that encodes task-relevant information (c.g.. past interaction history in AD or transition sequences in DPT). During testing. the trained policy evaluated on unseen tasks sampled from ((M) through direct interaction with the environment. With fixed policy parameters 0, all adaptations occur through the prompt/context 3 that is initially empty and gradually constructed from history interactions. As / evolves, the model refines its predictions, analogous policy updates in conventional RL. The objective is to learn an in-context policy that maximizes the expected episodic return over test tasks as J = EMMP(MM)[M (T)].
3.2 SCALABLE MODEL ARCHITECTURE
World Modeling, As shown in Figure (a). we design a lightweight prompt structure capable of encoding precise information about decision tasks. The world model, representing environment dynamics p(s', r|s, a) (Hafner et al. 2025), fully characterizes the underlying task and remains invariant to behavior policies or the datasets collected. Inspired by this fundamental property of RI,,
3

---

Published as a conference paper at ICLR 2026
Previous ststep Experiencee
The h-horizon Context Prompt
B
RA
Contert Rrooder
Contest Eauzsder &
Prompt
Causal Transformer
Representation
Cousal Transforerr
Duery
Dynutuic Decocrey o
brau
Q Fiead
V Vhead
Query %
To
(et (sg,8)
Qs (v.
Vig (s::
(a) World Model Pretraining
(b) In-textext Policy Pretraining
(c) Online Testing
Figure 1: The overview of S-ICQL. (a) We pretrain generalized world model to accurately capture task-relevant information from the multi-task offline dataset as in Eq. (3). and use the context encoder transform a small number of raw transitions into a precise and lightweight prompt /3 as in Eq. ( (b) We design prompt-based multi-head transformer model that simultaneously predicts the optimal policy Te (a |s; (8, the state value function (8; 3), and Q-function Qo a; (3) using separate heads, given the task prompt [ and corresponding query inputs (s or s, s,a.. We learn Vo by expectile regression as in Eq. (5) and use it to compute Bellman backups for training Qe as in Eq. (4). The in-context value functions
distilled into policy extraction using advantage-weigee in Eq. (6). (c) Online testing by interacting with the environment. The prompt is initially empty and gradually constructed from history interactions using the pretrained context encoder.
regression as
we pretrain a generalized world model to acquire task-relevant information from the multi-task offline dataset. The world model contains a context encoder Eg that abstracts recent k-step experiences Wi (st-k, dt k; A k ... 8t-1, 0t-1: 1: Tr-1, 8t, a²) into task representation z E (nf), and a dynamics decoder D. that predict the instant reward and next state conditioned on the task representation as [A,, & = (s,, ax; 2/). Then, for training the in-context policy, we sample a short h-step trajectory randomly from dataset D', and use the pretrained Es to transform raw transitions in this trajectory into compact task representations
8 = 12i,
2 = E.
A task i.
Then, we construct a lighteeight prompt 33 using the h-step task representation for fast and precise in-context inference, as opposed to AD algorithms that can require long training histories as context. Sec. 3.3 presents the world model and its pretraining process in detail.
Dynamic Programming. Classical RL algorithms learn value function to backpropagate expected returns using dynamic programming updates, showcasing high stitching capacity that allows for not only imitating collected data but also achieving substantial improvement beyond it (Sutton et al.) 1998). Inspired by this, we harness the stitching property to offer a promising avenue for unlcashing ICRL's potential toward explicit reward maximization. Following common practice in advanced offline RL. algorithms (Kostrikov et al. 2022; Snell et al.) 2023), we learn both a state value function V and an action value function Q. As shown in Figure (b),, we design prompt-based multi-head transformer architecture to simultancously predict the optimal and value functions. Let 0 denote parameters of the integrated transformer model. The model outputs the policy (a(a|s; B), the state value (s; 5), and the action value Qo (3) using three separate heads, given the task prompt B and corresponding query inputs (s or s, a). Sec. 3.4 presents the detailed learning process.
This flexible design maintains two appealing properties. One is parameter efficiency. Our model only introduces two additional lightweight heads compared to traditional ICRL methods such as DPT, and the resulting increase in parameters is negligible relative to the foundational transformer backbone. Another is model scalability. Our central component remains an end-to-end causal transformer architecture that holds the promise for training RL models at scale. We learn the in-context policy by advantage- -weatagg--eiighted regression. a supervised learning style with a simple and convergent maximum likelihood loss (refer to Eq. (6)). It ensures that our method preserves the scalability and stability inherent to the supervised pretraining paradigm. Moreover, we train the whole multi-heat transformer policy end-to-end (refer to Sec. 3.5. releasing the transformer's scalability with simplified pipeline.
4

---

Published as a conference paper at ICLR 2026
3.3 WORLD MODEL PRETRAINING FOR PROMPT CONSTRUCTION
The environment dynamics, i.e., the reward and state transition functions p(s', rls, a), can share some common structure across the task distribution. For each task M', we approximate its dynamics by generalized world model P that is shared across tasks, defined as
(r(re: 8141 St, ax) P(r:: 914 at; zt), V task i.
As the true task identity is unknown, we infer task representation zf from the agent's k-step experience within task M as ny. [st-k= Ot-k: Pt-k: St-1, at-1, "t-,,, at)². The intuition is that the true task belief can be inferred from the agent's ' history interactions, similar to recent meta-RL RL studies (Zintgraf et al. 202T Ni et al., 2023). We use context encoder Es, to abstract recent k-step experiences into task representation as 2 (m))). which augmented into the input of a dynamies decoder Dup to predict the instant reward and next state [fr, [t+1]=1,, (St, (st, az z4). Based on the assumption that tasks with similar contexts will behave similarly (Lee et al. (020), our generalized
world model can extrapolate meta-level knowledge across tasks by precisely capturing task-relevant information. The context encoder and dynamics decoder are jointly trained by minimizing the reward and state transition prediction error as
L( pp) = Bag [I][re, s141] - Dp(s,, az; 2 2 = Es (n:)]
4 task i.
(3)
After proper pretraining, we freeze the generalized world model for prompt construction in Eq.
3.4 IN-CONTEXT Q-LEARNING
In-Context Value Functions. The Q-function is trained to minimize the Bellman as
(o) = at 1>D (r(s), at) + + (s142: 8*) - Qo(s), at d")" A task i.
(4)
The in-context state value function Vi aims fit an upper-expectile of the O-function, and is trained to minimize an expectile regression loss as
(v = E [L (Qg(s) 8+) - (o(s); V taask i,
(5)
where (z( = |ا - 1m < 0) . u² is an asymmetric loss function with a expectile parameter E (0.5. 1). This form of expectile regression reduces the influence of Q V predictions by factor of 1 w- while assigning more importance to Q V predictions by factor of w. In this way, we predict an upper-expectile of the temporal-difference target that approximates the maximum of r(s), at) + yQ2 (s141. 1: f3) over actions d:+, constrained to the dataset actions. More details can be found in Implicit Q-Learning (Kostrikov et al. 2022).
In-Context Policy Extraction. The value function learning procedure allows for stitching suboptimal trajectories to discover globally optimal behaviors. Then, we distill the in-context value functions into policy extraction with advantage- advantage-weighted weighted regression (Peng et al., 2019), supervised learning style that uses a simple and convergent maximum likelihood loss function as
(0) - (s))))) exp ( (Qa(s), at; (²) (s1:;; B))) logg To(a (: 4 task i, (6)
where A > 0 is a temperature parameter. Using weights from in-context value functions, the objective not merely to clone behaviors from the dataset but learn policies that maximize Q-values under a distribution constraint from dataset actions. This formulation aims to select and stitch optimal actions in the dataset while retaining the scalability and stability of the supervised pretraining paradigm.
3.5 OVERALL OPTIMIZATION
To unify policy learning and value updating within a single architecture, the transformer backbone and its three specific heads are jointly optimized. This joint training objective combines supervised training with dynamic programming updates, allowing the model to simultaneously learn in-context policies and value functions while retaining scalability and stability. The overall loss is defined as
[(0)= с₁ C. (0) + (Q) + c3 (V).
where we set coefficients (G,, C2, C3) to a balanced ratio of 11: in all experiments. A detailed analysis of how different coefficient choices affect performance is provided in Appendix HH This design ensures that S- S-ICQL integrates lightweight prompt construction with a unified optimization pipeline, achieving efficient reward maximization and robust generalization across tasks.
5

---

Published as a conference paper at ICLR 2026
DIT
IC/QL
S-IQQL
Darkroom
700
Push
Reach
800
Average Neturn
1500
Average Return
500
20
650
100
10
Cheetah-Vel
Walber-Param
Ant.Dir
-40
500
800
Avorage Return
-50
Average Return
Avorage Raturn
-60
-70-
300
4200
200
200
Episode
Episode
Episode
Figure 2: Few-shot evaluation return curves of S-ICQL. and baselines on Mixed datasets.
EXPERIMENTS
We comprehensively evaluate the in-context learning capacity of S-ICQL on popular benchmarking domains across different dataset types. general, we aim answer the following questions:
Can S-ICQL consistently outperform other strong baselines on unseen tasks? (Sec. 4.1
How do the prompt construction via world modeling and the reward maximization via dynamic programming affect the in-context learning performance, respectively? (Sec. 4.2)
Can S-ICQL also achieve well performance for out-of-distribution (OOD) tasks? (Sec.
4.3
Does S-ICQL really have the stitching capacity find globally optimal policy? (Sec. 4.4
S-ICQL robust to the quality of offline datasets and hyperparameters? (Sec. 4.5 and Appendix
Can S-ICQL encode task-relevant information for efficient prompt construction? (Sec. 4.6
Does S-ICQL maintain strong performance
complex and diverse environments? (Sec. 4.7
Environments. We conduct experiments three challenging benchmarks commonly used for evaluating ICRL algorithms: (i) DarkRoom Laskin et al. (2023), ) a 2D discrete environment where the agent must locate an unknown goal. (ii) MuJoCo Todorov al. 2012), standard testbed including tasks with varying reward functions and transition dynamics. (iii) Meta- World ML1 Yu et al. (2020), robotic benchmark with 50 manipulation tasks. Tasks are randomly sampled and split into training sets wann and test sets Mes Further environment details are provided in Appendix
Pretraining Datasets. For DarkRoom, suboptimal datasets are generated using noisy action selection strategy by combining optimal and random policies. For MuJoCo and are collected using single task RL policy for each task. We construct two qualities of offline datasets: Mixed and Medi um. More details on dataset construction are provided in Appendix D
datasets
Baselines. We compare to six competitive ICRL approaches and one offline meta-RI. method, including: 1) IC-IQL TTrasov et al. 2025). 2) DIT (Dong et al.) 2025), 3) DICP (Son al. 2025), 4) IDT (Huang et al. 2024). 5) AD (Laskin et al. (023). 6) DPT (Lee et al. (023). and 7)
UNICORN (Li 2004b). Details about these baselines provided in Appendix
To ensure a fair comparison, we conduct a few-shot evaluation for all methods. During testing, each policy directly interacts with the environment for few episodes using fixed parameters, conditioned on a prompt sampled from past interactions, Results are reported as the mean of 10 trials with 95% bootstrapped confidence intervals, and standard emror also provided.

---

Published as a conference paper at ICLR 2026
Table 1:
Few-shot evaluation returns of S-ICOL and baselines on Mixed datasets, i.c., numerical results of converged performance from Figure 2
Methods UNICORN
Darkroom
Push
Reach
Cheetah- Vel
Walker-Param
Ant-Dir
127.69: 44.04
646,52 13.60
-55.29+ 7.90
232.23 21.73
416.35# 18.29
DPT
22.12 1.09
362.74= 1.91
736.72+ 39.42
-78.35 4.50
257.11 16.36
591.31 45.54
AD
42.72 2.14
604,50= 15.52
738,
96# 39.38
-67.37 3.53
424.82+ 19.23
215.01 46.61
IDT
40.70:: 1.44
621.58:
11.16
790.68:: 13.01
-59.46: 3.10
343.01 7.95
631.831 35.29
DICP
59.76: 2.80
487.28
66.33
706.46:: 14.32
-66.534 2.58
403.90:1 8.09
745.05: 24,12
DIT
30.90 1.96
633.584
19.22
758.92+
19.13
-74.50+
3.24
253.94 9:04
723.494: 27.36
IC-IQL
60.12 1.33
646.08
30.34
773.334
11.79
-56,534
1.98
391.38 13.97
713.264 27.23
S-ICQL
66.05 2.37
653.04 31.22
806.97 6.35
-35,48 1.33
466.72:1 24.06
813.34.1 14.12
Table Offline evaluation returns of S-ICQL and baselines on Mixed datasets.
Methods
Darkroom
Push
Reach
Cheetah- Vel
Walker-Paral
Ant-Dir
DPT
29.02+9998 13.22:
783.494222.3 110.16:107.03
-65.92:1-6.74
221 103.76
667.344:243.76
DIT
805.57+13.10 136.41 575.07.1118.11
-61.11+33.99
259.07=+0915
744.24=144.69
S-ICQL
76.63
818.04.18.66
-37.87.110.6
414.65
835.98 117.25
4.1
MAIN RESULTS
The baselines typically work under the few-shot setting, since they rely on task prompts or warm-start data to infer task representations. We evaluate S-ICQL against these baselines in an aligned few-shot setting, where all methods utilize the same number interaction trajectories for task inference. Figure 2 illustrates the evaluation return curves across various environments using Mixed datasets, and Table [ summarizes the numerical results of converged performance. Across environments with varying reward functions and transition dynamics, S-ICQL consistently demonstrates superior data efficiency and higher asymptotic performance. In more complex enviromments such as Ant-Dir and HalfCheetah- Vel, the advantage of S-ICQL more pronounced, highlighting its promising in-context learning abilities on challenging tasks. Moreover, S-ICQL generally exhibits lower variance during the testing phase, indicating both improved data efficiency and better training stability. Notably, DPT yields suboptimal performance across most environments, underscoring the critical role of incorporating fundamental reward maximization into our framework. The significant improvement over DIT and IC-IQL also verifies that our method provide more efficient way for policy extraction with diviaattag-weigghte regression or incorporating Q-learning objectives.
We also conduct offline evaluation to assess the robustness of S-ICQL trajectories collected by external behavior policies. Specifically, we use trajectories generated by non- On-S--CQLL policies and construct offline test sets by randomly sampling ten trajectories per task. We then evaluate DtT-style baselines on the same sampled trajectories to ensure a fair comparison. As shown in Table S-ICQL maintains strong performance in the offline setting, indicating that its advantages persist even when interaction is restricted and evaluation relies solely on externally collected data.
4.2 ABLATION STUDY
To assess the respective contribution of each component, we compare S-ICQL with three ab- lations: (i) w/o_c, which removes the world modeling component and directly uses a trajec- tory of raw transitions to construct prompts; (ii) w/o_q, Q. which removes the Q-learning compo- nent and learn in-context policies in pure su- eq, which
Table 3: Few-shot converged returns of S-ICQL and its ablations on Mixed datasets.
Ablation
Reach
Cheetah-Vel
Ant-Dir
w/o_oq 736.72 # 39.42 -78.35 = 4.50 591.31 + 45.54
w/o_c 792.09 6.66 -56.19 1.94 693.87 30.41
w/o_q 752,41 1 24.15 -63.66 1 8.26 784.07 1: 24.88
pervised paradigm; and (iii) w/ o/ w/o_cq, removes both world modeling and Q-learning reducing the model to the original DPT. In all ablations, the remaining structural components kept identical to those in the full S-ICQL.
S-ICQL 806.97 6.35 -35.48 1.33 813.34 14.12
7

---

Published as a conference paper at ICLR 2026
Table 4: Few-shot evaluation returns of S-ICQL and baselines for OOD tasks on Mixed datasets, i.e.,
numerical results of converged performance from Figure
Method UNICORN
DPT
IDT
DICP
DIT
IC-IQL
S-ICQL
Cheetah-Ver Vel 258. 394 43,67 1167-1377 137.26 16.28-112.54;
31,92 -103.25 1.74 16.533
266 113.18
2333 -101 89 3.50
-83.45: 2.58
Ant-Dir 367.461 15.43 205.29
295.29 1 40.00 158.781 5.17 519.80 67,10577.66111 611 51.25 51,2545..999 38.93 540.20 1: 47.22 664.95 28.95
Figure 3 shows the few-shot evaluation return curves of S-ICQL and its ablations on Mixed datasets the numerical results of converged perfor- mance. First, ablating the world model causes notable performance drop, especially in complex tasks like Ant-Dir, emphasizing its importance for precise in-context inference. Second, removing Q-learning reduces performance, underscoring its role in refining policies using suboptimal data. Inte- grating Q-learning enables S-ICQL to improve policies from noisy or inferior trajectories, enhancing its stitching ability to maximize rewards. Finally, removing both components results in the greatest decline, reducing the model to simple behavioral cloning that fails to generalize in complex tasks. Overall, the ablation results validate S-ICQL's effectiveness in harnessing both world modeling and Q-learning for efficient reward maximization and generalization across tasks.
representative environments. Taale
w/o .c
w/O.c
SHIOOL
Reach
Cheetah-Vel
Ant-Dir
Average Retumm
Average Return
Average Retum
-80
-100
700
10 Episode
-120
400
10 Episode
Episode
Figure 3: Few-shot evaluation return curves of S-ICQL and its ablations on Mixed datasets. w/o_c removes world modeling, w/o_q removes Q-learning, and w/ w/o_cq removes both components.
4.3 GENERALIZATION TO OUT-OF-DISTRIBUTTI (OOD) TASKS
Acknowledging the importance of evaluating ICRL algorithms under distribution shifts, we test S-ICQL against all baselines OOD tasks across representative environments. The target velocities for Cheetah- Vel range from [0.1, 3.0]. and the target directions for Ant-Dir range from
UNICORN
AD
DICP
ICH0L
Art.Brr
-75
Avcrogo Recum
-100
Arerape Retur
800
-125
-150
[0, 2m). We split tasks by target values con- struct an OOD setting. The model is trained on tasks with a lower range of target veloe- ities/directions and tested on the remaining higher range, ensuring that training and testing tasks come from different distributions. This setup allows us to assess S-ICQL's ability to generalize beyond the set of training tasks. As shown in Figure 4 and Table S-ICQL consistently outperforms baselines on OOD tasks. We extrapolate the meta-level knowledge across tasks by the extrapolation ability of the world model. which is more accurate and robust since the world model is intrinsically invariant to behavior policies or collected datasets. The world model shares some common structure across the task distribution (even for OOD tasks), e.g., the kinematics principle or locomotion skills. Hence, the extrapolation of the world model also works for OOD test tasks in this case. This observation is also consistent with the visualization results 5.6
175
200
200
1.0
100
Figure 4: Few-shot evaluation curves of S-ICQL and baselines for OOD tasks on Mixed datasets.
4.4 STITCHING CAPABILITY
8

---

Published as a conference paper at ICLR 2026
UNICORN
DPT
AD
IDT
DICP
DIT
IC
S-KQLL
800
Reach
-50
Cheetah-Vel
700
Ant.Dir
600
Avergee Return
Avarago Reburn
Averge Rehurm
800
-125
300
-150
500
175
400
10 Episade
15
20
-000
10 Episode
15
20
0-
10 Episode
15
20
Figure 6: Few-shot evaluation return curves of S-ICQL. and baselines on Medium datasets.
We conducted an in-depth analysis to examine whether S-ICQL truly exhibits stitching capa- bility. Specifically, we pretrained the model on offline datasets containing only suboptimal trajectories and evaluated on the correspond- ing training tasks. Figure [ compares the best returns in the offline datasets (OPT) with the performance of DPT, w/o_oo (S-ICQL without Q-learning), and S-ICQL, including both eval- uation curves and final returns. As w/ q and DPT follow the supervised pretraining paradigm, their performance never exceeds the dataset best, underscoring the limitations of purely super- vised methods. contrast, S-ICQL achieves performance exceeding the best in-dataset re- turn on training tasks, providing strong evidence of genuine stitching by composing suboptimal trajectory segments into globally superior policies.
OPT
S.ICQL
Cheetah-Vel
Average Return
-59.45
-80
77.688
-71.67
-100
10
20
Epicade
Ant-Dir
Averrag Return
Coverage Return
620.69
500
497.01
1334.0
514.67
0
20
Episode
Figure 5: Comparison of best dataset returns with
DPT, w/o_q, and S-ICQL on training tasks.
4.5 ROBUSTNESS TO THE QUALITY OF OFFLINE DATASETS
To evaluate S-ICQL's robustness to data quality, conduct experiments Medium datasets containing only suboptimal data. As shown in Figure and Table S-ICQL outperforms all baselines, particularly in complex environments like Cheetah- Vel and Ant-Dir, showing its su- periority when learning from suboptimal data. Notably, many baselines converge to suboptimal policies using imperfect data, especially for pure imitation learning methods like DPT. This again highlights our advantage of using Q-learning to not only simply imitate collected data, but also combine parts of suboptimal trajectories for finding globally optimal behaviors. These
Table 5: Few-shot converged results of S-ICQL and baselines on Medium datasets.
Methods
Reach
Cheetah- Vel
Ant-Dir
UNICORN
366.53+ 1.98 -103.23 19.35331.234 25.75
DPT
714.70= 32.51-161.88 12.61 73.72 34.71
AD
626.24 19.08 -77.30+ 2.06 183.51 11.82
IDT
604.49+ -103.23+
19.35
DICP
719.544: 10.20 -80.341
5.82 416.58 16.80
DIT
709.07
-82.61
3.15
169.22
IC-IQL
726.68= 15.61
-74,88+
2.78 500,
200,00+ 31,49
S-ICQL
743.72 22.15
-58.48
2.32 600.18 21.46
results, along with those in Sec. 4.1 confirm the robustness of method varying dataset quality.
4.6 VISUALIZATION INSIGHTS
We gain deep insights into the prompt construction process through t-SNE visualization on Cheetah- Vel and Ant-Dir tasks, as shown in Figure 8 We use continuous color spectrum to indicate task similarity. For each task, randomly sample 200 transitions and encode them into task representations using the pretrained world model. The initially entangled transitions are transformed into well- separate clusters in the prompt space, where points from different tasks clearly distinguished and similar tasks are grouped more closely. Further, representations of Cheetah- Vel form clear rectilinear distribution from blue (low velocity) to red (high velocity), exactly aligning

---

Published as a conference paper at ICLR 2026
Transition Representation
Promptt Representation
Transition Representation
Prompt Representation
Velocity
Direction
(a) Cheetah-Vel
(b) Ant-Dir
Figure 8: t-SNE visualization on Cheetah-Va Vel and Ant-Dir, where tasks differ in target velocities of 3.00 and target directions of [0, 2m). Data representations of raw transitions (8, a, r, s') and precise prompts 3 from distribution of tasks are mapped into rainbow-colored points.
with the rectilinear spectrum target velocities a physical sense. Similarly, representations of Ant-Dir follow cyclic spectrum that matches the periodicity of angular directions in physical sense. This finding highlights S-ICQL's ability to harness world modeling to distill meaningful task-specific information from raw transitions, enabling precise task inference to facilitate in-context learning.
4.7 ADDITIONAL RESULTS ON MORE COMPLEX ENVIRONMENT
To further assess the scalability and robustness of S-ICQL in demanding scenarios, extend our empirical evaluation to two challenging envi- ronments: (I) sparse-reward. hard-exploration task (PickPlace) and (2) a high-dimensional con- trol task (Humanoid-Dae
UNICORN
AD
DIEP
KAGL
DFT
1DT
DIT
S-KQU
5.00
PicPPlace
800
Hamenoidkär
300
further increasing the diversity and difficulty of our evaluation. As shown in Figure 7 and 6 S-ICQL consis- tently outperforms all baselines, demonstrating robustness and generalization in these environ- ments. Notably, the performance gap becomes even more pronounced in the high-dimensional setting, highlighting the particular strength of S-ICQL when dealing with large-scale problems, further underscoring its broader applicability.
200
SC0
20 Eptsode
20
400
10
Eppode
Figure 7: Few-shot evaluation return curves of all methods on Mixed datasets for complex tasks.
Table Few-shot converged returns S-ICQL and baselines on Mixed datasets for complex tasks.
Method UNICORN
DPT
AD
DDT
DICP
DIT
IC-IQI
S-ICQI.
PickPlace 234.78= auss 370.09 33.45 390.62: 42.99
295 786 47.55 425.30 13.1660. 35.00 12: 13.1772..6
Humamoid 575.00000 509.99 = 58.4088.088 19.6111.177 20.23
5 CONCLUSIONS, LIMITATIONS, AND FUTURE WORK
In the paper, we propose S-ICQL, an innovative framework that introduces dynamic programming and world modeling enable fundamental reward maximization and efficient task generalization in ICRL. S-ICQL employs multi-head transformer to jointly predict optimal policies and in- no-cottext value functions, guided by a pretrained world model that encodes precise task-relevant information for efficient prompt construction. Policy improvement is achieved by fitting in-context value functions with expectile regression and extracting policies via advantage- advantage-weighted regression, enabling reward maximization while preserving the scalability and stability of supervised pretraining. Extensive evaluations verify the consistent superiority of S-ICQL over a range of competitive baselines.
Though, our prompt length matches the sampled transitions, which may be too long for long-horizon interactive problems. Future work can explore using more compact task representations, e.g., encoding an episode a skill into single token. Another promising step is to leverage natural language Levee task prompt for ICLR. We also plan to utilize the dynamics decoder at test time to detect distributional shifts, facilitating adaptive learning through continual updates.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was supported in part by the National Natural Science Foundation of China (Nos. 62376122 and 72471 109), in part by the National Key Research and Development Program of China (No. 2025YTA 1016904)) and in part by the National Natural Science Foundation of China (Nos. 92370132 and 72471109).
ETHICS STATEMENT
We are not aware of any major ethical concerns arising from our work. Our study is conducted entirely within RL benchmarking environments, using only publicly available models and datasets for training and evaluation. No human subjects involved, and our research does not introduce sensitive or potentially harmful insights.
REPRODUCIBILITY STATEMENT
We provide detailed descriptions of our experimental settings in the Appendix, including environ- ments, datasets, implementation details, and hyperparameters. addition, we provide the original code in supplementary materials to facilitate reproducibility and verification of our results.
REFERENCES
David Ackley and Michael Littman. Generalization and scaling in reinforcement learning, In Advances in Neural Information Processing Systems, volume 2, 1989.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong. Luisa Zintgraf, Chelsea Finn, Shimon Whiteson, et al. A tutorial on meta-reinforcemmnt learning. Foundations and Trends@ in Machine Learning, 18(2-3):224 384, 2025.
Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In Proceedings of International Conference on Learning Representations, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020.
Zhenwen Dai, Federico Tomasi, and Sina Ghiassian. In-context exploration-exploitttio for rein- forcement learning. In Proceedingy of International Conference on Learning Representations, 2024.
Juncheng Dong, Moyang Guo, Ethan X Fang. Zhuoran Yang, and Vahid Tarokh. In-context reinforce- ment learning without optimal action labels. In ICML 2024 Workshop on In-Context Learning, 2024. URL URLLAttts://ppeeeeaaa
Juncheng Dong, Moyang Guo, Ethan X Fang, Zhuoran Yang, and Vahid Tarokh. In-context reinforce- ment learning from suboptimal historical data. In Proceedings of International Conference on Machine Learning. 2025.
Yan Duan, John Schulman, Xi Chen, Peter Bartlett, Ilya Sutskever, and Pieter Abbeel. RL?: Fast reinforcement learning via slow reinforcement learning, arXiv preprint arXiv:1611.027799 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of International Conference on Machine Learning, pp. 1126-1135, 2017.
Yunkai Gao, Rui Zhang, Jiaming Guo, Fan Wu, Qi Yi, Shaohui Peng. Siming Lan, Ruizhi Chen, Zidong Du, Xing Hu, et al. Context shift reduction for offline meta-reinforeemmett learning. In Advances in Neural Information Processing Systems, volume 36, pp. 80024-80045 2023.
11

---

Published as a conference paper at ICLR 2026
Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in in-context reinforcement learning for adaptive agents. In Proceedings of International Conference on Learning Representations, 2024a.
Jake Grigsby, Justin Sasek, Samyak Parajuli, Daniel Adebi, Amy Zhang, and Yuke Zhu. AMAGO-2: Breaking the multi-task barrier in meta-reinforeement learning with transformers. In Advances in Neural Information Processing Systems, volume 37, pp. 87473-87508, 2024b.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy max- imum entropy deep reinforcement learning with a stochastic actor. In Proceedings of International Conference Machine Learning, pp. 1861-1870, 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, pp. 1-7, 2025.
Zican Hu, Wei Liu, Xiaoye Qu, Xiangyu Yue, Chunlin Chen, Zhi Wang. and Yu Cheng. Divide and conquer: Grounding LLMs as efficient decision-making agents via offline hierarchical reinforce- ment learning. In Proceedings of International Conference on Machine Learning, 2025.
Zican Hu, Shilin Zhang, Yafu Li, Jianhao Yan, Xuyang Hu, Leyang Cui, Xiaoye Qu Chunlin Chen, Yu Cheng, and Zhi Wang. Diversity-inceniiized exploration for versatile reasoning. In Proceedings of International Conference on Learning Representations, 2026.
Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision transformer: Re- inforcement learning via hierarchical chain-of-thought, In Proceedings of International Conference on Machine Learning, 2024.
Joannis Kakogeorgiou, Spyros Gidaris, Bill Psomas, Yannis Avrithis, Andrei Bursue, Konstantinos Karantzalos, and Nikos Komodakis. What to hide from your students: Attention-guided masked image modeling. In Proceedings of European Conference on Computer Vision, pp. 300-318, 2022.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. In Journal of Artificial Intelligence Research, volume 76, pp. 201-264, 2023.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. In Proceedings of International Conference on Learning Representations, 2022.
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. Proceedings of International Conference on 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. In Advances in Neural Information Processing Systems, volume 36, pp. 43057-43083 2023.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization model-based reinforcement learning. In Proceedings of International Conference on Machine Learning, pp. 5757-5766, 2020.
Feng Li, Qing Jiang, Hao Zhang, Tianhe Ren, Shilong Liu, Xueyan Zou, Huaizhe Xu, Hongyang Li, Jianwei Yang, Chunyuan Li, et al. Visual in-context prompting. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12861-12871, 2024a.
Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: Efficient fully-offline reta-reinforeement learning via distance metric learning and behavior regularization. In Proceedings of International Conference on Learning Representations, 2021.
Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Yang YU, Junqiao Zhao, and Pheng- Ann Heng. Towards an information theoretic framework context-based offline renforeeiforreementt learning. In Advances in Neural Information Processing Systems, 2024b.
12

---

Published as a conference paper at ICLR 2026
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining. In Proceedings of International Conference on Learning Representations, 2024.
Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. In Proceedings of International Conference Machine Learning, volume 202, pp. 21362-21374, 2023.
Jinmei Liu, Haoru Li, Zhenhong Sun, Chaofeng Chen, Yatao Bian, Bo Wang. Daoyi Dong, Chunlin Chen, and Zhi Wang. Beyond the dirac delta: Mitigating diversity collapse in reinforcement fine-tuning for versatile image generation. arXiv preprint arXiv:2601. 12401, 2026.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foenster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, volume 36, pp.
2023.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting, In Proceedings of International Conference on Machine Learning, pp. 7780-7791, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529-333. 2015.
Amir Moeini, Jiuqi Wang, Jacob Beck, Ethan Blaser, Shimon Whiteson, Rohan Chandra, and Shangtong Zhang. A survey of in-context reinforcement learning, arXiv preprint arXiv25502.07978, 2025.
Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. MetaDiffuser: Diffusion model as conditional planner for offline meta-rl. In Proceedings of International Conference on Machine Learning, 2023.
Alexander Nikulin, Ilya Zisman, Alexey Zemtsov, and Vladislav Kurenkov. XLand-100B: A large- scale multi-task dataset for in-context reinforcement learning. In Proceedings of International Conference on Learning Representations, 2025.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. wdyantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint 1rXiv:1910.00117,.00 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
David Silver, Satinder Singh, Doina Precup, and Richard S Sutton. Reward is enough. Artificial Intelligence, 103555 2021.
Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. In Proceedings of International Conference on Machine Learning, pp. 45773-45793, 2024.
Charlie Victor Snell, Ilya Kostrikov, Yi Su, Sherry Yang, and Sergey Levine. Offline RL for natural language generation with implicit language Q-learning. In Proceedings of International Conference Learning Representations, 2023.
Jaehyeon Son, Soochan Lee, and Gunhee Kim. Distilling reinforcement learning algorithms for in-context model-based planning. In Proceedings of International Conference Learning Repre- sentations, 2025.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Denis Tarasov, Alexander Nikulin, Ilya Zisman, Albina Klepach, Andrei Polubarov, Lyubaykin Nikita, Alexander Derevyagin, Igor Kiselev, and Vladislav Kurenkov. Yes, Q-learning helps offline in-context RL. ICLR 2025 Workshop on Scalable Optimization for Efficient and Adaptive Foundation Models, 2025.
13

---

Published as a conference paper at ICLR 2026
Emanuel Todorov, Tom Erez, and Yuval Tassa. MuloCo: A physics engine for model-based control. In Proceedings of IEEE/RSJ International Conference on Intelligent Robots and Systems, pp- 5026-5033, 2012.
Zhendong Wang, Yifan Jiang, Yadong Lu, Pengcheng He, Weizhu Chen, Zhangyang Wang, Mingyuan Zhou, et al. In-context learning unlocked for diffusion models. In Advances in Neural Information Processing Systems, volume 36. pp. 8542-8562, 2023a.
Zhi Wang, Chunlin Chen, and Daoyi Dong. Lifelong incremental reinforcement learning with online Bayesian inference. IEEE Transactions Neural Networks and Learning Systems, 33(8); 4003-4016, 2022.
Zhi Wang, Chunlin Chen, and Daoyi Dong. A Dirichlet process mixture of robust task models for scalable lifelong reinforcement learning. IEEE Transactions Cybernetics, 53(12):7509-7520 2023b.
Zhi Wang, Li Zhang, Wenhao Wu, Yuanheng Zhu, Dongbin Zhao, and Chunlin Chen. Meta-DT: Offline meta-RL. as conditional sequence modeling with world model disentanglement In Advances in Neural Information Processing Systems, volume 37, pp. 44845 44870, 2024.
Wenhao Wu, Fuhong Liu, Haoru Li, Zican Hu, Daoyi Dong, Chunlin Chen, and Zhi Wang. Mixture- of-experts meets in-context reinforcement learning. In Advances in Neural Information Processing Systems, 2025.
Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguez, Q-learning decision transformer: Lever- aging dynamic programming for conditional sequence modelling in offline RL. In Proceedings of International Conference on Machine Learning, pp. 38989-39007, 2023.
Jianhao Yan, Yafu Li, Zican Hu, Zhi Wang, Ganqu Cui, Xiaoye Qu, Yu Cheng, and Yue Zhang. Learning to reason under off-policy guidance. In Advances in Neural Information Processing Systems, 2025.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-World: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, 2020.
Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In Proceedings of International Conference on Machine Learning, pp.
25747-25759, 2022.
Runzhe Zhan, Yafu Li, Zhi Wang, Xiaoye Qu, Dongrui Liu, Jing Shao, Derek F Wong, and Yu Cheng. ExGRPO: Learning to reason from experience. In Proceedings of International Conference on Learning Representations, 2026.
Shilin Zhang, Zican Hu, Wenhao Wu, Xinyi Xie, Jianxiang Tang, Chunlin Chen, Daoyi Dong, Yu Cheng, Zhenhong Sun, and Zhi Wang. Text-to-decision agent: Offline meta-reinforcemenn learning from natural language supervision. Advances Neural Information Processing Systems, 2025.
Luisa Zintgraf, Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. VariBAD: Variational Bayes-adaptive deep RL via meta-learning. The Journal of Machine Learning Research,
2021.
Ilya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence of in-context reinforcement learning from noise distillation. In Proceedings of the International Conference Machine Learning, pp. 62832-62846, 2024.
14

---

Published as a conference paper at ICLR 2026
A LIMITATIONS
Despite promising results, our approach has several limitations that warrant further investigation:
Prompt Length and Efficiency: Although the prompt length matches the sampled transitions, it may not be compact enough for long-horizon tasks, where increased complexity can cause inefficiencies. Future work should focus condensing the prompt without losing key context.
Task Information Extraction: The current approach uses prompts equal in size to the sampled transitions, which works for short tasks but may not scale well for more complex ones. Using compact representations, like encoding episodes or skills into single tokens, could improve scalability.
Leveraging Natural Language for S-ICQL: A promising direction is using natural language as a higher-level prompt for S-ICQL. Pretrained language models could enhance knowledge transfer, improving generalization and adaptability, offering valuable opportunities for future research.
These limitations highlight critical areas for improvement, particularly in reducing prompt size and enhancing scalability for long-horizon tasks, offering valuable avenues for future research.
B ALGORITHM PSEUDOCODES
Based on the implementations presented in Sec. this section provides an overview of the procedural steps of our method. Initially, Algorithm [ outlines the pretraining process for the world model. Subsequently, Algorithm describes the pipeline for training and testing S-ICQL.
Algorithm 1: Pretraining the World Model
Input: Training tasks Mttrain and corresponding offline datasets Dirainn Context encoder Eg:dynamics decoder decoder D,; Experience step k;;
for each iteration do
Sample task M M and obtain the corresponding dataset D from Dirain Sample a transition tuple (st, (t2, Tt, 81+1) with randomly selected
Obtain its h-step history ni (st-k: at-ki Tr-ks .... St-1- at Sz)
Compute the context 2 = Es(r
Compute the predicted reward and next state|fz, 8++1] 1= Dip Gri =/) Update Es and Do
using the loss as (( sp) = Mr [I][re, 8641] (8t, az = # = Es(ni)]
C THE DETAILS OF ENVIRONMENTS
DarkRoom:
The agent is randomly placed in a 10 x10 grid room, and the goal occupies one of the 100 grid cells. Thus, there are 100 possible goals. The agent's observation its current grid cell, i.e.. =10 x [10]. At each step, the agent selects one of five actions: move up, down, left, right, or remain stationary. The agent receives a reward only upon reaching the goal and 0 otherwise. The episode horizon for Dark Room is 100. Consistent with (Lee et al. 2023), , use 80 of the 100 goals for pretraining and reserve the remaining 20 goals to test our model's ''s in-context RL capability on unseen tasks.
MuJoCo: The multi-task MuloCo control testbed a classical benchmark commonly used in continual RL, multi-task RL, and meta-RL. This testbed concludes two environments with reward function changes and one environment with transition dynamics changes as
Cheetah- Vel: A planar cheetah must run forward at a specified target velocity along the positive x-axis. Each task is defined by distinct reward function that penalizes the absolute deviation between the cheetah's 's instantaneous velocity and its goal velocity. Goal velocities are drawn uniformly from 3.0), yielding suite of tasks with varying targets.
Walker-Param: A planar walker robot needs to move forward as fast as possible. Tasks differ in transition dynamics. For each task. the physical parameters of body mass, inertia, damping, and friction randomized. The reward function is proportional the running velocity in the positive
15

---

Published as a conference paper at ICLR 2026
Algorithm 2: Scalable In-Context Q-Learning
Input: Training tasks M and corresponding offline datasets Dirain Trained context encoder Egg Transformer model parameterized by 8: Prompt horizon h;
M P Prrerrainnng model
for each iteration do
Sample task M Muain and obtain the corresponding dataset De from Dirann
Sample (si. aj si+ and a h-step trajectory [$1, a1, Q1, "1:..., $/: Ob: TA]" randomly from dataset from DD'
Use the trained context encoder Eg to transform the h-step trajectory into lightweight prompt 3 := 2x]" = Esg ([51, (1, ([,1, "i,... 8b, Un,
Tn,T rul")
II In-context va vues learning Update V8 +- Vg - vV Toy (0) using Eq.
Update Qe + Qo pV (0) using Eq
Update i + (1-a)8
at
II In-connext policy extraction Update To 4 mg - PV (0) using Eq.
I Online test-time
deployment
Sample unknown task M M and initialize empty B = 0.
for each episode in max_ eps do
Deploy Ty by choosing ag ~ nT(- st, 3²) at step I.
Store (8t: Gt, Tt; 8t+1) and use trained context encoder Ess to transform the nearest h-step interactions into lightweight prompt Br.
x-direction, which remains consistent for different tasks. The agent must therefore adapt its policy to diverse dynamics to achieve optimal performance.
Ant-Dir: A quadrupedal ant robot must move in a specified direction. Each task defines a distinct target angle, and the reward is given by the cosine similarity between the agent's velocity vector and the unit vector in the target direction. Target directions are drawn uniformly from [0, 2n].
Humanoid- Dir: A high-dimensional humanoid agent must move in specified direction while maintaining balance and coordinated whole-body control. Each task defines target heading angle, and the reward is computed as the cosine similarity between the agent's 's velocity vector and the unit vector of the target direction. Target directions are sampled uniformly from [0, 27].
For all MuloCo domains, we allocate 45 tasks for training and reserve the remaining 5 for evaluation, with each episode capped at 200 timesteps.
The Meta-Learning (MLI) suite in Meta-World is a minimalist, single-task bench- mark for few-shot meta-reinforreemett learning. For each task, a Sawyer robotic arm in MuloCo must reach a randomly Selected surget-whose coordinates are withheld from foservations-forcing the agent to infer the goal by trial and error. MLL's sparse information and clear generalization challenge make it one of Meta-War World's most popular robotic manipulation testheds for evaluating rapid adaptation in meta-RI..
Meta-Workd:
Reach: a Sawyer robotic arm must reach a randomly assigned goal position in 3D space, with each Reach task differing in goal location and reward function. The objective is to learn an optimal policy that efficiently generates the action sequence required to reach the specified target.
Push: A Sawyer arm must push block to a randomly placed target on tabletop, with each task varying the block's start position and corresponding reward function. Agents receive only sparse distance-based feedback and must infer the goal through interaction. This setup evaluates the agent's ss ability to explore and adapt its pushing strategy under sparse supervision.
PickPlace: A Sawyer arm must pick up object and place it at target position sampled for each task. The object's start pose and goal location vary across tasks. Rewards are sparse and reflect successful grasping and placement, making this typical multi-step manipulation benchmark.
For all Meta- World domains, we allocate 45 tasks for training and reserve the remaining for evaluation, with each episode capped at 100 timesteps.
16

---

Published as a conference paper at ICLR 2026
D THE DETAILS OF DATASET CONSTRUCTION
Pretraining Datasets for Darkroom: In the Dark Room environment, the horizon of each trajectory set to 100 steps. At each step, follow the optimal policy with probability and random policy with probability e. We choose so that the mean return of the pretraining datasets is below 40% of the optimal policy return, reflecting the challenging yet common scenarios.
For each evaluation domain, we choose 45 tasks to construct the training datasets and train single-task policy independently for each task. We use soft actor-critic (SAC) (Haarnoja et al. (0188) for the MuJoCo domains. We use Proximal Policy Optimization (PPO) [Schulman et 2017) for the Meta-World domain. We collect two types of offline datasets for each evaluation domain as
Pretraining Datasets for MuJoCo and Meta-World:
Mixed: The dataset is constructed by mixing data from various policies, including those saved throughout the entire training process, offering a diverse range of experiences for training.
Medium: The dataset is constructed using data from medium-quality policies, which, while suboptimal compared high-quality policies, still offer valuable experiences for training.
Table 7 and Table 8 list the main hyperparameters for the SAC and PPO algorithms during offline data collection in all evaluation domains, respectively.
Table 7: Hyperparameters of SAC used to collect multi-task datasets.
Training
Warmup steps
Save frequency
Learning
Batch
Soft update
Discount
Entropy
Environments
steps
rate
size
factor
ratio
Cheetah- Vel
500000
2000
10000
3e-4
256
0.005
0.99
0.2
Walker-Param
1000000
2000
10000
3e-4
256
0.005
0.99
0.2
Ant-dir
500000
2000
10000
3e-4
256
0.005
0.99
0.2
Humanoid-De
500000
2000
10000
le-4
256
0.005
0.99
0.2
Table 8: Hyperparameters of PPO used to collect multi-task datasets.
Environments Total tineettpsstsps
n_steps Learning rate
Batch Bizeh_size n_epochs
Discount factor
Reach
400000
2048
3e-4
64
10
0.99
Push
1000000
2048
3e-4
64
10
0.99
PickPlace 1000000
2048
3e-4
64
10
0.99
E
THE DETAILS OF BASELINES
This section presents seven representative baselines addressing the meta-task generalization problem, including one context-based offline rentforeiifr learning (COMRL) method and six ICRL approaches. These basclines are thoughtfully selected to span the major domains of current offline meta-task research. Furthermore, since proposed S-ICQL method belongs the ICRL category, we incorporate more methods from this class as baselines for a comprehensive comparison. The detailed descriptions of these baselines
follows:
IC-IQL TTarsovv et al., 2025) extends standard ICRL by explicitly optimizing reinforcement learning objectives instead of relying solely supervised losses as in AD. It augments AD with Q-learning loss, both the policy and value are optimized by updating only their heads, with no gradient propagation to the shared transformer backbone.
DIT (Dong al. 2024) addresses the limitations of standard autoregressive imitation learning when trained on suboptimal trajectories. Instead treating trajectory prediction purely as supervised learning, DIT' emulates actor-critic algorithm in-context by applying weighted
17

---

Published as a conference paper at ICLR 2026
maximum likelihood estimation (WMLE) loss, where the weights directly computed from observed rewards.
DICP (Son et al. 2025). combines model-based reinforcement learning with prior in-context RI. approaches such AD and IDT. DICP leverages pre-trained Transformer not only to condition past transitions but also to simulate future trajectories and estimate long-term returns, enabling proactive decision-makin making without parameter updates. DICP jointly models environment dynamics and policy improvements in context, resulting in greater adaptability and sample efficiency, particularly in long-horizon tasks.
IDT (Huang et al. 2024), employs hierarchical learning framework that decomposes decision- making across temporal scales address the high computational cost of prior in-context RL methods on long-houe monghorizon tasks. Its architecture consists of three modules: (i) Secision-MMaknn which predicts high-level decisions; (ii) Decision-to-Go, which decodes these decisions into low-level actions; and (iii) Decision-Review, which maps low-level actions back to high-level representations. Built upon transformer architecture similar to that of AD, IDT addresses complex recision-making in long-horizon tasks while reducing computational costs through hierarchical modeling.
AD LLsskin al. (2023), casts in-context RL as a supervised sequence- sequence-modeling task: a causal Transformer is pretrained on across-episodic trajectories covering the entire RL learning to predict subsequent actions, thereby emulating standard RL update dynamics without explicit gradient updates. In AD, trajectories gathered across episodes organized into fixed-length sequences of length H, with each trajectory tokenized as an interleaved series of states, actions, and rewards. This allows the model learn purely from contextual information, enabling in-context adaptation of policies to new tasks.
DPT (Lee et al. 2023). adopts supervised training paradigm to enable in-context learning for RL tasks. The core idea train a transformer to predict the provided optimal action for a given query state purely based on context, using interaction histories from diverse tasks as in-context datasets. DPT treats each transition tuple (s, o, s'.r) r) as a single token, rather than decomposing into separate embeddings. This allows the attention mechanism to directly model relationships between full transitions, preserving the structural and semantic integrity of the interaction data.
UNICORN (Li et al. 2024b), integrates representative methods such as FOCAL (Li et al. 2021), CORRO Yuan Lu 2022), and CSRO (Gao al. (023), this work proposes unified framework that interprets these approaches as optimizing different approxi- mation bounds of the mutual information between task variables and their latent representations. Building on the information bottleneck principle, it further derives a general and unified objective for task representation learning, facilitating the extraction robust and transferable task represen- tations. This method provides unified information information-theoretti theoretic perspective that summarizes and connects several mainstream context-bart based offline meta-RL approaches. As it integrates key ideas from representative COMRL algorithms, it can serve strong and generalizable baseline.
Mainstream ICRL algorithms predominantly rely imitation learning paradigms, which tend learn suboptimal behaviors when pretraining data is limited in size or quality. To ensure fair comparison, we standardize the dataset size and quality across methods during pretraining. In our task setting, AD struggles to extract effective policy improvement operators due to insufficient offline data. Therefore. we adopt an AD variant that incorporates a reward-basee trajectory sorting mechanism from AT (Liu) & Abbeel] 2023) to better distill policy improvements. Additionally, since UNICORN under the COMRL setting requires warm-up data for task representation inference, we align it with the ICRL setup by initializing task representations using randomly sampled trajectories and updating them after each episode.
F IMPLEMENTATION DETAILS OF S-ICQL
World Model. In this paper, we adopt streamlined architectures for each component of the world model: the context encoder, and the dynamics decoder. The context encoder first em- ploys fully connected multi- perceptron (MLP) followed by gated recurrent unit (GRU) network, both using ReLU activations. The GRU processes the agent's k step history nf = (St-k, k; Ot- &t St-1, 4t-1, Tt-1, St, ).)". and produces 128-dimensional hidden vector. This
18

---

Published as a conference paper at ICLR 2026
Table 9: The network configurations used for S-ICQL.
World Model
Value
Causal Transformer
Value
GRU hidden dim
128
Layers num
4
Prompt representation dim
16
Attention heads num
1
Decoder hidden dim
128
Activation function
ReLU
Decoder hidden layers num
2
Activation function
ReLU
Table 10: Hyperparameters of S-ICQL various domains.
Hyperparameterss
Darkroom Push Reach Cheetah Vel Walker-Parann Ant- PickPlace Humanoid-Di
Training steps
2e5
4e5
4e5
4e5
4e5
4e5
4e5
8e5
Learning rate
3e-4
le-4
lc-4
3e-4
3e-4
3e-4
le-4
3e-4
Prompt horizon h
100
100
100
200
200
200
200
200
Emhedding dim
32
128
128
128
128
128
128
128
Output layers num
2
2
2
2
Temperature parameter
0.01
0.001
0.001
0.001
0.1
0.01
0.01
Expectile parameter
0.5
0.7
0.5
Soft update
0.005
0.005
0.005
0.005
0.005
0.005
0.005
0.005
Discount factor
0.99
0.99
0.99
0.99
0.99
0.99
0.99
0.99
vector is then projected via the MLP into a 16-dimensional task representation z/. The reward decoder an MIP that receives the tuple (s); at, ((1; and passes through two hidden layers of size 128 to predict the scalar reward ff. Analogously, the state transition decoder an MLP that takes (st, zf) as input and uses two 128-dimensional hidden layers to predict the next state 6/+1
Causal Transformer. We implement S-ICQL on top of the official DPT codebase released by DPT (Lee et al.) 2023) hnttps:/github..om///oo commjoonleecddeciionn--reee pretrained muetrained-transfommnee transformer). We ad- here to their architectural design, and construct the embeddings for the GPT-2 backbone as fol- lows. Specifically, given task dataset Dí, we sample (st, 0t, s(41)r)) and h-step trajectory (s1, 11, al, :.,..., 5h: ah, ra]' randomly from dataset from Di. Then encoder Eg to transform raw transitions in this trajectory into precise and lightweight prompt 3 = [Z ] = Es ([s1, 01, Ti, Sh, an rh.'). We form state vectors 5:, = (s), 0), next state vectors Ex = (s)+1,0) 0) and state-action vectors S(sa). = (s at, 0) by concatenating the relevant quantities and padding with zeros so that each ç¹ has dimension de := 2dg + d + 1. The (h (h+1) length sequence is given by X = (E* .f..... 21,) We first apply linear projection (inear(:) Then the (h + 1)-length tokens are
the pretrained context
to each vector and outputs the sequence Y = (No. .I..... Un.) fed into the transformer and predict output autoregressively using a causal self-attention mask. In the output layer, employ separate linear layers to produce the actions, the staue- and the Q-values, respectively. In summary, Table I the details of network structures.
Algorithm Hyperparameters. We evaluate the proposed S-ICQL algorithm eight environ- ments: Darkroom, Push, Reach, Cheetah Vel, Walker Param, Ant Ant-Dir: -iir, PickP lace and Humanoid-Dirr For all experiments, gradient-norm clipping at 10, an experience step of 4, and a batch size 128. Table I0 summarizes the detailed hyperparameters of S-ICQL in each domain.
the Adam optimizer with weight decay of le-4,
Compute. We train our models on NVIDIA RTX3090 GPUs paired with an AMD EPYC 9654 CPU and 512GB of RAM. Pretraining the world model takes approximately 1-2 hours, while pretraining the causal transformer requires about 2-22 hours, depending on the environment's complexity.
19

---

Published as a conference paper at ICLR 2026
G HYPERPARAMETER ANALYSIS
The temperature parameter À in Eq. 6 plays cru- cial role in balancing the trade-off between behavior cloning and the greedy pursuit of high Q-values. We conduct experiments to analyze the influence of X on S-ICQL's performance. Figure 9 and Ta- ble present the ablation results across represen-
Table 1I: Few-shot converged results of S- ICQL with varying À on Mixed datasets.
Reach
Cheetah- Vel
Ant-Dir
1.0 765 99: 17.93 -99.16= 4.00 486.23+ 30.98
0.1
796.42:: 7.20
-46.54
749.00 +99.00+
tative domains with varying values of X. A smaller
0.01 799.
799.29::20 2.77
-46.561
3.08 813.34 14.12
will make the distribution of advantage weights exp(+)(( V) less uniform, leading ploitation of Q-learning, As À decreases from 1.0,
0.001 806.97: 6.35 -35.48 1.33 767.55+: 19.89
the performance of S-ICQL can be greatly improved, highlighting the essentiality of harnessing Q- learning to steer ICRL architectures toward fundamental reward maximization. Practically, À = 0.01 or À = 0.001 leads to a satisfactory performance. The result further confirms the effectiveness and necessity of our in-context Q-learning
In prompt construction, the task representation =j is inferred from the agent's s k-step experience v/ as
Table 12:
Few-shot converged results of S- ICQL with varying k on Mixed datasets.
=f ~ (r))) The step k is crucial for capturing task-relevant information. We conduct experiments to investigate the impact of on S-ICQL's perfor- mance. Figure T0 and Table 12 present the few-shot evaluation returns of S-ICQL across representative environments with varying values of k. The results show that S-ICQL's performance is not sensitive to k. with a moderate experience step of k 4 performing the best in all evaluated domains. A small may provide insufficient task-relevant information, limiting precise in-context inference. A large k may introduce redundancy and noise leading overfitting during world model pretraining. Also, large k will increase the computation load in pretraining the world model and decrease the speed in constructing the prompt for in-context task inference. practice, moderate i achieves fast and precise in-context inference, highlighting the superiority of our prompt design via world modeling.
K
Reach
Cheetah- Vel
Ant-Dir
802.00 = 10.80
-37.51 1.80
717.95 37.81
4
806.97 = 6,35
-35.48
1.33
813.34 = 14.12
794.19 7.93
-40.36 1.36 648.70 = 56.75
A 1.0
0.1
2 0.01
A- 0.001
820
Reach
Cheetah-Vel
Ant-Dir
800
800
Averge Netum
Average Retur
Average Retum
700
780
-60
600
-80
500
-100
300
Epicade
Eplsode
Episode
Figure 9: Few-shot evaluation return curves of S-ICQL on Mixed datasets with varying values of X.
kG
820
Reach
-30
Cheetah-Vel
Ant-Diir
-35
800
Avengge Retum
Averege Reetrrn
-40
700
790
600
760
400
Episode
Episode
Episade
Figure 10: Few-shot evaluation return curves of S-ICQL on Mixed datasets with varying values of k.
20

---

Published as a conference paper at ICLR 2026
H ANALYSIS OF LOss WEIGHT COEFFICIENTS
We conduct a hyperparameter analysis of Eq. by evaluating multiple reweighting configurations for its three loss components, Cx (0), and (0). We vary the relative coefficients emphasize different aspects of the learning objective, including configurations that up-weight the policy loss to prioritize policy fitting and configurations that increase the weight of the value or Q-value losses to strengthen critic learning, while keeping all other training settings fixed. As shown in Table [3] the performance degrades substantially when the policy loss down-weighted (i.e., Cı = 0.1), while remaining insensitive to coefficients on value losses (i.e., C2/C3 = 0.1). The equal-weight configuration (1, 1) consistently offers a stable and straightforward choice, supporting the use of balanced weighting scheme for the three loss components in our main experiments.
Table 13: Analysis of loss weight coefficients in S-ICQL.
Coefficients (cI, C2, C3)
(0.1,
01,
(1, 0.1)
1,)
Cheetah- Vel
-3.294+3.26
.39.56+2.02
12.98
-33.48:t1.33
Ant-Dir
588.81
843.77+8.70
772.04+19.111
813.34+11.12
Reach
591.88 + 88.81
789.29 + 9.17
730.37 #38.46
806.97 6.35
I
THE USE OF LARGE LANGUAGE MODELS (LLMS)
We utilize Large Language Models (LLMs) to assist with polishing the writing and improving text readability. Specifically, LLMs are employed for proofreading, enhancing grammar, and refining sentence structure. The LLM was used solely for editorial purposes to improve clarity and did not contribute to research ideation, experimental design, implementation, analysis, or scientific conclusions. All core research contributions, experiments, and analyses were conducted independently by the authors without LLM assistance.
21


Reviews:
[{"review_id": "ABGelSMa3s", "reviewer": "Reviewer_72Wd", "summary": "This work proposes Scalable In-Context Q-Learning (S-ICQL), a novel in-context reinforcement learning (ICRL) framework that harnesses dynamic programming and world modeling to steer ICRL toward efficient reward maximization and task generalization. In particular, S-ICQL incorporating RL objectives into the supervised pretraining paradigm with a transformer architecture that simultaneously predicts optimal policies and in-context value functions. In addition, S-ICQL pretrains a world model to better capture the task information for more efficient in-context inference. The authors prove the effectiveness of S-ICQL through extension experiments on standard ICRL benchmarks.", "questions": "- Is it necessary to learn the in-context value functions together with the transformer policy? I think if you can first learn the in-context value functions, use the learned values functions to compute the weights for policy extraction, and lastly do the policy extraction. In this way, you don't have a moving target, i.e., the weights for AWR is fixed and this can improve training stability.\n- I think it is also standard to test offline deployment performance [4], that is, performance of the performance policy conditioned on a given trajectories collected by other behavior policies. \n- In Eq (7), there are no weights for the three losses? This is counter-intuitive and can be problematic if the three losses are not of the same scales.\n\n\n[4] Supervised Pretraining Can Learn In-Context Reinforcement Learning", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The manuscript is well-written and easy-to-follow. \n- The experiments are comprehensive, including an array of strong baselines. \n- The performance of S-ICQL is strong.", "weaknesses": "My biggest concern is about the novelty of the proposed framework S-ICQL. While it has many components, all of its components and insights have appeared in prior works. As a consequence, S-ICQL feels like a unification of ideas from prior works. Specifically, there are two core technical contributions, \n- **Section 3.3.** The world modeling is **identical** to [1]. Both the model design and pretraining process are almost identical to the task-embedding model of [1]. \n- **Section 3.4.** S-ICQL also proposes to estimate in-context (action) value functions and then extract the policy with advantage weighted regressions. The same idea has been implemented in [2]. The **only differences** are that (1) S-ICQL use a shared transformer backbone for the policy and value transformers, and (2) S-ICQL uses a different loss function, which is the same loss from IQL[3].\n\n[1] Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement.\n\n[2] In-Context Reinforcement Learning From Suboptimal Historical Data.\n\n[3] Offline Reinforcement Learning with Implicit Q-Learning", "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": "6LiKyd98YV", "reviewer": "Reviewer_gW7J", "summary": "This paper proposes S-ICQL, an in-context RL framework which encodes small chunks of experience into a compact task representation and then uses that task representation for transformer-based policy learning. Separate heads are used for the policy and value functions. Experiments show improvements over a number of ICRL/meta-RL baselines.", "questions": "See weaknesses", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Good experimental results which consider a number of relevant baselines\n- Ablations show that each of the components contributes meaningfully to performance\n- Stitching section shows improvements over best returns in the dataset", "weaknesses": "- My main issue is that it is very hard to understand the details of the method. Figure 1 is very complicated; there are many arrows going in different directions and it is not obvious what exactly “precise task representation” and “precise lightweight prompt” are. \n- The method is also quite complicated and not particularly novel. It seems to be a combination of a number of existing components (e.g., AWR, world model, transformer-based policies)", "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": "pJk0ZU0sA5", "reviewer": "Reviewer_qpLp", "summary": "The paper proposes a new In-Context Reinforcement Learning (ICRL) framework that integrates world modeling and dynamic programming, achieving state-of-the-art results under few-shot settings on DarkRoom, MuJoCo, and Meta-World benchmarks.  \nThe world model is used to learn task-relevant representations, which can be utilized by the learned policy.  \nThe dynamic programming component enables the model to derive improved policies from suboptimal supervised datasets.", "questions": "1.  After training the world model, only the encoder is used in subsequent stages.  \n    Are there possible ways to further utilize the world model during later training phases?\n    \n2.  The current setting trains on offline datasets and tests for generalization to OOD tasks.  \n    Would incorporating online learning with environment interaction in  in-distribution task improve the generalization performance on OOD tasks?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.  The proposed world model effectively captures task-relevant information, providing a more compact and informative representation compared to previous ICRL approaches.\n    \n2.  The integration of  Q-function learning and advantage-weighted regression (AWR) allows the model to optimize its policy from suboptimal data and obtain better-performing behaviors.\n    \n3.  The experimental results are relatively comprehensive, with comparisons to multiple baselines demonstrating consistent improvements.", "weaknesses": "1.  The experimental environments are relatively  simple. \n    It would strengthen the work to include results on more challenging or diverse environments.", "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": "SqBSguENne", "reviewer": "Reviewer_93W6", "summary": "This work proposes scalable in-context Q-learning (S-ICQL) that adapts implicit Q-learning (IQL) to the in-context reinforcement learning (ICRL) domain in the offline setting. S-ICRL pretrains a world model to capture the dynamics of the environments, which S-ICQL leverages to encode the task context into a more compact form. The empirical study across three benchmarks confirms the effectiveness of S-ICQL.", "questions": "- Why is pretraining a world model necessary for obtaining a good encoder $E_\\phi$? What prevents it from learning the encoder directly during the IQL?\n- If the main benefit of training the encoder is to save computation by projecting the raw transitions onto a lower-dimensional space, then why do the authors choose GRU instead of a simple linear layer?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- S-ICQL outperforms the other ICRL baselines consistently.\n- The paper is overall easy to follow.\n- The empirical study looks rigorous and includes the study of S-ICQL under out-of-distribution tasks and suboptimal data.", "weaknesses": "- I feel the comparison with other baselines is not entirely fair. For instance, algorithm distillation (AD) is more of an online method than an offline one. It requires a stream of experience generated by a consistently improving policy. In addition, the decision pre-trained transformer (DPT) requires the actions in the dataset to be optimal, which may not be satisfied in this work.\n- I think there is a considerable overhead of pre-training the world model. From what I understand, S-ICQL only uses the encoder $E_\\phi$ during the IQL process, and the dynamics decoder was discarded after pretraining. In addition, unlike pre-training a language model on a very general corpus, here the world model is only fitted on a particular class of tasks. Therefore, when the task class changes, the world model needs to be retrained.\n- Some equations are ill-written. For example, why is there an inner expectation in equation (3)? If the encoder $E_\\phi$ is a function, where does the stochasticity come from? Also, in equation (3), how is the squared norm of the tuple $[r, s]$ defined?\n- In the ablation study, I don't think removing the Q-learning component recovers DPT. Without optimal actions, the model is likely regressing on the offline dataset and will naturally learn a suboptimal policy if the data quality is poor.\n\nMinor concerns:\n- Though not significantly affecting reading, the writing has several typos and grammatical errors.\n- Since S-ICQL is claimed to be an ICRL algorithm, then the meta RL survey by Beck et al. (2025) and the ICRL survey by Moeini et al. (2025) should be cited for a richer context.", "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)