Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Authors: Long Ma
Fangwei Zhong
Yizhou Wang
Venue: ICML 2025 poster (2025)
Keywords: In-context Reinforcement Learning
Distribution Shift
Meta-Reinforcement Learning
Task Inference
Offline Reinforcement Learning
Primary area: reinforcement_learning->deep_rl
TL;DR: A model-based task inference method that is robust to changes in context behavior.

Abstract: The ability to adapt to new environments with noisy dynamics and unseen objectives is crucial for AI agents. In-context reinforcement learning (ICRL) has emerged as a paradigm to build adaptive policies, employing a **context** trajectory of the test-time interactions to infer the true task and the corresponding optimal policy efficiently without gradient updates. However, ICRL policies heavily rely on context trajectories, making them vulnerable to distribution shifts from training to testing and degrading performance, particularly in offline settings where the training data is static. In this paper, we highlight that most existing offline ICRL methods are trained for approximate Bayesian inference based on the training distribution, rendering them vulnerable to distribution shifts at test time and resulting in poor generalization. To address this, we introduce Behavior-agnostic Task Inference (BATI) for ICRL, a model-based maximum-likelihood solution to infer the task representation robustly. In contrast to previous methods that rely on a learned encoder as the approximate posterior, BATI focuses purely on dynamics, thus insulating itself against the behavior of the context collection policy. Experiments on MuJoCo environments demonstrate that BATI effectively interprets out-of-distribution contexts and outperforms other methods, even in the presence of significant environmental noise.

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Long Ma 12
Fangwei Zhong 32 Yizhou Wang 456
Abstract
People don't wear space suits on Earth. be the Moon!
The jump is not very high. could be training on Earth!
The ability to adapt to new environments with noisy dynamics and unseen objectives is crucial for Al agents. In-context reinforcement learn- ing (ICRL) has emerged as paradigm to build adaptive policies, employing a context trajectory of the test-time interactions to infer the true task and the corresponding optimal policy efficiently without gradient updates. However, ICRL poli- cies heavily rely context trajectories, making them vulnerable to distribution shifts from train- ing to testing and degrading performance, par- ticularly offline settings where the training data is static. this paper, we highlight that most existing offline ICRL methods are trained for approximate Bayesian inference based on the training distribution, rendering them vulnerable to distribution shifts at test time and resulting in poor generalization. To address this, we intro- duce -ehavicr-agnostic Task Inference (BATI) maximum-likelihood
Figure I. Illustration of the core idea behind BATI. The boy (left) infers the astronaut's location using behavioral cues (e.g., wearing space suits), while the girl (right) relies on physical dynamics (e.g., jump height), Although the former behavioral correlation casier to spot, the latter is more robust to distribution shifts and leads to the correct answer in this example. BATI mirrors the latter strategy. prioritizing dynamics over behavior.
1. Introduction
The ability of AT agents to adapt to new environments with noisy dynamics and unknown objectives is becoming in- creasingly important as we push the boundaries of artifi- cial intelligence applications. Meta-reinforcement learn- ing (Finn et al., 2017; Duan et al., 2016; Beck al., 2023) has emerged as a promising paradigm for developing adap- tive policies. This approach leverages the concept of learn- ing to learn, enabling agents to generalize from previous experiences and effectively tackle novel tasks quickly. Re- cently, the marriage of in-context learning (Brown et al., 2020; Min et al., 2022: Hendel et al., 2023) and meta- RRL has attracted attention from the community, referred as in-context reinforcement learning (ICRL) (Laskin et al., 2023; Grigsby et al., 2024; Ma et al., 2024). It leverages a context trajectory of interactions during testing to infer the true task and determine the corresponding optimal policy without any gradient update, showing significant potential for effective generalization across diverse and unknown en- vironments. However, such context-conditioned policies can be expensive and time-consuming to train with online interactions due to the sample inefficiency of RL algorithms and the costs associated with online data collection (Yu, 2018; Gu et al., 2024). People thus turn to offline ICRL to extract in-context policies from offline data without online interactions (Laskin al., 2023; Li et al., 2024).
for ICRL, a model-based maximum likelihood so- lution infer the task representation robustly. In contrast to previous methods that rely on a learned encoder the approximate posterior, BATI fo- cuses purely on dynamics, thus insulating itself against the behavior of the context collection policy. Experiments on MuJoCo environments demonstrate that BATI effectively interprets out- of-distribution contexts and outperforms other methods, even in the presence of significant envi- ronmental noise.
'Center for Data Science. Academy for Advanced Interdisci- plinary Studies, Peking University, Beijing, China Beijing Insti- tute for General Artificial Intelligence, Beijing, China 'School of Artificial Intelligence, Beijing Normal University, Beijing, China "School of Computer Science, Peking University. Beijing, China 'Institute for Artificial Intelligence, Peking University, Bei- jing. China "tate Key Laboratory of General Artificial Intelli- gence, Beijing, China. Correspondence to: Fangwei Zhong <tang- weizhong @ bnu.edu.en>.
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Despite its promise, ICRL faces substantial challenges, as the performance of ICRL agents be sensitive the be- havior shown in the context. During testing, any shifts in context distribution can lead to a marked decrease in perfor- mance (Gao et al., 2024). This issue is exacerbated when learning from offline data, where the distribution of train- ing data is likely to differ from the conditions encountered during deployment. As a result, the robustness and gen- eralization capabilities of offline ICRL methods are often compromised by context shifts.
2. Related Works
In-Context Rentfoeeinforemment Learning. In-context reinforcement learning (ICRL) aims to train agents that can generalize to solve new tasks using test-time interac- tions and reward signals, contexts (Duan et al., 2016; Laskin et al., 2023; Grigsby et al., 2024; Gao et al., 2024; Li et al., 2024). ICRL falls into the broader category of meta-reinforcement learning (Beck et al., 2023), which en- compasses both gradient- gradient-based methods (Finn et al., 2017; Song et al., 2020; Yoon et al., 2018) and context- based meth- ods (where ICRL belongs) for learning new skills at test time. Critically, to improve the eis-timme efficiency et al., 2024), ICRL policies need to acquire new capabilities with- out any gradient updates, resembling the in-context learning phenomenon of large language models (Brown et al., 2020; Min et al., 2022; Hendel et al., 2023). In this paper, we focus on the ICRL problem to obtain agents that can efficiently adapt to new tasks without parameter updates. While ICRL has many appealing properties, an in-context policy can be very expensive time-consuming to train due to the sample inefficiency of online RL algorithms and various costs of col- lecting online interactions in real-world scenarios (Grigsby et al., 2024; Yu, 2018; Gu et al., 2024). Offline ICRL (Li et al., 2024; Gao et al., 2024) has emerged to harness the advantages of both offline RL (Kostrikov et al., 2022; Wang et al., 2024; Zhong et al., 2025) and ICRL.. Our work seeks to address the context shift problem in this offline setting.
We identify a critical limitation in existing offline ICRL methods (Li et al., 2021b; Yuan & Lu, 2022; Gao et al., 2024; Li et al., 2024): they are typically trained for approxi- mate Bayesian inference based on the training distribution. This approach makes them vulnerable distribution shifts at test time, undermining their ability to generalize to new, unseen contexts. Addressing this limitation is crucial for ad- vancing the practical applicability of ICRL in real- real-world sce- narios. Figure| I shows an example of the effect of context shift when inferring the location of the jumping astronaut. In this case, the commonsense inference using behavioral characteristics leads to a wrong answer under distribution shifts, while the correct answer is given by an analysis of environmental dynamics.
In this paper, we propose new Behavior-agnostic Task Inference' (BATI) framework to enhance the robustness of ICRL. BATI is a model- proach that infers task representations without being in- fluenced by the behavior of the context collection policy. Unlike previous methods that depend on a leareed encoder to approximate the posterior, BATI focuses exclusively the dynamics of the environment. This focus allows BATI to remain insulated from the variability introduced by dif- ferent behavior policies, thereby improving its resilience to distribution shifts. We conduct extensive experiments sev- eral MuJoCo environments to evaluate the efffcctiveness of BATI. Our results demonstrate that BATI not only interprets out-of-distribution contexts more effectively than existing methods but also outperforms them in environments with significant noise. These findings highlight the potential of BATI to enhance the adaptability and reliability of Al agents operating in complex and dynamic settings.
moxiiuu--lkeelhoon ap-
Task Inference. Task inference methods (Humplik et al., 2019; Liu et al., 2021; Rakelly et al., 2019; Zintgraf et al., 2020; 2021) cast ICRL as a two-stage problem, where a latent representation of the true task is first inferred from the context and an in-context policy is conditioned on this latent to execute the corresponding optimal behavior. Previous works use supervision (Humplik al., 2019), contrastive- like objectives (Li et al., 2021b), or RL losses (Rakelly et al., 2019) to guide the task inference. However, the context distribution may shift between training and test time, posing a great challenge to the generalization capabilities of task inference methods (Lin et al., 2020; Yuan & Lu, 2022: Li et al., 2024; Gao et al., 2024; Xu et al., 2024), especially in the offline setting. To address this context shift, CSRO (Gao et al., 2024) proposes to minimize the mutual information between the task latent and the context collection policy to promote the true correlation between the latent and the task. Furthermore, UNICORN (Li et al., 2024) provides a unified for understanding task inference methods and proposes a tighter approximation to the true objective. In this paper, we analyze flaws in previous works and propose a maximum- likelihood based robust task inference method.
Our contributions are three-fold: 1) We identify and address key vulnerability in existing ICRL methods related to context distribution shifts; 2) We introduce BATI, robust task inference framework that enhances generalization by focusing on environmental dynamics; and 3) We validate the generalization of BATI through comprehensive experi- mental evaluations, setting a new benchmark for ICRL in noisy and unpredictable environments.
ieformmtioo-thhoretti framework
Project page: https: //tt://///eeeeeeee/eeeeeee www.eeoo........eee.ee......o... e. cww/vicw. bat - icrl
2

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
3. Preliminary
2019 Peng al., 2021; Kamienny et al., 2020). For exam- ple, Humplik et al. (2019) assumes that supervision of task parameters is available during training, while FOCAL (Li et al., 2021b) takes a distance metric learning approach.
3.1. Problem Formulation
We define the ICRL problem on a distribution of Markov decision process (MDP) M = (S, A. P. R. Po: Po: where S is the state space, A is the action space, P S x A + As is the transition function, R S x A - AR is the reward func- tion, po E As is the initial state distribution, and 7 is the dis- count factor. Hereafter we refer to each MDP as a task. We assume that all the tasks in support of the task distribution share the same S, A. Po, and y so that we can identify a task with PMr, RM- Define context X = {(ss, ae, ri, + C- as set of C' transition tuples in a task M: further denote its components X {(s;, 0;)) );),;; X {(ri, k. Such context could be generated by rolling out a (plain) pol- icy S + A in M, so 80 po; 0; ~ (((:),,r ri ~ RM a;), s/ = 5++1 ~ (s; a.) Vi) in which case we write the context random variables as XM X
However, in an offline learning setting, where the ICRL policy needs to be extracted from a fixed dataset without any online interactions, context shift becomes a major con- cemn The context distribution may dramatically change between the offline training and the online testing, causing the context encoder to go out of distribution and produce incorrect task encodings, which leads to unsatisfactory pol- icy performance. Among works that have sought to address the context shift problem (Lin et al., 2020; Li et al., 2024; Gao et al., 2024; Yuan & Lu, 2022; Li et al., 2021b;a), re- cently, UNICORN (Li et al., 2024) proposed an information- theoretic framework for offline ICRL, arguing that the task representation learning should be done by optimizing the following mutual information objective:
Our objective is to learn an in-context policy, or meta- policy ne S X X' + AA parameterized by optimize the following discounted objective:
max T (Z:M) +
(2)
max E ( Naa
Σ It' 120
(1)
where Z := foX) is the context encoding, fo is the context encoder parameterized by ø. Intuitively, Z should capture the part of X that describes the true task M. From this perspective, UNICORN proved that the objectives of sev- eral prior works can be streamlined as approximations or bounds of the mutual information objective. It proposed an alternative objective to achieve tighter and more robust approximation. A classification-lik loss can also be used to directly optimize Eq. 2 using task indices.
where r, is generated by rolling out TT(- I xas,p) in M, X = (S x A x E. x S)* is the context space, As the task M is sampled from a distribution and not directly revealed to Tp. To may only learn about the properties of the current task it needs to solve via the sampled context M,
Finally, to characterize the context shift problem, note that the above objective is an expectation over the joint distri- bution of task M and context collection policy μ.. Denote the training distribution as Porain (Mr ()) and the testing dis- tribution as (M, For the same task M, the training context M. induced by Prrain ( M) may have very different distribution from the testing context X XM Meee Prea Prea( (Jx) M), forming context shift.
However, we observe that this framework fails to formally characterize the change of / (Z; M) under context shift. In partieular, its modeling of the ICRL problem does not take account of the distribution of the context-collection policy μ... Expanding Eq. 2 under our formulation in Sec. 3.1 yields
I(Z;M)
H(M) H(M Z)
3.2. Background & Analysis
= HMM + EM,z p(M 2)
= (M) + +EM,, Exx., p(M (((())))))
In recent years, as interest in ICRL and meta-RL increased, many methods have been developed for training effective in- context policies. Among the existing works, task inference approaches consider ICRL to be the problem of inferring a latent task representation from given test-time interac- tion context. Once the test-time task is known, the optimal policy can be determined and exccuted to achieve a good performance efficiently without time-consuming and unsta- ble test-time gradient updates. Early works thus focus on effectively extracting task information from the context us- ing a context encoder, with the main differences being the learning objectives used (Humplik et al., 2019; Ren et al., 2022; Sohn et al., 2020; Zhang& & Kan, 2022; Lee et al.,
This derivation reveals that the mutual information objective depends on the joint distribution p(M, pa), which may shift between training and testing time. As a result, even if we were able to optimize the true objective on the training data, the testing MI is still not guaranteed to be large. Furthermore, the posterior estimate p(M (M,)) is also dependent on p( p(M, 11) and may give incorrect estimates about M in an icy performance as shown in our experiments.
scenario, hurting the pol-
To address the discrepancy, CSRO (Gao et al., 2024) ob- served that context shifts are introduced when the testing
3

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
(Fig. 2), where the inference of M is disrupted by the shift- ing distribution of j4, which subsequently contaminates the distributions of X and X
x
To circumvent this failure mode, we make further obser- vation from the graphical model that we can remove the influence of j by blocking X Once conditioned on X and the variable to be inferred M. all paths going from μ to X' are blocked while μ and X become independent, so we can correctly infer M using X'. Taking advantage of this observation, we propose behavior-agnostie task inference (BATI), a maximum- sikiimum-likelihoodd likelihood- solution to replace the Bayesian posterior inference of p(M X):
M
X
Figure 2. The graphical model of our proposed formulation. The dotted red arrow indicates the joint distribution of the task M and the context collection policy 14, which may change hetween training and testing, causing the context X", X' to shift. However, after conditioning on X" (blue node), the effect of i is blocked (solid red arrow), we can safely infer M (green node) via X'.
arg max log X M) M
(5)
This term, corresponding to the environment dynamics of M and irrelevant of μu, can now be safely estimated from offline data and transferred to online inference without worrying about the shifted distribution of j..
contexts are collected by different policy from the train- ing one, echoing our formulation. Subsequently, CSRO proposed to maximize the true correlation (Z; M) while minimizing the spurious correlation (Z: x to decor- relate Z with the behavior of the context collection policy contained in the state-action pair X However, as UNI- CORN noted, this mixed objective poses a trade-off' because the two components are sometimes contradictory. For ex- ample, consider a bandit-like like environment where every task corresponds to a specific action and training contexts always take that action. In this case, the mutual information be- tween the task M and the behavior X is already high, and it's not possible to find suitable Z with high I (Z; M) yet low I [Z; X*). More formally, we have
We can give another interpretation to this solution by view- ing it as robust version of the full likelihood p(X M):
p(X M)
S: plu M) M,
=
Sn plu M)p(X M, ()(()(X X M,
= p(xe 1x,, M) fl M)p(X*) M, u)dd
~ p(X* 1x,
The third equality stems from the properties of MDP, where the reward and the next state depend only on the state-action pair in any given MDP and are independent of the overall policy. In this expansion, only the integral over μ affected by the distribution shift. To estimate it, we would need information about the test-time joint distribution p( (,, (1) and the behaviors of every in every M, a tall order to fulfill. We thus assume that the integral is approximately the same over different tasks and ignore this term. Empirically, we find this approximation to have satisfactory performance with extensive experiments in Sec. 5.
Theorem 3.1. For arbitrary task representation Z,
[(Z M) IZ S H(M - X(M:X*)
(3)
See App. A for the proof. Consequently, when the training context collection policy is highly correlated with the task, the competing objectives of CSRO cannot be achieved si- multaneously, and a good encoder (in terms of the CSRO objective) does not exist.
Given our analysis above, direct Bayesian inference is problematic, why do existing methods still work to some extent? Expanding Eq. 4 in similar manner as above yields
4. Method
4.1. Behavior-agnostic Task Inference
In the analysis of previous works, we note that all of them fit context encoder fo to estimate the posterior of the task variable p(M [ | X) using different objectives and regulariz- ers. We argue that this direct Bayesian inference approach is inherently flawed, since
p(M I XX
x S P(M, p)p(X M, P(M ())(X)) M, H) M, p, xxXx)p
IM, ((M M. p)dd
which also includes the robust term p( X M, X*). Further- more, we observe that the evaluation environments of previ- ous works have deterministic d dynamics, which means that p(X' M, X*) is close to a delta function. Consequently, the true correlation may nevertheless dominate the Bayesian
[ ((M, p)
pMM X)
(4)
which inevitably depends on the joint distribution p(M, p)). This is illustrated in the graphical model of our formulation
4

---

Behavior-agnestic Task Inference for Robust Offline In-context Reinforcement Learning
Behavior-Agnostic Task Inference (BATI)
Baselines
Offline Training
Context ((s au No si
Selected by Task Index
State St
Context X
State
Z Emb.
Policy
Encoder
Policy
Dynamics
Estimator
Lrncom
Action at
Action at
Task Latent Z
Online Testing
Z Emb.
Z'
State St
Decoder
Task Index
arrg min frecon
Policy
Task
Dynamics
Inference
Estimator
(Frecon)
Loss
Gradient
Action at
Flow
Figure 3. Illustrations of Bchavior-Agnosiic Task Inference (BATI, left) and baselines (right). Baselines use the same encoder to encode the context X into task representation Z for both offline training and online testing, in effect performing Bayesian inference over the training distribution. In contrast, BATI uses the task index during training and searches for optimal task latent during online evaluation, avoiding the influence of context shifts.
inference and overwhelm the shifted integral. However, when the dynamics are noisy, as is common in real-world scenarios, the discriminative power of p(x* M, x²) might be weakened, leading the fitted posterior to rely on the spu- rious correlation p(X* M, μ4) as a shorteut (Geirhos al.,
and the value functions receive the same task embedding as the dynamics estimator, and all modules are trained si- multaneously. However, gradients from the policy and value functions are detached from the task embeddings, which are supervised exclusively with Creon See Alg. 1 for the pseudocode of the full training pipeline. Note that both the dynamics estimator and the policy are trained with purely offline data without any online interactions.
2020). We demonstrate this sensitivity noise with our ablations in Sec. 5.4.
4.2. Offline Training Pipeline
4.3. Online Evaluation Procedure
We now describe the instantiation of BATI and our of- fline training pipeline, as shown in Fig. 3. To estimate p(X X².M) without ground-truth task parameteriza- tions, we build a table of task embedding distributions (ZM) Yas (Fig. 3, stacked purple rectangles) containing an entry for each training task. We parameterize the Gaus- sian distributions with c. replacing the learned encoder fo in previous works. Both the task embeddings and the dy- namies estimator (orange rectangle) supervised by the following:
During the online evaluation, given contexts X collected by unknown policy for an unknown task sampled from Pest (M, μ). we perform the proposed maximum maximum-likelibood likelihood optimization in Eq. 5. Specifically, we compute
arg min Leeeon
(8)
with Z" sampled from M J for fixed number of times N. The in-context policy To takes Z" as input and interacts with the test task to evaluate its performance. See Alg. 2 for the online evaluation procedure.
min E M- (M), (M)x Zwzz CX (Ø, v) 9:3)
(6)
where
5. Experiments
(x+ - (xb 2)) + 1x
CX2 (d, v)
exp hpph Z)
We empirically validate BATT's performance several MuJoCo-based enviromments commonly used for offline ICRL. With the experiments in this section, we aim an- swer the following questions: 1) Can BATI achieve robust performance in the presence of significant context shifts? 2) How does BATT's inference procedure compare with a learned encoder? 3) How does the noise level impact the performance of BATI and baselines? 4) Can BATI outper- form baselines in noiseless environments? 5) How does
the dynamics reconstruction loss as the negative log- likelihood of N (X",,,, Z); exp (X*, Z)). Purain (M) is the training task distribution, DM is the offline training dataset for task M, and ge he represent the dynamics esti- mator.
To train the in-context policy ne. we use IQL (Kostrikov al., 2022) as the offline RI. algorithm. During training,
5

---

Bchavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
ACT
C2R0
C3RO
23
UNCDAN
Oookle
202
FOOAL
ROCA
250
Racon
Ceacl
100
352
-42
-452
20000
00020
#3200
100000
62000
82000
20000
00002
100033
(a) AntDir 0.4)
HalfChectah Vel (e 0.5)
(c) HopperParam 0.3)
1411
FOCA.
C510
300
CSRO
FOCAL
520
Racen
Recon
UNICORN
UNICONN
--. Oradi
420
120
-200
200
23000
42000
80302
202000
20020
40000
60302
50000
110000
25200
52000
50002
100023
Thaning State
Thanning Staps
Braring Siip
(d) HaliCheetahDir (c 0.7)
(e) WalkerParam (c = 0.2)
(L) HopperParam (c 0)
Figure 4. Online evaluation episodic return curves of BATI and baselines on contexts from Pes H) during training in our evaluation environments. BATI consistently achieves the best performance in environments and settings, converging faster and more stably.
BATI scale with more context data? 6) Can BATI handle OOD contexts and tasks simultaneously?
as the primary context collection policy associated with each task M and set
0.9,
H PM
Pwran (M, p)) = Puraia (M)
(9)
5.1. Setup
0.1 Perm
1 * PM
(M)
We choose five representative robot locomotion environ- ments based on the MuloCo simulator (Todorov et al., 2012) with varying properties and levels of difficulty. Among the environments, AntDir, HalfCheetah Vel. and HalfChee- tahDir have different reward functions RM, while Hopper- Param and WalkerParam have parameterized dynamics Par- To simulate real-worke scenarios more closely, which are often nonemeieriiiistic and noisy, we perturb the dy- namics of the environments with various noises of scale c. As noted in Sec. 4.1, dynamics noise weakens the true task correlation and adds to the challenge of task inference. The details of the environments are in App. C.I. We also con- duct ablative study in Sec. 5.4 demonstrating the effect of noise on baselines, compared to BATT which remains unscathed in all scenarios.
Prest Pres (M, p1) Prest Il PMr
(10)
where M is task most different from M. This creates an extremely challenging testing distribution, the in-context policy will be misguided by the context to execute behavior that is good for an adversarial task but quite bad for the true task. Unless otherwise stated, all evaluation results in this section are on Preat Prst(MM p). See App. C.2 for details on dataset construction.
We compare BATT with representative baselines from the field of offline ICRL: a) UNICORN (Li et al., 2024), re- cent state-of- ofee--of-thee-- tet-hhe-rtt art method using an cooommation-thooree objective, we use the UNICORN-SS variant reported to have better performance; b) CSRO (Gao et al., 2024), a strong baseline that promotes robustness by minimizing mutual information between the task representation and the context behavior; e) FOCAL (Li al., 2021b), classic method that uses distance metrie learning for self-supervised task inference; d) Recon, a model-based method that uses the same reconstruction loss as BATI but with a learned encoder for inference. This is also called UNICORN-SS-O or GEN- TLE (Zhou et al., 2024) and serves as an ablative baseline to isolate the effects of our task inference procedure. AntDir
In each evaluation environment, we randomly sample 20 tasks as PMrinn(M) and another 20 tasks as Pres( accord- ing to its task parameterization, e-g. target directions in AntDir or physics parameters in HopperParam (see App. C. for details). During the online evaluation, the true task in- dices or parameters are not directly provided to the in- context policy and must be inferred from the context. To construct contexts with large shift, we create a policy JaM
6

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
123
C350
Beoon
JNICOON
UMECORIN
baroed 11dmm
Crokho
Drudi
20
22
tpoed
balted
23-
20
20
-20-
-20
30000
30200
20000
Thitiin Stops
Taniin Supa
Tuaving
(a)
(b) 0.2
(c) 0.4 (Default)
Figure 5. Online evaluation episodic return curves during training for different noise levels on AntDir. With the noise level rising, baselines increasingly rely the behavior of the context collection policy μr to determine the true task, causing their OOD performance to become (b) or even decrease throughout training (c). At the same time, BATI consistently improves during training and performs similarly to or better than the baselines for all noise levels.
and HalfCheetah Vel, we also compare with DPT (Lee et al., 2024), a recent transformer-bes based method for ICRL. We implement BATI and the baselines on the same codebase with IQL (Kostrikov et al., 2022) as the base offline RL algorithm. See App. C.3 for implementation details.
trajectory X to produce Z, which is then used with X to decode X In contrast, BATI uses an embedding during training and maximum -ikeliio-ike inference during testing for Z. shown in Tab. 4 and Fig. 4, Recon underperforms BATI in every environment and have very high vari- ances in certain cases (c.g. HalfCheetahDir), In Fig. 4a, the performance of Recon decreases over time, indicating the capture of spurious correlations. This result reinforces our argument Sec. 3.2 that a learned encoder performs ap- proximate Bayesian posterior inference and cannot handle context shifts well.
5.2. Can BATI achieve robust performance in the presence of significant context shifts?
In this section, we report the online evaluation results on (M, p.) described in the previous section. Each online evaluation rollouts the in-context policy in all 20 testing envi- ronments with their respective contexts (Alg. 2) and returns the average episodic return. Fig. 4 shows the episodic return curves during training, while Tab. 4 contains the numerical results at convergence, computed as the average of the final 5 evaluations of each training run. Standard deviations are reported over 5 training seeds. Across all environments, BATI outperforms the baselines, usually by a large margin. In AntDir and HalfCheetahDir, we observe that the evalua- tion performances of several baselines decrease over time, indicating the gradual learning of spurious correlation. As a comparatively strong baseline, CSRO performs somewhat similarly to BATI in AntDir and HalfCheetah Vel but con- verges slowly and is less stable. It is dramatically outperformed in the more difficult environments. The rea- son for this underperformance and slow learning is that the competing objectives of CSRO cannot be fulfilled simultane- ously when the context collection policy is highly correlated with the true task (see Sec. 3.2).
5.4. How does the noise level impact the performance of BATI and baselines?
We now show the effect of dynamics noise on BATI and baselines in AntDir and HalfCheetahDir, two representa- tive environments with continuous (AntDir) and discrete (HalfCheetahDir) task parameterizations. We rerun the ex- periments with lower or no dynamics noise. As shown in Figures 5a to Se above and Figures 7a to 7c in App. D, increasing dynamics noise progressively destabilizes base- lines, while BATI maintains robust performance. The phenomenon of spurious correlation capture for baselines is especially evident in AntDir. While most baselines can achieve a decent performance in the noiseless setting (Fig. 5a), they fail to learn with a medium level of noise (Fig. 5b) and even grow steadily worse as learning pro- gresses in the default high-noise setting (Fig. 5c). This find- ing supports our analysis in Sec. 4.1 that noisy dynamics weaken the true correlation, forcing the learned Bayesian posterior to rely more on the spurious correlation of the context collection policy shorteut.
5.3. How does BATI's inference procedure compare with a learned encoder?
To demonstrate the effect of BATT's maximum-likelihen inference procedure, we compare BATI with the Recon baseline (also known GENTLE and UNICORN-SS-O) which shares the same reconstruction objective Lreson with BATI. The key difference is that Recon encodes the full
5.5. Can BATI outperform baselines in noiseless environments?
With deeper understanding of the role of noise, we now fur- ther show that baselines may fail even without any noise in

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
20
-100
Braokk 14lam
-00
-250
-55
34T
Rccen
0.2
0.5
-10.
-7.5
-50
0.0 (a)
-25
255
5.0 7.5
-10
10
Sarb Tuating Daca
BATI
(b) Recon
(c) Episodic Returns
Figure 6. t-SNE visualizations of inferred task latents of BATI (a) and Recon (b) HalfCheetahVel. The dots denote the latents of all tasks on Pran. and colors indicate the target velocities. Black dots linked by blue arrows are those inferred from interpolated contexts of a specific task. The corresponding online evaluation episodic returns are shown in (c). As the ratio of outto---istiibutiinn piest contexts rises, the latents inferred by Recon drift off' in the direction of the spurious task (from lower-right upper-left of (b)), At the same time, BATI consistently produces correct and in-support latent representations (lower-right of (a)), resulting in better performance.
Table 1. Online evaluation episodic returns of the final checkpoints of BATI and baselines on HalfCheetahDir under various context lengths, All policies trained with C 50 (*) and 0.9.
time horizon, this setting presents an enormous challenge to ICRL algorithms. Tab. contains the performances of BATI and baselines trained in the setting described above and tested with different context lengths. BATI outperforms all baselines with the shortest context length C 50 and improves continuously when provided with longer contexts. As the context grows, BATI can estimate log- logeliicls with lower variance, leading to improved quality of the in- ferred task latent and better policy performance. However, the baselines do not exhibit such a pattern and fluctuate around the initial performance, even getting slightly worse.
C = 50 (*)
C = 200
C 400
BATI
-4.7 1.2
1.8 1 0.5
4.7 1.5
CSRO
-33. 1.0
-24.7 0.7
25.3 1.1
FOCAL
-25.6 0.2
-2.66
Recon
022
0.3
0.2
UNICORN
2.3
24.0 1.9
2.44
hard scenarios. Fig. 4f shows the performances of BATI and baselines in HopperParam with no dynamics noise. Com- pared with the noisy results in Fig. 4c, baselines still under- perform BATI with only modest gains for the UNICORN baseline and little change for the others. As explained in Sec. 4. baselines fail when the discriminative power of the true correlation p(X* M, is weakened. The dynamics sufficiently complicated and nonlinear that v(X' M, x²) ) is already hard to estimate even under noiseless setting. Baselines thus opt to learn the shortcut of p(X* M, μ1) instead and fail when the context collection policy changes.
5.7. Can BATI handle OOD contexts and tasks simultaneously?
In addition to generalizing to OOD contexts, we demon- strate BATT's OOD task generalization capability through experiments in AntDir. OOD task generalization coupled with context generalization is even more challenging, since the policy is forced to adapt to an unfamiliar task that differs dramatically from those seen during training and requires very different behaviors to solve. We sample target direc- tions from [0, π)) during training while using directions from [r, 2.rr for testing. Results are presented in Tab. 2. While the performance of BATI decreases compared with the in- distribution case, it still outperforms baselines by large margin and is the only method to achieve positive returns. This further demonstrates the generalization capability of BATI to both OOD contexts and tasks.
of HopperParam
5.6. How does BATI scale with more context data?
A crucial desideratum for ICRL is the ability to keep im- proving with more context data, In this section, we demon- strate this property for BATI an extremely challenging setting. We choose HalfCheetahDir, an environment with two discrete task variants (going left right), and conduct an experiment with very high noise (c 0.9) and short train- ing context length (C = 50). Under this setting, the reward for each time step is computed in the true direction with probability 0.55 and in the other direction with 0.45, so the ICRL policy may only obtain an exceedingly weak signal of the true direction with each time step. Coupled with a short
5.8. t-SNE Visualization with Interpolated Contexts
In this section, we show t-SNE visualizations of task latents inferred by BATI and Recon in HalfCheetah Vel to demon- strate the impact of context shifts more elearly. We construct interpolated contexts of various ratios r consisting of N steps from Pies and N (1 - r) steps from Pusin- Visual- izations and episodic returns are shown in Fig. 6. While
8

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Table Final online evaluation episodic returns of BATI and baselines on Preal (4) with OOD tasks AntDir. BATT outperforms baselines and is the only method to achieve positive returns, further demonstrating its generalization potential.
Method
BATI
CSRO
FOCAL
Recon
UNICORN
Episodic Return
12.1 + 0.5 -45.8 + 14.5 -74.7 -51.8 + 37.0 -104.6 + 9.1
both methods produce reasonable latent patterns on Pinins the latents inferred by Recon drift off when the ratio of Ptest rises (Fig. 6b), causing its performance to drop (Fig. 6c). In contrast, the latents inferred by BATI consistently stay in the correct region (Fig. 6a) and yield performant policies.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
6. Conclusion
References
Intelligent agents in the real world must adapt to unseen tasks and environments in noisy conditions based on their experience. As in-context reinforcement learning emerges to train context-adaptive policies, the context shift prob- lem presents a great challenge to the task inference process and generalization capabilities. In this paper, we introduce Behavior-Agnostic Task Inference (BATI), a maximum- likeiihodd-base approach for robust task inference in ICRL. In contrast to the previous works, which predominantly use learned encoder to perform direct Bayesian inference, BATI carefully analyzes the problem and infers task latents using dynamics estimators only, thereby freeing itself from the influence of context shifts. Extensive experiments in various benchmark environments validate the performance of BATI over strong baseline methods.
Beck, J., Vuorio, R., Liu, E. Z., Xiong. Z., Zintgraf, L., Finn, C., and Whiteson, S. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.0002 2023.
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 in Neural Information Processing Systems, 33: 1877-1901, 2020.
Chen, Y., Geng, Y., Zhong, F., Ji, J., Jiang, J., Lu, Z.,, Dong, H., and Yang, Y. Bi-dexhands: Towards human-level bimanual dexterous manipulation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 26(5):8004- 2818, 2023.
BATI still has certain limitations. We conducted experi- ments in several state-based MuJoCo environments com- monly used in the field of offline ICRL. Going forward, we hope scale BATI vision-based embodied scenarios (Wu al., 2022; Wang et al., 2023; Chen et al., 2023; Zhong et al., 2024) enhance its applicability in the real world. The counterfactual estimation of p( ( M,, (4) may fur- ther improve the performance of BATI. We may also scale BATI to multi-agent scenarios (Wang et al., 2022; Pan et al., 2022; Ci et al., 2023; Long et al., 2024) the future, where the tasks involve interactions with or are defined by other agents.
Ci, H., Liu, M., Pan, X., Zhong, F., and Wang, Y. Proactive multi multi-camera collaboration for 3d human pose estimation. arXiv preprint arXiv:2303.03766 2023.
Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I., and Abbeel, P. RI 2: Fast reinforcement learn- ing via slow reinforcement learning. arXiv preprint orXiv:1611.000
2016.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional Conference Machine Learning, pp. 1 126--1133 PMLR, 2017.
Gao, Y., Zhang, R., Guo, J., Wu, F., Yi, Q.. Peng. S., Lan, S., Chen, R., Du, 7., Hu, X., et al. Context shift reduction for offline meta-reinforremenn learning. Advances in Neural Information Processing Systems, 36, 2024.
Acknowledgements
This work was supported by the National Science and Technology Major Project (2022ZD01 14904), NSFC- 6247070125, NSFC-62406010, the State Key Lab of Gen- eral Artificial Intelligence at Peking University, the Funda- mental Research Funds for the Central Universities, Qual- comm University Research Grant.
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Bren- del, W., Bethge, M., and Wichmann, F. A. Shortcut learn ing in deep neural networks. Nature Machine Intelligence, 2(1 1):665-672 2020.
Grigsby, J., Fan, L., and Zhu, Y. AMAGO: Scalable in- context reinforcement learning for adaptive agents. In The

---

Bchavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Twelfth International Conference on Learning Represen- fations, 2024. URL https: /ttl:///pppnnt htwp//pppereee..eeeeeeeeeeee...m net./ Barum?id-M6WooEB
Li, L.., Yang, R., and Luo, D. FOCAL: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. In International Conference on Learning Representations, 2021b. URL https: openrevi New net / forum? ? d=8cpHT
Gu, S., Yang, L., Du, Y., Chen, G., Walter, F., Wang. J., and Knoll, A. A review of safe reinforcement learning: Methods, theories and applications. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Li, L., Zhang, H., Zhang, X., Zhu, S., YU, Y., Zhao, J., and Heng, P.-.. Towards information theoretic frame- work of context-based offline meta-reinforeemeen learn- ing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / openreview net/ www.mrrriiiiiooo Searrmmid=OFsSSSVV 2ww9xx
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor- critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Con- ference on Machine Learning, pp. 1861-1870. PMLR, 2018.
Lin, Z., Thomas, G., Yang, G., and Ma, T. Model-based Adversarial Meta-Reinforcemeen Learning. In Advances in Neural Information Processing Systems, volume 33, pp. 10161-10173, Curran Associates, Inc., 2020.
Hendel, R., Geva, M., and Globerson, A. In-context learn- ing creates task vectors. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 9318- 9333, 2023.
Liu, E. Z., Raghunathan, A., Liang, P., and Finn, C. Decou- pling exploration and exploitation for meta-reinforreemnn learning without sacrifices. In International Conference Machine Learning, pp. 6925-6935. PMLR, 2021.
Humplik, J., Galashov, A., Hasenclever, L., Ortega, P. A., Teh, Y. W., and Heess, N. Meta reinforcement learning as task inference. arXiv preprint arXiv:/905.066 1905.0644,, 2019.
Long, Q. Zhong, F., Wu, M., Wang. Y., and Zhu, S.-C. Socialgfs: Learning social gradient fields for multi-agent reinforcement learning, arXiv preprint arXiv:2405.01839, 2024.
Kamienny, P.-A., Pirotta, M., Lazaric, A., Lavril, T., Usunier, N., and Denoyer, L. Learning adap- tive exploration strategies in dynamic environments through informed policy regularization. arXiv preprint arXiv:2005.029:: 029344, 2020.
Ma, L., Wang, Y., Zhong, F., Zhu, S.-C., and Wang, Y. Fast peer adaptation with context-aware exploration. In Pruceedings of the 41st International Conference on Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 33963-33982, PMLR, 21-27 Jul 2024. URL https: https///poocceeeeeee mlr press/ v235 m 24 html.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit q-learning. In International Confer- ence on Learning Representations, 2022. URL https: //openreview net. /tt///oormmmiie in (i..dd.6nn2222... wtt.:/wwwee//eeueeeeeeee
Laskin, M., Wang, L., Oh, J., Parisotto, E., Spencer, S., Steigerwald, R., Strouse, D., Hansen, S. S., Filos, A., Brooks, E., maxime gazeau, Sahni, H., Singh, S., and Mnih, V. In-context reinforcement learning with algo- rithm distillation. In The Eleventh International Confer- ence on Learning Representatiinss, 2023. URL https: //appnreview.nnee http://eee///tteeeeeeeea/eee forum? hih Uv.
Min, S., Lyu. X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H.,, and Zettlemoyer, L. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pp. 11048- 1 11064, 2022.
Lee, G., Hou, B., Mandalika, A., Lee, J., and Srinivasa, S. S. Bayesian policy optimization for model uncertainty. In International Conference on Learning Representations, 2019. URL https:///// http//ppwwweewwccccece wtw.eeewwwee//orur forum? id=SJGvnsOcK7.
Pan, X., Liu, M., Zhong, F., Yang, Y., Zhu, S.-C., and Wang. Y. Mate: Benchmarking multi-agent reinforcement learning in distributed target coverage control. Advances in Neural Information Processing Systems, 35:27862- 27879, 2022.
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.
Peng, M., Zhu, B., and Jiao, J. Linear representation meta-reinforrcement learning for instant adaptation. arXiv preprint arXiv:2 2021.
Li, L., Huang, Y., Chen, M., Luo, S., Luo, D., and Huang, J. Provably improved context-based offline meta-rl with attention and contrastive learning. arXiv preprint 1xXv:2/02.11074, 2021a.
Rakelly, K., Zhou, A., Finn, C., Levine. S., and Quillen, D. Efficient off-policy reif-reenforeement learning via prob- abilistic context variables. In International Conference on Machine Learning, pp. 5331-5340 PMLR, 2019.
10

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Ren, Z.,, Liu, A., Liang, Y., Peng, J., and Ma, J. Efficient meta reinforcement learning for preference-based fast adaptation. Advances in Neural Information Processing Systems, 35:15502-15515 2022.
Yu, Y. Towards sample efficient reinforcement learning. In IJCAI, pp- 5739-5743, 2018.
Yuan, H. and Lu, Z. Robust Task Representations for Offline Beta-Reeifforemmet Learning via Contrastive Learning. In Proceedings of the 39th International Conference on Machine Learning, pp. 25747-25759. PMLR, June 2022. http: httpssss/ppccceedeeeeeee mmlr mllr. press/ v1 V162/ 22w/22an22a...m html. ISSN: 2640-3498.
Sohn, S., Woo, H., Choi, J., and Lee, H. Meta reinforce- ment learning with autonomous inference of subtask de- pendencies. In International Conference on Learning Representations, 2020. URL https: oppnrevven.. net/ Eorumiid Wtg.wwkkt...............
Zhang, H. and Kan, Z. Temporal logic guided meta q learning of multiple tasks. IEEE Robotics and Automation Letters, 7(3):8194-8201 2022.
Song, X., Gao, W., Yang, Y., Choromanski, K., Pacchi- ano, A., and Tang. Es-maml: Simple hessian-free free meta learning. International Conference on Learning Representations, 2020. URL https: /eppneevee.eeeee net. Eorum? wwt.mmiiddsseexxnnn.mmm
Zhong, F., Wu, K., Wang. C., Chen, H., Ci, H., Li, Z.,, and Wang, Y. Unrealzoo: Enriching photo-realistic virtual worlds for embodied ai. arXiv preprint arX1v:2412.20997, 2024.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ Inter- national Conference on Intelligent Robots and Systems, pp. 5026-5033 IEEE, 2012.
Zhong, F., Wu, K., Ci, H.,, Wang, C., and Chen, H. Empow- ering embodied visual tracking with visual foundation models and offline rl. In European Conference on Com- puter Vision, pp. 139-155 Springer, 2025.
Wang, D., Zhong, F., Li, M., Wen, M., Peng, Y., Li, T., and Yang. A. Romat: Role-based multi-agent transformer for generalizable heterogeneous cooperation. Neural Nei- works, 174:106129, 2024.
Zhou, R., Gao, C.-X., Zhang, Z., and Yu, Y. Gener- alizable task representation learning for offline meta- reinforcement learning with data limitations. In Proceed- ings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 17132-17140, 2024.
Wang, H., Wang, Y., Zhong, F., Wu, M., Zhang, J., Wang, Y., and Dong, H. Learning semantic- aemantic-agnostic and spatial- aware representation for generalizable visual-audio nav- igation. IEEE Robotics and Automation Letters, 8(6): 39000 2023.
Zintgraf, L., Shiarlis, K., Igl. M., Schulze, S., Gal, Y.,
Hofmann, K., and Whiteson, S. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. In International Conference on Learning Representations, 2020. URL /Rttttps:///ppneeeeeee /wppnrevieweeeeeeeeeeeeeeeeeeeee net net/forun? 1 dd=kk1 9J1BYvr.
Wang, Y. Zhong, F., Xu, J., and Wang, Tom2c: Target-oriented multi-agent communication and coop- eration with theory of mind. International Confer- ence on Learning Representations, 2022. URL https: //openrevies net forum? wd BB.
Zintgraf, L.,, Devlin, S., Ciosek, K., Whiteson, S., and Hof- mann, K. Deep interactive bayesian reinforcement learn- ing via meta-learning. arXiv preprint 10Xiv:2101.3300 2021.
Wu, T., Zhong, F., Geng, Y., Wang, H., Zhu, Y., Wang, Y., and Dong, H. Grasparl: Dynamic grasping via adversarial reinforcement learning. arXiv preprint arXiv:200.002119 2022.
Wu, Y., Tucker, G., and Nachum, O. Behavior regu- larized offline reinforcement learning. arXiv preprint 19X1v.1911.11366 2019.
Xu, T., Li, Z., and Ren, Q Meta-reinforcement learning robust to distributional shift via performing lifelong in- context learnin. In Proceedings of the 4|st International Conference Machine Learning, pp. 55112-55125. PMLR, July 2024.
Yoon, J., Kim, T., Dia, O., Kim, S., Bengio, Y., and Ahn, S. Bayesian aggostic meta-learning. Advances in Neural Information Processing Systems, 31, 2018.

---

Bchavior-agnestic Task Inference for Robust Offline In-context Reinforcement Learning
A. Proof of Theorem 3.1
To prove this inequality, we use the following simple lemma:
Lemma A.l. For arbitrary random variables X, Y. Z.
H(X Y) + H(Y Z) H(X 2).
(11)
Proof.
H(X Y) H(Y Z)
H(X Y. Z) + H(Y Z)
H Y Z)
N H(X Z)
Now, for the main proof:
Proof. Applying Lemma A.I and rearranging, we have
H(Z X H(Z M) & H(M x²)
Now,
M) I(Z
H(Z) H(Z M) (H(Z) (Z xx))
H(Z X²) - H(Z M)
S H(M X*)
H(M) - (M( - H(M x*))
= H(M) I(M;
We now that the CSRO objective (at À upper bounded by H(M) (MM; X*). Consequently, when 7(M; high, e.g. M can be determined fully from X have / (M; X²) H (M) and ((Z IZz As result, cannot find Z that simultaneously captures M and is not correlated with X
B. Pseudocode of the Offline Training and Online Testing Procedure
Algorithm Offfine Training Pipeline of BATI
Require: Training task distribution Pusin (M) and associated offline datasets {DM}
1: Randomly initialize 0 o. vb
2: while Maximum training step not reached
3: Sample batch of tasks { (M,) ~ Prrain(M] and corresponding training contexts (XM,) ~ (DM,)
4: Sample {2M,} (Z } with reparameterization
5: Update o, db with (XM.) } and (ZM,) } using the dynamics reconstruction loss Eq.
6: Update 0 with (XM,) M, and (detached) {2M,) using IQI. losses
7: end while
8: Return final 0, o. v
C. Experiment Details
C.1. Environments
We conduct our experiments in several commonly used environments in the field of offline ICRL, Based on MuJoCo (Todorov et al., 2012), a popular physics simulator, the environments feature several robot locomotion tasks where the robots must move accomplish certain task:
12

---

Bchavior-agnestcc Task Inference for Robust Offline In-context Reinforcement Learning
Algorithm 2 Online Evaluation Procedure of BATI
Require: Testing task distribution (M) and associated context datasets {C evaluated parameters 6, Φ, v. number of
latent samples per inference N
1: R -00
2: for M E Prest (M) do
3: Compute Zt with X ~ CM using the dynamics reconstruction loss Eq. 8
4: Sample initial state ~ Pa of M
5: while Episode not done do
6: Sample action ss, Z²)
7: Execute action in task M and get the next state reward r, done flag d
8: R. 4 R +
9: end while
10: end for
11: Return R/ Prest [M)
A four-legged ant-like robot needs to go along the direction specified by the task. The task is goal direction uniformly sampled from [0, 2Tr, and the reward is the inner product between the position delta and the goal direction. At every time step f, we set the goal direction M; to the sum of the true task A/* and Gaussian noise mean 0 and standard deviation e, such that the reward function stochastic with respect the true task.
AntDir:
HalfCheetah Vel: bipedal cheetah-like like robot needs to run with the velocity specified by the task. The task is a target velocity uniformly sampled from [0, 2). and the reward is the negative absolute difference between the current and the goal velocity minus control cost. We add a similar Gaussian noise of mean 0 and standard deviation to the target velocity at every time step.
DalrCheetahDir: A bipedal cheetah-like robot needs to run along the direction specified by the task. The task is target direction uniformly sampled from {left, right), and the reward is the velocity in the target direction minus a control cost. We use the true target direction to compute the reward with probability 1 - e, and a uniformly random direction with probability e.
HopperParam, WalkerParam: A single-leg hopper or a bipedal walker needs to move forward. In these two environments, tasks differ by the transition (physics) Par of the MDP, instead of the reward function Ry in the environments above. The task is a log friction coefficient uniformly sampled from [-3,30 3), and the reward is the velocity in the forward direction minus control cost. Furthermore, the friction coefficients are also used to perform affine projection to the action ensure that different tasks have different optimal behaviors. We add Gaussian noise of mean 0 and standard deviation to the log friction coefficient at every time step.
C.2. Datasets Construction
To construct the datasets used for training and testing, we train an expert policy μM using SAC (Haarnoja et al., 2018) for each training task M. For the environments differing rewards, we use = 0 when training the context-colleetion policies. For the testing task M,, we use the expert policy for the most similar training task as its primary context collection policy μM, and that for the most different training task the testing context collection policy P Furthermore, to ensure that Phest remains supported by the training data Prrin, we mix 90% data from Ma with 10% data from μaM to construct Piest for HopperParam and WalkerParam. Distance between tasks is defined as follows for each environment:
AntDir: The absolute difference of the goal direction, measured radians modulo 2x: min (aimmmx(MM) M2) min(M,, M2), min(M1, M2) + 2 max(M,, M2)).
HalfCheetah Vel: The absolute difference in target velocities.
HalfCheetahDir: The target direction is the same (0) or different (1).
HopperParam, WalkerParam: The absolute difference in log friction coefficients.
13

---

Bchavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
Table 3. Hyperparameters used in each of evaluation environments.
Environments
Parameter Name
AntDir
HalfCheetah
HalfCheetahDir
HopperParam
WalkerParam
Learning Rate
3 * 10-4
3 10
3 * 10
3 * 10-4
3 10
Batch Size
4096
4096
4096
4096
4096
Task Contrastive Batch Size
16
16
16
16
16
IQL T
0.8
0.8
0.8
0.8
IQL
0.05
0.05
0.05
0.05
0.05
IQL Exp. Adv. Clip
100
100
100
100
100
Gradient Steps
105
105
105
105
105
Episode Length
200
200
200
200
200
Dataset Size
105
2 10
2 * 10
3 * 10
3 * 10
Task Latent Dim
5
5
5
40
40
BATI # Latent Samples N
40
40
40
40
UNICORN Weight
0.15
0.15
0.15
1.5
CSRO CLUB Weight
5.0
1.0
1.0
2.5
2.5
CLUB Encoder Hidden Dims
[200]
[200] *
[200] # 3
[200] * 3
[200] *
Encoder Hidden Dims
[64, 64]
[64, 64]
[64, 64]
[128, 128)
[128, 128]
Decoder Hidden Dims
[64, 64]
[64, 64]
[64, 64]
[128, 128]
[128, 128]
RL Hidden Dims
[256, 256]
[256, 256]
[256, 256]
[256, 256]
[256, 256]
Table Final online evaluation episodic returns of BATI and baselines Prea(M, p)). BATI achieves superior performance in of the benchmark environments, dramatically outperforming baselines. IQL provides additional optimizations the basis of oracle data and may lead to methods with even better performance.
AntDir
HalfCheetah
HalfCheetahDDr
HopperParam
WalkerParam
Oracle (Pre-IQL.)
58.7
-139.3
291.2
290.2
549.2
BATI
46.4
2.9
-122.8 1.6
286.1 + 13.8
285.9 + 5.5
565.2 : 7.9
CSRO
3.0
-134.8 8.9
-340.0 17.2
144.4 21.3
279.5 + 34.0
FOCAL
-46.5 2.2
-279.4 + 9.4
-336.4 8.5
162.6 44.5
317.0 + 37.0
Recon
-18.3 13.7
7.4
127.1
157.8 = 16.7
292.6 + 38.4 38.4
UNICORN
-32.0 + 4.2
267..3 20.5
-362.9 + 6.8
144.7 + 21.4
319.0 + 11.5
C.3. Implementation Details
Building on the codebase of the official implementation of UNICORN (Li et al., 2024), we implement BATI and all our baselines. All methods share the same offline RL algorithm and relevant architectures. The networks are implemented as MLPs with ReLU activations, with hidden dimensions specified Tab. 3. The CLUB model used in the CSRO baseline ported from the official implementation CSRO (Gao al., 2024). We found BRAC (Wu al., 2019), base offline RL algorithm used in several prior works, to behave unstably in some cases, and switched to IQL (Kostrikov al., 2022) instead for all methods, which yields consistent and satisfactory performance. See Tab. for hyperparameters of the main experiments. Note that some hyperparameters apply only to the methods that require them, e.g. Task Contrastive Batch Size applies only to baselines with FOCAL-like losses (Li et al., 2021b), and Encoder Hidden Dims apply only to baselines and not BATI (which does not have an encoder network). The task latent embedding distributions of BATI M } M are parameterized as Gaussian distributions with learnable mean and fixed log variance of -4. An optional KL divergence between the latent distributions and (0 I) can be used to regularize the task latent distributions. We empirically find the unregularized version have good performance and not use this regularization for all methods, echoing UNICORN.
14

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
C20
CSRO
FOCAL
TOCAL
RAEST
302
Drace
220
Orode
266
232
-292
-250
-1500
Talning Staps
Tatn Steps
Ta Tanging Stepe
(a) 0
(b) c 0.3
(c) E 0.7 (Default)
Figure 7. Online evaluation episodic return curves during training for different noise levels on HalfCheetahDir. The performance of baselines mostly decreases relative BATT the noise level increases.
C200
FOCAL
FOCAL
Reson URICO 1M
-152
250
USBCORNN
52332h 14tm
1utum
200
202
22
220
202
-22-
20000
Thitiin Strpe
Tni no Steps
Thiving Siice
(a) AntDir
(b) HalfCheetah
(c) HopperParam
BNTI
RKTI
CS10
CERO
TOCAAL
FOCAL
5:0
Eakr04 Nolm
200
Spraod. (utum
420
320
100
220
25000
41200
6023E
80302
105000
20020
40000
80202
$1000
120000
20200
400003
62000
080932
100093
Pariing
(d) HaifChectahDir
(c) WalkerParam
(f) HopperParamn 0)
Figure 8. Online evaluation episodic of BATT and baselines contexts from Parain M, 11) during training in evaluation environments.
D. Additional Figures and Experiments
D.1. Additional Figures and Tables Referred by the Main Text
Due to space constraints, we moved some figures and tables used in the main text here. Tab. 4 contains the final online evaluation episodic returns of the main experiment, while Fig. 7 contains noise level ablations in HalfCheetahDir.
D.2. In-distribution Performances
We show the online evaluation episodic return curves on Prain here. In Fig. 8, most methods can achieve similar performance when evaluated on the training Prain contexts, echoing our argument that learned encoders can be trained to perform in- distribution Bayesian inference. The performance of CSRO is somewhat unstable, potentially due to the conflicts among its objectives. Similar phenomena can observed the in-distribution performances of the noise level ablations in AntDir (Fig. 9 and HalfCheetahDirr (Fig. 10).
15

---

Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning
120
BNT
AAT CSRD LMBCORN
142
Recsn
120
Rocan UNIDOW
A02
LNKCOA
02-
92-
Exkced:
00000
00202
102020
20020
40000
20105
D0032
(b) c
(c) c 0.4 (Default)
Figure 9. Online evaluation episodic return curves Purin (trin(M, (M, during training for different noise levels AntDir.
1250
IATI
CIRO
C2-10
FOCAL
CLRO COMM
SFFT
1000
Pauo
JNICORK
6:0
750
420-
250
220
-500
22000
42000
00202 202000
20020
40000
60302
82000
110000
23020
43020
60002
102020
Parsin Stere
Trann Stepe
1a Stepo
(8) 0
(b) 0.3
(c) c 0.7 (Default)
Figure 10. Online evaluation episodic return curves on contexts from Pirai P (M, μ)) during training for different noise levels on HalfChee- tahDir. All methods perform similarly while CSRO and Recon somewhat unstable.
D.3. Adaptation Different Numbers of Training Tasks
We demonstrate the robustness of BATI with respect to the task embedding table size (corresponding to the number of training tasks). We split the 40 evaluation tasks AntDir into different training and testing splits and reran the experiments. As shown in Tab. across all splits, BATI achieves the best performance uniformly and highly stable.
D.4. Adaptation Multi-agent Scenarios
We conduct preliminary experiments to showcase the general applicability of BATI multi-agent domains. We choose Kuhn Poker, a two-player card game with discrete state and action spaces, differing from the continuous MuJoCo environments used in our paper. We generate different player-2 (opponent) policies "tasks" and learn adaptive policy for player-1 I over 10 episodes (20 steps) of contexts. As shown in Tab. 6, BATI maintains superior performance over all baselines, further showcasing its capabilities and generalization.
Table 5. Final online evaluation episodic returns of BATI and basclines on Prea M, n1) with different train/test splits in AntDir, BATI outperforms baselines and consistently good performance across different numbers of training tasks.
Train/Tess Split
BATI
CSRO
FOCAL
Recon
UNICORN
10/30
45.0 - 6.0
4.7 4.2
-33.2 + 3.4
2.8
-22.2 4 2.5
20/20 (Main)
46.4 + 2.9
28.9 3.0
-46.5 + 2.2
-18.3 13.7
-32.0 + 4.2
30/10
49.9 + 4.6
46.2 4.8
-33.9 + 1.4
-4.3 5.0
7.6477.6
Table 6. Final online evaluation episodic returns of BATT and baselines on Peest (M, (1) in KuhnPoker.
Method
Oracle
BATI
CSRO
FOCAL
Recon
UNICORN
Episodic Return
0.0734
0.049 + 0.025
-0.180 + 0.032
0.191 + 0.020
00.185 + 0.052
0.243 + 0.042
16


Reviews:
[{"review_id": "Z6sUIZhYmf", "reviewer": "Reviewer_2Aoe", "summary": "The authors propose a modification to the way offline context-based meta-RL methods supervise task identification, which they call BATI. The core idea is to remove the correlation between the task estimate and the behavior of the policy collecting the context. That way, the test-time context policy can be significantly different from the policy used to collect training context, and the method will still accurately select a task representation that leads to high returns when used as an augmented state feature for a standard offline-RL policy. The authors evaluate BATI in the standard meta-RL extensions of gym locomotion tasks under a setup that highlights robustness to stochastic dynamics and adversarial test-time context.\n\n## Update After Rebuttal\nI appreciate the authors' reply and have no further concerns. I maintain my positive 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 kind words and insightful review! We are happy to answer your questions below.\n\n> Q1. Training contexts generated by adaptive policies.\n\nR1: The result would likely depend on the exact behavior of the adaptive policy. For example, if the adaptive policy is bad, it may randomly walk around the state space in any task and reveal little information about the true task identity, in which case all approaches would fail due to the uninformative nature of the context. At the other end of the spectrum, if the adaptive policy is good and finds the optimal policy quickly, its behavior may diverge significantly between tasks, resembling the single-task expert case where BATI works and baselines don't. In general, we expect BATI to outperform or at least remain competitive with the baselines for all cases, and the baselines to work only when the adaptive policy behaves similarly in all tasks yet still reveals a lot of task-related information.\n\n> Q2. Challenges brought by extremely large task spaces.\n\nR2: A good question! Indeed, this is one of the potential future works we'd like to explore. The task inference procedure of BATI has a time complexity of $\\mathcal{O}(NL)$ where $N$ is the number of samples (lower-bounded by the number of training tasks) and $L$ is the length of the context and can be fully vectorized. This allows the inference procedure to run reasonably fast. However, if the task space is extremely large over which we couldn't practically enumerate, more sophisticated techniques could be employed to perform the inference. For example, we can use [diffusion posterior sampling](https://arxiv.org/abs/2209.14687) to sample from the conditional posterior $p(x \\mid y, c)$ where $x$ is the task latent, $y$ is $X^t$, and $c$ is $X^b$. Under this formulation, the task prior $p(x)$ is represented by a diffusion model instead of a table of embeddings, and (robust) likelihood $p(y \\mid x, c)$ is provided by the dynamics model whose gradient $\\nabla\\_x p(y \\mid x, c)$ serves as a form of classifier guidance. Leveraging the powerful expressitivity of diffusion sampling procedures, we can efficiently and robustly sample from a much larger task space. However, to demonstrate the advantage of such an approach, a benchmark with a sufficiently complicated task space is required, which is still emerging, as pointed out in your review, and we will explore this further in our future work.\n\n> Q3. Connection with BOReL (Dorfman et al., 2022).\n\nR3. Thank you for bringing our attention to this related work! We note that the training data of BOReL are actually *augmented* via their reward relabelling (RR) procedure, in effect rolling out all context collection policies in all tasks and making $M$ and $\\mu$ independent. This avoids the spurious correlation issue during training, but requires additional properties (e.g. tasks differing only by rewards) and ground-truth reward functions. Indeed, as noted in BOReL, *\"With reward relabelling ablated ... the agent believes the reward is at the point it first visited on the semi-circle\"* (Page 8, lower left), which is a consequence of the spurious correlation phenomenon. Our BATI can be applied to BOReL to replace RR without any additional requirements.\n\n> Q4. Oracle references and number of seeds.\n\nR4: We have updated the main results to include oracle references by single-task experts, available anonymously [here (Fig. 1 and Tab. 4)](https://sites.google.com/view/bati-icrl). BATI is able to match or even slightly outperform the expert data in most cases due to optimizations performed by IQL. Due to limited time and computational resources, we have run two additional seeds for the main experiments, bringing the total to `5 seeds`; the changes are also reflected in the figure and table above.\n\n> Q5. Usage of \"in-context RL\".\n\nR5.: Thank you for the detailed and insightful comments! In BATI, we treat ICRL as mostly synonymous with context-based meta-RL, i.e. learning policies that can ***efficiently adapt to the context and improve on-the-fly*** without laborious fine-tuning, regardless of transformer or implicitness. We argue that this \"functional\" definition better captures the essence of ICRL. For example, there are [state-space model approaches](https://arxiv.org/abs/2303.03982) that produce a fixed-length representation of the context without transformers, and methods like [DPT](https://arxiv.org/abs/2306.14892) that are trained to perform explicit Bayesian inference of optimal actions. However, we stress that **we greatly value inputs from the community about terminological choices, and are open to changes if deemed necessary**.\n\n> Q6. Additional hyperparameter and implementation details.\n\nR6: Thank you for your advice! As the ICML 2025 reviewing policy prevents authors from changing the PDF during rebuttal, we'll add relevant details to the appendix for a more self-contained reading experience when we get to update the paper.", "claims_and_evidence": "Yes, the main claims are well studied by the experiments.", "methods_and_evaluation_criteria": "The authors evaluate their method in a set of mujoco locomotion benchmarks with a small set of randomized objectives and dynamics that have been staples since the first wave of deep meta-RL papers. These tasks play to the strengths of methods that can learn precise task estimates from short context over a narrow range of behavior. These benchmarks are quite saturated at this point. However, diverse and affordable alternatives are still starting to emerge, and this choice isn’t meaningfully impacting my review. The authors go to great lengths to try and make these tasks more interesting by adding noise to their dynamics and creating adversarial context sequences at test-time (mainly consisting of actions chosen from the expert data collection policy for the least similar task). Still, I wonder whether the method would scale to diverse training distributions that might require large M (the size of the task set and number of embeddings).", "theoretical_claims": "I reviewed the appendix including the argument of Theorem 3.1", "experimental_designs_or_analyses": "The experiments seem sound and implement all the baselines fairly on top of the same codebase. Transferable improvements (such as a change in the base offline RL algorithm) are applied across multiple baselines. The authors are relying on a nearly worst-case-scenario setup to evaluate their method on the tasks chosen, but this may reflect real applications. \n\nIt may be useful to report some oracle (task-conditioned or single-task expert) references scores for the experiments where the dynamics of the locomotion envs (like AntDir) are given significant noise.\n\nThe low seed counts are a bit of an issue given the very thin margins between most of the baseline results.", "relation_to_broader_scientific_literature": "The paper’s use of the term “in-context RL” (ICRL) is confusing — if this term even means anything unique at this point (and maybe it doesn’t). To review, “in-context learning” is primarily used to describe sequence models’ (almost always Transformers/LLMs) ability to understand a task at test-time based on a small sample of input/output pairs. This ability is often attributed to *implicit* Bayesian inference of the task inside the activations of the model, though there are some other explanations of this effect in NLP. \n\nModern Meta-RL and Meta-IL have begun to borrow the term when using Transformers to *implicitly* learn the ability to improve based on their input sequence. It is broadly used to rebrand traditional subsets of black-box meta-learning by using the connection to LLMs to highlight the flexibility of implicit task inference outside of standard few-shot meta-RL benchmarks (e.g., to few-shot prompting, opponent modeling). Or, in the case of Algorithm Distillation, to the implicit ability to learn the RL gradient update itself. BATI’s explicit task inference by conditioning the online policy on one of M task embeddings that minimizes a reconstruction objective on data provided by another behavior policy fits in neither category. There is no task estimate emerging implicitly in the context length of a sequence model. I had never seen the term “in-context RL” used this way. The authors write “We identify a critical limitation in existing offline ICRL methods” and then cite several sources, but the only one of these sources that actually call it “in-context RL” is the most recent UNICORN (Li et al., 2024), which I assume is how this came to be. This might be a losing battle, as the vocab and taxonomy of meta-RL has been messy for quite some time, but I really don’t think the term “in-context” applies here at all.", "essential_references_not_discussed": "The literature is well covered, but I am interested in the authors’ thoughts on the connection to [Dorfman et al., 2022](https://arxiv.org/abs/2008.02598). My understanding is that the evaluation setting in BATI is a similar worst-case scenario to what is discussed there, in that the context dataset is collected by an expert in a single task and therefore its behavior heavily implies the task identity.", "other_strengths_and_weaknesses": "The main argument of behavior-agnostic task identification from context sets generated by single-task experts is made clearly and given plenty of motivation.", "other_comments_or_suggestions": "I think the paper would benefit from some more details in the Appendix related to the network architecture and overall implementation. Currently, the authors are relying on UNICORN for much of this information, which prevents the paper from standing alone. That’s a shame because the core argument for the change in inference objective is given enough space in the main text that someone would not necessarily have to read the baseline papers to follow it.", "questions_for_authors": "1. How would the authors expect results to change when the train-time context data is generated by policies that are adapting over some distribution of tasks (rather than single-task experts trained by SAC)?\n2. What would be the challenges involved in scaling BATI to domains with extremely large task spaces?", "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": "3eGvsi71A8", "reviewer": "Reviewer_bfNC", "summary": "The paper introduces Behavior-Agnostic Task Inference (BATI) to improve offline in-context reinforcement learning (ICRL) under distribution shifts. BATI, a model-based maximum-likelihood approach, infers task representations robustly by focusing on environmental dynamics. Results show BATI outperforms existing methods, especially with context shifts and noise.", "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 review. We're encouraged by the overall positive assessment that our paper `\"is generally well-written\"` and `\"addresses a critical challenge in ICRL\"`. Below, we address the main concerns raised regarding the generalizability of our method beyond MuJoCo environments and the theoretical connections between our analysis and the proposed BATI approach. As demonstrated in our preliminary multi-agent experiments and theoretical explanations, BATI effectively addresses the identified limitations of previous methods while maintaining strong performance across diverse domains. The following are our detailed responses:\n\n> Q1. Results in other domains.\n\nR1: Due to the tight rebuttal schedule, we have conducted a preliminary multi-agent experiment to demonstrate the general applicability of BATI in other domains. We choose Kuhn Poker, a two-player card game with discrete state and action spaces, differing from the continuous MuJoCo environments used in our paper. We generate different player-2 (opponent) policies as \"tasks\" and learn an adaptive policy for player-1 over 10 episodes (20 steps) of contexts. As shown in the table below, BATI maintains superior performance over all baselines, further showcasing its capabilities and generalization.\n\n| Method | Oracle (pre-IQL) | BATI | CSRO | FOCAL | Recon | UNICORN |\n| -------- | --- | -------- | -------- | --- | --- | --- |\n| Episodic Return     | $0.0734$ | $\\mathbf{-0.049 \\pm 0.025}$     | $-0.180 \\pm 0.032$     | $-0.191 \\pm 0.020$ | $-0.185 \\pm 0.052$ | $-0.243 \\pm 0.042$ |\n\n\n> Q2. Theoretical analysis of our method.\n\nR2: We apologize for any potential confusion. We establish the preliminaries and analyze the prior methods in Section 3. Building on these insights, Section 4.1 is primarily dedicated to our method that circumvents the prior failure modes. In Section 4.1, we provide two perspectives on the reason why BATI works: 1) the graphical model perspective (Fig. 2) where BATI achieves robustness by blocking $X^b$; 2) the robust likelihood perspective (page 4 right, line 179-219) where the core objective of BATI can be understood as a robust version of the full likelihood that does not depend on $\\mu$ through derivations. We will communicate these points more clearly in the revision and greatly appreciate any suggestions to further improve the writing.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "The proposed method, BATI, makes sense for addressing the identified problem of distribution shifts in offline ICRL. It effectively shifts the focus from Bayesian posterior inference to a maximum-likelihood estimation of environmental dynamics, thus being more robust to context shifts. The evaluation criteria, using MuJoCo environments and varying noise levels, are appropriate for assessing BATI's performance and robustness in different scenarios.", "theoretical_claims": "The theoretical result is standard.", "experimental_designs_or_analyses": "The experimental designs and analyses are generally sound, using MuJoCo environments, relevant baselines, ablation studies, and appropriate evaluation metrics.", "relation_to_broader_scientific_literature": "The paper improves offline in-context reinforcement learning (ICRL) under distribution shifts.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n\nThe paper is generally well-written and easy to follow. The authors clearly explain the problem, the proposed solution, and the experimental setup.\n\nThe paper addresses a critical challenge in ICRL, which is the vulnerability to distribution shifts. Overcoming this limitation is essential for the broader applicability of ICRL in real-world scenarios. \n\nWeaknesses\n\n The experiments are primarily conducted in MuJoCo environments. While these are standard benchmarks, such as MetaWorld. it would be beneficial to see results in other domains to further demonstrate the generalizability of BATI.\n\nThe theoretical analysis in Section 3 and Section 4.1 seems to have less connection to the proposed method. Why the proposed method can address the theoretical issue of the previous works shown in Section 3 and Section 4.1? Is there any theoretical analysis for the proposed method?", "other_comments_or_suggestions": "The paper spends several pages to explain the issue of the previous method. Some issues, such as spurious correlation of the context learner. The paper can be strengthened by the analysis of the proposed method on the issue.", "questions_for_authors": "Please address the above comments.", "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": "BstlK3d99A", "reviewer": "Reviewer_4HP8", "summary": "Authors propose Behavior-agnostic Task Inference (BATI) approach for meta RL problems which is claimed to be more robust to noisy dynamics compared to previous methods like UNICORN or CSRO and which works at the same level in noise-free cases.", "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 efforts in reviewing our paper and the detailed review! We address your concerns below. **Due to space limits, all the figures and tables referred below are available anonymously [on this website](https://sites.google.com/view/bati-icrl).**\n\n> Q1. Experiment protocol of BATI and baselines.\n\nR1: We clarify that all methods, including CSRO, are evaluated on the same context distribution $p_\\text{test}(M, \\mu)$ during test time, as described in Section 5.1. We design our testing protocol to be as challenging as possible since this aligns better with real-world scenarios where dynamics are noisy and the policy may not have control over the context (e.g. provided by a human operator). Other baselines like UNICORN and FOCAL also use a similar protocol. To address your concerns more thoroughly, we evaluate CSRO on HalfCheetahDir using their self-collection protocol and obtain a performance of $-16.2 \\pm 25.3$. While the stronger protocol improves performance, CSRO still lags behind BATI by a lot.\n\n> Q2. Choice of base offline RL algorithm and hyperparameter tuning.\n\nR2: We switched from BRAC to IQL because we observed in some cases that BRAC simply failed to learn; see website Fig. 2 for results in WalkerRandParams. We also performed little hyperparameter tuning and directly adopted hyperparameters from previous works. The IQL- and CSRO-related ones are borrowed from their respective papers (CSRO CLUB weights are divided by 10 since their implementation has a weight of 10 for the FOCAL loss while ours has 1) while the rest comes from UNICORN, including those of data-collection runs (with learning steps and dataset sizes slightly adjusted to accelerate training).\n\nFurthermore, we would like to remark that as argued in Sec. 3 and 4, the core problem with baselines lies in task inference and not policy learning; the performances of baselines may even worsen with hyperparameter tuning, since the meta-policy would then get better at executing a ***wrong*** policy.\n\n> Q3. OOD task experiments.\n\nR3: OOD generalization for tasks is extremely challenging, and to the best of our knowledge, the primary experiments of UNICORN and CSRO focus on **generalization to OOD contexts on in-distribution tasks**, same as BATI. Only UNICORN attempted an OOD task setting with a model-based RL approach, which is out of scope for our paper.\n\nNevertheless, we acknowledge the importance of OOD generalization and conduct the experiment you requested in AntDir with training goal directions sampled from $[0, \\pi)$ and testing directions from $[\\pi, 2\\pi)$. The results are shown on the website (Tab. 1). It can be seen that while the performance decreases compared to the in-distribution case, BATI still outperforms baselines by a large margin and is the only method to achieve positive returns. This further validates the generalization capability of BATI and its task embedding designs.\n\n> Q4. Task embedding table size; sensitivity to hyperparameters.\n\nR4: The size of the task embedding table is the same as the number of training tasks, as described in Section 4.2. To demonstrate the robustness of BATI with respect to the number of training tasks, we conduct an ablation in AntDir that splits the 40 tasks into different train/eval splits. The results are shown on the website (Tab. 2). It can be seen that BATI outperforms baselines in all settings and is highly stable.\n\nWe conduct another ablation in WalkerRandParams to validate the robustness of BATI with respect to task embedding size. Compared with the main result of $565.2 \\pm 7.9$ with embedding size 32, size 16 and 64 yield performances of $556.4 \\pm 1.7$ and $536.2 \\pm 6.4$, respectively, demonstrating the robustness of BATI.\n\n> Q5. Discussion and comparison with transformer-based methods.\n\nR5: Thank you for the additional related works! We remark that all of the three methods belong to the same paradigm where a transformer is trained to predict actions conditioned on the context, an act of Bayesian inference similarly flawed as the context encoders of our baselines.\n\nTo support this argument, we add a new baseline [DPT](https://arxiv.org/abs/2306.14892) that also uses a transformer to encode the context and predict the optimal action, outperforming AD. Due to limited time, we implement DPT and evaluate it in AntDir and HalfCheetahVel, as shown on the website (Tab. 3). DPT is inferior to BATI in both environments due to spurious correlations.\n\n> Q6. Dynamics model training objective $\\mathcal{L}\\_\\text{recon}$.\n\nR6: Thank you for pointing this out! We apologize for the confusion, the correct objective for Eq. 7 should be $\\mathcal{L}\\_\\text{recon}^{X, Z}(\\phi, \\psi) := \\frac{(X^t - g\\_\\psi(X^b, Z))^2}{\\exp h\\_\\psi(X^b, Z)} + h\\_\\psi(X^b, Z)$. We design $p(X^t \\mid X^b, Z)$ to be a Gaussian distribution with mean and (log) variance parameterized by neural networks g,h, and L_recon is its negative log likelihood (up to a constant). This will be fixed in the revision.", "claims_and_evidence": "Claims are supported by the evidences. But experimental design is not clear for me which makes evidences questionable. E.g. CSRO which is the strongest baseline was not designed for the evaluation under the adversarial context and it is not clear whether it is tested under the same protocol with BATI or not. More on this in following sections. Replacing BRAC with IQL in baselines without demonstrating the difference is not reliable. Also, I haven’t found details on hyperparameters tuning for BATI and baselines which is crucial for RL.", "methods_and_evaluation_criteria": "Evaluation criteria and used benchmarks (meta MuJoCo environments) make sense and common in the field. \n\nBut I would like to offer additional set of experiments from the side of tasks. How would performances change if we have different distributions for train and test tasks? E.g. in AntDir what if train goals are sample from $[0, \\pi)$ and test goals from $[\\pi, 2\\pi)$? UNICORN and CSRO should be able to handle this as demonstrated by the corresponding works and it is not clear for BATI because of it’s design (the finite codebook of latent representations for tasks Z).", "theoretical_claims": "I did not carefully checked theoretical claims and did not observe issues during reading. I did not read theoretical part in Appendix.", "experimental_designs_or_analyses": "Authors provide valid experimental design to support claims that their method is robust to noisy dynamics and irrelevant contexts. \nHowever, it is not clear whether the comparison against CSRO is correct. I haven’t seen the test protocol for CSRO which is designed to infer task by collecting the context on it’s own. Was it provided with the same context as BATI? It is not supposed to work with adversarial contexts and might perform better if the context is collected from scratch.\n\nBased on the previous remark, I thing it is important to check how the algorithms will behave without adversarial contexts, i.e. collecting entire contexts needed for task inferences on their own (zero-shot).\n\nIf I understood the BATI correctly, there is a finite “code-book” (table) for latent representations of the tasks build based on the training instances. First, I did not understand whether the size is equal to amount of training tasks or not. Anyway, intuitively, BATI might heavily depend on this table size and there is no experiment which demonstrate the sensitivity. In case the table size matches training tasks, it is important to compare BATI and baselines under various sizes of training dataset in terms of number of tasks. \n\nThere is also a set of experiments about OOD environments instances I mentioned in “Methods And Evaluation Criteria” which I find lacking.\n\nAuthors replaced BRAC which was used in CSRO and UNICORN as offline RL algorithm with IQL. While in Appendix it is stated with text that IQL is more stable, I consider it as a major change for baselines and wish to see the exact performance difference between approaches backed up with BRAC and IQL. Now I can imagine that BRAC was not working for BATI while working for other baselines. It would be also nice to make a reference that IQL and BRAC-like algorithms demonstrate competitive performance (https://arxiv.org/abs/2210.07105). \n\nI did not found any details on the hyperparameters tuning for BATI or baselines. Given the fact that authors changed offline RL algorithm from BRAC to IQL and collected their own datasets, it is important to show that BATI and baselines (or at least the most competitive CSRO) were tuned with similar hyperparameters tuning budget.\n\nIt would be good to also see BATI sensitivity to it’s hyperparameters, e.g. # Latent Samples N or task latent dimensionality.", "relation_to_broader_scientific_literature": "Authors present approach which is more robust to noisy dynamics and adversarial contexts -- two aspects which seem essential for meta RL. They demonstrate that prior powerful offline meta RL methods are not able to handle such shifts easily.", "essential_references_not_discussed": "While mentioning Algorithm Distillation, authors do not discuss follow-up works based on it’s findings or similar approaches based on scalable Transformer architecture. I believe those works should be at least mentioned in the context of the work and ideally compared against (but I do not require it now as an important change due to the computational costs). Transformers are known for being able to adopt to completely novel tasks after a certain scale. Here is a list of works I would recommend discussing:\n\n1. Prompt DT (https://arxiv.org/abs/2206.13499) which solves offline meta RL tasks in-context by training transformer to predict next action given the context of interactions with environment.\n2. https://arxiv.org/abs/2312.03801 where authors demonstrate that transformers are able to adapt to completely novel environments in-context after particular offline pre-training.\n3. Headless AD (https://arxiv.org/abs/2312.13327) proposes an AD modification which is able to adapt to completely novel action space, i.e. novel environment dynamics.", "other_strengths_and_weaknesses": "The latent Z table might be a strong limitation when compared to previous methods which are based on context encoder models. The experiments I’ve asked for in this regard should reveal whether it is true.", "other_comments_or_suggestions": "I did not understand where from comes the $L_{recon}$ training objective. While numerator seems intuitive, I do not understand why there is denominator. Wouldn’t $h_{\\psi}$ just try to produce large positive numbers to reduce the loss? Why is it there at all?", "questions_for_authors": "All of my questions were asked in previous sections of the review.", "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": "1ScMjO19GA", "reviewer": "Reviewer_7pDq", "summary": "This work analyzes the shortcomings of existing In-context Reinforcement Learning (ICRL) methods, pointing out their inability to handle context shift scenarios. The authors theoretically analyze the necessity of maximizing the true mutual information between context representation and task indices. Building on this foundation, they propose Behavior-agnostic Task Inference (BATI), which ensures that the context representation focuses solely on the environmental dynamics. Finally, experiments conducted in environments with noisy dynamics demonstrate the effectiveness of BATI.\n## update after rebuttal\nThanks for the author's reply, I have no further questions.", "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 review! We answer your questions below.\n\n> Q1. Types of environments.\n\nR1: We use a set of MuJoCo environments that is standard in the field of meta-RL and consistent with our baselines. To further demonstrate the generality of BATI, we have additionally conducted a preliminary multi-agent experiment. We choose Kuhn Poker, a two-player card game with discrete state and action spaces, differing from the continuous MuJoCo environments used in our paper. We generate different player-2 (opponent) policies as \"tasks\" and learn an adaptive policy for player-1 over 10 episodes (20 steps) of contexts. As shown in the table below, BATI maintains superior performance over all baselines, further showcasing its capabilities and generalization.\n\n| Method | Oracle (pre-IQL) | BATI | CSRO | FOCAL | Recon | UNICORN |\n| -- | --- | --- | --- | --- | --- | --- |\n| Episodic Return     | $0.0734$ | $\\mathbf{-0.049 \\pm 0.025}$     | $-0.180 \\pm 0.032$     | $-0.191 \\pm 0.020$ | $-0.185 \\pm 0.052$ | $-0.243 \\pm 0.042$ |\n\n> Q2. Comparison with single-model, transformer-based approaches.\n\nR2: We note that Algorithm Distillation imposes additional requirements on the training data (i.e. generated by the learning histories of an RL algorithm), making it less general than BATI and unsuitable for our evaluation setup. In general, we remark that the family of single-model approaches performs Bayesian inference to deduce the posterior given a context trajectory, and is thus similarly flawed as the context encoders of our baselines, as argued in Sec. 3 and 4 in our paper.\n\nTo support this argument, we add a new, recent baseline [DPT](https://arxiv.org/abs/2306.14892) that also makes use of a transformer to encode the context and predict the optimal action, outperforming Algorithm Distillation. Due to the computational requirements of training transformers and the tight rebuttal schedule, we implement DPT and evaluate it in AntDir and HalfCheetahVel, as shown below. It can be seen that DPT is outperformed in both environments.\n\n| Environments        | Oracle (pre-IQL) | BATI             | DPT              | CSRO             | FOCAL            | Recon            | UNICORN           |\n|----------------|---|------------------|------------------|------------------|------------------|------------------|-------------------|\n| AntDir         | $58.7$ | $46.4 \\pm 2.9$   | $-8.8 \\pm 0.9$   | $28.9 \\pm 3.0$   | $-46.5 \\pm 2.2$  | $-18.3 \\pm 13.7$ | $-32.0 \\pm 4.2$   |\n| HalfCheetahVel | $-139.3$ | $-122.8 \\pm 1.6$ | $-408.9 \\pm 3.3$ | $-134.8 \\pm 8.9$ | $-279.4 \\pm 9.4$ | $-201.8 \\pm 7.4$ | $-267.3 \\pm 20.5$ |\n\n\n>Q3. Design and capacity of task embeddings.\n\nR3: We demonstrate the expressitivity of our task embedding design with extensive experiments in our paper, showing that BATI outperforms the baselines across the board. To further support the capabilities of BATI, we conduct a new OOD task generalization experiment in AntDir, which is even more demanding for the task embedding, requiring generalization to both OOD contexts and tasks. We sample training goal directions from $[0, \\pi)$ and testing directions from $[\\pi, 2\\pi)$ to build disjoint distributions of training and testing tasks. The results are shown in the table below. It can be seen that while the performance decreases compared to the in-distribution case, BATI still outperforms baselines by a large margin and is the only method to achieve positive returns. This further validates the generalization capability of BATI and its task embedding designs.\n\n|    Method         | BATI           | CSRO             | FOCAL           | Recon            | UNICORN          |\n|-------------|----------------|------------------|-----------------|------------------|------------------|\n| Episodic Return | $12.1 \\pm 0.5$ | $-45.8 \\pm 14.5$ | $-74.7 \\pm 6.1$ | $-51.8 \\pm 37.0$ | $-104.6 \\pm 9.1$ |", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I have reviewed Theorem 3.1 proposed by the authors, but I cannot guarantee that the proof is entirely correct.", "experimental_designs_or_analyses": "The experimental design is basically reasonable", "relation_to_broader_scientific_literature": "The authors' discussion on context shift in ICRL is quite intriguing and provides new insights for future work on generalization in ICRL.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**\n\n1.\tThe authors' analysis of context shift appears to be insightful. Particularly, the example in Figure 1 aptly illustrates their motivation.\n2.\tThe analysis of the UNICORN and CSRO works theoretically explains the motivation behind BATI's focus on environmental dynamics.\n3.\tThe experimental results demonstrate that BATI's performance is competitive, especially in scenarios where the environmental dynamics noise increases.\n\n**Weaknesses**\n\n1.\tAlthough BATI avoids the impact of behavioral policies on context shift by predicting environmental dynamics, it seems to also limit the information capacity of the task inference encoder.\n2.\tDespite the experimental results are promising, the types of experimental environments are somewhat limited. In addition, the comparisons do not include baselines like Algorithm Distillation, which use Transformer as a backbone, thereby reducing the persuasiveness of BATI.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "From my perspective, methods like Algorithm Distillation handle both task inference and context-conditioned policy within a single model, such as a Transformer. Compared to these methods, what are the advantages and disadvantages of the authors' approach, which uses two separate models for learning?", "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 work proposes a new method for task identification in offline context-based meta RL, enhancing the agent's robustness to distribution shift in the test time. Experiments are conducted in MuJoCo domains, demonstrating reasonable performance improvements. Theoretical analysis are provided to complement the algorithmic design. 

3 out of 4 reviewers appreciate the contribution of the work and are in favour of acceptance. The only negative review is relatively short and the reviewer (bfNC) did not engage well in author reviewer discussion. I therefore recommend acceptance. That being said, I do agree with bfNC that the work can benefit from experiments in more challenging domains such as Meta World. The author response does provide extra experiments and I hope the authors could provide more in the camera ready version.
Decision: Accept (poster)