CERTAIN: Context Uncertainty-aware One-Shot Adaptation for Context-based Offline Meta Reinforcement Learning
Authors: Hongtu Zhou
Ruiling Yang
Yakun Zhu
Haoqi Zhao
Hai Zhang
Di Zhang
Junqiao Zhao
Chen Ye
Changjun Jiang
Venue: ICML 2025 poster (2025)
Keywords: Offline Meta Reinforcement Learning
Context-based
One-Shot Adaptation
OOD
Ambiguity
Uncertainty
Primary area: reinforcement_learning->batchoffline

Abstract: Existing context-based offline meta-reinforcement learning (COMRL) methods primarily focus on task representation learning and given-context adaptation performance. They often assume that the adaptation context is collected using task-specific behavior policies or through multiple rounds of collection. However, in real applications, the context should be collected by a policy in a one-shot manner to ensure efficiency and safety. We find that intrinsic context ambiguity across multiple tasks and out-of-distribution (OOD) issues due to distribution shift significantly affect the performance of one-shot adaptation, which has been largely overlooked in most COMRL research. To address this problem, we propose using heteroscedastic uncertainty in representation learning to identify ambiguous and OOD contexts, and train an uncertainty-aware context collecting policy for effective one-shot online adaptation. The proposed method can be integrated into various COMRL frameworks, including classifier-based, reconstrution-based and contrastive learning-based approaches. Empirical evaluations on benchmark tasks show that our method can improve one-shot adaptation performance by up to 36% and zero-shot adaptation performance by up to 34% compared to existing baseline COMRL methods.

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for Context-based Offline Meta Reinforcement Learning
Hongtu Zhou
Ruiling Yang"
Yakun Zhu'
Haoqi Zhao
Hai Zhang' Changjun Jiang'
Di Zhang'
Junqiao Zhao
Chen Ye
Abstract
ing(Mnih, 2013; Lample & Chaplot, 2017), robotic con- trol(Nguyen & La, 2019), and recommendation systems (Zheng et al., 2018). However, RL faces a critical challenge of low sample efficiency due to its reliance on extensive interactions with the environment, limitation that becomes more pronounced in multi-task settings (Levine et al., 2020; Li et al., 2020a).
Existing context-based offline meta- reinforcement learning (COMRL) methods primarily focus on task representation learning and given-context adaptation performance. They often assume that the adaptation context is collected using task-specific behavior policies through multiple rounds of collection. However, in real applications, the context should be collected by a policy in a one-shot manner to ensure efficiency and safety. find that inherent context ambiguity across multiple tasks and out-of cuttof-diisributtinn (OOD) issues due distri- bution shift significantly affect the performance of one-shot adaptation, which has been largely overlooked in most COMRL research. To address this problem, we propose using heteroscedastic uncertainty representation learning identify ambiguous and OOD contexts, and train an uncertainty-aware context collecting policy for effective one-shot online adaptation. The proposed method can be integrated into various COMRL frameworks, including classifier-based, reconstrution-basee and contrastive learning- based approaches. Empirical evaluations on benchmark tasks show that our method can improve one-shot adaptation performance by up to 36% and zero-shot adaptation performance by up to 34% compared to existing baseline COMRL. methods.
Context-based offline meta reinforcement learning (COMRL) provides a promising solution to improve sample efficiency multi-task environments by leveraging offline datasets for training and enabling fast adaptation to new tasks with minimal online samples (Li al., 2020b). COMRL treats task-specific trajectories
transitions context information and operates in two key phases: the offline meta-training and the online adaptation. In the first phase, a context encoder is learned through task representation learning, alongside a meta-policy conditioned on the learned task representations. During the second phase, the meta-policy adapts new tasks by conditioning on the task representation inferred from the online-collected context.
Despite its promise, COMRL encounters two major chal- lenges, as depicted in Figure left: context ambiguity and -istiiiiiiiittionn (OOD). Context ambiguity arises when context fails to uniquely infer task because it appears across multiple tasks (Dorfman al., 2021). Context OOD, the other hand, occurs when context falls outside the task distribution of the offline dataset. Both issues can lead the context encoder to infer unreliable task representations, resulting in suboptimal policy performance.
Most existing COMRL methods predominantly focus on task representation learning (Li et al., 2020b; Gao al., 2024: Li al.. 2024). These methods often assume that contexts either in-distribution, thereby avoiding OOD issues, can be collected through multiple rounds, which mitigates the impact of ambiguity. While CORRO (Yuan & Lu, 2022) attempts to address OOD challenges, its applica- bility is limited to (s, a) pairs observed in the dataset, and fails to generalize to truly OOD (s, a) scenarios. However, in many practical scenarios, context collection is constrained to one-shot manner due to safety and efficiency considera-
1. Introduction
Reinforcement learning (RL) has achieved remarkable suc- cess in a wide range of domains, including game play-
Eqqual contribution 'School of Computer Science and Tech- nology, Tongji University, Shanghai, China. Correspondence to: Junqiao Zhao #zhaojuuniia@ttongii.c........
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

