In-Context Reinforcement Learning From Suboptimal Historical Data
Authors: Juncheng Dong
Moyang Guo
Ethan X Fang
Zhuoran Yang
Vahid Tarokh
Venue: ICML 2025 poster (2025)
Keywords: In-Context Learning; Reinforcement Learning; Transformers; In-Context Reinforcement Learning
Primary area: reinforcement_learning
TL;DR: We propose a supervised pretraining framework for in-context RL without optimal action labels

Abstract: Transformer models have achieved remarkable empirical successes, largely due to their in-context learning capabilities. Inspired by this, we explore training an autoregressive transformer for in-context reinforcement learning (ICRL). In this setting, we initially train a transformer on an offline dataset consisting of trajectories collected from various RL tasks, and then fix and use this transformer to create an action policy for new RL tasks. Notably, we consider the setting where the offline dataset contains trajectories sampled from suboptimal behavioral policies. In this case, standard autoregressive training corresponds to imitation learning and results in suboptimal performance. To address this, we propose the *Decision Importance Transformer* (DIT) framework, which emulates the actor-critic algorithm in an in-context manner. In particular, we first train a transformer-based value function that estimates the advantage functions of the behavior policies that collected the suboptimal trajectories. Then we train a transformer-based policy via a weighted maximum likelihood estimation loss, where the weights are constructed based on the trained value function to steer the suboptimal policies to the optimal ones. We conduct extensive experiments to test the performance of DIT on both bandit and Markov Decision Process problems. Our results show that DIT achieves superior performance, particularly when the offline dataset contains suboptimal historical data.

In-Context Reinforcement Learning From Suboptimal Historical Data
Juncheng Dong'
Moyang Guo
Ethan X. Fang?
Zhuoran Yang Vahid Tarokh
Abstract
1. Introduction
Transformer models (TMs) have achieved remarkable em- pirical successes (Radford et al., 2019; OpenAI, 2024). In partieular, TMs trained on vast amount of data have shown remarkable in-context learning (ICL) capabilities, solving new supervised learning tasks only with a few demonstra- tions and without requiring any parameter updates (Brown et al., 2020a; Akyürek et al., 2022). Meanwhile, substan- tial evidence demonstrates that autoregressive TMs excel at solving individual reinforcement learning (RI.) tasks, where TM-based policy is trained and tested on the same RL task (Li et al., 2023b). Inspired by these, recent research has explored the use of TMs for in-context RL (ICRL). In this setting, we pretrain TMs on an offline dataset consisting of trajectories collected from a family of different RL tasks. Af- ter pretraining, we deploy the pretrained TMs to solve new and unseen RI. tasks (Laskin et al., 2022; Lee et al., 2024). See Figure (a) and (c) for comparisons between standard offline RL and ICRL. When presented with context dataset containing environment interactions collected by unknown and often suboptimal policies, pretrained TMs predict the optimal actions for current states from the environmental information within the context datasct. See Figure 1 for a visual illustration. Two recent works, Algorithm Distil- lation (AD) (Laskin et al., 2022) and Decision Pretrained Transformer (DPT) (Lee al., 2024), have demonstrated impressive ICRL abilities, inferring near-optimal policies for new RL tasks.
Transformer models have achieved remarkable empirical successes, largely due to their in- context learning capabilities. Inspired by this, we explore training an autoregressive transformer for in-context reinforcement learning (ICRL). In this setting, initially train a transformer on offline dataset consisting of trajectories col- lected from various RL tasks, and then fix and use this transformer to create an action policy for new RL tasks. Notably, we consider the setting where the offline dataset contains trajectories sam- pled from suboptimal behavioral policies. In this case, standard autoregressive training corresponds to imitation learning and results in suboptimal performance. To address this, we propose the Decision Importance Transformer (DIT) frame- work, which emulates the actor-critic algorithm in an in-context manner. In particular, we first train transformer-basee value function that es- timates the advantage functions of the behavior policies that collected the suboptimal trajectories. Then we train wansformer-basee policy via a weighted maximum likelihood estimation loss, where the weights are constructed based on the trained value function to steer the suboptimal poli- cies to the optimal ones. We conduct extensive experiments to test the performance of DIT on both bandit and Markov Decision Process prob- lems. Our results show that DIT achieves superior performance, particularly when the offline dataset contains suboptimal historical data.
Challenges. However, existing supervised pretraining ap- proaches focus on training TMs to imitate the actions in the pretraining datasets and thus have stringent require- ments on the pretraining datasets. For example, AD requires the pretraining dataset to capture the learning process of RL algorithms episodes generated by randomly ini- tialized policies to those collected by nearly optimal poli- cies accsssccros a wide range of RL tasks; DPT requires access to optimal policies to generate a set of optimal action labels for its supervised pretraining of TMs. To overcome these limitations, this work considers pretraining TMs for ICRL using only suboptimal historical data. While this presents significant challenges, it also offers substantial potential benefits by significantly improving the feasibility of ICRI., as suboptimal trajectories are far easier to gather. For in- stance, large companies often maintain extensive databases
'Department of Electrical and Computer Engineering, Duke University. Durham, US Department of Biostatistics and Bioin- formaties, Duke University, Durham, US Department of Statistics and Data Science, Yale University, New Haven, US. Correspon- dence to: Juncheng Dong <juncheng.deee dong @ duke.edu'..
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

In-Context Reinforcement Learning From Suboptimal Historical Data
(a) Offline RL
(b) Supervised Pretraining for In-context RL
Training
(A,1S)
ILV
Testing
X
Oftine Datases
Bny M
(<) In-context optimal action labels
ICRL Deploymmntt
A
R
Xnv 2
Bupprrvtad Fretraining Datacet
Pretreining
To (A
Ray M-1
Optinal aotion Jabal
Rlatorios) Datenos
(d) Decision Importance Transformer (DIT)
Squery
Causal Transtormer To
a
Phase In- In-Context Advantage Estimation
Phase 2: Weighted Supervised Pretraining w/o optimal sction labels
Fav
Bav
Baz
Sx
Causal Transfeemer
In-context Adv Est
C
W
query
Cauasl Trenstormer To
à
(S, Agg 1)
Figure I: (a) and (c) Comparison between Offline RL and ICRL. Standard offline RL trains and tests a policy π in the same task (Env Oy; ICRL pretrains TMs on trajectories collected from family of different RL tasks (Env Env 2..... Env M), and deploys the pretrained TMs unseen tasks (Env M+/). ICRL Deployment. The pretrained TMs generate actions conditioned on the current states and context datasets consisting of offline trajectories collected by (suboptimal) behavioral policies from the unseen tasks. (b) Supervised Pretraining. Presented with offline trajectories and optimal action labels, TMs are pretrained to predict the optimal actions for query states across RL tasks. (c) ICRL from Suboptimal Historical Data. This work addresses the challenging problem of ICRL without optimal action labels. (d) Schematic Overview of the Proposed Framework DIT. Lack of the optimal action labels, the proposed framework employs in-trajectory state-action pairs as query states and pseudo-optime action labels, and weighted pretraining objective, where the weights are based on the optimality of actions, estimated by TM-based in-context advantage function estimator.
of historical trajectories from non-expert users.
that actions with high advantage values receive more weights during pretraining, leading to guaranteed pol- icy improvements over the behavior policies.
Contributions. In pursuit of this goal, we introduce Deci- sion Importance Transformer (DIT), supervised pretrain- ing framework for ICRL using only historical trajectories collected by suboptimal behavioral policies across distinct RL tasks. When the pretraining datasets contain only sub- optimal trajectories, existing approaches correspond to imi- tation learning and thus result in suboptimal performance. DIT overcomes this challenge through several techniques:
Notably, although advantage weighted regression has been studied in standard RL (Wang et al., 2018; Peng al., 2019), it remains unclear how generalize this approach ICRL. The primary challenge is that the weighting function in ICRL must be task-dependent, thus requiring the estimation of advantage functions for all RI tasks in the pretraining dataset. To this end, the most significant technical difficulty arises from the unknown souree tasks of pretraining trajectories, preventing from grouping trajectories from the same RI tasks to improve estimation. As result, must estimate the advantage functions individually for each trajectory in the pretraining dataset.
DIT learns to infer near-optimal actions from sub- optimal trajectories through an exponential reweight- ing technique that assigns good actions in the offline dataset with more weights during supervised pretrain- ing. These assigned weights guide the suboptimal poli- cies toward the optimal ones.
In particular, the assigned weights are constructed from the advantage functions of the behavior policies such
To address this formidable challenge, DIT trains a TM--baed advantage estimator that interpolates across
2

---

In-Context Reinforcement Learning From Suboptimal Historical Data
trajectories from different tasks for an in-contexi es- timation of the advantage functions to facilitate the weighted supervised pretraining framework. See Fig- ure (d) for a visualization.
Transformer (DT) (Chen et al., 2021; Zheng et al., 2022; Liu et al., 2023; Yamagata al., 2023) uses return-conditioned supervised learning to tackle offline RL. Although salable to multi-task settings (i.e., one model for multiple RL prob- lems), DT' is commonly criticised for its inability to improve upon the offline datasets and provably sub-optimal in certain scenarios, e.g., environment with high stochasticity (Brand- fonbrener et al., 2022: Yang et al., 2022: Yamagata et al., 2023). To this end, AD (Laskin et al., 2022) uses sequential modeling to emulate the learning process of RL algorithms, i.e., meta-learning (Vilalta & Drissi, 2002). The work most closely related to ours is DPT, supervised pretraining ap- proach for in-context decision making (Lee et al., 2024). DPT trains transformers to predict the optimal action given query state and set of transitions. Both AD and DPT have stringent assumptions on the pretraining datasets. Our work overcomes those drawbacks and does not not require query to optimal policies nor the learning histories of RI algo- rithms (Laskin et al., 2022; Lee et al., 2024).
Empirical Results. Through extensive experiments vari- ous bandit and Markov Decision Process (MDP) problems, we demonstrate that pretrained DIT models generalize to unseen recisson-mmainn problems. On bandit problems, the performance of DIT models matches that of the theoretically optimal bandit algorithms (e.g.. Thompson Sampling (Russo al., 2018)). In four challenging MDP problems including two navigating tasks with sparse rewards and two complex continuous control tasks, DIT models achieve superior per- formance, particularly when the pretraining dataset contains suboptimal trajectories. Notably, in many scenarios, DIT is comparable to DPT in both online and offline testings, despite being pretrained without optimal action labels.
2. Related Work
3. Preliminary
Offline Reinforcement Learning. Since we consider pretraining with historical data, our work falls within the broader field of offline RL. While online RL algo- et al., 2018)
Markov Decision Process. Sequential decision problems can be formulated Markov Decision Processes (MDPs). An MDP is described by the tuple (S,A, P.. Rr, 7, Pr) where S is the set of all possible states, A is the set of all possible actions, P, S x A - ASS) is the transition function that describes the distribution of the next state, R2 S x A IK is the reward function, 7 E (0, 1) is the discounting factor for cumulative rewards, and pr A(S) is the initial state distribution. An agent interacts with the environment as follows. At the initial step h = 1, an initial state $ E S is sampled according to Pr. At each time step h, the agent chooses action a₁ E AA and receives reward r4 = R, (sh, an). Then the next state Sh+1 is gener- ated following Pr (r(s,, an). A policy π : + A(A) maps the current state to an action distribution. Let G, E >"" " ward of π for task т. The goal of an agent is to learn the optimal policy x that maximizes (TT).
rithms (Kaelbling et al., 1996; François- learn optimal policies by interacting with the environments through trial and error, offline RL (Levine et al., 2020; Matsushima et al., 2020; Prudencio et al., 2023) aims to infer optimal policies from historical data collected by (sub- optimal) behavioral policies. One of the most substantial challenges for offline RL is the distribution shift caused by the mismatch between behavioral policies and optimal policies (Levine et al., 2020; Kostrikov et al., 2021). To this end, offline RL algorithms learn pessimistically by ei- ther policy regularization underestimating the policy re- turns (Wu et al., 2019; Kidambi et al., 2020; Kumar et al., 2020; Rashidinejad et al., 2021; Yin & Wang, 2021; Jin et al., 2021; Fujimoto& & Gu, 2021; Dong et al., 2023). While the goal of offline RL is to solve the same RL tasks from where the offline datasets are collected, the goal of ICRL is to efficiently generalize unseen tasks after pretraining with offline datasets from diverse RL tasks.
/T, 7] denote the expected cumulative re-
Transformer: Our proposed approach builds upon the model architecture of DPT, which is a su- pervised pretraining method for TMs to have ICRL capa- bilities (see Figure 1(b) for its architecture). DPT assumes set of tasks (r1) sampled independently from task distribution pr. with each r' as an instance of MDP. For each task r", a context dataset De is sampled, consisting of interactions between a behavioral policy and r'. That is, D' = {(s), ai. fh+t:rr)h where of is chosen by a behavioral policy. Additionally, for each task 7', a query state si sqpar E S is sampled, and an associated optimal ac- tion label of is sampled from (8qurry) ), where nt, is the optimal policy for T¹. The complete pretraining dataset is
Transformer Models and Autoregressive Decision Mak- ing. Large language models and autoregressive mod- els (Radford et al., 2019; Brown et al., 2020b; Wu et al., 2023b; Touvron et al., 2023; OpenAI, 2024) have achieved astonishing empirical successes in wide range of ap- plication areas, including medicine (Singhal al., 2023; Thirunavukarasu et al., 2023), education (Kasneci et al., 2023), finance (Wu et al., 2023a; Yang et al., 2023), etc. As it is natural to use autoregressive models for sequential decision-maker making, transformer models have demonstrated superior performance in both bandit and MDP problems (Li et al., 2023a; Yuan et al., 2023). In particular, Decision
3