CERTAIN: Context Uncertainty-awarc One-Shot Adaptation for COMRL
Uncertainty Estimation Network
Task1
Task2
Task3
ODD
Ambiguous
Full Transition Space
Uncertain
1
Task
Uncertainty-aware Collecting Policy
Uncertain Context Restraining
OOD
Otfine Detenet
Aoabiguons
Task
Certain
Deplay
Trajectory
Context
Task3!
goe
Taski
aak
Sampling
Task
Transition Uncertainty (Ambiguity/o0n
Offline Meta Training
Online One-shot Context Collecting
Task Inferring
Figure 1. Left: Contexts consist of certain task-specific transitions and uncertain transitions that are either ambiguous or OOD. Right: Our approach involves training uncertainty estimation network and uncertainty-awwi aware context collecting policy restrain the detrimental influence of uncertain contexts on task inference during online adaptation.
tions, making context ambiguity and OOD issues even more challenging.
2. Related Work
2.1. Offline Reinforcement Learning
To address these challenges, we propose CERTAIN: Context encertainty-awarc One-Shot Adaptation. The core insight behind CERTAIN is that both ambiguous and OOD contexts contribute to high task representation uncertainty, whereas contexts that reliably distinguish tasks exhibit lower uncer- tainty (Figure 1, left).
Offline reinforcement learning (Offline RL) aims to learn a policy from a static dataset without requiring interactions with the environment, significantly improving sample effi- ciency (Levine al., 2020). Recent research has focused on addressing key challenges such value overestimation (Kumar et al., 2020; Li et al., 2022; Singh et al., 2023; Mao et al., 2024) and distributional shift (Fujimoto et al., 2019; Kumar et al., 2019; Kostrikov et al., 2021; Brandfonbrener et al., 2021; Ran et al., 2023; Zhang Tan, 2024). Our work adheres to the offline RL setting, targeting the learning of the task representation and the meta-policy from multi-task offline datasets.
Building this observation, introduce an uncertainty estimation network during the offline meta-training phase to estimate context uncertainty alongside task representation learning (Figure 1, right). By leveraging these uncertainty estimates, CERTAIN effectively identifies ambiguous and OOD contexts, mitigating their negative impact on task infer- ence during online adaptation. To enable fast and accurate task inference, we further propose training uncertainty- aware context collecting policy that prioritizes the collection of low-uncertainty contexts during the online adaptation phase.
2.2. Meta Reinforcement Learning
Meta reinforcement learning (Meta RL) seeks to train meta-policy that generalizes across multiple tasks, enabling rapid adaptation to new tasks. Context-based methods (Duan et al., 2016; Rakelly et al., 2019; Fakoor et al., 2019; Zintgraf et al., 2019; Humplik et al., 2019; Wang et al., 2023b; Chu al., 2024) frame Meta RL as a partially ob- servable Markov decision process (POMDP), where histori- cal trajectories serve the context. These methods extract task-specific information from the context and treat it as hidden state to guide decision-making. Gradient- based methods (Finn et al., 2017; Finn & Levine, 2017; Xu et al., 2018; Rothfuss et al., 2018; Xu et al., 2020; Liu et al., 2022; Xu & Zhu, 2024) aim to learn meta-policy that provides strong initialization for task task-fpec policies. Adaptation to new tasks is achieved by fine-tuning the meta-policy us- ing meta- getaiggaadeets computed from online few-shot samples. Our work builds on the context-based Meta RL framework, which is better suited for real-world applications due to its practicality and scalability.
CERTAIN designed as plug-in framework that seam- lessly integrates into existing COMRL methods, includ- ing classifier-based, reconstruction-baste based, and contrastive learning-basse bbaeed approaches, to enhance robustness against context ambiguity and OOD issues during online adaptation. We validate CERTAIN through extensive experiments on a range of benchmark tasks. The results show that CER- TAIN significantly improves the performance of existing COMRL methods in the one-shot adaptation setting, achiev- ing up to a 36% increase in average return. Furthermore, in the zero-shot adaptation setting-where the meta-polly policy is directly applied to new tasks without any online adapta- tion CERTAIN delivers up to a 34% improvement in aver- age return, outperforming state-of-the-aa (SoTA) COMRL methods.
2

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
2.3. Offline Meta Reinforcement Learning
maximize the expected return J across tasks:
Offline meta reinforcement learning (OMRL) focuses on learning meta-policy from multi-task offline datasets to able rapid adaptation to new tasks (Li et al., 2020a). Similar to Meta RL, OMRL methods are categorized into gradient- based (Mitchell et al., 2021) and context-based (Li et al., 2020b) approaches. The key distinction is that OMRL oper- ates under an offline RL setting.
max E Mirp((M [ew-D 0,0
(4)
here, 2 is the latent task representation inferred from the context e, specifically:
K Σ
1
(0) K
(5)
Our work is grounded in context-based OMRL methods. Recent COMRL approaches (Li al., 2020b; Yuan & Lu, 2022; Gao et al., 2024) emphasize task representation learn- ing during the meta- traittttaiinnggggg phase, which can be unified under the perspective of mutual information maximization (Li et al., 2024). Other works (Gao et al., 2024; Wang et al., 2023a) address the challenge of context collection during the online adaptation phase. Our method aligns with the latter category, focusing on improving context collection for robust one-shot adaptation.
In the online adaptation phase, trajectories are collected for the unseen task M' using a context collecting policy ne to form the adaptation context d':
e' ((,)),,,,,
(M', xe)
(6)
Subsequently, the adapted policy no(als, 2/) is derived based on the inferred task representation 2' go(2)c)
3.2. Context Uncertainty in Task Inference
2.4. Uncertainty Estimation
In COMRL, the task representation is inferred from the context using context encoder (=(z))) Assuming reasonable qá. the reliability of task inference based on = heavily depends on the context c. Two major factors that lead to unreliable task inference are context ambiguity and OOD contexts. Before formally defining ambiguity and OOD, we first define the probability of context ((f Tjryy, ( o' given task M and a policy и as:
Uncertainty can be categorized into epistemic uncertainty and alcatoric uncertainty. Prior work has tackled epistemic uncertainty estimation using ensemble methods (Lakshmi- narayanan et al., 2017) and Bayesian neural networks (Gal & Ghahramani, 2016), while aleatoric uncertainty is typi- cally modeled via heteroscedasticity by parameterizing the variance of the output distribution (Kendall Gal, 2017). However, the heteroscedastic uncertainty obtained with the above method can also be affected by the model's 's capacity and the training dynamics. Building upon the heteroscedas- tic uncertainty framework, we propose an approach to esti- mate context uncertainty in the setting of COMRL.
K
p(c.M T) = II j=1
p(s)
where p(s) denotes the initial state distribution, and
PM r(s,a) is the model of the task M.
3. Problem Formulation
With this formulation, we define context ambiguity and OOD contexts as follows:
3.1. Preliminaries
Definition 3.1. Context Ambiguity: Context ambiguity occurs when a given context e corresponds at least two distinct tasks M, and M,, such that p(c].Ms) > 0 and P(cl-M)) > 0, where p(c]MM := p(c]M, T)dT.
A task in RL is formally represented Markov Decision Process (MDP) MM = (8,A, P, R), where S denotes the state space, A the action space, P the transition dynam- ics, and R the reward function. In COMRL, each task is assumed to be sampled from a task distribution:
Definition 3.2. OOD Context: A context e is considered
OOD if it cannot be sampled by the behavior policy Th within the offline dataset task i.e.e., p(c|M, TT) = 0.
M; (S,A,P., Ra) p(M)
(1)
Definition 3.3. Context Uncertainty: Contexts affected by either ambiguity or OOD issues are likely to produce unreliable task representations for task inference, which we refer to as Uncertain contexts. Conversely, a context Certain
The context e for task M₁ is defined as a set of transition tuples:
c= {(}} Kn Tj; k
Given an offline dataset D containing M tasks:
that is neither ambiguous nor OOD is defined context, capable of reliably inferring a task based on z. To quantify the reliability of a context for task inference, we define Context Uncertainty as a, which captures the degree of uncertainty in task inference for a given context.
D = (D, {(5}, ag. y,, $1) M
the objective of offline meta-training in COMRL is to learn a context encoder (s(2)c) and a meta-policy (a 2) that
3

---

CERTAIN: Context Uncertainty-awar One-Shot Adaptation for COMRL
Off shehl--seeffee representation learning
(1,0,1;)))
Tol- +0
Softmax
(s,,,,
(1
8,7
(ale)
Uncertain Context Restraining
(7 s,)
Ma 5,
Stop gradient
beett
loggo
Nu(- Context Collecting Training
(a) Uncertainty-awarc representation learning
(b) Uncertainty-aware meta-policy learning
(c) Uncertainty-aware adaptation
Figure 2. Overview of the proposed method: (a) combining off-the-shelf task representation learning with heteroscedastic uncertainty framework, (b) integrating Uncertainty in meta-policy learning and context collecting policy learning, and (c) uncertainty-awar one-shot adaptation in unseen task.
3.3. One-shot Adaptation Challenges
jectives.
In the online adaptation phase, most existing COMRL meth- ods assume that the adaptation context is collected by the task-specific behavior policy ωy for the unseen task M', following distribution similar to the offline dataset D. Al- ternatively, these methods assume that e' can be collected over multiple rounds using an exploratory policy, e.g. random policy, to obtain more certain contexts, thereby mit- igating the impact of uncertain contexts. However, under the one-shot adaptation setting, the adaptation context d' derived from a single trajectory collected by the context collecting policy Te- This restriction substantially increases the likelihood of encountering uncertain contexts, which can impair task inference based on the task representation, thereby degrading policy performance.
Classifier-based methods (Zhang et al., 2025) train the con- text encoder by predicting the task label y using classifier Po (vlz). The objective to minimize the cross-entropy loss between the predicted task label and the ground truth:
Lats ""ssaa""" [+(=))
(8)
Reconstruction -basdd methods (Li et al., 2024) train the context encoder by reconstructing the dynamics and reward model with decoder (r, s'|s, 2). The objective is to maximize the likelihood of the reconstructed dynamics and reward model:
free B(s.ana") (onaoa log 1o
(r, a,, a)
(9)
4. Method
Contrastive learning- based methods (Le-Khac al., 2020)
train the context encoder using contrastive losses, includ- ing InfoNCE (Oord et al., 201 8), FOCAL (Li et al., 2020b) etc. These objectives minimize the distance between query and positive samples while maximizing the distance be- tween query and negative samples. One commonly used contrastive loss (Li al., 2020b)
We now present our method, CERTAIN, which consists of three key components: (a) uncertainty-ave task represen- tion learning, (b) uncertainty-awart meta-policy learning, and (c) uncertainty- adaptation. Figure provides an overview of our method.
4.1. Uncertainty-aware Task Representation Learning
Lcont = 1/y yı 3i - 22
Building on the heteroscedastic uncertainty framework pro- posed in (Kendall Gal, 2017), we introduce a method to evaluate context uncertainty task representation learning by slightly modifying the heteroscedastic loss. This for- mulation allows the context uncertainty be seamlessly integrated into the three primary paradigms of task represen- tation learning in COMRL: (1) classifier-based methods, (2) ieconstruction-based methods, and (3) contrastive learning- based methods. All these methods share the context encoder (2(c) as a core component but differ their training ob-
B
(10)
+ 1 + vz)
|| - 22
To integrate context uncertainty, the above learning objec- tives are modified by incorporating uncertainty estima- tion network h (alc). As shown in Figure 2 (a), given an input context c, the task representation is inferred using the context encoder (=(c), and task representation learn- ing objective Cs is computed. Simultancously, a context uncertainty ow predicted by byp (a(e). The heteroscedastic

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
loss Loyy is then formulated by combining Ls with og:
4.2.2. CONTEXT COLLECTING POLICY LEARNING
Le + log d
In the one-shot adaptation setting, collecting an effective context for task inference is crucial. Existing methods not applicable they either assume that adaptation contexts are either collected under the same behavior policy as the offline dataset or allow for multiple collection attempts using an exploratory policy.
(11)
The context uncertainty dg is positively correlated with the task representation learning loss La: the detailed proof is provided in Appendix A. This correlation implies that v can serve as a reliable measure of context uncertainty, as a higher task representation learning loss Ls typically in- dicates that the context is more likely to be ambiguous or OOD. Moreover, the heteroscedastic loss formulation pro- vides an additional benefit: it adaptively adjusts the effective learning signal for L. allowing the model to place less em- phasis on uncertain contexts and more on confident ones, thereby improving the overall robustness of task representa- tion learning.
To address these limitations, we propose training a con- text collecting policy that actively explore low-uncertainty contexts, thereby avoiding ambiguous and OOD contexts.
As illustrated in Figure 2 (b), the context collecting pol- icy is defined as To (a)s,, 20), where 20 represents zero- initialized task representation. The policy no Ta(a) 20) is trained to maximize the action value function Qa(s, a, 20). while incorporating an uncertainty penalty. This encourages the collection of informative and reliable contexts that im- prove task inference and enhance downstream adaptation performance:
4.2. Uncertainty-aware Policy Learning
4.2. 1. META-POLICY LEARNING
[To((ass 20)) = max Blsannd)-- [rr aa)
The meta-policy no (als, z) in COMRL is conditioned on the task representation 2 to maximize the expected return across tasks. Its theoretical foundation lies in the framework of a partially observable Markov decision process (POMDP). Once the task representation is inferred from the context, the tuple 2) is treated as a fully observable state, effec- tively transforming the POMDP into a Markov decision process (MDP). This enables the meta-policy in COMRL to be optimized using standard offline RL algorithms, such as CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2021), BRAC (Wu et al., 2019), and TD3BC (Fujimoto & Gu, 2021).
+ 7Be 20)
(13)
where a a hyperparameter to balance the reward and the uncertainty penalty and 7 is the discount factor.
4.3. Uncertainty-aware One-shot Adaptation
As illustrated in Figure 2 (c), during the online adapta- tion, we employ the proposed context collecting policy no (a|s,, =o) to collect one trajectory as adaptation context e in the unseen M'':
The task representation in the meta-policy is typically derived as the mean of task representations of the context in existing COMRL methods. Nevertheless, assigning equal importance to all transitions within the context is vulnerable to ambiguous and OOD contexts. To address this limitation, we propose leveraging the uncertainty a estimated for each transition in the context c, and computing weighted sum of each the task representation. This approach reduces the influence of uncertain transitions, thereby improving the accuracy of task inference.
é' d(=((s,,,,, (1,1,, 82), (52, 02, r2, s2, (sT,
$z)..... .... (M', sa(sss,=)))
(14)
Since the ambiguous and OOD context cannot be fully pre- vented, we then utilize the context encoder 4a and the uncer- tainty estimation network hp infer the task representation z while restraining the influence of uncertain contexts in d':
As illustrated in Figure 2 (b), given a context e with n transition tuples, the context encoder (2(c) is used to extract task representations for each transition, yielding {zili 1, n,. n). The uncertainties for each transition {a.li = 1,...,n, n are then predicted by the uncertainty estimation network (a)z). These uncertainties are then transformed into weights (wy i = 1,....,"! n.} via a softmax function. Finally, zu is computed as the weighted sum of {zili = n..
Σ w
where
21 go(ze|se, de, 8141)
(15)
expl h (o
Wi
T Ei (orlze)
Finally, the adapted policy M' obtained by conditioning on the weighted task repre- sentation z'.
for the unseen task
2m = En W;
(12)
5

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
Second Test
Second Episode Test
Second Episode Test
-530
Aovespe Banum
-020
-200
Rocal
focal
cassimer
Fecal-our
-750
-200
fecaleturs
clasitirr
clasifter oara
cuss 1eus
claoffirrout
un com
urecom
uncam
com-oure
-030
urccom-our
uncem-ouuss
Mag
40
40
Stops
Stoosile31
(a)Point-Robo
(b) Ant-Gooa
(h)eett Vel
Second Episade Test
Second Episode Test
300
200
250
150
200
300
Roal
Feealooure
150
(ocal nours clockficr chooffermmiu
Feral
cansther
100
50
Caccifor-outt
unicom OUTS
out
Mao
40
60
100
Tain Stepo)le3)
(d)Walker-Fe) Rand--arand Params
HeHHopper-Rande -aaee--Rnndd-aaaaae
Figure Average return of one-shot in five environments, Fach curve represents the average return of four seeds.
5. Experiments
CLASSIFIER (classifier-basee) learns context repre- sentation by predicting task labels from the context using a cross-entropy loss.
In our experiments, we aim to demonstrate the effectiveness of CERTAIN in online one-shot adaptation. Additionally, we conduct ablation study to assess the necessity of both the uncertainty-awant context collecting policy and the uncertain context restraining mechanism.
FOCAL (contrastive learning- -baeed) learns context rep- resentations using a contrastive metric loss.
UNICORN (reconstruction-baas extends FOCAL by incorporating reconstruction loss.
5.1. Setup
Environments. We evaluate CERTAIN on one toy environ- ment, Point-Robot, and four MuJoCo environments: Ant- Goal, Half-Cheetah-Vel. Walker-Rand-Panee Params, and Hopper- Rand-Params.
Additionally, we compare with IDAQ (Wang al., 2023a), which builds upon FOCAL and explicitly consider the OOD problem in online adaptation. The complete experimental configurations are summarized Appendix B.
The Point-Robot environment consists of an agent starting (0, 0) and goal positioned on semicircle with radius of 1. Tasks are distinguished by different goal positions. This environment enables intuitive visualization of agent tra- jectories and facilitates analysis of our method's underlying mechanism.
Evaluation Metries. We evaluate CERTAIN in the one- shot online adaptation setting, Each evaluation consists of collecting two episode trajectories:
First episode (context collection phase): The trajectory is collected using the context collecting policy, serving as the one-shot context.
The MuJoCo environments present more complex and chal- lenging scenarios. Ant-Goal and Half-Cheetah- Vel, tasks distinguished by different reward functions. In Walker- Rand-Params and Hopper- Ranner--and-Paaaas tasks vary based on differences in environment dynamics parameters.
Second episode (adaptation phase): The trajectory is collected using the adapted policy, conditioned on the task representation inferred from the first episode's context.
Baselines. We integrate CERTAIN with three baseline meth- ods: CLASSIFIER (Zhang et al., 2025), FOCAL (Li et al., 2020b), and UNICORN (Li et al., 2024).
For each evaluation, we run experiments across four random seeds and report the average return from the second episode as the final performance metric.
6

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
First Episode Test
First Episode
First Episode
150
630
-200
Aovesppe Bantrr
-250
-750
Rocal
Auveape
foo
focal
freal oure clasiiiie
fecaleturs
-820
cussifer
-250
clusilirr
clasifie tara
closifir
1945
chsoffirro.u rours
un com
uricom
-000
uncam
com
-910
unccmm-uur
uncemmeurr
Ide
-50
30
40
100
22
60
Stopsilee)
Stepor3e3)
Stoosile31
(a)Point-Robo
(b) Ant-Goal
(c)eheeah- Vel
First Episode Test
First Episode Test
250
150
Average Betum
115
200
Feral
Inesl+ oure
100
locall our clockfiir cheollierouus
canstter
Cascffor-ourt
unicom
unicom oUTS
out
Mao
0
40
60
100
Tain Stepogle3)
(d)Walker-Ree -aaleer---n----ee Rand-Params
(e) Rand-Params
Figure 4. Average return of zero-shot in five environments.
5.2. Main Results
Moreover, observe that the zero-shot performance of baseline methods tends to degrade as training progresses. We hypothesize that this is because the policy TØ(s, z0) in the baselines becomes increasingly conservative over time. To support this claim, we provide an interpretable visualization of this performance decline in the Point-Robot environment in Appendix C.
5.2.1. COMPARISON WITH BASELINES
We incorporate CERTAIN into three baseline methods and evaluate their one-shot adaptation performance against the original baselines. As shown in Figure 3, CERTAIN consis- tently matches or outperforms the baseline methods across all enviromments. Notably, CERTAIN+FOCAL consistently outperforms IDAQ (Wang et al., 2023a). The most signifi- cant improvement is observed in the Point-Robot environ- achieves 36.9%
Additional results on online adaptation in training tasks are presented in Appendix D, and results using a fixed offline context are provided in Appendix E. While our primary fo- cus is on one-shot adaptation- the most challenging case in the few-shot setting also evaluate CERTAIN in a few- shot setting, where the agent can collect multiple context trajectories and update the task representation z after each episode. As shown in Appendix F, CERTAIN consistently outperforms the baselines, demonstrating its effectiveness in few-shot adaptation as well.
ment, where CERTAIN+CLASSI performance improvement over the original CLASSIFIER.
Additionally, we evaluate zero-shot performance based on the return of the first episode trajectory, as shown in Figure 4. The results demonstrate that CERTAIN tends to outperform baseline methods in most environments. We attribute this to CERRTIN'S ss context collecting policy, which gathers an initial trajectory with lower uncertainty, often resulting in higher returns.
5.2.2. VISUALIZATION
Interestingly, FOCAL achieves higher zero-shot returns in the Point-Robot and Ant-Goal environments. We hypothe- size that this is due to contrastive learning enabling FOCAL to learn a more balanced initial latent representation 2u, al- lowing the policy (8 =0) explore diverse goals more effectively. In contrast, CERTAIN's 2n policy prioritizes collecting low-uncertainty transitions, which does not neces- sarily correlate with higher returns in the sero-shot setting,
To better understand the rationale behind CERTAIN, we visualize the uncertainty of each context transition in the Point-Robot offline dataset. As shown in the left of Figure 5, transition uncertainty is notably higher near the starting point (shaded in red), where transitions are likely to oc- cur across multiple tasks. This overlap leads to context ambiguity, making it challenging for the model to reliably distinguish between different tasks. In contrast, transitions farther from the starting point exhibit more pronounced re-
7

---

CERTAIN: Context Uncertainty-awan One-Shot Adaptation for COMRL
2.2
0.4
1986
(a) CLASSIFIER
FOCAL
UNICORN
Figure 5. Context Uncertainty of the Robot offline dataset. The left shows the uncertainty distribution across the state space, while the right presents the histogram of uncertainty values.
ward differences, facilitating task identification and resulting in lower uncertainty (shaded in blue). To characterize the uncertainty distribution, we present histogram of all transi- tions on the right of Figure 5, showing that most transitions have low uncertainty, while highly uncertain transitions be- come progressively less frequent.
(d) CLASSIFIER+
(e) FOCAL+
(I) UNICORN+
Figure 7. Visualization of two episode trajectories in the Point- Robot during evaluation. '+' denotes CERTAIN.
the goal is likely located, whereas baseline methods tend to explore around the starting point. In the second episode (orange). CERTAIN's trajectory more direct and accu- rate compared to the baselines. Moreover, the uncertainties along CERTAIN's trajectories are lower than those of the baseline methods, suggesting that CERTAIN effectively re- duces context uncertainty.
Ood Taak
5.3. Ablation Study
To assess the necessity of the uncertainty-awart collecting policy and the uncertain context restraining mechanism in CERTAIN, we conduct ablation studies on CLASSIFIER- CERTAIN. For the ablation of the uncertainty-an lecting policy, we compare CERTAIN with two variants: VERTAIN-Randon and CERTAIN-Naives
0.2
Figure 6. Context Uncertainty of OOD context and OOD task the Roint-Robot. The left shows the OOD context under in- distribution tasks, while the right presents the context uncertainty under OOD-task collected with behavior policy.
col-
CERTAIN- Random collects the adaptation context trajectory using random policy, whereas CERTAIN-Naive-Cn employs the same collecting policy as CERTAIN but without training during the meta-training phase. For the ablation of the uncertain context restraining mechanism, we introduce an- other variant, CERTAIN- which computes the mean of the representations inferred from the context trajectory. As shown in Table 1, CERTAIN consistently outperforms CERTAIN-Randone CERTAIN-Naive-a and CERTAIN- Mean across all environments. These results highlight the necessity of both the uncertainty-aware collecting policy and the uncertain context restraining mechanism CERTAIN. The ablation study for FOCAL-CERTAIN and UNICORN- CERTAIN is provided Appendix G.I.
We evaluate CERTAIN's 's OOD- awareness by visualizing context uncertainty in the Point-Robot under two settings: OOD context and OOD task. In the OOD context setting, tasks remain the same as in training, but transitions are col- lected from the bottom semicircle, while the offline dataset comes from the top. In the OOD task setting, transitions are collected from the top semicircle using the same behav- ior policy in offline dataset, but with goals located in the bottom comicircle-unsedd during training. As shown in Figure 6, both OOD settings exhibit significantly higher con- text uncertainty than the in-distribution case, highlighting CERTAIN's ability to detect distribution shifts-cruce for robust one-shot adaptation.
To evaluate the impact of the uncertainty penalty coeffi- cient in Equation (13), we conduct an ablation study. We observe no clear correlation between performance and the value of a, and hypothesize that the penalty should be scaled to match the magnitudes of both the reward and context un- certainty. Detailed results are provided in Appendix G.2.
To further demonstrate the effectiveness of CERTAIN, we visualize the online trajectories over two episodes for dif- ferent methods in the Point-Robot environment. As shown in Figure 7, the first episode trajectory (green) collected by CERTAIN moves directly toward the semicircle where
8

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
Table Average return and standard deviation of the second episode trajectory collected with CLASSIFIER-CFRTAN and its variants.
Env
CERTAIN
Naive-
Random
Mean
Point-Robot
-9.27 = 0.69
-13.30 + 1.00
-10.04 = 0.78
-9.89 + 0.35
Ant-Goal
-539.37 = 40.83
-576.40 + 52.04
-583.9 = 72.63
-555.67 + 68.17
Cheetah-Vel
-110.02 15.21
-173.90 + 7.34
-126.23 = 15.63
-133.31 + 33.38
Walker
144.40 + 26.19
105.66 + 13.26
100.99 + 44.19
113.16 + 37.07
Hopper
223.45 + 28.23
202.22 + 13.64
213.47 + 20.01
206.6 + 36.15
We also study the effect of offline dataset quality on CER- TAIN's performance in the Point-Robot environment. Re- sults show that performance improves most when the dataset of medium quality. We hypothesize that extreme dataset qualities impair uncertainty estimation, resulting in subopti- mal performance. Full results are presented in Appendix H.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
6. Conclusion and Limitations
Brandfonbrener, D., Whitney, W., Ranganath, R., and Bruna, J. Offline without off-policy evaluation. Advances in neural information processing systems, 14493344466, 2021.
In this paper, we address the challenges of context ambigu- ity and OOD issues in one-shot adaptation for COMRL. By leveraging context uncertainty learning, we effectively iden- tify ambiguous and OOD contexts. To further enhance task inference accuracy during one-shot online adaptation, context collecting policy and an uncertain context restraining mechanism. Experimental results demonstrate that our method outperforms baseline approaches both task inference accuracy and its ability to handle ambiguous and OOD contexts.
Chu, Z., Cai, R., and Wang, H. Meta-reinnforcemen learning via exploratory task clustering. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 11633-11641, 2024.
design uncertainty-awee
Dorfman, R., Shenfeld, I., and Tamar, A. Offline meta reinforcement beaming-identifiabiity challenges and ef- fective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-4618, 2021.
Despite its effectiveness, our method has certain limitations. First, it does address the issue of spurious correlations, common challenge in COMRL. However, it can be in- tegrated with existing approaches, such relabeling the context with task-specific models, to mitigate this problem. Second, while our method exhibits zero-shot adaptation ca- pabilities, we attribute this primarily to the fact that our context collecting policy tends to gather low-uncertainty context trajectories, which often yield higher returns. Third, the context uncertainty estimated by method can capture the presence of both OOD and ambiguous contexts, but it cannot explicitly distinguish between them. In future work, we plan to investigate techniques for disentangling OOD and ambiguous contexts and explore integrating our method with approaches designed handle spurious correlations.
Duan, Y., Schulman, J., Chen, X., Bartlett, L., Sutskever, I., and Abbeel, ing via slow reinforcement learning. arXiv preprint arXiv: 161 1.02779, 2016.
R12: Fast reinforcement learn-
Fakoor, R., Chaudhari, P., Soatto, S., and Smola, A. J. Meta- g-lcarning. arXiv preprint 17000:191..00125. 2019.
Finn, C. and Levine, S. Meta-learning and universal- ity: Deep representations and gradient descent can approximate any learning algorithm. arXiv preprint 2017.
Finn, C., Abbeel, P., and Levine, S. Model-agnostie meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
Acknowledgments
The work is supported by the National Key Research and Development Program of China (No. 2020YFA071 1402).
Fujimoto, S. and Gu, S. S. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20145 2021.

---

CERTAIN: Context Uncertainty-awarc One-Shot Adaptation for COMRL
Fujimoto, S., Meger, D., and Precup, D. Off-policy deep reinforcement learning without exploration. In Interna- tional conference machine learning, pp. 2052-2062. PMLR. 2019.
Li, J., Tang, C., Tomizuka, M., and Zhan, W. Dealing with the unknown: Pessimistic offline reinforcement learning. In Conference on Robot Learning, pp. 1455-1464. PMLR, 2022.
Gal, Y. and Ghahramani, Dropout as a bayesian approx- imation: Representing model uncertainty in deep learn- ing. In international conference machine learning, pp- 1050-1059. PMLR, 2016.
Li, L., Yang, R., and Luo, D. Focal: Efficient fully- offline meta-reenforcemeett learning via distance met- ric learning and behavior regularization. arXiv preprint arXiv:2010.011 100000000111.... 2020b.
Gao, Y., Zhang, R., Guo, J., Wu, F., Yi, Q.. Peng. S., Lan, S., Chen, R., Du, Z., Hu, X., et al. Context shift reduction for offline meta-reeinforeement learning. Advances in Neural Information Processing Systems, 36, 2024.
Li, L., Zhang, H., Zhang, X., Zhu, S., Zhao, J., and Heng. P.-A. Towards an information theoretic framework of context-based offline meta-reinforcement learning. arXiv preprint arXiv:2402.02433) 2024.
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 2019.
Liu, B., Feng, X., Ren, J., Mai, L.,, Zhu, R., Zhang, H., Wang, J., and Yang. Y. A theoretical understanding of gradient bias in meta-reinforcement learning. Advances Neural Information Processing Systems, 35:31059-31072, 2022.
Kendall, A. and Gal, Y. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Mao, Y., Zhang, H., Chen, C., Xu, Y., and Ji, X. Supported value regularization for offline reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Kostrikov, Nair, A., and Levine, S. Offline reinforce- ment learning with implicit q-learning. arXiv preprint arXiv:21 10.06/69, 2021.
Mitchell, E., Rafailov, R., Peng. X. B., Levine, S., and Finn, C. Offline meta-reinforcement learning with advan- tage weighting. In International Conference on Machine Learning, pp. 7780-7791, PMLR, 2021.
Kumar, A., Fu, J., Soh, M., Tucker, G., and Levine, S. Stabilizing ooo-po q-learning via bootstrapping error reduction. Advances in neural information processing systems, 32, 2019.
Mnih, V. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5600 2013.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33: 1179-1191, 2020.
Nguyen, H. and La, H. Review of deep reinforcement learning for robot manipulation. In 2019 Third IEEE international conference on robotic computing (IRC), pp. 590-595. IEEE, 2019.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learn- ing with contrastive predictive coding. arXiv preprint arXiv: 807.03778..0. 2018
Lample, G. and Chaplot, D. S. Playing fps games with deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
Rakelly, K., Zhou, A., Finn, C., Levine, S., and Quillen, D. Efficient off off-policy meta-reinforcemmnt learning via prob- abilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
Le-Khac, H., Healy, G., and Smeaton, A. F. Contrastive representation learning: A framework and review. leee Access, 1939944 2020.
Ran, Y., Li, Y.-C., Zhang, F., Zhang, Z., and Yu, Y. Policy regularization with dataset constraint for offline reinforce- ment learning. In International Conference on Machine Learning, pp. 28701-28717. PMLR, 2023.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0/44 00600. 2020.
Rothfuss, J., Lee, D., Clavera, I., Asfour, T., and Abbeel, P. Promp: Proximal meta-policy search. arXiv preprint arXiv: 1810.06784, 2018.
Li, J., Vuong, Q., Liu, S., Liu, M., Ciosek, K., Christensen, H., and Su, H. Multi-task batch reinforcement learning with metric learning. Advances in neural information processing systems, 33:6197-6211 2020a.
Singh, A., Kumar, A., Vuong, Q. Chebotar, Y., and Levine, Reds: offline reinforcement learning with
10

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRI
heteroskedastic datasets via support constraints. In Pro- ceedings of the 37th International Conference Neural Information Processing Systems, pp. 23921-23933, 2023.
for bayes-adaptive deep rl via meta-learning. arXiv preprint arXiv: 2019.
Wang, J., Zhang. J., Jiang, H., Zhang. J., Wang, L.,, and Zhang, C. Offline meta reinforcement learning with in- distribution online adaptation. In International Confer- Machine Learning, pp. 36626-36669. PMLR, 2023a.
Wang, M., Bing, Z., Yao, X., Wang, S., Kai, H., Su, H., Yang, C., and Knoll, A. Meta-reinforremment learning based self-supervised task representation learning. In Proceedings of the AAAI Conference on Artificial Inteli- gence, volume 37, pp. 10157-10165, 2023b.
Wu, Y., Tucker, G., and Nachum, O. Behavior regu- larized offline reinforcement learning. arXiv preprint 19X1v 19t.:/991.111111,111111111111111 2019.
Xu, S. and Zhu, M. Meta-reinforcement learning with universal policy adaptation: Provable near-optimality under all-task optimum comparator. arXiv preprint
arXiv:2410.09728, 2024.
Xu, Z., van Hasselt, H. P., and Silver, D. Meta-gradient reinforcement learning. Advances neural information processing systems, 31, 2018.
Xu, Z., van Hasselt, H. P., Hessel, M., Oh, J., Singh, S., and Silver, D. Meta-gradient reinforcement learning with an objective discovered online. Advances in Neural Infor- mation Processing Systems, 33:15254-15524 2020.
Yuan, H. and Lu, Z. Robust task representations for of- fline meta-reinforcement learning via contrastive learning. In International Conference on Machine Learning, pp. 25747-25759. PMLR, 2022.
Zhang, H., Zheng, B., Ji, T., Liu, J., Guo, A., Zhao, J., and Li, L. Scrutinize what we ignore: Reining in task rep- resentation shift of context-based offline meta reinforce- ment learning. The Thirteenth International Confer- on Learning Representations, 2025. URL https: // /oppnneviiew..c../oo net httpi/ewwweee/ooreeeeeeeeeee forum www.iidrrriiii@@....
Zhang, and Tan, X. An implicit trust region approach to behavior regularized offline reinforcement learning. In Proceedings of the AAAI Conference Artificial Intelli- gence, volume 38, pp. 16944-16952, 2024.
Zheng, G., Zhang, F., Zheng, 7.., Xiang, Y., Yuan, N. J., Xie, X., and Li, Z. Drn: A deep reinforcement learning framework for news recommendation. In Proceedings of the 2018 world wide web conference, pp. 167-176, 2018.
Zintgraf, L.. Shiarlis, K., Igl, M.. Schulze, S., Gal, Y., Hof- mann, K., and Whiteson, Varibad: very good method

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRI
A. The Proof of Context Uncertainty
The modified heteroscedastic loss is defined as:
Le
(o.v) d + logo
(16)
where Ls denotes the task representation learning loss which is non non-negative, and dΦ represents the context uncertainty predicted by the network.
Since the neural network outputs the log-variance Og = log Jih we can reparameterize the loss as:
L + 02
(17)
exp(20g)
We then compute the first- and second-order derivatives of the loss with respect og!
aL (
2L.
(18)
Dog
exp(20
0 (o.v)
4L0
2 0
(19)
do
exp(200)
Since L& 2 0, the second derivative is non-negative, implying that the loss is convex with respect to og. Setting the first derivative to zero yields the minimizer o7:
o log(2Tbp)
(20)
and consequently, the optimal context uncertainty is:
2T
(21)
This derivation demonstrates that the optimal context uncertainty o is positively correlated with the task representation learning loss Lg. thus justifying the use of σw as proxy for context uncertainty.
12

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRI
B. Experimental Details
Table 2. Contigurations and hyper-parameters used in the training process.
Configurations
Point-Robot
Ant-Goal
Cheetah Vel
Walker-Rand Rand-PPram
Hopper- Rand-Param
dataset size training steps test time context size
1c3
1c6
60k
200k
trajectory (20 steps)
trajectory (200 steps)
unicorn weight = punish (eight(focal) punish weight(classifier) punish (nigho()
0.5
0.1
0.5
0.5
10
0.1
5
0.1
10
5
10
25
10
2
2
2
context batch size
1024
512
100
256
RL batch size
256
task representation dimension learning rate
20
3e-4
RL network width
256
RI. network depth
encoder width
200
encoder depth
3
uncertainty network width uncertainty network depth
200
13

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
C. Explanation of Baseline Decline in Figure 4(a) with Training Steps
To further investigate generalization under the zero-shot condition, we examine how baseline performance evolves as training progresses. The following figures visualize the average return of FOCAL, CLASSIFIER, and UNICORN at various training steps (x10%). In each subfigure, the trajectory labeled "Episode Line I" corresponds to zero-shot evaluation. We observe that the average return of baselines tends to decline as training continues, suggesting possible overfitting or memorization of training context.
Beten Priceet I2A, Rrloen Doboel AT
ee Tyncud 132 Artoe Dolond
10
22
14
144
(a) Focal- - Step O
(b) Focal Step 10
(c) Focal - Step 20
(d) Focal Step 40
Ioaer of Eptedde b. 101, Naure 2
Ia feiedde L... Milre of Saladd
Iatar of Eytide t. 1800 Motce 2atuade
Tatan Saaicde ITA Milurs of Splaado
Urcotarn
Urcontants
Uncortants
Urcontants
(e) Classifier- - Step 0
(f) Classifier- Step 10
(g) Classifier - Step 20
(h) Classifier - Step 40
Houre Lo to 4.sk save ot AELY
Iotrr craaaae a couu matanuu splocde
Piturn crupedde cna retarial spleode
Prttrn cripedee a con ratam
Upccertartyy
Uecertaitty
(i) Unicorn- Step 0
Unicomn- Step 10
(k) Unicorn- - Step 20
(I) Unicorn- Step 40
Figure& 8. Average return of baseline methods under shot conditions at different training steps (x10"). Each trajectory labeled "Episode Line 1" corresponds to the zero-shot rollout.
14

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
D. Online Adaptation Results on Training Tasks
D.I. One-shot Results
We present the online adaptation results on the training tasks in Figure 9. outperforms the baselines across almost all the environments, mirroring its performance on the testing tasks (Figure 3).
The results indicate that CERTAIN consistently
Second Eptsode
Second Episade
Second Episode
Aecope toum
-550
Sereous
1490
Tra Suppatkud
Trin Stepar
Trin Stopar230))
(a) Point-Robot
(b) Ant-Goal
(c) Cheetah-Vel
Second Ephsode Train
Second Episode Train
375
335
75
25
333
Thain Supplia31
(d) Walker-Ran Land-Params
(c) Hopper- Rand-Params
Figure 9. Average return of one-shot in five environments on training tasks. Each curve represents the average return of four seeds.
D.2. Zero-shot Results
We also present the zero-shot results on the training tasks in Figure 10, which are consistent with those observed on the testing tasks in Figure 4.
First Eplcode Pain
First Eplsode Train
First Episade Train
-800
-2.0
Averge fatmm 8
-250
Roce
candier
-320
for
unkca
unkamesus
unkamssas
-50
300
(a) Point-Robot
(b) Ant-Goal
(c) Cheetah-VV Vel
First Epicode
First Episade Train
303
Ioc
clatteer
untariouu 504
Taai Stpp
Trah Stope
(d) Birkker-RandPPaaamen
(c) Hopper-Rand Caapper-andd-Paane Params
Figure 10. Average return of zero-shot in five environments on training tasks.
15

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
E. Adaptation Results with Given Offline Context
Similar to other COMRL methods, we present adaptation results with given offline context Figure and Figure 12.
E.1. Results on Training Tasks
Office Return Thin
Ofline Return Train
Ofline Return Train
-000
-200
Ceatteer
for
urkam
-40
10
Stopsclest
43
CG
168
40
C6
100
Point-Rohot
(h) Ant-Goal
(c) Cheetah-Vel
Offine Return Train
Di ine Return Thain
300
430
2:3
case byr
LEcom--u
unicmmau
46
503
(d) Walker-Rand RRand-Paramss
(c) Hlopper-aand-Prrmmee
Figure 11. Average return on the training tasks with given offline context in five environments. Fach curve represents the average return of four seeds.
F.2. Results on Testing Tasks
Offinne Return Test
Offinne Retur Test
Office Return Test
-400
-300
-300
-02
unkkam
-14
10
309
23
40
CD
300
Tunn Staps Daid
(a) Point-Robot
(b) Ant-Goal
(c) Cheetah-Vel
Offinne Retur Test
Office Retur Test
300
230
Acsrage Mahm
cenlie
unikam
4B
153
(d) Walker-Rand
(c) Hooper-Rannd-Pante
Figure 12. Average return testing tasks with given offfine context in five environments.
16

---

CERTAIN: Context Uncertainty-awarc One-Shot Adaptation for COMRL
F. Few-shot Performance
In the main paper, a one-shot setting is employed; here we supplement results under a few-shot setting. Under few-shot setting, we randomly sample 5 episodes during evaluation. The table below presents the average return (= standard deviation) for each episode, with the best return in each episode highlighted in bold.
Table Few-shot average return and standard deviation across five sampled episodes in the Point-Robot environment. Bold values indicate the best performance per episode across all methods.
Episode
FOCAL
FOCAL+Ours
CLASSIFIER
CLASSIFIER+Ours
UNICORN
UNICORN+Ours
Episode
16.744+0.08
-17.40-0.28
-19.13+0.26
-17.66++0.1
-19.49+40.47
-18.52+1.73
Episode
17.82:10.33
111.97=1.11
-1.251
-8.92.0.81
-11.46::..07
-11.00+2.17
Episode
-18.13-10.24
1.2.34
-10.79
-9.3
-13.26::1.95
-9.53+L0.79
Episode
-18.3110.26
1.0.36=1.05
-10.69
-9.6440.66
-14.66+2.19
-9.94
Episode
-18.48-0.23
-12.35=0.98
-13.14+1001
-9.69+1..00
-15.07
-10.53+0.61
G. Ablation Study
G.1. Ablation Results on FOCAL-CERTAIN and UNICORN-CERTAIN
In this section, we provide additional ablation results FOCAL-CERTAIN and UNICORN-CHRTAIS We compare the and UNICORN-CERTAN with two variants: Naive-2g and Random. For detailed explanations, please refer Section 5.3.
performance
Table 4, Average return and standard deviation of the sccond episode trajectory collected with FOCAL-CERTAIN and its variants in the point-robot environment.
Env
CERTAIN
Naive-2u
Random
Mean
Point-Roue
-11.20 0.77
-16.38 0.54
-12.52 1.28
-13.80 + 2.63
Ant-Goal
-432.83 12.98
-531.40 31.06
-464.78 23.65
-459.33 +'2 6.08
Cheetah- Vel
-117.14 = 31.13
-137.00 + 14.40
-90.37 = 16.54
-146.70 + 43.84
Walker
139.98 + 16.44
104.71 + 18.89
112.27 + 22.70
136.65 34.17
Hopper
231.2 22.51
234.45 + 29.55
186.20 14.79
192.70 + 9.90
Table 5. Average return and standard deviation of the second episode trajectory collected with UNICORN-CERTAIN and its variants in the point-robot environment.
Env
CERTAIN
Naive-s0
Random
Mean
Point-Robot
-10.06 1.27
-15.98 2.20
-15.81 1.20
-15.98 + 2.20
Ant-Goal
-637.00 66.39
-618.93 59.30
-630.73
25.91
-660.45 + 29.03
Cheetah- Vel
-139.10 39.70
-14..38 + 16.89
-117.02 18.78
-169.53 + 22.23
Walker
186.35 + 37.75
129.91 + 60.48
148.63 + 25.42
109.63 37.53
Hopper
188.93 + 10.11
197.83 + 21.95
148.57 = 4.38
225.47 13.51
17

---

CERTAIN: Context Uncertainty-aware One-Shot Adaptation for COMRL
G.2. Impact of Uncertainty Penalty on Performance
In this section, we provide ablation results for the hyperparameter of uncertainty penalty a, comparing the results for 1,
Table Average return and standard deviation of the second episode trajectory under different methods and uncertainty penalty settings
Env
Classifier
Focal
Unioom
Pain--Robet
-9.53-2-194
-0 -0.58220.70
-9.27=00.99
-12.54+11.71
-154745455
-11.200:0.77
-13.09:--483
-13,07=5.69
-9482 1.35
An--Goal
-527.4434.86
-519.68--488GG
-330.6=240080
-451,23++18.9
-500.28--:43.77
(889-13-661...
465.48-34.17
Cheesah Vel
1131.184-33.15
-164.33-4-99166
-110.02-15.22 15.21
-114.98 - 40.86
-14..744:55800
-106.9-18.94
-139.10-4-39.70
-175.95 51
-195 185-55-4-40:24
Walker Hupper
1:3001
(06.02::41.56
83:23-4:1..22
100.32:1:25.87
186.23-1:39.77
132.95 16.20
168.32-1:4.76
117.5] 2:27772
91.77-1.4.44
232.48:125.77
19339
223:45:1:28.23
231.2:1:2251
194.30:1 101.60
194.77111299
174.17-156.26
183.991180.66
(4..03:555.8
Table 7. Average return and standard deviation of the first episode trajectory under different methods and uncertainty penalty settings
Emn
Classifee
Focal
Unicurn
Point-Robat
-17.6::::91
-17.63:
-17.68:20.49
-17.32=0.13 461.700=38.44
-17.70:20.50
-11.57::0.6
-17.5M=20566
-17.05:1.322
1.5544:1.57
Ass-Goal
38-53384533385
-639.47:--0077 20.72
462.135220.72
-671.254:6699
-615.525-+ 16.73
.004.7::3.49
-08.4=61.56
Cheetab Vel
-183.68---400.6
-173:73-- 19.21
-172.05 17.91
-1807884:2000
-186.88-4-17.26
-210.13 17.55
-18183-
219.18--31.25
Walker
131.65 9:3669
288.75
122.04:6.3.397
101.57-1:99.4
185.08 1-499-
137,25 :5:445
179.73/:61.36
148.0125 4:4978
37.74
Hooper
200.28:113.20
219.85::11.26
199.95119.89
14.1
165 165-222:50000
180.03:18.25
140,7:154.38
166.3:1:44.2
34.05-133.00
H. Performance Results under Low, Mid, and Expert Offline Data
While the main paper adopts the mix offline dataset, this section complements the results by evaluating our method under low, mid, and expert quality offline datasets in the Point-Robol environment.
We visualize the zero-shot (left) and one-shot (right) performance across different offline data quality levels. The solid lines represent our method, and the dashed lines denote the corresponding baselines. All results evaluated every 1000 training steps.
Data Quality Definitions:
Low-quality data contains high levels of noise and inaccuracies, often due to environment randomness or poor dynamics modeling, leading to unstable training and higher uncertainty.
Mid-quality data contains moderate noise and some useful information for generalization, providing a balanced trade-off between reliability and diversity.
Expert-level data is high-quality data with minimal noise, curated by expert policies. It is reliable but may contain spurious correlations due lack of exploration diversity.
H.1. Low Offline Dataset
Our method improves over FOCAL and UNICORN baselines, indicating robustness to noisy data. However, it slightly underperforms the CLASSIFIER baseline, suggesting that noise and incompleteness can hinder effectiveness. This highlights the need for additional noise handling strategies in future work.
H.2. Mid Offfine Dataset
Across all baselines, our method consistently improves second episode performance, demonstrating its effectiveness in moderately reliable settings and achieving ofate-of-the-art results under mid-quality data.
H.3. Expert Offfine Dataset
Despite high data quality, both baselines and our method perform worse than on mid or mix datasets. Visualizations suggest that expert data may introduce spurious correlations due to limited exploration, which hinders robust task inference.
18

---

CERTAIN: Context Uncertainty-awar One-Shot Adaptation for COMRL
First Episode Test
Second Episode Test
-16
-7.
-18
-10.0
Averagi Return
-1255
-15.0
Tocal
-17.
Tocal
-24
classifier
-00.0
classifier
tocall
classifiers
-22.5
classfter
unicom
unicorn
unicom++ours
-25.0
unicorm+ours
-28
20
30
40
50
30
50
60
Train Steps(1e3)
Train Steps(le3)
(a) Zero-shot on Low-quality Dataset
(b) One-shot on Low-quality Dataset
Figure 13. Performance under low-quality offline data.
First Episode Test
Second Episode Test
-16
-8
-17
Average Return
Average
ours
-14
focal+
classifier
classifler
clessifier+ouurs
classifier +ours
-21
unicom
unicern
unicorn+
-22
60
Steps(1e3)
Steps(1e3)
(a) Zero-shot on Mid-quality Dataset
(b) One-shot on Mid-quality Dataset
Figure Performance under mid-quality offline data.
First Episode Test
Second Episode Test
-1500
-75
-17.5
-20.0
Average Retur
-22.5
-15.0
Average
-25.0
Tocal
focal
-27.5
focal+ours ours classifier
-17.5
focal+ours
classifier
-30.0
classfier+ ours
-20.0
classiter+
unicorn
-325
unicom+our
unicorn+ours
10
20
30
40
50
60
10
20
30
40
50
60
Train Steas(1e3)
Train Steps/1e3)
(a) shro on Expert-quality Dataset
(b) One-shot Expert-qualip Dataset
Figure 15. Performance under expert-level offline data.
19


Reviews:
[{"review_id": "tGaKXKW9HZ", "reviewer": "Reviewer_NvXZ", "summary": "This paper deals with the context ambiguity problem in context-based offline meta-reinforcement learning. The authors propose an uncertainty-aware context-collection algorithm to produce in-distribution, unambiguous contexts using heteroscedastic uncertainty estimates as rewards. Experiments are conducted in several MuJoCo environments.\n\n## Update after rebuttal\n\nThis paper introduces an interesting idea into the field of COMRL and has the potential to make significant contributions, but probably needs more work regarding theoretical justifications and empirical validations. The authors are encouraged to revise the paper in these aspects.", "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": "We sincerely appreciate the reviewers’ constructive feedback. Below, we respond to each concern point by point. We will incorporate all the reviewers’ suggestions in the final version of the paper. \n\n> Reviewer: Definitions 3.1-3.3 are not very clear ... How is $\\sigma$ defined and computed?\n\nResponse: $p(c|M_i)$ is defined as the probability of context $c$ occurring under the $i_{th}$ MDP $M_i$. The equation $p(c|M,\\pi_\\beta)=0$ holds for all tasks in the offline dataset. We refer to all collection strategies used during the dataset collection process as the behavior policy $\\pi_\\beta$. It can be a mixture of multiple policies, such as human experts, other agents, etc. We define the degree of context uncertainty as $\\sigma$, which is a value greater than 0, calculated by a heteroscedastic network with the loss $\\frac{L}{\\sigma^2}+\\log \\sigma$. It should be noted that $\\sigma$ captured the loss $L$ for the training dataset and, therefore, not only represents the heteroscedastic uncertainty but also the confidence of a specific encoder model (a.k.a the epistemic uncertainty).\n\n> Reviewer: The uncertainty estimator $h_\\psi$ may struggle to produce correct estimates for OOD inputs ... more evidence in this scenario is needed to support these claims.\n\nResponse: Yes, we acknowledge that $h_\\psi$ does struggle to produce correct estimates for OOD inputs in theory. However, in many cases, it generalizes well to OOD context. We have added experiments in the point robot environment under OOD context conditions, where the context used for inferring task representations consists of 25% in-distribution samples (upper semicircle) and 75%  OOD samples (lower semicircle) at <https://anonymous.4open.science/r/CERTAIN-6073/experiment1.md>.\n\n> Reviewer: Furthermore, CERTAIN chooses to use a deterministic context encoder with an additional network for uncertainty estimation, while a more natural choice could be to directly use a probabilistic context encoder for capturing this uncertainty.\n\nResponse: A probabilistic context encoder can indeed capture uncertainty within the training dataset, but it is typically restricted to VAE-based methods. In contrast, heteroscedastic uncertainty captures the magnitude of the loss, making it adaptable to classification and contrastive learning. This allows for a more generalizable representation of context uncertainty.\n\n> Reviewer: \n> 1. The empirical performance improvement of CERTAIN seems a bit marginal. In Fig. 3 ...\n> 2. What is the context collection policy in the first episode ...\n> 3. (Minor) Both figures are a bit hard to read ...\n\nResponse: \n1. Since our method does not directly improve the policy but provides a mechanism for collecting and weighting contexts during adaptation, it is reasonable that performance does not significantly improve when the uncertainty in the contexts is already low. To clarify, we have provided performance improvement percentages in tabular <https://anonymous.4open.science/r/CERTAIN-6073/experiment6.md>. The One-shot experiment in Figure 3 and the Zero-shot experiment in Figure 4 are not directly related in terms of performance. Instead, this result suggests that the quality of the collected context is not necessarily correlated with the return of consequent episodes. Therefore, it is crucial to prioritize high-quality contexts with lower uncertainty. To further investigate the decline in baseline performance observed in Figure 4a during training, we visualized zero-shot trajectories (episode line 1) at different training steps at at <https://anonymous.4open.science/r/CERTAIN-6073/experiment7.md>, and found that as training progresses, the zero-shot trajectories gradually become more conservative, leading to a performance drop.\n\n2. The first trajectory for all methods is collected using $\\pi_\\theta(s, z_0)$. We sincerely appreciate the reviewer pointing out the incorrect labeling of the “X-axis”—it should indeed be “training step.”\n3. Thank you for the reviewer’s suggestions regarding the figures. We have redrawn the experimental plots in <https://anonymous.4open.science/r/CERTAIN-6073/experiment5.md>.\n\n> Reviewer: The paper misses comparison and discussion about a relevant prior work BOReL ...\n\nResponse: BOReL defines MDP ambiguity, which is different from our concept of context ambiguity. BOReL ultimately provides guidance on how to collect training data and how to use an oracle model to relabel existing datasets for correction. However, CERTAIN aims to identify uncertain contexts and improve the few-shot (one-shot) adaptation performance. Moreover, our method is a plug-in approach that can be applied to any context encoders.", "claims_and_evidence": "Definitions 3.1-3.3 are not very clear. How are the conditional probabilities $p(c \\mid \\mathcal{M}\\_i)$, $p(c \\mid \\mathcal{M}, \\pi\\_\\beta)$ defined? Does $p(c \\mid \\mathcal{M}, \\pi\\_\\beta)=0$ need to hold for a single $\\mathcal{M}\\_i$ or all of the tasks in the offline training dataset? What is the relationship between $\\pi_\\beta$ and the offline datasets $\\mathcal{D}$? How is $\\sigma$ defined and computed?\nFurthermore, the empirical evidence is not very convincing (see the method and experiment parts below).", "methods_and_evaluation_criteria": "While the use of heteroscedastic uncertainty estimates is common among existing works, they are primarily used to estimate aleatoric uncertainty (e.g., in [1]). In this paper, they are also used to estimate epistemic uncertainty for OOD cases. The uncertainty estimator $h\\_\\psi$ may struggle to produce correct estimates for OOD inputs, in which case the epistemic uncertainty estimation approach of [1] could be better suited. As the authors make specific claims about OOD contexts, more evidence in this scenario is needed to support these claims.\n\nFurthermore, CERTAIN chooses to use a deterministic context encoder with an additional network $h\\_\\psi$ for uncertainty estimation, while a more natural choice could be to directly use a probabilistic context encoder for capturing this uncertainty (see the reference discussion below).\n\n[1] Kendall, A. and Gal, Y., 2017. What uncertainties do we need in bayesian deep learning for computer vision?. Advances in neural information processing systems, 30.", "theoretical_claims": "N/A as there are no theorems or proofs.", "experimental_designs_or_analyses": "1. The empirical performance improvement of CERTAIN seems a bit marginal. In Fig. 3, CERTAIN variants only display a relatively clear advantage in Point-Robot and Walker-Rand-Params while being outperformed in Point-Robot and Ant-Goal in Fig. 4. The authors' explanations for this fail to convince. Also, the performance of baselines seems to decrease over time in Fig. 4(a) which is confusing.\n\n2. What is the context collection policy in the first episode for the baseline methods? What is the x-axis \"Test Steps\" in Fig. 3 and 4? \n\n3. (Minor) Both figures are a bit hard to read; it may help with readability to distinguish between baselines and CERTAIN+baselines using e.g. solid/dotted lines of the same color.", "relation_to_broader_scientific_literature": "This paper proposes a method to promote exploration and reduce task uncertainty in COMRL settings, which is a well-explored topic in the field of meta-RL. See the next section for specific discussions about the relevant literature.", "essential_references_not_discussed": "The paper misses comparison and discussion about a relevant prior work BOReL [1], which is cited but not thoroughly analyzed. BOReL seeks to address the ambiguous context problem by learning a meta-policy conditioned on the distribution of the latent posterior, making it also uncertainty-aware as CERTAIN. The idea of representing uncertainty as a latent distribution is common among related works, e.g. VariBAD [2] and PEARL [3], which are also cited but not discussed. Since CERTAIN takes another approach and estimates the uncertainty with the deterministic latent embedding, the paper could benefit greatly from more detailed comparisons and discussions.\n\n[1] Dorfman, R., Shenfeld, I. and Tamar, A., 2021. Offline Meta Reinforcement Learning--Identifiability Challenges and Effective Data Collection Strategies. Advances in Neural Information Processing Systems, 34, pp.4607-4618.\n\n[2] 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.\n\n[3] Rakelly, K., Zhou, A., Finn, C., Levine, S. and Quillen, D., 2019, May. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning (pp. 5331-5340). PMLR.", "other_strengths_and_weaknesses": "The paper is well-motivated and the method makes intuitive sense. Fig. 1 and 2 are good illustrations and provide great clarity about the method. Visualizations are provided.", "other_comments_or_suggestions": "Typo: Definition 3.1 missing subscript in $p(c \\mid \\mathcal{M}\\_i)$; table captions mention the point-robot environment but include results from multiple environments (e.g. Tab. 1); Eq. (12) could be missing something, e.g. discount factor and max operators.", "questions_for_authors": "1. Can CERTAIN reliably identify OOD contexts?\n\n2. How does CERTAIN compare with Bayes Adaptive methods, e.g. BOReL?", "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": "UH0Lv1l6v5", "reviewer": "Reviewer_a856", "summary": "This paper studies the task representation problem in context-based offline meta-reinforcement learning (COMRL). It first identifies the problem of task uncertainty in a context, including task ambiguity and out-of-distribution problems. Then, the paper proposes a training method that learns both context representation and uncertainty and trains a policy that maximizes return while minimizing the estimated uncertainty. In experiments on MuJoCo tasks, the method outperforms baselines in both one-shot adaptation settings and zero-shot settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewers’ constructive feedback. Below, we respond to each concern point by point. We will incorporate all the reviewers’ suggestions in the final version of the paper. \n\n> Reviewer: \n> - (line 40, right) ... However, CORRO [1] focuses on addressing OOD contexts in OMRL. \n> - Definitions 3.1, 3.2, 3.3 are more like an intuitive explanation and the definitions are either informal or not used in the following writing.\n\nResponse: CORRO primarily focuses on improving representation robustness through data augmentation. It uses a CVAE-based approach to augment the dataset, enhancing the encoder’s robustness in OOD scenarios. However, its data augmentation relies on (s, a) pairs that have appeared in the dataset, inherently limiting its OOD generalizability. In contrast, our method identifies and evaluates ambiguous and potentially OOD contexts for adaptation purposes and is designed as a plug-in approach. Thus, we consider CORRO and our method orthogonal, meaning they can be combined to further improve robustness in OOD tasks.\n\nRegarding Definitions 3.1, 3.2, and 3.3, they are introduced to formally define key concepts in the paper.\n\n> Reviewer: \n> I have check the experiments in detail. Here are some issues:\n> - In Figure 3 and 4, ...\n> - In the zero-shot setting, the returns of which episodes are reported?\n> - How are the behavior policies trained and the offline datasets collected for these experiments? How is the data quality, compared with the performance of the trained OMRL policies?\n\nResponse:\nYes, the x-axis in Figures 3 and 4 should be “training step.”\nIn the Zero-shot setting, we collect a trajectory using the context collection policy $\\pi_\\theta(s, z_0)$, and its return is reported as the Zero-shot performance. To further investigate the decline in baseline performance observed in Figure 4a during training, we visualized zero-shot trajectories (episode line 1) at different training steps at at https://anonymous.4open.science/r/CERTAIN-6073/experiment7.md, and found that as training progresses, the zero-shot trajectories gradually become more conservative, leading to a performance drop.\nRegarding the dataset, we strictly follow the data collection procedure used in baselines such as FOCAL. Specifically, in each task, we train an agent from scratch using SAC and collect data at training steps. Additionally, we have provided results on the performance of different methods under varying dataset qualities in <https://anonymous.4open.science/r/CERTAIN-6073/experiment2.md>. We found that dataset quality significantly affects all methods, particularly FOCAL. We hypothesize that FOCAL is more susceptible to spurious correlations compared to other methods.\n\n> Reviewer: I hope the authors address the issues in the writing of experiments and some related works.\n\nResponse: We have conducted experiments on OOD contexts, where the context used for inferring task representations consists of 25% in-distribution samples (upper semicircle) and 75% out-of-distribution samples (lower semicircle) in <https://anonymous.4open.science/r/CERTAIN-6073/experiment1.md> and OOD tasks, where the goal is at  lower semicircle in <https://anonymous.4open.science/r/CERTAIN-6073/experiment4.md>. \nIn the final version of the paper, we will include a more detailed discussion on OOD.", "claims_and_evidence": "Most claims made in this paper are clear. But there are two issues:\n- (line 40, right) \"Prior methods often assume that contexts are either in-distribution or can be collected through multiple rounds.\" However, CORRO [1]  focuses on addressing OOD contexts in OMRL. \n- Definitions 3.1, 3.2, 3.3 are more like an intuitive explanation and the definitions are either informal or not used in the following writing. Though, this can be a minor issue because the paper does not claim some theoretical contributions.\n\n[1] Robust task representations for offline meta-reinforcement learning via contrastive learning. ICML 2022.", "methods_and_evaluation_criteria": "The proposed method makes sense for the identified problems. The experimental settings properly follow prior works.", "theoretical_claims": "The paper makes no theoretical claims.", "experimental_designs_or_analyses": "I have check the experiments in detail. Here are some issues:\n- In Figure 3 and 4, I guess horizontal axes should denote training steps, not test steps, since the number goes to 1e5. Is this a typo? If so, why do returns decrease in Point-Robot and Ant-Goal in Figure 4?\n- In the zero-shot setting, the returns of which episodes are reported?\n- How are the behavior policies trained and the offline datasets collected for these experiments? How is the data quality, compared with the performance of the trained OMRL policies?", "relation_to_broader_scientific_literature": "The notion of context uncertainty has been well explored in online meta-RL, such as VariBAD [1]. But in OMRL, it is less explored. The proposed method is an intuitive and good solution to this problem in OMRL and I think it is a good contribution to the community.\n\n[1] Varibad: A very good method for bayes-adaptive deep rl via meta-learning. 2019.", "essential_references_not_discussed": "I did not find a missing citation. But I think some related works should be further discussed. For example, some prior works, such as CORRO, explores the problem of out-of-distribution context in COMRL.", "other_strengths_and_weaknesses": "Strengths:\n- The paper unifies the problem of OOD context and task ambiguity into the problem of task uncertainty, which is important in the literature of COMRL. The proposed method makes sense.\n- The proposed method can seamlessly integrate various approaches, including classification-based methods, reconstruction-based methods, and contrastive learning methods.\n\nWeaknesses:\n- As I mentioned above, some experimental details are confusing. Related works about OOD context and task uncertainty can be further discussed.", "other_comments_or_suggestions": "Please see issues above.", "questions_for_authors": "I hope the authors address the issues in the writing of experiments and some related works.", "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": "q1xEzlfHZv", "reviewer": "Reviewer_eruf", "summary": "This paper presents the CERTAIN method for Offline Meta Reinforcement Learning. The CERTAIN method models uncertainty for each transition sample explicitly, and performs task representation by selecting samples with lower uncertainty, thereby achieving more accurate task identification. Experimental results in MuJoCo environments show that CERTAIN, when combined with three different Offline Meta RL baselines, demonstrates superior reward performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewers’ constructive feedback. Below, we respond to each concern point by point. We will incorporate all the reviewers’ suggestions in the final version of the paper.\n\n> Reviewer: Uncertainty Estimation: The paper uses the Heteroscedastic Uncertainty loss function from regression tasks ...\n\nResponse:\nIn this paper, the heteroscedastic loss is formulated as $\\frac{L}{\\sigma^2} + \\log \\sigma$, where the heteroscedastic term $\\sigma$ captures the magnitude of the original loss L. Thus, as long as L is a positive-valued function—such as MSE, cross-entropy, or contrastive loss—the learning process of heteroscedasticity $\\sigma$ is guaranteed. Moreover, heteroscedasticity $\\sigma$ has already been applied to contrastive losses, as demonstrated in [1], where it was incorporated into the triplet contrastive loss.\nRegarding the role of heteroscedasticity $\\sigma$ in contrastive loss, it can be interpreted as the separability of samples from different tasks given a specific encoder model.\n\n[1] Unsupervised Data Uncertainty Learning in Visual Retrieval Systems\n\n> Reviewer: Context Collecting Policy Learning: In the learning of the context collecting policy, the policy still needs to be learned through the reward r. But ...\n\nResponse:\n1. We consider only cases where the offline dataset includes rewards, meaning all transitions in the dataset are complete in the form of  $(s, a, r, s')$. For scenarios where the dataset lacks rewards, we acknowledge this as an important and interesting problem and plan to explore it in future work.\n2. The training process of the context collection policy $\\pi_\\theta(s, z_0)$ is task-agnostic. Consequently, $\\pi_\\theta(s, z_0)$ focuses on overall uncertainty across the dataset distribution and cannot ensure that collected samples for a specific task have low uncertainty. To mitigate the impact of high-uncertainty samples, we employ an uncertainty-restraining method, which helps reduce task inference uncertainty and ultimately improves meta-policy performance.\n\n> Reviewer: Data Augmentation Methods: The paper lacks a discussion or comparison with data augmentation methods in COMRL, such as MBML [1], COSTA [2], or ReDA [3]. Data augmentation techniques are likely to be effective in reducing uncertainty.\n\nResponse: Data augmentation can improve the robustness of the context encoder in OOD scenarios. However, our method focuses on identifying and evaluating ambiguous and potentially OOD contexts for adaptation purposes. Furthermore, since our approach is plug-in compatible, it can be integrated with methods that utilize data augmentation.\n\n> Reviewer: Comparison with SOTA Methods: While the paper mentions Mutual Information-based methods like CORRO [4] and CSRO [5] in contrastive learning, it does not compare CERTAIN with these methods ...\n\nResponse: In our adopted baseline [2], a complete theoretical derivation of CORRO and CSRO, along with their performance evaluations, is provided within a unified framework, UNICORN, from a mutual information perspective. Therefore, we believe that conducting experiments based on UNICORN effectively represents CORRO and CSRO.\n\n[2] Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning\n\n> Reviewer: Narrow Focus on One-Shot In-Distribution Task Adaptation: The paper primarily focuses on one-shot in-distribution task adaptation ...\n\nResponse: We consider one-shot learning the most challenging case within few-shot settings. Achieving high performance with as few shots as possible better demonstrates the effectiveness of the adaptation algorithm. We acknowledge the reviewer’s suggestion regarding broader few-shot settings. Therefore, we have supplemented few-shot (5 episodes) experiments on the Point Robot task at <https://anonymous.4open.science/r/CERTAIN-6073/experiment3.md> and added experiments on OOD tasks (where goal is at lower semicircle) at <https://anonymous.4open.science/r/CERTAIN-6073/experiment4.md>.", "claims_and_evidence": "See Strengths and Weaknesses.", "methods_and_evaluation_criteria": "See Strength and Weaknesses.", "theoretical_claims": "No theoretical claims in this paper.", "experimental_designs_or_analyses": "See Strength and Weaknesses.", "relation_to_broader_scientific_literature": "See Strength and Weaknesses.", "essential_references_not_discussed": "See Strength and Weaknesses.", "other_strengths_and_weaknesses": "This paper demonstrates certain advantages in addressing the problem of Offline Meta Reinforcement Learning:\n\n\n\n- The paper clearly outlines two reasons for task identification failures in COMRL: Context Ambiguity and OOD Context.\n\n- The proposed method is simple and straightforward.\n\n- Experiments in the MuJoCo environment show that the CERTAIN method provides an improvement in reward performance.\n\n\n\nHowever, there are several issues with the paper:\n\n\n\n1. **Uncertainty Estimation**: The paper uses the Heteroscedastic Uncertainty loss function from regression tasks. While this loss function is theoretically well-supported for maximizing likelihood in regression tasks, it may not necessarily have the same properties in the context of representation learning as presented in this paper. Therefore, I am skeptical about the choice of this uncertainty estimation method. For example, when combining CERTAIN with classification methods, would it be more reasonable to use a classification-specific Heteroscedastic Uncertainty loss function? Additionally, what is the significance of such a loss function in contrastive learning?\n\n2. **Context Collecting Policy Learning**: In the learning of the context collecting policy, the policy still needs to be learned through the reward r. But if there is no prior task, what is the reward r in this case? Furthermore, the uncertainty estimation also requires input from (s,a,s′,r), meaning that uncertainty should be related to the reward, which corresponds to the task itself. However, the context collecting policy should be a task-independent policy. How then can uncertainty be applied to the context collecting policy? For example, suppose there are three tasks with significantly different distributions. A given transition s,a,s′ may have low uncertainty when the reward r1 from task 1 is given, but be an out-of-distribution (OOD) sample when the rewards r2 or r3 from tasks 2 or 3 are provided. In this case, a sample that is valid for task 1 may become invalid for tasks 2 or 3. Hence, I don't find the uncertainty-based context collecting policy to be fully reasonable.\n\n3. **Data Augmentation Methods**: The paper lacks a discussion or comparison with data augmentation methods in COMRL, such as MBML [1], COSTA [2], or ReDA [3]. Data augmentation techniques are likely to be effective in reducing uncertainty.\n\n4. **Comparison with SOTA Methods**: While the paper mentions Mutual Information-based methods like CORRO [4] and CSRO [5] in contrastive learning, it does not compare CERTAIN with these methods. Given that these methods are likely more state-of-the-art than FOCAL, such a comparison would strengthen the paper.\n\n5. **Narrow Focus on One-Shot In-Distribution Task Adaptation**: The paper primarily focuses on one-shot in-distribution task adaptation, which is a rather narrow setting. One-shot adaptation requires only a single trajectory to adapt to the task, but there are many few-shot methods, such as Thompson Sampling. Expanding the one-shot scenario to a few-shot setting and comparing it with few-shot adaptation methods could make the paper more solid. Additionally, many previous COMRL algorithms, including the baseline FOCAL and methods like CORRO, deal with OOD task adaptation, but this paper only experiments with in-distribution tasks.\n\n\n\nIf the issues above are effectively addressed, I would consider re-reviewing the paper.\n\n\n\n[1] Multi-task Batch Reinforcement Learning with Metric Learning\n\n\n\n[2] Cost-aware Offline Safe Meta Reinforcement Learning with Robust In-Distribution Online Task Adaptation\n\n\n\n[3] Disentangling Policy from Offline Task Representation Learning via Adversarial Data Augmentation\n\n\n\n[4] Robust task representations for offline meta-reinforcement learning via contrastive learning\n\n\n\n[5] Context Shift Reduction for Offline Meta-Reinforcement Learning", "other_comments_or_suggestions": "None", "questions_for_authors": "See Strength and Weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fzxOda0hy1", "reviewer": "Reviewer_eS1K", "summary": "The paper presents CERTAIN, a novel framework designed to address challenges in context-based offline meta-reinforcement learning (COMRL), particularly context ambiguity and out-of-distribution (OOD) issues, in one-shot adaptation settings. The authors propose leveraging heteroscedastic uncertainty in task representation learning to identify and mitigate the negative impact of ambiguous and OOD contexts on task inference during online adaptation.\n## update after rebuttal\nThanks for the author's reply, I have no further questions.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewers’ constructive feedback. Below, we respond to each concern point by point. We will incorporate all the reviewers’ suggestions in the final version of the paper.\n\n> Reviewer: \n> \n> Weaknesses\n> \n> 1.While the paper provides strong empirical results ...\n> \n> 2.While the paper presents ablation studies on certain components ...\n\nResponse:\n1. We acknowledge the lack of a theoretical analysis and plan to explore this aspect further in future work. However, our primary contribution lies in addressing uncertainty in few-shot (specifically one-shot) adaptation, which we hope will open new avenues for research in the OMRL field.\n2. To further validate our approach, we have added multiple experiments in the anonymous link, demonstrating the ability of the uncertainty estimation network to distinguish OOD contexts at <https://anonymous.4open.science/r/CERTAIN-6073/experiment1.md> and highlighting the importance of the context collection policy for few-shot performance at <https://anonymous.4open.science/r/CERTAIN-6073/experiment3.md>.\n\n\n> Reviewer: 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 multiple models for learning?\n\nResponse:\nUsing a single model for both task inference and policy execution typically requires a massive dataset and a large number of parameters. Moreover, these two fundamentally different objectives can interfere with each other’s learning, making optimization more challenging. In contrast, our approach separates these processes into multiple models, reducing learning difficulty while enhancing interpretability and controllability.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No proofs for theoretical claims.", "experimental_designs_or_analyses": "The experimental design is basically reasonable.", "relation_to_broader_scientific_literature": "The key contributions of this paper build on the context-based offline meta-reinforcement learning (COMRL) framework, addressing the limitations of context ambiguity and OOD issues, which have been largely overlooked in prior COMRL methods. The integration of heteroscedastic uncertainty in task representation learning and context collection policies extends existing approaches, such as classifier-based and contrastive learning methods, by improving adaptation robustness in one-shot settings.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**\n\n1.The proposed integration of heteroscedastic uncertainty into task representation learning for COMRL is a novel approach that addresses both context ambiguity and OOD issues, which have been insufficiently explored in prior research.\n\n2. The paper effectively identifies a practical problem in one-shot adaptation settings and introduces a solution with direct real-world relevance, improving sample efficiency and safety in reinforcement learning tasks.\n\n3. Extensive empirical evaluations across multiple environments, including toy tasks and complex MuJoCo environments, demonstrate the robustness and effectiveness of the CERTAIN framework.\n\n4. CERTAIN is designed as a plug-in framework that can be easily integrated into existing COMRL methods (e.g., classifier-based, reconstruction-based), making it adaptable to various contexts and methods.\n\n**Weaknesses**\n\n1.While the paper provides strong empirical results, it lacks formal theoretical analysis or guarantees regarding the performance of the uncertainty-aware components, which would strengthen the theoretical foundation.\n\n2.While the paper presents ablation studies on certain components, further detailed analysis on the interplay between the different uncertainty mechanisms (e.g., uncertainty estimation network and context collection policy) could provide a deeper understanding of their individual contributions.", "other_comments_or_suggestions": "No other suggestions.", "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 multiple 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 paper proposes CERTAIN, an uncertainty-aware context collection method for one-shot adaptation in context-based offline meta reinforcement learning. Leveraging heteroscedastic uncertainty estimates, CERTAIN can be used with prior COMRL methods to learn an uncertainty-aware task representation. The estimated uncertainty can further be used as penalties in context collection policy learning. Experiments are conducted in several commonly used MuJoCo benchmark environments and show improved performance.
The reviews are generally mixed with initial scores of [2, 2, 3, 3] and post-rebuttal [2, 3, 3, 3]. Overall, the reviewers appreciated the generality of CERTAIN, which can be integrated with various existing COMRL methods to improve performance, as demonstrated by the experiments. The problem of OOD context identification and mitigation is important to the broader COMRL community and real-world applications, for which CERTAIN provides a simple method. Existing concerns mainly involve the lack of theoretical analysis of CERTAIN, which the authors acknowledged during rebuttal. Empirical evidence in OOD scenarios is also weak, potentially due to the tight rebuttal timeline.

Overall, this is a borderline paper, with the benefits (generality of method, importance of problem) outweighing the downsides (theoretical foundations, OOD evaluation). So, I recommend a weak accept while encouraging the authors to improve the paper accordingly in the final version.
Decision: Accept (poster)