---

In-Context Reinforcement Learning From Suboptimal Historical Data
Dpre (Di, S Squery 0 m Let To denote causal trans- former with parameters 0 (Radford et al., 2019). The pre- training objective of DPT is defined as
index T is also known and can be provided as input meta-policy. Then in Section 4.2, introduce solutions for scenarios where A (s, a) and need to be estimated. All proofs of the theoretical results in this section are deferred to Appendix C.
min
log To (a sunys Dt).
(1)
4.1. Weighted Maximum Likelihood Estimation
ICRL Deployment. After pretraining, the pretrained au- toregressive TM To can be deployed as both an online and offline agent. During deployment, an unseen testing task 7 is sampled from p.- For offline deployment, a dataset Den is first sampled from 7, e.g.. Dor contains trajectories gathered from a behavioral policy in r, then DPT follows the policy To(-Ish, Dog) after observing the state S/ at time step h For online deployment. DPT initiates with an empty dataset Don In each episode, DPT follows the policy To(-|sh, Dan) to collect trajectory {51, { "l,,,,,,........................ 88, an, ] which will be appended into D This process repeats for pre- defined number of episodes. See Algorithm 2 in appendix for the pseudocodes of both offline and online deployments.
Motivation. To motivate DIT, we first consider the set- ting of imitation learning where the agent is trained and tested on the same task. Given a dataset of transitions D = {(S,, an; 9h41,7h)3 collected by a behavior pol- icy (a(ss) with advantage function (s, a). Wang et al. (2018) proposes to optimize weighted objective:
arg max
Σ
exp(A"(s,, -log
क (anl(s))
(SKDAL)E
The rationale is that the good actions in the offline dataset, that is, ag with high advantage value (sh, an); should be given more weights during imitation learning. These weights essentially work as importance sampling ratios so that the action distribution is closer to the optimal one.
4. Decision Importance Transformer
Here we introduce our proposed framework Decision Im- portance Transformer (DIT).
Weighted Pretraining for ICRL. In contrast to imitation learning that focuses on individual RL tasks, the objective of DIT is to learn a task-conditioned policy r[a|s;)) with the task index as input. In particular, [[aas; r) should perform well for 7 Pr.
Pretraining with Suboptimal Data. Similar to DPT, DIT assumes a family of datasets D= (Dim [i=] where D' con- sists of H transitions {(s), "/, s/+1 r)) the (suboptimal) behavioral policy T.. in task 7' which it- self is independently sampled from the task distribution Pr- In contrast to DPT, however, DIT does not require the set of paired query states and optimal action labels (squery ot) 1m Jim1.
h=1 collected by
Motivated by the aforementioned weighted imitation learn- ing objective, DIT has the following weighted maximum likelihood estimation (WMLE) loss for pretraining:
A (s.a) r
which are often difficult to obtain in practice.
L(T)- = -Br,3 exp
log log (2)
Notations. In the sequel, for any task r, assume that it has index (parameter) also denoted by such that the task information 7 can be an explicit input to meta- veta-policy x (s|a;7) 7) which can generate distinct policies based on the received task т. For example, in robotic control tasks, т may represent the physical parameters of the robots such as robot mass or the environmental parameters such as ground friction. We use T/ (a|s) ) denote the behavioral policy for task Denote
Pr. ds),, and ~ <=(ass) where ds(s) is the dis- counted visiting frequencies of क. (als) defined as dy (s) (1 - |Σ 7 1(s; = s) s)|T, = The effectiveness of the objective in Equation (2) is demonstrated by the following result which states that the optimizer to DIT's pretraining objective is also the solution to another policy optimization problem that is easier interpret.
The expectation in Equation (2) is with respect
80
Proposition 4.1. Consider the following optimization prob- lem where Er,, is defined as in Equation (2) except that a <[a|s; T). i.e. the action is sampled from the task- conditioned policy rather than the behavioral policies:
(=(() = E
h
rh 81 = S,
h=1
8
(s a) = E
h- rn = 0,01 a, 7, n
max ((T) = -Er,a,a A - (ru "
as its value and action-value functions respectively, and let (s, = (s) - (s)) be its advantage function.
For presentation clarity, in Section 4.1, we first consider the scenarios where (i) A4 (8, a) is known and (ii) the task
where DKL is the Kullback Leibler (KL) divergence, and let n E arg max, g J(T) be its optimizer. Then we have for
4

---

In-Context Reinforcement Learning From Suboptimal Historical Data
any policy T)
4.2. In-context Task Identification and Advantage Function Estimation
Ey Brap(r),owdd (0) [KLL (T (1s;;
However, two key challenges remain: (i) During pretrain- ing, the advantage function A (s, a) is not accessible for generating the weights required by the WMLE loss; (ii) During deployment, the task index T is not accessible as only a context dataset D, is presented. In other words, the true identity of the testing task T is unknown.
= --E,,s,a
(s) exp (8, log x (a)
+ C,
(4)
where Br, is defined as in (2), C is a constant independent of v and (s) Ea (als) exp A (s,a) /n).
In Equation (3), the objective is to find policy x² that improves over the behavior policy (by maximizing term (I) and does not stray too far from the behavior policy (by minimizing term (II).. When the behavioral policy n (a)a) near-optimal. 77 should set to large value so that we can have safe improvements over the behavioral policy. On the other hand, when the behavioral policy is highly sub-optimal, 1) should set to small value so that we have more freedom for policy improvement to decrease the sub- optimality. Note that the DKL constraint (term (II) in Equa- tion (3)) is critical for pretraining large transformer models to prevent policy collapse (Schulman, 2015).
In-context Task Identification. To address the second prob- lem, we follow DPT to instantiate ι (a]s; 7) with autore- gressive transformer Ta parameterized by 0. Conditioned on given context dataset D, consisting of environment inter- actions collected by a behavioral policy in т, the T'-based policy To(a|s. s, D2) first implicitly extracts task information about from the context D, and chooses an action based on the extracted task information (see Lee et al. (2024) for de- tailed discussion). During pretraining, To learns extract useful task information for the pretraining tasks { fm conditioned on the pretraining context datasets {D*}m and generalizes to unseen tasks during testing.
RED
Comparing Equation (4) with the pretraining objective of DIT in Equation (2), we observe that DIT aims to identify a policy that is closest to x² by setting (s) = 1 (we pro- vide a brief discussion for why this is valid in Appendix C.3). When AC a) is known, the pretraining objective of DIT can be estimated with the given pretraining dataset 'D by minimizing the following loss function
In-context Advantage Function Estimation. The first problem more critical. Given that during pretraining the context dataset D may contain up to several trajectories for each task 7² in the setting of ICRL, estimation of AP, (s, a) bused on D' alone can be unreliable. To this end, in the same spirit of ICRL, we propose to use an in-context advan- tage function estimator As(sh, "f|7²)) to estimate the advan- tage value of any state- pair (s/, of.) in the pretraining dataset D. Specifically, As(s,, of (r') is implemented by two TMs:
m
H
(TT = mH
of log ग (ah1s))7)) (5)
1=1
1=1
where w/ = exp (4 (s), a/)nn) Next we establish that DIT can provably achieve policy improvement.
at As ohl-') Db) (a(s) (7)
Proposition 4.2 (Policy Improvement). Let n* be the pol- icy that optimizes (3). For any task and policy n, let (T) E TT r] represent the expected cumu- lative reward of π for T. Let n,; denote xt" (a)s:T). Then we have
where V. and Qc are two transformers, parameterized by & and (. acting as the in-context value and action value estimators that interpolate across tasks to have an improved estimation.
Model Architecture. Let C E# Nhyh in-trajectory discounted cumulative reward starting from step h. For any observed state-action pair (s%, of)) in the pretraining dataset, Qc(s), af (Dgt) and Ta(s/ID es- timate the action-value function Q (sh, ai)) and value function V (h.(s) respectively, conditioned on the histo- ries of transitions D2y = {(s),, az. G)))==1 and D&A {(s). G) where we employ (ittiie ch as the noisy la- bels for value functions to facilitate in-context learning. See Figure 6 for their visual representations.
"h be the
Erepe (T - Gr (T++))
2 1-7
17 T
E, Br-p. ICP1
(6)
- (1 71 Byupr [Aa VCP
27
where CP Bond (e) DKI (-|s; (+100) and CAA = max, FFawx"(aeen A (s)
In particular, when the magnitude of the advantage func- tion A is small, the right-hand side of Equation (6) is nonnegative. In this case, the policy x¹ obtained by solv- ing Equation (3) is strictly better than the behavior policy. Equivalently, adding the exponential weights in Equation (5) is strictly better than vanilla imitation learning, when the total number of pretraining tasks m is large.
Training. We train V/ and Qc with the following objective function:
min (0) = Iune (e. (0.5++ +A (IE(0) + LO(3)), (8) 0.C
5

---

In-Context Reinforcement Learning From Suboptimal Historical Data
where A > 0 is hyperparameter to balance
across tasks where d denotes the dimension of linear ban- dit problems. The reward of bandit a E A in task r is r'(a) ~ .N (u o²) where nnر E[rla, + = (0', σ(a)) and x² = 0.3. Here, 8¹ is the task -ask-ipecific parameter that defines task r². We conduct experiments on LB problems where K = 20, d = 10 and // = 200. The pretraining dataset for DIT are generated as follows.
1 x
(Vo(s)ID()) )i)²
(og((.)) ()) == mH
+ ( ai P2 Gi)²
LQ(()= = mH- Σ (0h(() - @c( 1 (Di
Pretraining Dataset. For LB problems, we generate the feature function Φ A + R by sampling bandit fea- tures from independent Gaussian distributions, i.e., <(a) Na (0, La/d) for all a E AA To generate the pretraining tasks {r'}, we sample their parameters (0*) independently fol- lowing 0 ~ Nu Ia/d). To generate context dataset D we randomly generate behavioral policy by mixing (i) a probability distribution samples a Dirichlet distribution and (ii) point-mass distribution on random arm. The mix- ing weights are uniform sampled from (0.0, 0. 1.0). At every time step h, the behavioral policy samples an ac- tion aj, and receives ri. We do not enforce extra coverage of the optimal actions for bandit problems. Following the setting of DPT (Lee et al., 2024), we collect 100k context datasets for LB problems.
i
where Qh(1)== = r + ( (h Di"" and
LE(0) = mH
1 Σ
(v)(0) 1o
2,6
where ()(0)= rr + Vo D
Here, LB and L regularize the transformer models with the Bellman equations for value functions.
DIT with In-context Advantage Estimator. After training, with As 1') defined in Equation (7) as an estimation of the true advantage function, we can now optimize the objective function of DIT to have the pretrained TM To- for ICRL, i.e.
0" E arg min OE
Σ
wh log To (ahlsh) D'), 2,1
(9)
Comparisons. We compare to the following baselines (see Appendix B for more details): Empirical Mean (EMP) se- lects the bandit with the highest average reward; Upper Con- fidence Bound (UCB) (Auer, 2002) builds upper confidence bounds for all bandits and selects the bandit with the highest upper bound: Lower Confidence Bound (LCB) (Xiao et al., 2021) builds lower confidence bounds for all bandits and selects the bandit with the highest lower bound; Thomp- son Sampling (TS) (Russo et al., 2018) builds posterior distribution for the rewards of all bandits and selects the bandit with the highest sampled mean. In terms of metrics, for offiin learning, we follow the convention to use the subeptimality defined as (Ma. paa) where (la* is the mean reward of the optimal bandit and pa the mean reward of the chosen bandit; for online learning we use the cumulative regrel defined Σn (Ma Maa) where a, is the chosen action at time step h.
mH
where w/ exp((A As of l/))) complete procedure of DIT in Algorithm I.
We summarize the
5. Experiments
We empirically demonstrate the efficacy of DIT through ex- periments on various bandit and MDP problems. In bandit problems, DIT showcases matching performance to that of the theoretically optimal bandit algorithms in both online and offline settings. In MDP problems, we corroborate that DIT can infer close-to-optimal policies from suboptimal pretraining datasets. Notably, albeit without optimal action labels during pretraining, DIT models demonstrate perfor- mance as strong as that of DPT, which has access optimal action labels during pretraining.
Empirical Results. In Figure 2, DIT models demonstrate superior performance in the online setting to those of the theoretically optimal bandit algorithms, i.e., UCB and TS. Deployed for unseen bandit problems, DIT models quickly identify the optimal bandits at the beginning and maintain low regrets over the horizon. In the offline setting, DIT can infer near-optimal bandits from trajectories collected by sub- optimal policies. When the behavioral policies (captioned as BEH) are randomly generated policies, DIT significantly outperforms both TS and LCB, the theoretically optimal al- gorithm for bandit problems. When the context is collected by expert policies, DIT models further improve upon their performance. We also observe that DIT is slightly outper-
Implementation. We follow Lee et al. (2024) to choose GPT-2 (Radford et al., 2019) as the backbone for To. Q5- and Va due to limited computation resource, and note that the performance may be further improved with larger mod- els. We set 7 = 0.8 for all tasks. We choose 7) = 1 for all tasks. Due to space constraint, see Appendix G for more details.
5.1. Bandit Problems
We consider linear bandit (LB) problems with an underly- ing structure shared among tasks. Specifically, there ex- ists a bandit feature function A -> Rd that is fixed
6

---

In-Context Reinforcement Learning From Suboptimal Historical Data
6
BEH
TS (BEH)
Cumulative Regret
UCB
DPT
DPT
EMP
Suboptimallty
0.1
DIT
DIT
DPT
LCB
LCB
DIT
TS
20
82
25
20
75
100
123
350
175
300
25
52
75
100
125
150
175
200
23
50
100
125
150
175
200
Horizon
Data
Data
Figure 2: Results for Linear Bandits (lower values indicate better performance). Left: Online testing. Middle: Offline testing conditioned trajectories gathered by highly suboptimal, randomly generated policies. Right: Offline testing condtioned on trajectories gathered by experts.
formed by DPT. This is expected given that DPT uses the optimal bandit information during pretraining. However, the loss curves of DPT and DIT demonstrate similar trend, showcasing the effectiveness of DIT's weighted pretraining.
Comparisons. We compare DIT to other in-context algo- rithms as well as RI. algorithms without pretraining. The baseline algorithms are briefly described next (see their im- plementation details in Appendix B).
5.2. MDP Problems
Soft Actor Critic (SAC) (Haarnoja et al., 2018): SAC an online RL algorithm that trains an agent from scratch every environment.
Environments. We conduct experiments on four chal- lenging MDP environments: two navigating tasks with sparse reward Dark Room (Laskin et al., 2022) and Mini- world (Chevalier-ooisvertt al., 2023), well two plex continuous-contro) tasks Meta-world (reach-v2) (Yu et al., 2020) and Half-Cheetah (velocity) (Todorov et al., 2012). In Dark Room, the agent is randomly placed in a room of 10x x 10 grids with an unknown goal location on one of the grid. The agent needs to move to the goal location by choosing from actions in 100 steps. Miniworld, the agent is placed in image and its direction as input. It can choose from four possible actions to reach target box, out of four boxes of different colors. In Meta- World, the task is to control a robot hand to reach a target position in 3D space. In Half- Cheetah, the agent controls robot to reach target velocity, which is uniformly sampled from the interval [0, 3). and is penalized based how far its current velocity is from the target velocity. See Appendix D for more details.
Algorithm Distillation (AD): AD is a sequence modeling-based approach for ICRL that emulates the learning process of RL algorithms (Laskin al., 2022). To this end, AD requires the pretraining dataset con- sist of complete learning histories of an RL algorithm ffom episodes generated by randomly initialized poli- cies to those collected by nearly optimal policies- across a wide range of RL tasks. In this work we use SAC as the RL. algorithm for AD to emulate.
and receives (25 25 x 3) color
Decision Pretrained Transformer (DPT): DPT and DIT use the same context datasets for pretraining! However, DPT requires query states and their associ- ated optimal action labels across different tasks.
Prompt-DT (PDT): PDT is Decision Transformer- based approach (Xu et al., 2022), which leverages the transformer's prompt framework for few-shot adapta- tion. PDT uses the same pretraining dataset as DIT. Thus, the performance gain of DIT over PDT high- lights the effectiveness of DIT's design.
Pretraining Datasets. For Dark Room and Miniworld, to ensure coverage of optimal actions (so that optimal policies can inferred), at every step, with probability p (respec- tively p) we use optimal policy (respectively random policy) to choose action. We choose p so that the average re- ward of the trajectories in the pretraining dataset is less than 30% of that of the optimal trajectories. For Meta- World and Half-Cheetah, we construct the pretraining datasets using historical trajectories generated by Soft Actor Critic (SAC). Specifically, SAC is trained until convergence for each task, then we sample from its learning trajectories to build the dataset. Our SAC model training follows the settings out- lined in Haarnoja et al. (2018). See Appendix E for details.
Behavior Cloning (BC): We include an variation of DIT without the exponential reweighting. This ap- proach closely imitates BC, with the following pre-
'Note that the context datasets in our experiments are collected using suboptimal policies. as opposed to the uniformly random policies used by DPT in Lee et al. (2024). As result, the reported performance of DPT in the Dark Room environment differs from that in the original paper.

---

In-Context Reinforcement Learning From Suboptimal Historical Data
Optinal
AD
80
80
DPT
BC
Return
60
SAC
Return
Retum
Prompt-DT
40
40
40
20
5
10
15
20 Episodes
25
30
35
40
Optmal DPT
DIT PDT
AD
BC
Optimal DPT
DIT PDT
AD
BC
Methods
Methods
(a) Online Testing
(b) Offline Random
(c) Offline Expert
Figure 3: Results on Dark Room (higher values indicate better performance). (a): Change in return of policies with additional online episodes for (in-context) learning, (b) and (c): Offline evaluations with context trajectories sampled from random and expert policies.
training objective:
ilarly, for Half-Cheetali out of the 40 total tasks, we use 35 tasks train and to test. Due to space constraints, the results for Meta-World and Half-Cheetah are presented in Figure 5 in the Appendix. We observe that DIT outperforms PDT and BC in all testing scenarios. Moreover, DIT consis- tently outperforms AD despite with less information used for pretraining. It can also be observed that the performance gap between DPT and DIT is larger in the Meta- World en- vironment compared
H
min 0 mH
log To 19, D*).
In particular, AD and DPT require extra information during pretraining: AD requires the complete learning history of RL algorithms while DPT requires optimal action labels. Given that DIT only relies on suboptimal historical data, the comparison is inherently unfair. Notably, despite these disadvantages, DIT outperforms AD and matches with DPT in most scenarios. In terms of metrics, we follow the con- vention to use the episode cumulative return Σ The
Salf-Cheetahh We believe this is is more challenging environment than Half-Cheetah. As result, the additional set of optimal action labels for out-of-trajectory query states used by DPT has a greater impact on performance, while DIT can only utilize in in-trajectory states and actions as query states with pseudo-optimal labels.
because
In-context Decision-man making for Navigating Tasks. We explore how our method generalizes to unseen RL tasks, using the Dark Room environment (Laskin et al., 2022). Following the evaluation protocal of DPT (Lee al., 2024), use 80 goals for training and evaluate the remaining 20 unseen goals. For SAC, since it is an online learning method, we directly train from scratch on the 20 goals to benchmark the returns of ICRL. Figure 3a shows the online evaluation over 40 episodes. After 40 episodes, SAC gains little in return, demonstrating the difficulty of the RL tasks for testing. Restricted by their capability efficiently plore in new tasks, BC also perform poorly. Although our method (DIT) initially has lower returns than DPT and AD, it quickly surpasses them and continues to improve. Fig- ures 3b and 3c show the results for offline evaluations with expert (high-reward) trajectories and random (low-reward) trajectories. Despite being pretrained without the optimal action labels, DIT models demonstrate competitive perfor- mance to that of DPT.
Ablation Study on Weighted Supervised Pretraining. While DIT's significantly improved performance over BC (the unweighted version of DIT) already demonstrates the effectiveness of the proposed weighted pretraining objective, we now conduct experiments in the Miniworld (Chevalier- Boisvert et al., 2023) environment to explore whether DIT reaches the limits of the weighted pretraining framework.
To this end, we compare our model to the DPT model that a pretraining dataset containing only query states that belong to the set of observed states in the pretraining dataset, along with their associated optimal action labels. In this scenario, the total number of pretraining context datasets and optimal action labels for DPT remains the same, but the query states are restricted. This restriction makes the DPT model function as an oracle upper bound for DIT, as all query states used by DIT in the weighted pretraining originate from the observed states.
In-context Continuous Control. We explore two complex continuous control tasks, Meta-World (Yu et al., 2020) and Half- Cheetah (Todorov et al., 2012). Meta- World has 20 tasks in total, to evaluate our approach's 's ability to generate new RL. tasks, we use 15 tasks train and to test. Sim-
The significant performance gain of DIT over BC (the un- weighted version of DIT) demonstrate the effectiveness of the weighted pretraining framework. Surprisingly, in the online setting, DPT struggles to perform, while DIT mod- els gradually improve their returns, as shown in Figure 4.
8

---

In-Context Reinforcement Learning From Suboptimal Historical Data
Optimal
BC
DPT
SAC
30
Return
Retum
Retum
20
20
20
10
10
10
10
15
20 Episodes
25
30
35
40
Opt
DPT
DIT
BC
Opt
DPT
DIT
BC
Methods
Methods
Figure 4: Ablation Study on Miniworld. From left to right: online testing, offline random, offline expert.
In the offinn setting, DIT again demonstrates competitive performance with DPT. These results indicate that DIT has effectively leveraged the pretraining dataset to a significant extent.
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learn- ing? investigations with linear models. arXiv preprinl arXiv:2211.1566 2022.
Auer, P. Using confidence bounds for exploitation- exploration trade-offs. Journal of Machine Learning Research, 3(Nov):397 2002.
6. Discussion
We have proposed framework DIT for pretraining TMs from suboptimal historical data for ICRL. DIT has guar- anteed policy improvements over the suboptimal behavior policies and demonstrated superior empirical performance on a comprehensive set of ICRL benchmarks. Despite these strengths, DIT still requires the behavior policies to have reasonable rewards. Most historical data typically adheres to this constraint. That said, it is highly unlikely infer near-optimal actions solely from random trajectories with- out any information about optimal policies. To this end, we will further explore the limits of the proposed weighted pretraining framework in future work.
Brandfonbrener, D., Bietti, A., Buckman, J., Laroche, R., and Bruna, J. When does return-conditioned supervised learning work for offline reinforcement learning? Ad- vances in Neurul Information Processing Systems, 35: 1542-1553, 2022.
Brown, T.. Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances neural information processing systems, 33: 1877-1901, 2020a.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Acknowledgements
Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020b.
Juncheng Dong and Vahid Tarokh were supported in part by the National Science Foundation (NSF) under the National AI Institute for Edge Computing Leveraging Next Gener- ation Wireless Networks Grant #21 12562. Ethan Fang is partially supported by NSF grants DMMS2346299 and DMS- 2434666. Zhuoran Yang acknowledges support from NSF DMS 2413243.
Canonne, C. L. A short note an inequality between kl and tv. arXiv preprint arXiv:2202.07 198, 2022.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34: 10044-1555977 2021.
Impact Statement
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences our work, none of which we feel must specifically highlighted here.
Chevalier-Boisvert, M., Dai, B., Towers, M., de Lazcano, R., Willems, L., Lahlou, S., Pal, S., Castro, P. S., and Terry, J. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
References
Dong, J., Mo. W. Qi, Z., Shi, C., Fang, E. X., and Tarokh, V. Pasta: pessimistic assortment optimization. In Interna- tional Conference on Machine Learning, pp. 8276 8295. PMLR, 2023.
Achiam, J.. Held, D., Tamar, A., and Abbeel, Constrained policy optimization. In International conference ma- chine learning, pp. 22-31. PMLR, 2017.

---

In-Context Reinforcement Learning From Suboptimal Historical Data
François-Laver, V., Henderson, P., Islam, R., Bellemare, M. G., Pineau, J., al. An introduction deep reinforce- ment learning. Foundations and Trendy@ in Machine Learning, 11(3-4):11-335- 2018.
Li, W., Luo, H., Lin, Z.., Zhang, C., Lu, Z.,, and Ye, D. A sur- vey on transformers in reinforcement learning. Transac- tions Machine Learning Research, 2023a. ISSN 2835- 8856. URL https: /oppnreviei net forum? i 1d=r30yuDPvf2. Survey Certification.
Fujimoto, S. and Gu, S. S. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20141 2021.
Li, W., Luo, H., Lin, Z., Zhang, C., Lu, Z., and Ye, D. A survey on transformers in reinforcement learning, 2023b. URL httthhtt///// arxiv org/ abs 1000.30110000000 03044.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off- off--ooicy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference machine learning, pp. 1861-1870. PMLR, 2018.
Liu, Z., Guo, Z., Yao, Y. Cen, Z., Yu, W., Zhang, T., and Zhao, D. Constrained decision transformer for offline safe reinforcement learning. In International Conference Machine Learning, pp. 21611-11122222 16611-21630. PMLR, 2023.
Jin, Y., Yang, Z., and Wang. Z. Is pessimism provably efficient for offline rI? In International Conference on Machine Learning, pp.
Matsushima, T., Furuta, H., Matsuo, Y., Nachum, O., and Gu, S. Deployment-efficicee reinforcement learning via model-based offline optimization. arXiv preprint arXiv:2006.03: 03647, 2020.
PMLR, 2021.
Kaelbling, L. P., Littman, M. L., and Moore, A. W. Re- inforcement learning: survey. Journal of artificial intelligence research, 4:237-285, 1996.
OpenAl. Gpt-4 technical report, 2024.
Peng, X. B., Kumar, A., Zhang, G., and Levine, Advantage-weighted regression: Simple and scalable off off-policy reinforcement learning. arXiv preprint arXiv: urXiv:9910.00177 2019.
Kasneci, E., SeBler, K., Küchemann, S., Bannert, M., Dementieva, D., Fischer, F., Gasser, U.. Groh, G., Günnemann, S., Hüllermeier, E., et al. Chatgpt for good? on opportunities and challenges of large language models for education. Learning and individual differences, 103: 102274, 2023.
Prudencio, R. F., Maximo, M. R., and Colombini, A survey on offline reinforcement learning: Taxonomy, review, and open problems. IEEE Transactions Neural Networks and Learning Systems, 2023.
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems, 33: 11810-21820 2020.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforce- ment learning with implicit q-learning. arXiv preprint arXiv:21 10.06169, 2021.
Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., and Dormann, N. Stable-baselines3: Reliable rein- forcement learning implementations. Journal of Machine Learning Research, 22(268): 2021. URL http: html.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020.
www.:opppppppppppp.p..otop...eep //jmlr. .rg. papppsaapssssss//peep wttpsr//v/2////22
Rashidinejad, P., Zhu, B., Ma. C., Jiao, J., and Russell, S. Bridging offline reinforcement learning and imita- tion learning: A tale of pessimism. Advances in Neural Information Processing Systems, 34:11702-11716, 2021.
Laskin, M., Wang, L., Oh, J., Parisotto, E., Spencer, S., Steigerwald, R., Strouse, D., Hansen, S.. Filos, A., Brooks, E., et In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14211 2022.
Russo, D. J., Van Roy, B., Kazerouni, A., Osband, I., Wen, Z., et al. A tutorial on thompson sampling. Foundations and Trends@ in Machine Learning, 11(1):1-96, 2018.
Lee, J.. Xie, A., Pacchiano, A., Chandak, Y., Finn, C., Nachum, O., and Brunskill, E. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Schulman, J. Trust region policy optimization. arXiv preprint arXi::150..05477 2015.
Singhal, K., Azizi, S., Tu, T., Mahdavi, S. S., Wei, J., Chung, H. W., Scales, N., Tanwani, A., Cole-Lewis, H., Pfohl, S., et al. Large language models encode clinical knowledge. Nature, 620(7972): 2023.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.00 00000000.0/64,.... 2020.
10

---

In-Context Reinforcement Learning From Suboptimal Historical Data
Thirunavukarasu, A. J., Ting, D. S.J., Elangovan, K., Gutier- rez, L.,, Tan, T. F., and Ting, D. S. W. Large language models in medicine. Nature medicine, 193))1930-, 2023.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., and Gan, C. Prompting decision transformer for few- shot policy generalization. In international conference on machine learning, pp. 24631-24645. PMLR. 2022.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ inter- national conference on intelligent robots and systems, pp- 5026-5033. IEEE, 2012.
Yamagata, T., Khalil, A., and Santos- Rodrige-o Q-
learning decision transformer: Leveraging dynamic pro- gramming for conditional sequence modelling in offline rl. In International Conference on Machine Learning. PP. 38989-39007. PMLR, 2023.
Touvron, H., Martin, L., Stone, K. R., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D. M., Blecher, L.., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshomn A. S., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I. M., Korenev, A. V., Koura, P. P.S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I,, Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Sub- ramanian, R., Tan, X., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P. Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang. S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Ilama 2: Open founda- tion and fine-tuned chat models. ArXiv, abs/2307.09288. 2023. URL https: /apii www.etiiascseeeoeoeeoo org/ htt//oorpusID:23955000
Yang, H., Liu, X.-Y, and Wang, C. D. Fingpt: Open-source financial large language models, 2023.
Yang, M., Schuurmans, D., Abbeel, P., and Nachum, O. Dichotomy of control: Separating what you can control from what you cannot. arXiv preprint arXiv:2210.1344 134355,,, 2022.
Yin, M. and Wang, Y.-.. Towards instance-optimal offline reinforcement learning with pessimism. Advances in neural information processing systems, 34:4065-4078. 2021.
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, Weea-world: benchmark and evalua- tion for multi-task and meta reinforcement learning, In Conference robot learning, pp. 1094-1100, PMLR, 2020.
Vilalta, R. and Drissi, Y. A perspective view and survey of meta-learning. Artificial intelligence review, 18:77-95, 2002.
Yuan, W., Chen, J., Chen, S., Lu, L., Hu, Z., Li, P., Feng. D., Liu, F., and Chen, J. Transformer in reinforcement learn- ing for decision-making: A survey. Authorea Preprints, 2023.
Wang, Q. Xiong, J., Han, L.,, Liu, H., Zhang, T., et al. Exponentially weighted imitation learning for batched historical data, Advances in Neural Information Process- ing Systems, 31, 2018.
Zheng, Q., Zhang, A., and Grover, Online decision trans- former. In international conference on machine learning, pp. 27042-27059. PMLR, 2022.
Wu, S., Irsoy, O., Lu, S., Dabravolski, V., Dredze, M.. Gehrmann, S., Kambadur, P., Rosenberg, D., and Mann, G. Bloomberggpt: A large lan- guage model for finance. ArXiv, abs/2303.17564, 2023a. URL https: http:/ppp...mmmancccccc....oeeee : 257833842.
Wu, T., He, S., Liu, J., Sun, S., Liu, K., Han, Q-LL,, and Tang, Y. brief overview of chatgpt: The history, status quo and potential future development. JEEE/CAA Journal of Automatica Sinica, 10(5):1 122-1136, 2023b.
Wu, Y., Tucker, G., and Nachum, O. Behavior regularized offline reinforcement learning, 2019.
Xiao, C., Wu, Y., Mei, J., Dai, B., Lattimore, T., Li, L.., Szepesvari, C., and Schuurmans, D. On the optimal- ity of batch policy optimization algorithms. Inter- national Conference on Machine Learning, pp. 11362- 11371. PMLR, 2021.

---

In-Context Reinforcement Learning From Suboptimal Historical Data
A. Results on Meta-World and Half- Caff-Cheetah
10
DIT
BC SAC
PromptOT
15
Return
6
Return
Return
10
10
5
10
15
20
25
30
35
40
DPT
DIT
PDT Methods
AD
BC
DPT
DIT
PDT Methods
AD
BC
Episodes
100
100
100
DIT
BC
Prompt-DT
75
75
Return
Retur
Return
40
25
DPT
Episades
Methods
Methods
(a) Online Testing
(b) Offline Random
(c) Offline Expert
Figure 5: Top row shows results on Deta:World; bottom row shows results on Half-Cheetah.
B. Baselines
B.1. Bandit Algorithms
Empirical Mean (EMP). We follow (Lee et al., 2024) to consider strengthened version of EMP which, in the offline setting, only chooses from actions that have been observed at least once in the offline dataset while, in the online setting, at least choosing every action once. At every time step, EMP chooses actions as
arg Pax aEA
where D. is the average observed reward for action
Upper Confidence Bound (UCB). Motivated by the Hoeffding's Inequality, at each time step, UCB chooses actions as
à E arg max pha + C V1/na.
aeA
where C' is a hyperparameter and n, is the number of times a has been chosen. For unseen actions, P. set to 0 and n. set to 1. We follow (Lee et al., 2024) to set to as demonstrates the best empirical performance.
Lower Confidence Bound (LCB). LCB is on the contrary of UCB. In the offline setting, LCB only chooses from observed actions in the offline dataset. Specifically, it chooses actions
p CC.//// v1/maa
E arg max asA
where C' is hyperparameter and n. is the number of times has been chosen. Similar to hyperparameter of UCB, the hyperparameter C' for LCB is also set to 1 due to its strong empirical performance.
Thompson Sampling (TS). We use Gaussian TS (Russo et al., 2018) with Gaussian prior. The mean and variance the prior are set to the true mean and variance of the pretraining tasks: for mean and for variance.
12

---

In-Context Reinforcement Learning From Suboptimal Historical Data
B.2. RL Baselines
Peedsion-Pretrannee Transformer (DPT). The Decision-Pretrainee Transformer (DPT) is designed to perform in-context learning for reinforcement learning (RL) tasks by leveraging a supervised pretraining approach. The core idea is to train transformer model to predict optimal actions given query state and corresponding in-context dataset, which contains interactions from a variety of tasks. These interactions are represented as transition tuples consisting of states, actions, and rewards, offering context for decision-making. During pretraining, DPT samples a distribution of tasks. For each task T; an in-context dataset D; is constructed to include sequences of setle-action-eeweeee with that task. Additionally, a query state s" is sampled from the MDP's state distribution, and the model is trained to predict the optimal action based on this query state and the context D,. Formally, the training objective is to minimize the expected loss over the sampled task distribution by predicting a distribution over actions given the state and context.
interactions that represent past experience
Prompt-based Decision Transformer (Prompt-DT). Prompt-DT arranges its data to facilitate few-shot policy general- ization by using trajectory prompts. For each task Ti, prompt Tę¹ of length K* is constructed from few-shot demonstration data Pt. containing tuples of state, action, and reward-to-go (s", a", G*). This prompt encodes task-specific context necessary for policy adaptation. Additionally, the recent trajectory history TÚ of length K, sampled from an offline dataset D,, is appended to the prompt to form the full input sequence Tingut- Formally, this input sequence is represented as .fx +K & *K* "K* This sequence Tinput (T,' T²) = (fI, si, ai....,, Kk, sK+, ax.; f KK- "1:1:.
contains 3 K* + K) tokens, following the state-actio-reeward format. The full sequence Tinput is then passed through a Transformer model, which autoregressively predicts actions at the heads corresponding to each state token. We follow Prompt-DT's setting and set = 20.
Algorithm Distillation (AD). Algorithm Distillation (AD) transforms the process of reinforcement learning (RL) into an in-context learning task by training a transformer model to predict optimal actions based on a cross-episodic trajectory. AD gathers trajectories from training episodes, where each trajectory T of length H encodes the states, actions, and rewards observed over multiple episodes. Instead of training via traditional gradient updates, AD models the training history to predict actions for subsequent episodes, effectively distilling the behavior of RI algorithms like SAC into the transformer. This enables the model to learn directly from context, facilitating quick adaptation to new tasks and improving learning efficiency.
Behavior Cloning (BC). Behavior Cloning (BC) is a supervised learning approach for imitation learning, where the goal is to learn to mimic the behavior of policy by mapping states to actions. Specifically, the objective is to minimize the discrepancy between the actions predicted by the learned policy wo and the target policy's actions, often through a loss function such as mean squared error or cross-entropy for continuous or discrete action spaces, respectively: J(O) [[((T)
B(s,aa)) [[(TT( (st), a.),, where D is the dataset of state-action pairs collected from the target policy's 's demonstrations, St the state at time step i, and a, the corresponding target action.
Soft Actor-Critic (SAC). Soft Actor-Critic (SAC) is an off-policy deep reinforcement learning algorithm that balances exploration and exploitation by maximizing a trade-off between expected reward and entropy. The core objective SAC is to learn a policy that not only maximizes cumulative rewards but also encourages exploration by maximizing the entropy of the policy's actions. SAC uses an actor network to predict actions, and two critic networks to estimate the Q-values of latteaaction pairs. The training objective involves learning the parameters of the policy to maximize a soft objective function: J(T) E, Be
(Q((= (S1,, - a log π (dr]s:)], where Q(s,, at) is the Q-value estimated by the critics, is temperature parameter controlling the trade-off between reward and entropy, and π (a, is the action probability distribution given the state. SAC is trained by sampling
of transitions from a replay buffer to update the policy (actor) and Q-value estimates (critics). For model and training settings, we use the default implementation from Stable Baselines3 (Raffin et al., 2021).
13

---

In-Context Reinforcement Learning From Suboptimal Historical Data
C. Theoretical Results
C.1. Proof of Proposition 4.1
Consider the following optimization problem where E,,, is defined as in Equation (2) except that a ~ T() ie. the action sampled from the task-conditioned policy rather than the behavioral policies:
max ( E,,,,0
A (s, -1 (K
(10)
(1)
where DKL, is the Kallback-Leiblees (KL) divergence, and let 'T" arg (TT be its optimizer. Then we have for any policy (;)s))),
E Brpp(r)) (s) [RL (T (1)=;
(11)
= -
(+(s) exp (s, a)/n) log x( ((als:T) + +C,
where 5,, is defined as (2), C a constant independent of and (+(s) Ea (a((s) (x(aa)))xpppp)) exp (8, /n).
Proof of Proposition 4.1. For any task and fixed state 8, have
max (a)s())) [ a) 77 (-
min E (alar))))) flog T() 1 2 (s)
(alss 17
<()];;T)
min E (alart))))
log (als) (lss)exp(444 (A4 (s, /n)
T())s;T)
min E (aler))
log ex(((A))xpAAA (s, (s)
log Z, (s)
<()(sT))
min E (alar)))
log (a(sss (1flls))xppAAAl (xp(A4 (s, a) (s)
(Z,(s) independent of x)
min (T(ss
where n (a|s) = (-1ss exp(A" a) /m) / / (s). Note that the optimum for fixed 8 and task is obtained at =; which is unique by the uniqueness property of KL divergence, i.e., (=l(T||=)) = Of and only if π (;(ss.. Thus, the optimal task-conditioned policy is
π (a|s;)) = क (xp() (s (s).
Thus, we further have
(xd (DIL. (TT (+1s; ()s; (+};s]];
x (als;;)) log x(a|siT)
By-p(r) swdd, (a)
By (s) Σ (als)) (s.a) /n)////Z) <a(ss:))
T(a|s:T)
-EETp(() (s),am)((al)) (exp(A) (s. (s) +C,
where C' = E, (t)) [exp(Ab (4)
C.2. Proof of Proposition 4.2
Let nt be the policy that optimizes Equation (3). For any task and policy let (T) = E "Tl/ represent the expected reward of for T. Let π; denote x (als; T). Then
Eyop, [G, G, (TT,))
>
n E, Erap, [C
2y Er
CA V CD/2
(12)
1
(1 (1-7)² 7)
14

---

In-Context Reinforcement Learning From Suboptimal Historical Data
where Cp (e) DRL( (x* (117)))))))))) and C, = maxx, (olss) a)].
Proof of Proposition 4.2. First consider any fixed task From Corollary in (Achiam al., 2017), we have
G+( -x
2
1 (s)
Σ क [ (als: (s.a)
24, E Band, (1
(13)
(1 (1-
where C.A = max, (alos) 1 a)l and || ||r is the total variation distance between two distributions. In the proof of Propsition 4.1, we observe that: for any and s,
n" 7-{=;; max [(x, s) (a-ar) [4 (s. 17. (I
Thus, L(m;, 2 (a s). which implies that
Eawx (ale:r) A (s, a) - 1 DKI. (1177)))1)))))))) 2 (a-n)( [4, (s, as = 0.
Hence, we have
Egnd (s), [A* (s, a)]
(14)
Moreover, from Pinsker's inequality (Canonne, 2022).
E Egdd, (6) (6)]l (0 (185;) (-1((1))]TV S Band, (0) ((ass
(15)
Eswd (s) (T+ (11=)))))
(16)
where the last inequality comes from Jensen's 's Inequality, Pluging (15) and (14) into (13), we have
-()t)),((T))) 2 " E Band, (1) DKLL ( (11s;))|
(17)
(1 (1- V / /1 (PL (1877))))
27Cr
(18)
Taking expectation with respect T concludes the proof:
E (T)) G, 1-7 Erppr 10
"
(1 11-7)1 7)2 Erapr OA, CD
2y
(19)
where CP = (s) (DKL(T* (-Iss) and C²A1 = max, [Bar (4 ())
C.3. Justification for the identity (s)
Assume that (s, a)/n << |log 1log= (a|s)]. Note that this can always be satisfied through reward normalization. Then
(x) Σ
(s. a) /n) (al (8, a) /v)]
(aln) + A (s, a) /n o((A (s, /n) (by Taylor expansion).
Moreover, by definition of the advantage function, we have
E (a) (s, a)] = Ea (+. -V) = 0.
Thus,
(+) = -1+ E ((L)))
15

---

In-Context Reinforcement Learning From Suboptimal Historical Data
D. MDP Environments
Dark Room. The agent is randomly placed in a room of 10 x 10 grids, and there is an unknown goal location on one of the grid. Thus, there are 10x10- 100 goals. The agent's observation is its current position/grid in the room, i... S [10] x [10]. The agent needs to move the goal location by choosing from 5 actions: move in one of the directions (up, down, left, right) or stay still. The agent receives reward of 1 only when it is at the goal; otherwise, it receives O. The horizon for Dark Room is 100. We follow (Lee et al., 2024) to use the tasks on 80 out of the 100 goals for pretraining, and reserve the rest 20 goals for testing our models' in-context RL capability for unseen tasks. The optimal actions are defined as: move up or down until the agent is on the same vertical position as the goal; otherwise move left or right until the agent reaches the goal.
Miniworld. The agent is placed a room with four boxes of different colors, one of which being the target box. The goal to reach box of a specific color in the room. The agent receives a (25 x 25 x 3) color image and its 2-D direction as input, and can choose from four possible actions: to turn left/right, move straight forward, stay still. Similar to Dark Room, it receives reward of 1 only when it is near the target box while the horizon is 50. The optimal actions are defined as follows: turn left/right towards the correct box if the agent's is front is not within 15 degrees of the correct box; otherwise move forward and stay if the agent is near the box.
Morl-W The agent needs to control a robotic arm to pick up an object and place it at a designated target location. In each task, the state space is in 39 dims including the gripper's 's position and state (open or closed), the 3D position of the object to be manipulated, and the coordinates of the target location. The agent operates in a continuous action space, where it can adjust the gripper's s 3D position and control the open/close state to enable successful grasping and releasing of the object. It provides partial rewards for moving the gripper towards the object, grasping it correctly, transporting it to the target location, and successfully releasing it there. The task goal is to learn an optimal policy that efficiently achieves the sequence of actions required to pick up and accurately place the object at the specified location. Each task has a different goal position. We train in 15 tasks and test in 5 tasks.
The agent needs to control 2D half-cheetah robot to achieve and maintain varying target velocities, which change across episodes. The state space contains the cheetah's 's motion, including joint angles, velocities, body velocity, and position. These observations enable the agent to learn intricate movement patterns and maintain balance while running. The action controls the torques applied to each joint of the cheetah, thus dictating its locomotion and stability. The reward is designed to align with the core task objective: matching the agent's 's velocity to the target velocity. Each task has different target velocity, and we use 35 tasks to train and 5 to test.
Half-Cheetah.
E. Pretraining Dataset
Pretraining Datasets for Dark Room and Miniworld. To ensure coverage of optimal actions (so that optimal policies can be inferred), at every step, with probability p (respectively p) we use optimal policy (respectively random policy) to choose action. We choose p that the average reward of the trajectories in the pretraining dataset is less than 30% of that of the optimal trajectories, reflecting the challenging yet common scenarios. For Dark Room, to test whether DIT models can generalize to unseen RL problems in context, we collect context datasets from only 80 out of the total 100 goals and reserves the rest 20 for testing. For each training goal, we follow the setting of DPT to collect 1 context datasets, leading to a total of 80k context datasets in the pretraining dataset (64k for training and 16k for validation). For Miniworld, we collect 40k context datasets (32k for training and 8k for validation), 10k datasets for each of the four tasks corresponding to four possible box colors.
Pretraining Datasets for Meta-World and Half-Cheetah. We construct the pretraining datasets using historical trajectories generated by agents trained with Soft Actor Critic (SAC). Specifically, SAC is trained until convergence for each task, then we sample from its learning trajectories to build the dataset. Our SAC model training follows the settings outlined in (Haarnoja et al., 2018). For the Meta-World environment, we use its built-in deterministic policy as the optimal policy; for ChlffChhettah we use the optimal SAC policy. In Meta-World, we used 15 tasks to train and 5 to test. Similarly, for Half-Cheetan we used 35 tasks to train and 5 to test.
16

---

In-Context Reinforcement Learning From Suboptimal Historical Data
01
6
Gh
Gi
6
Gp
auto-regressive model
euto-regressive model
aj
22
G2
.
s/
aj
G
52
Gh
Figure 6: Model structure of the in-context action-value transformer Q (left) and value transformer V of the i-th pretraining task.
(right) the trajectory
F. Training Parameters.
For all methods, we the A optimizer with a weight decay of le 4, learning rate of 1e 3, and batch size of 128.
G. Model Details
Decision Transformer Architecture. Our model is based on a causal GPT-2 architecture (Radford et al., 2019). It consists of attention layers, each with single attention head, and an embedding size of 256. To separately encode state, action, and reward pairs, we employ three fully connected layers. We use single fully connected layer decode from the transformer's output.
Value Function Transformer Architecture. transformer.
The architecture of the value function transformer mirrors that of the decision
H. Computation Requirements
Our experiments can be conducted on single A6000 GPU. typically takes less than one hour to generate the required dataset for training in parallel. For PPO, training usually takes less than 10 minutes per task. For the other methods, we observe that the transformer model converges within 50 epochs.
Pseudocodes
Algorithm Pretraining of Decision Importance Transformer
1: Input: Pretraining Dataset D = {Di}: transformer models To, Qc, V/g.
2: // In-context Estimation of Advantage Functions
3: Randomly initialize and train Q₂ and Vo by optimizing the loss in Equation (8).
4: Construct the in-context advantage estimator as:
As = Qc -Vo-
5: // Weighted Pretraining
6: Randomly initialize Tg.
With trained As and D, train Tp by optimizing the loss in Equation (9).
J. Estimation of Advantage function
Figure 7 illustrates the performance of our value function estimators. Notably, the ground truth labels represent the cumulative rewards empirically sampled using Monte Carlo, rather than the in-trajectory cumulative rewards. From the two graphs, we observe that our function estimator effectively learns the empirical distribution of cumulative rewards. Furthermore, the difference between the Q-function and V-function estimators provides the advantage function.
17

---

In-Context Reinforcement Learning From Suboptimal Historical Data
Algorithm 2 Deployment of In-Context RL Models
1: Input: Pretrained transformer Model Ta: Horizon of episodes H; Number of episodes N for online testing; Offline
dataset Dar = {(*,, 4/: Ph+1:Th))b h. consisting of transitions collected by a behavioral policy.
2: // Offline Testing
3: for every time step / E [l,..., H do
4: Observe state Sh
5: Sample action with Tg:
on To (11sh, Det)
6: Collect reward rh
7: end for
8: // Online Testing
9: Initialize an empry online data buffer Da = 0
10: for every online trial n {I...., ] do for every time step h E {1..... H) do
Observe state sh
13:
Sample action with To:
On To (-|sh, Da.)
14:
Collect reward
15:
end for
16:
Append the collected transitions {8h, (hi 8h++1 , rh)
into Dan
end for
Prediction Ground
Prediction vs Ground Truth
Prediction
Prediction
0.50
0.5
60
80
100
20
40
60
80
100
Time step
step
Q function
(b) function
Figure 7: Performance of Q and V function estimator. On the x-axis is time step of horizon; on the y-axis is the model predictions or ground truth values.
K. Effectiveness of in-context trajectory
Figure illustrates the effectiveness of the in-context trajectory for DIT. Since DIT predicts actions based on the current state and the historical states in the in-context trajectory, it is crucial to ensure that the task goal of the in-context trajectory aligns with the current task that DIT predicting. Here, "In Task" refers to cases where the in-context trajectory is sampled
18

---

In-Context Reinforcement Learning From Suboptimal Historical Data
80
60
Return
40
20
0
In Task
Out Task
Figure 8: Performance of DIT when the in-context trajectory is aligned (In Task) or misaligned (Out Task) with the current task goal.
from same task as the current task, while "Out Task" indicates that the in-context trajectory is sampled from a different task.
From Figure 8, observe that alignment between the in-context trajectory and the current task goal is critical for effective performance. This finding also validates that DIT relies heavily on the in-context trajectory for action prediction, as misalignment with the current task goal leads to significant decrease in performance.
19


Reviews:
[{"review_id": "ZYVqWAvJcB", "reviewer": "Reviewer_JvGg", "summary": "The paper proposes the Decision Importance Transformer (DIT), a novel framework for in-context reinforcement learning (ICRL) that is designed to work with historical datasets generated by suboptimal behavioral policies. Unlike previous approaches that require optimal action labels or complete learning histories, DIT uses only suboptimal data. Its key idea is to incorporate an exponential reweighting scheme in the supervised pretraining objective. This weighting is derived from an estimated advantage function that is computed via transformer-based value and action-value estimators. The overall system is built upon an autoregressive transformer (using a GPT-2 backbone) that, once pretrained on a diverse set of tasks, can generalize to unseen tasks by extracting task-specific information from context trajectories. The paper supports its contributions both theoretically—with propositions that connect the weighted maximum likelihood objective to policy improvement—and empirically, by demonstrating competitive or superior performance on a range of problems (from bandit settings to challenging MDPs including navigation and continuous control tasks).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments. We hope our response can address your concerns. \n\n### Presentation \n\n> **Style Improvement.** We will follow the reviewer’s suggestion to remove most of the colored text, keeping only minimal phrases with key importance. \n\n> **Clarification on the term “In-Context.”** We appreciate the reviewer highlighting this potential ambiguity. While “in-context learning” originated in the LLM literature, the term has recently been adopted by several works in reinforcement learning to describe similar ideas—namely, models that adapt to new tasks by conditioning on sequences of past transitions, without gradient updates. Our use of “in-context RL” follows this growing convention. \nNonetheless, we agree that the dual usage can be confusing. In the revised manuscript, **we will explicitly distinguish in-context RL from ICL in language models to ensure clarity and avoid potential misunderstandings**.\n\n### Methods.\n\n> **Zero-shot Adaptation.**  DIT builds on the supervised pretraining framework of DPT and can act as an online meta-policy in new environments. Specifically, **it can predict actions without conditioning on a context dataset, as the prediction is only based on the token corresponding to the query state**. In the zero-shot setting—i.e., with no context tokens—DIT and DPT default to the behavior learned during pretraining, effectively leveraging a strong prior. Cold-start performance is demonstrated in the online testing experiments (Figures 2, 3, and 4). These results show that **DIT quickly improves its performance in-context with more information (trajectories), adapting efficiently even in the cold-start setting**.\n\n### Learned Advantage Function and Distribution Shift\n\n> We appreciate the reviewer’s insightful comment. We would like to clarify that **DIT does not require robustness to distribution shift during the pretraining**. \n\n> This is because the learned advantage functions are only applied to actions taken by the same behavioral policies that generated the trajectories—meaning **all computations remain in-distribution with respect to those behavioral policies**. \n\n> This design ensures stability and avoids the challenges typically associated with off-policy corrections. We consider this to be a key strength of DIT, and we will make this point more explicit in the revised manuscript.\n\n### Clarification of Contributions\n> While similar theoretical results have been established in the standard RL setting, **our contribution lies in extending these insights to the in-context reinforcement learning (ICRL) framework**. Although theory is not the primary focus of this work, we include the analysis to better motivate and ground the design of DIT, **helping the reader understand why the method is effective**.\n\n> Our core insight is that if \"relatively good\" actions can be identified and emphasized during supervised pretraining, a transformer-based policy can match the performance of meta-policies trained with significantly more expensive data, such as DPT. Although individual weighted action labels may be noisy, our analysis and experiments show that **the weighted MLE objective—when averaged over diverse environments—can yield high-quality meta-policies capable of generalizing to unseen tasks**.\n\n> Furthermore, with the proposed in-context advantage estimator, we observe that transformer models can learn to generate reliable value function estimates across tasks and behavioral policies using only supervised learning. We believe these empirical insights are important to share, as they highlight the practicality and potential of ICRL.\n\n> On the practical side, **DIT is simple to implement and significantly improves the feasibility of in-context RL**, as suboptimal trajectories are much easier to collect in real-world systems, especially in industry settings where large amounts of historical data are already available. **This opens the door for broader adoption of ICRL methods in practical applications**.", "claims_and_evidence": "Core Claims:\n- Policy Improvement via Reweighting: The authors claim that by reweighting action labels according to an estimated advantage function, the transformer can learn to “steer” suboptimal behavioral data toward near-optimal policies.\n- Generalization from Suboptimal Data: DIT is posited to work well even when the training data do not contain optimal action labels, a scenario that is common in real-world settings where only historical data are available.\n\nEvidence:\n- Theoretical Analysis: The paper presents Proposition 4.1 and Proposition 4.2, which relate the weighted supervised objective to a policy improvement problem and provide conditions under which the learned policy strictly improves over the behavior policy. Although proofs are deferred to the appendix, the propositions outline a clear link between the exponential weighting scheme and performance guarantees.\n- Empirical Results: Extensive experiments are conducted on both bandit problems and various MDPs (including navigation tasks like Dark Room and Miniworld, as well as continuous control tasks such as Meta-World and Half-Cheetah). The results show that DIT can match or exceed the performance of baselines (including methods that have access to optimal action labels) in both online and offline settings.", "methods_and_evaluation_criteria": "Methodology:\n- Weighted Maximum Likelihood Pretraining: DIT is trained using a weighted maximum likelihood objective where the weights are an exponential function of an estimated advantage function. This aims to give higher importance to actions that are deemed “better” in the historical data.\n- In-Context Advantage Estimation: The paper introduces transformer-based modules to estimate the value and action-value functions in-context, thereby approximating the advantage for each state–action pair.\n- Task-Conditioned Policy: The transformer model is conditioned on the context (a set of transitions from the environment) so that it can adapt its decision-making to new, unseen tasks during deployment.\n\nEvaluation:\n- The experiments are carried out in both online (where the agent gathers additional data) and offline (using fixed historical trajectories) settings.\n- The evaluation metrics include cumulative return and regret (for bandit problems) and episode cumulative return (for MDPs), which are standard and appropriate for the problem domain.", "theoretical_claims": "I did not check the detailed derivations, but they seem reasonable as the conclusions are standard in RL literature.", "experimental_designs_or_analyses": "The paper evaluates DIT on both simple (linear bandit) and complex (MDP) environments.\nFor bandit problems, the method is compared against theoretically optimal algorithms (e.g., UCB and Thompson Sampling), showing that DIT quickly identifies the optimal bandit.\nIn MDP settings, the authors test on environments with both sparse rewards (Dark Room, Miniworld) and complex dynamics (Meta-World, Half-Cheetah), and compare against several baselines including DPT, AD, and behavior cloning variants.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Weaknesses\n- According to the proposed method, it is extremely important to learn a universal advantage function that is robust to distribution shifts, which may not be scalable.\n-  In context learning needs a few tokens to construct the context. What is the cold-start performance of ICRL? How to perform reliable zero-shot adaptation?\n-  There seems to be a reload of meaning between the \"In Context\" RL and \"In Context\" Learning of LLMs.\n- The theoretical results are less significant, considering that Prop. 4.1 is a well-known conclusion in RL, and the quadratic reliance on effective planning horizon in Prop. 4.2 makes the bound rather loose.", "other_comments_or_suggestions": "- Too many sentences are marked with red, which hinders a smooth reading process. I would recommend to mark very few phrases that are truly important.\n-  The Primary Area: Reinforcement Learning->Everything Else should be used sparsely. Is there a specific category that this paper falls into?", "questions_for_authors": "See weaknesses.", "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": "NNqfRXvSWs", "reviewer": "Reviewer_EYEv", "summary": "This work focuses on in-context reinforcement learning where the source data/policy is suboptimal. In this case, the traditional ICRL algorithm could perform bad. This work proposes Decision Importance Transformer (DIT), which emulates the actor-critic algorithm in an in-context manner. It achieves superior performance than the baselines when the offline data is suboptimal.\n\n\n## update after rebuttal\n\nI appreciate the authors' rebuttal. However, I think the authors should address my Question 7 more properly in the paper. I am not fully convinced by the authors' explanation. DIT uses the state from the context as the query state while DPT considers a more general way (use a random query state), shouldn't DPT work better in terms of learning a general understanding of an MDP? DIT will be limited by the size and richness of the context dataset. In any case, it remains unclear why DIT can outperform DPT.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful and constructive comments. We will incorporate the referenced work into the final manuscript. In addition, we will include citations supporting the availability of abundant historical data in real-world settings (e.g., using [1] as a standard reference). Please see below for our responses to your other concerns.\n\n### Why DPT and ICRL Work\n> As discussed in the original DPT paper, on a high level, DPT internally conducts posterior sampling (PS) for MDP, where the transformer model is pretrained to infer the target MDP from the given context dataset and to take action according to the optimal policy for the inferred target MDP. **As the context size increases, the inferred MDP becomes closer to the true target MDP, thus leading to improved performance**. Meanwhile, there are also theoretical works establishing the efficacy of the supervised pretraining approach taken by DPT, e.g., [2]. \n\n> **Prediction without a context dataset.** The DPT model can still predict an action using only the query state as input (without a context dataset), as the prediction is only based on the token corresponding to the query state. \n\n### Method \n\n> **Choice of hyperparameter.** As shown in our theoretical result, $\\eta$ represents the penalty for the KL divergence. Thus, if the historical data is less suboptimal (high quality), it should be high as we would like to stay close to the behavior policies; on the contrary, if the data is low quality, it should be chosen as a small value so that we can improve more over the behavior policies. \n\n> **Value transformer model structure.**  Our design follows the in-context learning setup, where transformers are trained for regression tasks $y = f(x)$ using sequences like $x_1, y_1, ..., x_n, y_n$. In our case, value estimations (for V and Q) can be seen as regression problems, e.g., set $x = s$ and $y = G$ (cumulative reward) when estimating V. Based on this, we design our value transformer as shown in Figure 6. Figure 7 provides empirical evidence that it can accurately estimate value functions in context.\n\n> **Cost of training in-context advantage estimator.** As detailed in the manuscript, we train them with standard supervised training objectives for transformers, which are stable and straightforward to optimize. While this moderately increases the total pretraining time, we emphasize that **the primary bottleneck in ICRL lies in collecting high-quality pretraining data, not in training time**. To this end, DIT reduces this data requirement considerably, trading off for a modest increase in training complexity. We believe **this tradeoff is worthwhile and offers strong practical benefits**, making ICRL more accessible and scalable in real-world applications.\n  \n### Experiments\n> **Historical Data Suboptimality.** In our experiments, we use behavioral policies with performance 30%-50% of the optimal policies to collect historical data, a commonly used setting for offline RL. We would like to clarify that we use the intermediate training checkpoints of SAC to collect trajectories so that they are indeed suboptimal. We will update the manuscript to make this clear in the introduction to avoid confusion. \n\n> **Experiment Setup.** For Meta-world, we follow the setting of Prompt-DT to use Meta-World reach v2. We use Mini-world mainly for ablation study to understand the importance of the proposed reweighting mechanism and the effect of lack of optimal action labels. Given this purpose, we compare DIT with DPT and a variation of DIT without reweighting. \n\n> **Performance of DPT.** The observed performance difference of DPT between this work and the original paper is due to the difference of the pretraining datasets. Due to space constraint, please refer to our response to Reviewer Gczk for more details. \n\n### DIT sometimes outperforms DPT\n> We appreciate the reviewer’s insightful observation. This is indeed a compelling research direction. \n\n> At a high level, DPT uses a single randomly sampled query state paired with an optimal action label for each trajectory during pretraining, whereas DIT leverages multiple reweighted suboptimal action labels. **Interestingly, these many reweighted suboptimal labels can collectively provide a stronger learning signal than a single optimal label, resulting in better pretraining objectives in some settings**. Of course, if DPT were given access to many optimal action labels per trajectory, it would likely surpass DIT, as DIT does not use any optimal supervision. \n\n[1] Levine, S., Kumar, A., Tucker, G., & Fu, J. (2020). Offline reinforcement learning: Tutorial, review. and Perspectives on Open Problems, 5.\n\n[2] Lin, Licong, Yu Bai, and Song Mei. \"Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining.\" arXiv preprint arXiv:2310.08566 (2023).", "claims_and_evidence": "Overall yes. But I still have some questions and recognize some weaknesses. Please see them in the corresponding section below.", "methods_and_evaluation_criteria": "Yes. The problem is useful and the benchmark environments are reasonable.", "theoretical_claims": "Yes, they look correct to the best of my knowledge.", "experimental_designs_or_analyses": "Yes I checked. Please see the weakness and my questions below.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "The following reference is missing, which also studies how the offline data affect the performance of ICRL.\n\nZisman, I., Kurenkov, V., Nikulin, A., Sinii, V., & Kolesnikov, S. (2023). Emergence of in-context reinforcement learning from noise distillation. arXiv preprint arXiv:2312.12275.", "other_strengths_and_weaknesses": "1. In the sentence: \"When presented with a context dataset containing environment interactions collected by unknown and often suboptimal policies, pretrained TMs predict the optimal actions for current states from the environmental information within the context datase\". I have two concerns. \n\na) It sounds like there needs to be a pre-collected context dataset that is presented to the TMs. However, the online deployment of DPT already shows that the pretrained TMs can collect its own contexts online. \n\nb) \"predict **the optimal actions**\" doesn't sound good to me, because AD does not predict the optimal actions, it predicts whatever in the source RL algorithm. Indeed, here the pretrained TMs predicts whatever it learns during the pretraining, which may not be the optimal actions, depending on which ICRL algorithm is considered. Maybe it is just a language thing, but make it more clear and accurate.\n\n2. \" For instance, large companies often maintain extensive databases of historical trajectories from non-expert users.\" Why? is there a reference? The authors can explain the motivation in a more prefessional way, in terms of the language and example.\n\n3. In Line 138, \"causal transformer\" may not be familiar for the readers. Explain it with more details and provide the reference.\n\n4. In section 3, the authors use $p_\\tau$ to denote both pretraining and test distributions. If DIT is deployed in an unseen new test task, in general, the test distribution should differ from the pretraining distribution?\n\n5. The sentence \"$D_{off}$ contains trajectories gathered from a random policy in $\\tau$\" is not correct right? The offline deployment in DPT paper considers two offline datasets. The offline dataset can be either collected by random policy or the expert policy. See Figures 4 and 5 in DPT paper.\n\n6. The weight in DIT is manually selected. Is there a way to automatically select/learn the weight? Or is there a hint that what weight should one choose given a problem?\n\n7. In DPT's online deployment, how does it work initially when $D_{on}$ is empty? I mean, it appends a trajectory every episode. In the whole first episode (not just the first step), the context is empty. Does TM still predict actions?\n\n8. The authors claim that \"DIT learns to infer near-optimal actions from suboptimal trajectories\" and \"DIT is comparable to DPT in both online and offline testings, despite being pretrained without optimal action labels.\" These are very strong claims. My question is: How \"suboptimal\" is your dataset? Should the suboptimal trajectories contain the near-optimal actions? If so, how much it needs? Please explain it more explicitly and with more contents in the Introduction, **not just one or two sentences in the experimental section**. Otherwise readers might think that DIT learning from a very very terrible offline dataset can achieve comparable performance with DPT, which is irrealistic. \n\n9. I am confused about the pretraining dataset for Meta-World and Half-Cheetah. If they are collected from SAC, how is that \"suboptimal\"?\n\n10. I appreciate that the authors provide the codes for validation. However, there is no instructions at all about how to install the dependencies and how to run and evaluate the algorithms.\n\n11. Why AD is not considered in Miniworld?", "other_comments_or_suggestions": "1. The authors are encouraged to not write equations in the text, which negatively impacts readability, e.g., the texts after Eq.(2).\n\n2. Immediately introduce $\\eta$ when first mentioning it in Eq.(2).", "questions_for_authors": "1. Why in-context reinforcement learning works? I mean, after the Transformers are pretrained, they are kept frozen when applying to the **new unseen tasks**. Why it can learns in-context? Specifically, why the performance can improve when there is more and more contexts while the Transformers parameters are frozen?\n\n\n2. This work is built upon DPT. But why DPT can learn in-context? As discussed in the original AD paper, the methods that learn an good policy like DT cannot improve in-context with frozen Transformers parameters, i.e., cannot reinforcement learn in context. In this sense, isn't DPT also learning optimal policies? Why DPT can show in-context reinforcement learning abilities?\n\n\n3. Comparing with AD/DPT, the DIT proposed in this work will need two extra transformer-based value functions. Does it cost a lot? I think the authors should report how much cost does it need compared with the baseline in terms of e.g., running time? More interestingly, does DIT still perform well compared to the baselines when **they have the same time budget**? This comparison is important.\n\n4. Why the authors want to select a expotential weight? What is the advantage?\n\n5. The context is genenrally just {s, a, r, s'}. But in Q and V estimators, the context is {s, a, G}. Why do we want to use this? Is it fair that the label for Q and V transformers should be the return G. But why the context needs to contain G?\n\n6. Which tasks do the authors consider in Meta-World?\n\n7. It is **extremely surprising** that DIT performs better than DPT in Figure 3(a) and Figure 4(a). Figure 5 looks more normal where the DPT is the oracle baseline. In addition, why DPT in Figures 3 and 4 does not have the comparable performance as that in the original DPT paper?", "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": "OfRNu2Inbi", "reviewer": "Reviewer_GczK", "summary": "The new method Decision Importance Transformer (DIT) was proposed in the paper. This method is an enhancement of existing Decision Pretrained Transformer (DPT). While DPT requires expert target actions for training, DIT can be trained on trajectories sampled from suboptimal behavioral policies.  Since there is no need in expert policies, the proposed approach is easier to use and more versatile than DPT and other ICRL methods. To learn near-optimal policies through sub-optimal historical data, exponential reweighting with advantage function technique was utilised. Authors propose to add weight into common next-action prediction objective, to force model prioritise actions with higher advantage values.\n\nDIT was tested on bandit and MDP problems. Dark Room and Miniworld was used as environments with discrete action spaces, Half-Cheetah and Meta-World was used as environments with continuous action spaces.  In all settings DIT showed performance competitive to DPT despite being pretrained without the optimal action labels.\n\n## Update after rebuttal\nI appreciate the authors’ rebuttal, which clarified my questions and addressed points that were previously unclear. The primary weakness of the paper was the insufficient explanation of the experiments, which left many aspects unclear to readers. This issue has now been resolved in the rebuttal. The idea of using the Advantage function to train DPT on suboptimal data is both novel and valuable for the ICRL area. In light of these improvements, I have raised my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments. Please see our responses below to address your concerns. \n\n### References\n> We appreciate these shared references. We will include a discussion of AD-\\epsilon, AWR, and AMAGO-2 into our manuscript. In particular, we highlight that AD-\\epsilon still requires sampling actions from the good policies to perform noise distillation. Moreover, to create noise distilled dataset, they need to actively sample trajectories following a noise schedule. **In comparison, our work assumes a setting of only historical data, which is more realistic and easier to satisfy**. \n\n### Theoretical Results\n> The authors deeply appreciate the reviewer’s detection of this typo, and we will correct it in our updated manuscript. Yes, the expectation in Equation (3) should be with respect to the meta-policy $\\pi(a|s\\tau)$ and the expectation in Equation (4) is with respect to the behavioral policy $\\pi^b_{\\tau}(a|s)$: we use data collected by the behavioral $\\pi^b_{\\tau}(a|s)$ to learn a meta-policy  $\\pi(a|s;\\tau)$. \n\n### Experiment Setup\n>**Comprehensive Benchmarks.** Our goal is to provide a more comprehensive empirical evaluation than prior works, such as DPT (which focuses on bandits and navigation tasks like Dark Room and Mini-world) and prompt-DT (which focuses solely on continuous control). **To that end, we include representative benchmarks from all three domains: bandits, navigation, and continuous control**.\nWhile we agree that Key-to-Door is a valuable benchmark, it shares significant structural similarities with tasks like Dark Room and MiniWorld. As such, we believe it would offer limited additional insight beyond the settings already covered. Due to the time constraints of the rebuttal phase, we are unable to include new results at this stage. However, **we are more than happy to run additional experiments on Key-to-Door and include them in the final manuscript to further strengthen the empirical evaluation**.\n\n> **Meta-World Benchmark.** We follow the same environment setting of Prompt-DT[1] to use 20 tasks from the Meta-World reach-v2 suit (specifically ML1-pick-place-v2), while we choose the task horizon to be 100. We choose the converged SAC policy as the optimal policy. The cumulative reward of a random policy is around 1.5 within a horizon of 100 and around 20 for the optimal policy. We will update the manuscript to make this clear.   \n\n### Extra Experiments\n> To provide additional insights, we conduct extra experiments to compare DIT with several policy reinforce [2] baselines where the reweighting is based on cumulative reward rather than the advantage function. Specifically, we consider reweighting with cumulative rewards and exponentiated cumulative rewards for offline testing from expert trajectory on Meta-world and summarize the key results in the following table. \n\n| Method                        | Cumulative Rewards | Exponentiated Cumulative Rewards | DIT |\n|------------------------------|--------------------|----------------------------------|-----|\n| Return (higher means better) | 5.0                | 4.7                              | 8.2 |\n\n> **The significantly improved performance of DIT over these two baselines further proves the effectiveness of our proposed method**.\n\n### Performance of DPT.\n\n> The difference in DPT's performance between our work and the original paper stems from the nature of the pretraining datasets. In the original DPT paper, pretraining trajectories were collected using **uniformly random policies**, which ensured broad coverage of the state-action space. In contrast, as noted in our manuscript, we use suboptimal policies (achieving only ~30% of optimal performance) to collect pretraining data. This choice naturally reduces coverage and impacts DPT’s performance. \n\n> However, we believe **it better reflects realistic scenarios**—especially in practical applications where historical data are typically collected by suboptimal or heuristic policies, not uniformly random ones.\n\n> To ensure a fair comparison, both DPT and our method (DIT) are trained on the same set of suboptimal trajectories. However, **we provide DPT with additional privileged information**: a set of randomly sampled out-of-trajectory query states with corresponding optimal action labels sampled from optimal policies. This positions DPT as a strong oracle-style upper bound relative to DIT, further underscoring the strength of DIT’s performance under more constrained assumptions.\n\n[1] Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., & Gan, C. (2022, June). Prompting decision transformer for few-shot policy generalization. In international conference on machine learning (pp. 24631-24645). PMLR.\n\n[2] Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8, 229-256.", "claims_and_evidence": "The main claims of the paper is following: \n\n1. DIT is able to learn in-context and adapt to unseen tasks while it is pertained on sub-optimal data *without the optimal action labels.*\n2. DIT models demonstrate competitive performance to that of DPT\n\nThe authors evaluated DIT on bandits and several MDP problems (Dark Room, Miniworld, Half-Cheetah and Meta-World). Bandits, Dark Room and Miniworld is a classic set of ICRL benchmarks and it is shown in the paper that DIT emerges in-context learning on this environments, but I have concerns about presented results. The scores of DPT on Dark Room and Miniworld  is much lower than it was reported in the original paper. Therefore, the statement that DIT has performance comparable to DPT is inaccurate. \n\nMoreover, in my opinion, it would be great to evaluate DIT on Key-to-Door and Watermaze environments as these problems are well-known and common benchmarks for ICRL. For example these benchmarks are used in AD (https://arxiv.org/pdf/2210.14215) and AD-epsilon (https://arxiv.org/pdf/2312.12275) papers. Without these established benchmarks, the set of evaluation tasks seem less comprehensive and insufficient.\n\nIt is good, that authors evaluated the method on continuous control tasks (Half-Cheetah and Meta-World), this makes the claims made in the submission stronger. But I spotted several inaccuracies. It is said that “Meta-World has 20 tasks in total, to evaluate our approach’s ability to generate to new RL tasks, we use 15 tasks to train and 5 to test.”, but Meta-World benchmark has 50 task in total. Why only 20 tasks were used and which task were used for training and evaluation, there are no information about it in the submission. Also I have concerns about presented scores on Meta-World tasks. The values of returns are in [0, 10] segment, but experts on these tasks can achieve much more higher results ( hundreds if the length of episode is 100 as it was in JAT or thousands if the length of episode is 250 as it was in GaTo (https://arxiv.org/abs/2205.06175)), so the value of returns in the segment of [0, 10] seem like totally random agents and does not support the main claims.", "methods_and_evaluation_criteria": "The proposed benchmarks is fully suitable for verifying the article's claims, but it could be enhanced by additional ones commonly-used in ICRL area. \n\nPresented results is questionable and the setup of experiments are not fully described. What Meta-World tasks were chosen for training and test sets, has hyperparameters search been performed for considered methods, why demonstrated scores for DPT don’t match with reported values in original work, what is the length of context in pretraining datasets for each benchmark. In my opinion this information should be provided for better understanding the experiment setup and results of the work.", "theoretical_claims": "There are a few theoretical claims in the submission. I have carefully reviewed them and would like to ask some questions as I found contradiction. In Proposition 4.1, the optimization problem  $J(\\pi)$ is considered. This objective is the expected value of the difference between the advantage function and the KL divergence.\n\nIn the mathematical expectation operator at the start (Expression 10), actions are sampled from the behavior policy for a certain task,  $a \\backsim \\pi_{\\tau}^b(a|s)$. However, in the “Proof” section, the mathematical expectation operator samples actions from the meta-policy,  $a \\backsim \\pi(a|s; \\tau)$.\n\nFrom my understanding, the second variant is correct because if the expectation is taken over the behavior policy  $\\pi_{\\tau}^b(a|s)$, then maximizing  $J(\\pi)$ becomes trivial. This happens since only the  $D_{KL}(\\pi(.|s;\\tau)||\\pi_{\\tau}^b(.|s))$ term would depend on $\\pi(a|s; \\tau)$.", "experimental_designs_or_analyses": "In my opinion experimental design and analysis is valid for this work.  All my concerns about experimental setup were considered in “Claims And Evidence” and “Methods And Evaluation Criteria” paragraphs.", "relation_to_broader_scientific_literature": "Decision Importance Transformer (DIT) is a novel approach in the field of In-Context Reinforcement Learning. Existing approaches, such as DPT(https://arxiv.org/pdf/2306.14892) and AD(https://arxiv.org/pdf/2210.14215), require expert data in their training datasets. However, collecting such data can be difficult, whereas suboptimal trajectories are only available for training agents. Many methods in offline RL, including CRR(https://arxiv.org/abs/2006.15134), IQL(https://arxiv.org/pdf/2110.06169), and CQL(https://arxiv.org/abs/2006.04779), utilize suboptimal data to learn better policies, yet no comparable techniques exist for ICRL. Advantage weighting is a well-known strategy in offline RL, and the authors propose a way to generalize it for ICRL.  Some works have begun exploring this—for example, AMAGO-2 (https://arxiv.org/abs/2411.11188) — but it remains an understudied area. In my opinion, this is an impactful contribution that will reduce the data requirements needed to train ICRL agents.", "essential_references_not_discussed": "Talking about “*stringent requirements on the pretraining datasets”* authors claim that AD need full learning trajectories of RL agents, but they don’t mention, that it is possible to collect dataset via noise distillation. This method is called $AD^{\\epsilon}$ (https://arxiv.org/pdf/2312.12275). This work should be mentioned.\n\nAlso advantage weighted regression (https://arxiv.org/abs/1910.00177) is mentioned in this work, but there are no reference.\n\nAuthors are talking about using AWR in In-Context RL setting, but don’t mention AMAGO-2 (https://arxiv.org/pdf/2411.11188), where AWR was used. In my opinion, this work should be mentioned too.", "other_strengths_and_weaknesses": "Strengths:\n\n- The paper is well written and it's easy to follow.\n- Detailed and clear description of the proposed method;\n- Applying method from Offline RL into In-Context RL area is good and novel idea;\n- Theoretical justification of the proposed method\n\nWeaknesses:\n\n- The experimental description is incomplete, leaving questions after reading the paragraph;\n- The claimed DPT scores differ from those in the original article", "other_comments_or_suggestions": "I do not have any other comments or suggestions.", "questions_for_authors": "1. Why do the DPT returns on Dark Room and Miniworld differ between the submission and the original paper?\n2. What Meta-World tasks were chosen for training and test sets?\n3. In the mathematical expectation operator at the start (Expression 10), actions are sampled from the behavior policy for a certain task,  $a \\backsim \\pi_{\\tau}^b(a|s)$. However, in the “Proof” section, the mathematical expectation operator samples actions from the meta-policy,  $a \\backsim \\pi(a|s; \\tau)$. Why?\n4. Did you check what DIT do on Meta-World tasks (render video) or what the expert return is on these tasks? According to my experience the good returns on Meta-World tasks (when agent actually solve the problem or almost solve the problem) are much higher than 10.\n5. What is the length of the context in the pertained datasets?\n\nI may change my evaluation and raise the score if the answers remove my concerns.", "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": "iOzlmg5c6h", "reviewer": "Reviewer_g8M9", "summary": "The paper introduces a method for multi-task RL meta-learning with suboptimal behavior policies. The goal is to train a common transformer model to imitate and improve upon the observed behaviors to maximize online rewards in new tasks. Towards this end, the paper introduces two methods: 1. a weighted policy imitation algorithm where the weights are based on exponentiated advantage function values and 2. a model architecture to transfer key dataset characteristics to infer the type of the new tasks. Experimental results on bandits and standard RL problems are included.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and constructive feedback. We hope our responses below address your concerns.\n\n### Experiments\n\n> **Bandit (Figure 2).** Our bandit experiments are designed to be comprehensive and aligned with established evaluation standards. As bandit problems are well-understood, **we follow the evaluation protocol used in DPT [1], comparing DIT against strong baselines such as the known optimal algorithms Thompson Sampling and UCB/LCB, as well as DPT itself** — which serves as an oracle for DIT due to its access to additional optimal bandit information during pretraining. These comparisons provide a rigorous evaluation of DIT in this setting.\n\n> **MDPs (Figures 3(a) & 4(a)).** The goal of these experiments is to evaluate adaptation speed — how quickly each method improves reward on new tasks. **While DIT has not fully converged to its optimal performance in Figure 3, the performances of all baselines have plateaued, and DIT already surpasses them in terms of performance**. We believe this clearly demonstrates the superior adaptability of DIT. \n\n> To provide additional empirical insights, we follow your advice to continue evaluation until the performance of DIT fully converges. As DIT’s performance in Figure 4(a) (Miniworld) has already converged, we focus on Darkroom and we run 10 additional episodes to see the final converged reward, it turns out that **DIT will continue advance the score a bit then keep stable (around 50), compared to 40 at episode 40**, expanding its performance gain over baseline methods. \n\n| Episode  | 5   | 10   | 15   | 20   | 25   | 30   | 35   | 40   | 41  | 42  | 43  | 44  | 45  | 46  | 47  | 48  | 49  | 50  |\n|----------|-----|------|------|------|------|------|------|------|-----|-----|-----|-----|-----|-----|-----|-----|-----|-----|\n| Reward   |3.87 |11.5 |16.9 |19.3 |23.8 |33.0 |36.9 |44.3 |45.5 |45.8 |46.2 |46.4 |46.8 |47.1 |47.7 |47.4 |47.2 |47.4 |\n\n### Method\n\n> **Normalizing Constant and Policy Reinforce Baselines.** We prove in Appendix D.3 that, when using the exponentiated advantage function for reweighting, the normalizing constant (partition function) can be safely ignored and set to 1. \n\n> To further illustrate the effectiveness of DIT, we follow the reviewer’s suggestion to conduct an extra ablation study to compare DIT to several policy-reinforced baselines with different reweighting schemes. We evaluated three weight alternatives:\n> 1. **Cumulative reward**\n> 2. **Exponentiated cumulative reward**\n> 3. **Batch-normalized DIT**, where advantage values are first estimated and then normalized across all trajectories in the pretraining dataset.\n\n> We consider offline testing from expert trajectory on Meta-world and summarize the key results in the following table. \n\n| Method                                      | Cumulative Rewards | Exponentiated Cumulative Rewards | DIT (Batch-Norm Advantage) | DIT |\n|---------------------------------------------|--------------------|----------------------------------|-----------------------------|-----|\n| Return (higher means better)                                  | 5.0                | 4.7                              | 4.1                         | 8.2 |\n\n\n> We observe that **DIT significantly outperforms all baselines**. In particular, batch-normalized DIT performs worse than standard DIT. This **aligns with our theoretical analysis**, which highlights the importance of using trajectory-specific weights. Normalizing across the entire dataset violates this principle, which likely degrades performance.\n\n> We will include these results and a discussion of the policy reinforcement methods into our experiment section to further improve the quality of our work. We thank the reviewer for this constructive suggestion.\n\n### Training Objective (on-policy value function)\n> Thank you for this insightful observation. We agree that the value function equals the expected Q-function only under on-policy rollouts. \n\n> Indeed, in our case, DIT performs pretraining using trajectories generated by behavioral policies, and the learned value functions are used to weight the actions taken by those same behavioral policies. **Thus, all training rollouts are on-policy with respect to the value functions being used, and there are no off-policy discrepancies in the objective**. \n\n> We believe this is one **key strength** of DIT and will clarify this point in the revised version. Additionally, we appreciate the note about the time indices in the Bellman equations and will correct them accordingly.\n\n[1] Lee, Jonathan, et al. \"Supervised pretraining can learn in-context reinforcement learning.\" Advances in Neural Information Processing Systems 36 (2023): 43057-43083.", "claims_and_evidence": "Yes. Claim 1 is supported by an equilibrium proof and Claim 2 is supported by a detailed model design. The authors provided empirical evidence to supplement the discussions.", "methods_and_evaluation_criteria": "Yes. The problem is typical of RL meta-learning and the datasets are standard for bandit / RL problems.", "theoretical_claims": "Yes. The main theoretical claim is in Proposition 4.1. I briefly thought about it and the conclusions seem reasonable based on standard analysis on exponential family models.", "experimental_designs_or_analyses": "Partially. The author provided detailed evaluation setups, but they missed to include a few baselines in bandit experiment. The authors should also consider policy reinforce algorithm in the RL experiments.", "relation_to_broader_scientific_literature": "The authors were exhaustive in the related work section. The authors surveyed decision transformers (PDT), algorithmic distillation, and general behavior cloning. The authors appear less focused on meta-learning for RL, though the literatures there could be less a bit old.", "essential_references_not_discussed": "Not to the best of my knowledge.", "other_strengths_and_weaknesses": "Strengths:\n1. I like Proposition 4.1, which forms the foundation of the proposed policy improvements. The conclusions seem intuitive.\n2. The follow-up algorithmic description is detailed, increasing my confidence in the empirical results.\n3. The experiments validates the author's claims.\n\nWeaknesses:\n1. In the discussions of Proposition 4.1, the authors discarded the partition function. This could introduce additional variance, leading to empirical instability. The authors should consider other normalization techniques such as batch-normalization. The authors should also include closely-related policy reinforce baselines where the exponentiated advantage function is replaced with plain reward function.\n2. A few results are missing from the bandit experiment, indicating a rush to complete the paper?\n3. In the training details, for L_req objective, the authors expected the average of the q function to be the value function. This is only true for on-policy roll-outs. For off-policy roll-outs, the authors should also include policy-likelihood ratios in the L_reg objective. Related, the time-indices for the Bellman equations are in backwards.", "other_comments_or_suggestions": "I like the paper a lot, but I cannot vouch for acceptance when there are still doubts in the experiments. To gain confidence, I would appreciate:\n1. Completing Figure 2\n2. Running Figure 3 & 4 until the proposed method meets optimal performance\n3. Introducing policy reinforce baselines to show that simpler alternatives would not be sufficient to model reward values", "questions_for_authors": "Regarding the three weaknesses:\n1. Can the authors include and discuss policy reinforce methods in the experiments?\n2. Can the authors include batch-normalization in the proposed methods - assuming that the authors left out normalization due to computational complexity?\n3. Please complete Figure 2. Also for Figure 3 & 4, the methods are still far from converging to the optimal. How would you close this gap?", "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:

Meta-review: This paper proposes a method to learn effectively from suboptimal historical data without requiring optimal action labels for in-context reinforcement learning (ICRL). Reviewers generally appreciated the idea of using readily available suboptimal data for ICRL, which is more practical than methods requiring optimal labels (like DPT) or full learning histories (like AD). They also appreciated theoretical grounding and strong empirical results showing DIT performing competitively, and sometimes surprisingly better, than baselines like DPT trained under similar data constraints. There was a concern regarding why DPT performed lower than in its original paper and why DIT sometimes surpassed it, which was partially explained by the authors during the rebuttal period. 

Overall, this paper presents a valuable and practical contribution to the ICRL field by enabling learning from more realistic, suboptimal datasets. During the post-rebuttal discussion period, the reviewers strongly suggested the authors to explain why DPT performs poorer than expected and make a more precise claim about DIT's performance relative to the "particular" implementation of DPT in the camera-ready version. Assuming that the authors will make this change, I recommend to accept this paper.
Decision: Accept (poster)