Context Shift Reduction for Offline Meta-Reinforcement Learning
Authors: Yunkai Gao
Rui Zhang
Jiaming Guo
Fan Wu
Qi Yi
Shaohui Peng
Siming Lan
Ruizhi Chen
Zidong Du
Xing Hu
Qi Guo
Ling Li
Yunji Chen
Venue: NeurIPS 2023 poster (2023)
Keywords: offline meta-reinforcement learning
offline reinforcement learning
meta-reinforcement learning
TL;DR: This paper studies the context shift problem in offline meta-reinforcement learning and proposes a novel to address this problem.

Abstract: Offline meta-reinforcement learning (OMRL) utilizes pre-collected offline datasets to enhance the agent's generalization ability on unseen tasks. However, the context shift problem arises due to the distribution discrepancy between the contexts used for training (from the behavior policy) and testing (from the exploration policy). The context shift problem leads to incorrect task inference and further deteriorates the generalization ability of the meta-policy. Existing OMRL methods either overlook this problem or attempt to mitigate it with additional information. In this paper, we propose a novel approach called Context Shift Reduction for OMRL (CSRO) to address the context shift problem with only offline datasets. The key insight of CSRO is to minimize the influence of policy in context during both the meta-training and meta-test phases.  During meta-training, we design a max-min mutual information representation learning mechanism to diminish the impact of the behavior policy on task representation. In the meta-test phase, we introduce the non-prior context collection strategy to reduce the effect of the exploration policy. Experimental results demonstrate that CSRO significantly reduces the context shift and improves the generalization ability, surpassing previous methods across various challenging domains.

Context Shift Reduction for Offline Meta-Reinforcement Learning
Yunkai Gao 12233
Rui Zhang
Jiaming Guo
Fan Wu2345
Qi Yi
Shaohui Peng
Siming Lan 1,2, Ruizhi Chen
Zidong Du
Xing Hu
Qi Guo² Ling Liit5
Yunji Chen
University of Science and Technology of China, USTC, Hefei, China
2 State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences, Beijing, China
Cambbicon Technologies
University Chinese Academy of Sciences, UCAS, Beijing, China
5 Intelligent Software Research Center, Institute of Software, CAS, Beijing, China
Shanghai Innovation Center for Processor Technologies, SHIC, Shanghai, China
{gyk314, yiqi, lansm)@mail. wustc..dd.... (wuf an2200, pengshachui, ruizhi, 1iling/@iscas. acc.ciacccsccc.c........ {zhangrui, guojiaming, duzidong, huxing, guoqi, syj]0ict.ac.co
Abstract
Offline meta-reinforcemment learning (OMRL) utilizes pre-collected offline datasets to enhance the agent's 's generalization ability on unseen tasks. However, the context shift problem arises due to the distribution discrepancy between the contexts used for training (from the behavior policy) and testing (from the exploration policy). The context shift problem leads incorrect task inference and further deteriorates the generalization ability of the meta- meta-policy. Existing OMRL methods either overlook this problem or attempt to mitigate it with additional information. In this paper, we propose novel approach called Context Shift Reduction for OMRL (CSRO) address the context shift problem with only offline datasets. The key insight of CSRO is to minimize the influence of policy in context during both the meta-training and meta-test phases. During tetannra we design max-mi min mutual information representation learning mechanism to diminish the impact of the behavior policy on task representation. In the meta-test phase, we introduce the non-prior context collection strategy to reduce the effect of the exploration policy. Experimental results demonstrate that CSRO significantly reduces the context shift and improves the generalization ability, surpassing previous methods across various challenging domains.
Introduction
Rentfoeiinoor Learning (RL) has made significant strides in enhancing agents' generalization capabilities for unseen tasks. Meta-RL methods [6] 5 27 leverage vast amounts of trajectories gathered across a task set during the meta-training phase to learn the shared structures of tasks. During meta-testing, these meta-RI. methods exhibit swift adaptation to new tasks with limited number of trials. Context-based meta-RI. methods approaches, showcasing remarkable performance and rapid task adaptation. These methods first acquire task information by learning a task representation from the context (i.e., the of transitions) and subsequently condition the policy the task representation to execute appropriate actions.
have emerged as the most popular
"Corresponding Author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Due to the requirement of a vast number of trajectories sampled from a task set, data collection in
meta-RL can be both expensive and time-consuming. Offline RL [T7 32 14] presents potential solution to mitigate the cost of data collection. These methods only utilize large datasets that were previously collected by some policies (referred to as the behavior policy), without any online interaction with the environment during the training phase.
By combining meta-RL and offline RL, offline meta-reinforcemeet learning (OMRL) 18 20 37
amalgamates the benefits of generalization ability on unseen tasks and training without requiring interaction with the environment. During meta-training, both the acquisition of task representation and the training of the policy utilize pre-collected offline trajectories obtained from the task-dependent behavior policy.
However, OMRL faces a context shift problem: a distribution discrepancy arises between the context obtained from the behavior policy during meta-training and the context collected from the exploration policy during meta-testing [26). Given that the task representation is learned from the context, this context shift problem hinders the accurate inference of task information by the task representation. Consequently, this leads to a decline in performance when adapting to unseen tasks. As shown in Table we compared the test results of FOCAL 20 and OffPearl p2 two different environments and found that context shift significantly reduced testing performance.
Existing OMRL methods primarily aim to incorporate the benefits of offline RL into meta-RL, RL. but they do not fully address the context shift problem. Some approaches [20] 21 19 utilize pre- collected offline datasets from the test tasks to acquire task representation, disregarding the context shift problem. However, this approach is impractical as obtaining pre-collected offline datasets for new and unseen tasks is extremely challenging. Other methods attempt to mitigate the context shift problem but require additional information beyond the use of offline datasets alone. For instance, [ assumes that the reward functions of all tasks are known, (26 requires the collection of substantial volume of online data post offline training.
In this paper, introduce novel approach called Context Shift Reduction for OMRL (CSRO) to address the context shift problem using only offline datasets. The key insight of CSRO is to minimize the impact of the policy in context during both the meta-training and meta-test phases. Firstly, the context is collected from the task-dependent behavior policy, as depicted in Figure| left. When encoding the context into the task representation during meta-training, the task representation tends to embed the characteristics of the behavior policy. This incorrect task inference during the meta-test phase stems from the disparity between the behavior policy and the exploration policy. Hence, our objective is to reduce the influence of the behavior policy on task representation during the meta-training phase. To achieve this, we design max-min mutual information representation learning mechanism that minimizes the mutual information between the task representation and the behavior policy, while maximizing the mutual information between the task representation and the task information. During the meta-test, the current exploration policy of OMRL [27] conditions on an initial task representation. The collected context corresponds the policy of the initial task representation, which significantly differs from the behavior policy, as shown in Figure I middle. This context leads to incorrect task inference and control, as depicted Figure | right. Therefore, we propose the non-prior context collection strategy to reduce the impact of the exploration strategy. In practice, the agent initially explores the task independently and randomly at each step, then the agent proceeds to explore based on the context. Through random exploration, the agent initially gains preliminary understanding of the task and gradually enhances the accuracy of task recognition during subsequent exploratory steps.
Experimentally. we compare the proposed CSRO with previous OMRL algorithms. For a fair comparison, we use the same offline backbone algorithm and the same offline data for environments including Point-Robot and MuloCo physics simulator 29 with reward or dynamic function changes. Experimental results demonstrate that the proposed CSRO can effectively mitigate the context shift problem and enhance the generalization capabilities, outperforming prior methods on a range of challenging domains.
2 Related Works
Meta Reinforcement Learning. Meta-RL leverages training task sets sampled from the task distribution to learn and generalize to new tasks. Gradient-based meta-RL [6 33] methods seek

---

Table 1: Context shift leads to testing performance decline. Context A refers to the use of offline data collected by behavior policy as the context during testing, with no context shift problem. Context B, on the other hand, involves using trained meta- meta-policy to explore and collect context based on an initial task representation, resulting in a context shift problem.
Env
Point-Robot
Half Salf-Cheetaah-e Vel
context A context
context A .5.77F2.77 -123.0 -162.8:228.8
context
FOCAL
-..
-14.9+1.1
-69.5+99.6
OffPearl
-5.1
-17.8+1.5
Offlinc Datasets of Meta-Training
Contexts of Meta-Test
Trajectorics of Mcta-Test
goal
goal
goal
Offline meta-RL on the Point-Robot: one of the tasks where the task objective is control the agent to reach the goal on the semicirele. Left: During meta- veta-training. the offline dataset of this task is collected by task-dependent behavior policy and comprised solely of trajectories leading towards the goal. Middle: The agent was trained using the OffPearl
Figure 1:
(7 method. During the
meta-test phase, the agent utilizes the current exploration policy of OMRI. 27 26 and employs the learned meta-policy, which is conditioned on initial task representation sampled from prior distribution, to collect contexts. Right: The agent navigates out trajectory based on the context depicted in the middle subfigure, the same color represents the corresponding context and trajectory.
policy parameter initialization that requires only few update steps to generalize well on the new tasks. Since the parameters that cause reward or dynamic changes across different tasks are unknown, context-based meta-RL (27 39 16 [1] address this problem PODMP. These methods employ context encoder to extract task representations from historical trajectories, enabling the identification of task domains, with policies conditioned on these acquired task representations. Our proposed CSRO is based on context- based meta-RL frameworks (27
Offfine Reinforcement learning. Unlike on-policy and off-policy RL, in offline RL, the agent does not interact with the environment during the training process and only interacts with the env environment during the test phase. Offline RL methods 15 policy for training, effectively reducing the cost of data collection. the distribution shift between the behavior policy and the learned policy introduces the challenge of Q overestimation [17] [32 [3]
g use pre-collected data generated by behavior
make the distribution shift bounded by constraining the differences between learned policy and
behavior policy. We employ BRAC [32 as our offline backbone algorithm.
Offline Meta Reinforcement learning. OMRL solves the generalization problem and expensive data collection problem by using meta-learnigg on pre-collected offline datasets. Existing OMRL 31 use historical trajectories from offline datasets to infer task representations during meta- training. In the meta-test phase, existing methods can be categorized into two types based on the source of the context: offline test and online test. Offline test methods [18] 20 use sampled context from the pre-collected offline data in the test task, which impractical and ignores the context shift problem. In contrast, online test methods [ 26] use the context collected by the agent the test task. While these online test methods can mitigate the context shift problem, they require additional information beyond just offline datasets, such reward functions of all tasks [4] and additional online data (26). We focus on improving the generalization performance of the online test by addressing the context shift problem using fully offline datasets.
methods 220 18 4
37
Mutual Information Optimization. In order to learn more about essential representation, mutual information optimization is often employed to enhance the correlation between the representation
3

---

Meta-Training
Test Phase
4,(2)e)
Liwann
Explore
(s.gns }
Φ
(-llaans)))
=
Link
S (,(l)) (ajss,
Luan
(=K5.a))
xp(a(s.2)
3 Leo
Figure 2: CSRO framework. Left: In meta-training, the context encoder 9s extracts each transition embedding from (s, a, r, r') transition tuple sample from offline datasets, the CLUB variational distribution estimator 96 is used to approximate p(2 (s, a)). We use the mean of transition embedding as task representation which conditions both the critic Qu and the actor TT. Right: In test phase. agent uses NN to explore new tasks.
and other properties. The greater the mutual information, the greater the correlation between the two variables. 2 22) maximizes the mutual information between variables by maximizing the estimated mutual information lower bound. [ minimizes the mutual information between variables by
minimizing the upper bound of the estimated mutual information. In this work, we utilize CLUB 3 to minimize the mutual information between task representation and behavior policy.
Preliminaries
Problem Formulation. In RL, we formalize each task Markov Decision Process (MDP), defined by M = (S, A, P.Poo Po, R2 7) with state space S, action space A, transition function P(s'ss,a), reward function R(s, a), initial state distribution Po(s) and discount factor 7 (0, 1). The distribution of policy is (a)s). We aim to maximize the expected cumulative rewards (()) Er 1r 7 ? at),, where is trajectory following So Po. dr T(aels)). P(5+41 St, a.). We define the value function and action-value function as:
8141
(s+) PenTT,nuiI
y (($142,
(s ax) = R(ss, ar) + yB Plasall (s++1)
(2)
In offline meta-reinforcemeee learning (OMRL). assume each task sampled from prior task distribution, MV,- = (S, Pi, 7) ~ p(M). Different tasks have the same state space and action space, but the reward function transition dynamic is different due different parameters, So the task distribution can be defined as p(R,, P). The offline datasets D, {(s..j Gi.j Tiys {{, N of task M; is composed of the state, action, reward, and next-state tuples collected by the task- dependent behavior policy T s). Since the parameters of different tasks are unknown, for each task, context-based OMRL samples mini-batch of historical data from the offline datasets as the context {(s) aj T. s}))) The context encoder transforms the context into task representation = go (=le). Policy, value function, and action-value function are all conditioned on We use offline data to train the meta policy To(a 2) with parameters and will not interact with the environment during the training. OMRL aims to maximize the expected cumulative rewards in the test task:
(TT) EM [M
4 Method
4.1 Description of Context Shift Problem
OMRL uses pre-collected offline data from task-dependent behavior policy to maximize expected cumulative reward. Generally, given multiple different tasks, each behavior policy is the standard RL policy trained on each task. Then each offline dataset is collected from the behavior policy interacting with this environment. For N MDPs sampled from p( P. R), the behavior policy of task M, for

---

collecting data is (a]ss, and the offline data is D, = {($L.) Paj 81,) ) of offline data is denoted by
Nais. The distribution
Naize
p(D.) - p(su)
II ==0
(4)
In the meta-test phase, the data is online collected by the exploration policy.
The different sources of offline datasets used in meta-training and online data used in the meta-test lead to the context shift problem: distribution discrepancy between context from behavior policy and context collected from exploration policy. As shown in Eq 4 the distribution of collected offline datasets is jointly determined by task and behavior policy, which means the context used to train the context encoder contains some characteristics specific the behavior policy. When the characteristic of behavior policy is correlated with the task, the context encoder is possible to mistakenly embed the characteristic of behavior policy as the task information. Since the characteristics of the exploration policy may be very different from the behavior policy, the task representation derived from the context encoder will deliver incorrect task information and cause performance degradation in the meta-test phase.
For example, as shown in Figure left, the task goal is reach the position on the semicircle, while the offline datasets only contain the trajectories of behavior policy that go towards the corresponding goal. Thus the direction of the behavior policy is its characteristic which is highly correlated with the task goal. Based on this, the context encoder tends to memorize the direction of the context as the task information. During the meta-lest phase, the common exploration policy is conditioned on an initial task representation 20 sampled from a prior p(z) to collect context. The initial task representation zo contains a specific direction, so the collected context will move to this specific direction which is different from behavior policy, as shown in Figure I middle. Affected by this context, the agent will go along the specific direction instead of going to the task goal, as shown in Figure right.
To address the context shift problem, we propose that CSRO reduce the effect of policy in the context in both the meta- meta-training and meta-test phases. Figure 2 gives an overview of our CSRO framework.
4.2 Max-min Mutual Information Representation Learning
OMRL uses offline datasets to train a context encoder Zo to distinguish tasks. We first sample the context c; = {(s..j ot.j, Tij. st 5 Ni from the offline datasets D, of task M,. And then the context encoder extracts each transition embedding from each transition of the context: 2i.j
(=(=((8i,), (i.jj (i.j. +(,)) s(,)).. Finally, we use the mean of transition embedding 2; Ey(zz,) as the task representation of task M,. Ideally, the agent should pay more attention to the reward function of the environment where the reward changes, and the dynamic function of the environment where the dynamic changes, and only infer the task from there. However, since the offline dataset is determined by the behavior policy, the task representation will contain the characteristics of the behavior policy, Due to the context shift problem, this context encoder will lead to incorrect task inference in the meta-test phase.
In order to ensure the task representation only focuses on representing the feature of the task we design the max-min mutual information representation learning mechanism. Since mutual information measure of the correlation between two random variables, we can use mutual information to guide the optimization direction of the context encoder. This mechanism consists of two parts, one is to maximize the mutual information between task representation z and the task, and the other is to minimize the mutual information between and the behavior policy IT..
Since the task representation needs represent the task information, we need to maximize the mutual information between task representation and the task. We use distance metric learning in 200 to obtain mutual information approximation with task:
3
A ( = {v = v,Hlı - 20112 + 1 + Vi || 23
(5)
1 +
where is context sampled from task M,. yı is task label, 2; = (0(2)c;) is a task represtantion. Context encoder minimizes the distance of task representations from the same task and maximizes the distance of task representations from different tasks.
5

---

Definition 1. In the test task M; = (S, R₂) p(M). context e as context collected by exploration policy Te. The expected return of test task M, which is evaluated by a learned meta-policy To(a|s, 2) is:
H
(To =
rel.
(6)
t=
For task Mj and its corresponding behavior policy T... when exploration policy Te and behavior policy T' are different, if and only if mutual information between task representations and policies me; 0, the expected test return for both the same Jar; (To, Te) Ja, (TT, Ti))-
Proposition 1.
As analyzed in Section 4.1 the context encoder is possible to contain the characteristics of the behavior policy. Proposition 1 states that we also need to minimize the mutual information between task repre- sentations and policies Th remove the influence of behavior policy in task representation. However, due to the collection strategy of offline datasets, each task has only one policy to collect data, which makes it difficult to obtain the representation of behavior policy. Consider that each transition tuple (s, r, s') is jointly determined by task and behavior policy p(a, r, s'18) = a ( s)p(s' 18, a)r(8, a), use (s, a) tuple to represent behavior policy. Following A. we introduce using CLUB min- imize mutual information between transition embedding 2 and (s, as, Mutual information upper bound estimation of CLUB defined as:
(x (s, a)) = (.llog((2:((00 a:)) - -E, E,,oog((z]((s, a:))]]-
However the conditional distribution p(=l(s, a)) unavailable on our tasks, therefore we use variational distribution (2l(s, a)) to approximate p(=[(s, by minimizing lig-liielihood:
(v) -E M) B [log qu (2) a))]..
(8)
We minimize mutual information between transition embedding and (s, a) by minimize mutual information upper bound estimation:
L (Ø) EM (2 (st, (s)) - [log gu (z)l(s;, a.))l,
which 2; represents task embedding obtained by (st, az Tt, s() through the context encoder. Here, the variational distribution estimator ds and context encoder Q6 are used for adversarial training. If each transition embedding obtained by the context encoder go (=)(s,,, a, s')) contains more (s, a) information, the gæ (=]((s,,)) has higher accuracy, but mutual information upper bound T will be higher. If the transition embedding does not contain (s, information, the ability predict z through (st, a.) is the same as that from other (sj, az), LminMI (Ø) is lower. In this adversarial training, the mutual information between behavior policy and task representation can be minimized, the context encoder will pay more attention to the information of tasks.
By combining the above two mutual information, can get the total loss of context encoder as:
Leneoder (Ø) LIaaxM + ALminMis
(10)
where À is hyperparameter to adjust the weight of the two mutual information. Based the max-min mutual information representation learning, we can reduce the influence of policy in task representation.
4.3 Non-prior Context Collection Strategy
In the online meta-test phase, the agent needs to explore the new task to collect context e and update the posterior distribution qo (2(c) according to context The current exploration strategy of OMRL [7 4. (26] follows the common exploration strategy of meta-RL (27). The agent uses the learned meta-policy which is conditioned on initial task representation 2e sampled from the prior distribution p(z) to explore the new task. After that, the agent updates posterior qo (2(2)c by collected context c and samples
Then the agent continues explore the task to collect context conditioned on The steps of sampling z and collecting context are performed iteratively. Finally, the agent executes task inference according to all collected contexts.
Because there only one policy for each task to collect data, the context distribution sampled from offline datasets is highly correlated with the behavior policy. Even if we propose the max-min

---

mutual information representation learning which minimizes the mutual information between task representation and behavior policy, we still cannot completely eliminate the influence of policy in task representation. If we use the common exploration strategy to collect context c, the distribution of context will correlate with the sampled initial Eg. In this way, the posterior (o(=)c) will be affected by =o. leading to incorrect task inference.
On the other hand, there are some methods in current online meta-reinforcement learning that specifically train an exploration policy to collect context 77 38 However, in offline settings, there is no way to interact with the environment, and using offline datasets to train the exploration policy, the exploration policy is limited to the vicinity of the datasets, and the conservatism of offline is in conflict with exploration, by training an exploration policy to solve this problem in offline settings is difficult.
We propose non-prior context collection strategy eliminate this problem. Our strategy is not conditioned on sampled initial 20 from prior. The agent first explores the task independently and randomly at each step. Then posterior p(z]c) will be updated after this exploration. Through this exploration, the agent preliminarily recognizes the task. After that, the agent continuously explores the task according the posterior 4s (=/e) and updates the posterior each step until the context c is collected. In this process, the agent gradually recognizes the task more accurately. Finally, the agent evaluates based on the collected context c. By adopting this approach, we mitigate the influence of the sampled initial 20.
4.4 Algorithm Summary
Finally, in order to avoid the problem of overestimation of Q caused by the distribution shift of the actions offline learning, we use BRAC [32] to constrain the similarity between learned policy mo and behavior policy Ts by punishing the distribution shift of learned policy Tp. The optimization objectives under the actor-critic framework become:
(ws =
r- -1Q (''d',2))) (11)
Luctor (0) -B,,,,, D,, D ( (t)) (D(TT (-1,, 2), TT(-13 2]]] (12)
where D is the divergence estimate of learned meta-policy Ta( (-1s, =) and behavior policy و. T, z).
and we use the KI. divergence.
We summarized our meta-training and meta-test method in the pseudo-code of Appendix A
5 Experiments
We propose max-min mutual information representation learning and the non-prior context collection strategy to mitigate the context shift between the meta-train train and meta-test. To demonstrate the performance of our method, we compare CSRO with the prior offline meta-RL methods on meta-RL tasks in which the reward function or dynamic function is variable. Code is available at hhtps git
//gihhub. wwt.:/e....oooee/ee.e
5.1 Experimental Settings
Environments. We evaluate our method on the Point-Robot and MuloCo 29 that often used as the offline meta-RI. benchmarks:(1) Point-Robot is a 2D navigation task of a point in continuous space, and the goal located on semicircle far from the origin. (2) Half-Cheetch- Vel is to control the agent to reach the target velocity of tasks. (3) Ant-Goal is a task modified on Ant, which needs to control the agent to reach different goals on the circle. (4) Humanoid-Dir is control the is to control the velocity of the agent as fast as possible. (6) Walker--Rand-Parans is also to control the biped agent to move forward as quickly as possible. For (1) and (3), the goal positions of different tasks are different. (2) and (4) are different in target velocity and direction respectively. These four types of environments are all reward functions changing environments. (5) and (6) are dynamic functions changing environments. Robots different environments have different physical parameters, including mass, inertia, stamping, and friction.
humanoid agent to move in the direction provided by the task. (5) Klopper-Rand-Paraan
7

---

Roint-Robot
Half- half-Cheetah-Me
Gool
-140
-315
-1.5
Average Beturn
-100
Aver age Return
Averagee Return
-02
-15.0-
-30.
1.7
-200
FOCG
CICC
HOCA
C2O
KOCA
-2.5
COR60
-07
DORRO
D0Te.
Dohe
D0NK
-300
Train Steps(1e5)
Stepatle5)
Humandid-Dir
Hopper-Rand-PPn Params
Whiker-Rand-Pn Params
350
290
Averege Return
Average Return
Average
300
032
CIRO CCSSO
CSno Otheat CORBO
CS0O
300
Othead
DORRC
50NA.
Train Steps(1e5)
Train Steps(1e5)
Train Steps(le5)
Figure The average return of the online test phase on unseen test tasks versus other OMRL methods. Point-Robot, Half- Chee-Chhethh--e Veel, Ant-Goal, and Humanoid- -ii are all reward functions changing environments. Hopper-Rane Seaper-Rand-Paramte Paramss and Miiter-Rand-Paatte are dynamic functions changing environments. The shaded region shows standard deviation across 8 seeds.
Offline Data Collections. For each environment, sampled 30 training tasks and 10 test tasks from the task distribution. We use SAC (10] on each training task to train an agent and save the policy different training times as behavior policy. We each policy roll out 50 trajectories in the corresponding environment as offline datasets.
Compared Methods. To demonstrate the performance of our method, we compare CSRO with the following methods:(1) OffPearl is an extension that applies Pearl [27] to offline RL. The context encoder outputs the distribution of task representation and updates by Q-function. The pre-collected offline datasets are used to replace the interaction with the environment. (2) FOCAL [20) uses metric distance learning train context encoder, reduce the task representation distance of the same task, and increase the task representation distance of different tasks. (3) CORRO (37 trains CVAE [28] or adds noise to produce negative samples, and uses InfoNCE 22 to train robust context encoder. builds VariBAD 39). For fair comparison, we use variant of BORel. that does not utilize oracle reward functions, as introduced in the original paper
(4) BOReL
For fairness comparison, all methods here use the same offline RL algorithm BRAC [32 and the same offline datasets. More details of the experimental setting and results available in the Appendix.
Table 2: Comparing CSRO with other baselines in the online test phase without and with the non- prior context collection strategy(Np) Point-Robot, Half-Cheetah- Vel, and Malker-Rand--ae Params environments.
Env
Point-Robot
Half-Cheetah- Vel
Walker-Rann Randr--randd
w/ Np
w/o Np
w/ Np
w/o Np
w/ Np
w/o Np
CSRO
-6.420.8
-9.220.6
-48.4=3.9
688.511.99
344.2=38.0
319.7438.4
FOCAL
11..8+1.6
-1..9+1.1
-60.9-5.7
-69.5++9.6
253.3+142.7
247.5++29.4
OffPear
117.0+1.6
-17.8+1.5
-133.7=18.9
-162.8+28.8
284.5++3099
262.0+22.5
CORRO
-7.8++1.9
-10.5+3.00
-65.6=9.3
-92.1+23.2
312.5+46..6
275.2+73.9
BOReL
221.613.99
-23.245.8
-90.1428.3
10..
260.6-140.2
245.8432.9

---

Point- Foint-Robot
Half-Chectah-vel
Ant-Goal
-312
Average Beturn
Average Return
-193
Return
-10
-125
-190
-200
-25
-20
CaRD
-22
UMR
CSRO hp
-25
CSRO
C313 mo9e
-300
C2RO wo k
CS10 mer.
Stepsile5)
Train Stepsíle5)
Steps(1e5)
Humanoid-Dii
Hopper-Rand-Fee Parppar--nn-Paaraam
Walkerr Seizer-Rand-Parrane
185
Average Return
Averoge
300
Average
230
100
200
300
C21D Alc
10
CSRO wr kp
150
CGRO ARt
Train Steps(le5)
Train Steps(le5)
Train Steps(lé5)
Figure 4: Ablation study in six environments, compare CSRO with methods that without CLUB minimize mutual information (minMI) and con-prior context collection strategy (Np) components.
5.2 Comparison Online Test Results.
To evaluate our performance in the online test phase, we compare CSRO with other methods across all six environments. In the online test phase, each method needs to explore new tasks and collect context. The four baseline methods follow the common exploration strategy, CSRO uses the non-prior context collected strategy that proposed.
We plot the mean and standard deviation curves of returns across 8 random seeds in Figure In the six environments, CRSO outperforms all the baselines, especially in Point-Robb Roboot and Ant-Goal. Experimental results prove the effectiveness of our method.
5.3 Ablation.
Max-min mutual information representation learning and non-prior context collection strategy are the key components of our proposed CSRO. Max-min mutual information representation learning uses metric distance learning and CLUB minimizes mutual information. We compare the methods with those without CLUB minimize mutual information(minMI) and non- non-prior context collection strategy (Np) components, to show the effect of each component. We conducted experiments on six environments, and the results are shown in Figure 4 In six environments, components are not used, the average return of the online test the lowest. After adding minMI, the average return is improved to varying degrees. In most environments, the average return is also increased after using Np. We noticed that in the Point-Robot and Ant-Goal, our component effect is more significant. Because in the Point-Robor and Ant-Goal the behavior policy of different tasks moves in different directions, which makes the influence context shift more serious.
notice that if both
In the online test phase, We also compare the average return of CSRO with other baselines without and with Np. shown in T the ouline test without Np. CSRO outperforms other baselines in most environments. This shows that CSRO has learned the context encoder with better generalization. Comparing online tests without Np and with Np. the average return of most methods is improved after using Np. this shows the effectiveness of Np. However, the performance of other methods with Np lower than CSRO, because others do not fully reduce the influence of behavior policy on context encoder.

---

CSRO
FOCAL
OffPear!
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.5
0.0
1.0
0.5
CORRO
BOReL
1.0
1.0
0.5
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
Figure 5: t-SNE visualization of the learned task representation space in Half-Cheetah- Vel.. Each point represents a context sampled by the non-prior context collection strategy from test tasks. Tasks of velocities from to are mapped rainbow colors, from red purple.
5.4 Visualization.
In order to analyze the task representation of the context encoder, we use t-SNE (30 to map the embedding vectors into 2D space and visualize the task representations. For each test task of Half- Cheetah- Vell task embedding. In Figure compared to the other methods, CSRO successfully clusters the task representations of the same task and distinguishes the task representations from different tasks better than other baselines.
the non-prior context collection strategy sample 100 contexts visualize
Conclusion
In this paper, we study the context shift problem in OMRL, which arises due to the distribution discrepancy between the context derived from the behavior policy during meta-training and the context obtained from the exploration policy during meta-test. Context shift problem will damage the generalization capability of the meta-policy by leading to incorrect task inference unseen tasks. We propose that CSRO address the context shift problem with only offline datasets with the key insight of reducing the effect of policy in context during both meta-training and meta-test. In meta-training, we design the max-min mutual information representation learning mechanism which reduces the influence of behavior policy in task representation. In the meta-test, we introduce the non-prior context collection strategy to reduce the effect of the exploration policy. We compared CSRO with previous OMRL methods range of challenging domains with reward or dynamic function changes. Experimental results show that CSRO can substantially reduce the context shift and outperform prior methods.
Acknowledgments
This work is partially supported by the NSF of China(under Grants 61925208, 62102399, 62222214, 62002338, U22A2028, U19B2019), Beijing Academy of Artificial Intelligence (BAAI), CAS Project for Young Scientists in Basic Research (YBncch(SSBRO222) Youth Innovation Promotion Association CAS and Xplore Prize.
10

---

References
[1] Alexander Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:/6/2.00410,, 2016.
[2] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint
2018.
[3] Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. Club: contrastive log-ratio upper bound of mutual information. In International conference on machine learning, pages 1779-1788. PMLR, 2020.
[4] Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning- identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 34:4607-400 4618, 2021.
[5] Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 1rXiv:1611.0211 2016.
[6] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. International conference on machine learning, pages 1126-1138 PMLR, 2017.
[7] Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for mnet-reeinforcementt learning: approach based contrastive learning. ArXiv, abs/2009.13891, 2020.
[8] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019.
[9] Jiaming Guo, Rui Zhang. Xishan Zhang, Shaohui Peng, Qi Yi, Zidong Du, Xing Hu, Qi Guo, and Yunji Chen. Hindsight value function for variance reduction in stochastic dynamic environment. arXiv preprint arXiv:2 107. 12216, 2021.
[10] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[11] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. ArXiv, abs/2107.02722 022299, 2021.
[12] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit g-learnin. arXiv preprint arXiv:2110.06169, 2021.
[13] Aviral Kumar, Justin Fu, G. Tucker, and Sergey Levine. Stabilizing off-policy q-larnning via bootstrapping error reduction. Neural Information Processing Systems, 2019.
[14] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179- 1191, 2020.
[15] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer, 2012.
[16] Kimin Lec, Younggyo Seo, Scunghyun Lee, Honglak Lee. and Jinwoo Shin. Context-aware dynamics model for generalization model-based reinforcement learning. In International Conference Machine Learning, pages 5757-5766. PMLR, 2020.
[17] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.016433 2020.

---

[18] Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Christensen, and Hao Su. Multi-task batch reinforcement learning with metric learning. Advances in Neural Information Processing Systems, 33:6197-6210, 2020.
[19] Lanqing Li, Yuanhao Huang, Mingzhe Chen, Siteng Luo, Dijun Luo, and Junzhou Huang. Provably improved context-based offline meta-rl with attention and contrastive learning. arXiv preprint arXiv: 102. 10774, 2021.
[20] Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforceemett learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.01112, 2020.
[21] Sen Lin, Jialin Wan, Tengyu Xu, Yingbin Liang, and Junshan Zhang. Model-based offline meta-reinforeemmnt learning with regularization. arXiv preprint arXiv: 02929, 2022.
[22] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation Iearning with contrastive predictive coding, arXiv preprint 1xXiv:187.00744 2018.
[23] Shaohui Peng, Xing Hu, Qi Yi, Rui Zhang, Jiaming Guo, Di Huang, Zikang Tian, Ruizhi Chen, Zidong Du, Qi Guo, al. Self-driven grounding: Large language model agents with automatical language-aligned skill learning. arXiv preprint arXiv:2309.0133. 2023.
[24] Shaohui Peng, Xing Hu, Rui Zhang. Jiaming Guo, Qi Yi. Ruizhi Chen, Zidong Du, Ling Li, Qi Guo, and Yunji Chen. Conceptual reinforcement learning for language-conditong tasks. arXiv preprint arXiv:2303.05069 2023.
[25] Shaohui Peng, Xing Hu, Rui Zhang, Ke Tang, Jiaming Guo, Qi Yi, Ruizhi Chen, Xishan Zhang, Zidong Du, Ling Li, et al. cousality-driven hierarchical structure discovery for reinforcement learning. arXiv preprint arXiv:2210.0664
2022.
[26] Vitchyr H Pong. Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reinforcement learning with online self-supervision. In International Conference on Machine Learning, pages 17811-17829. PMLR, 2022.
[27] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340. PMLR, 2019.
[28] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
[29] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 0026-5033. IEEE, 2012.
[30] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
[31] Jianhao Wang, Jin Zhang, Haozhe Jiang, Junyu Zhang, Liwei Wang, and Chongjie Zhang. Offline meta reinforcement learning with in-distribution online adaptation. arXiv preprint 2305 19529, 2023.
[32] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: 1361, 2019.
[33] Zhongwen Xu, H. V. Hasselt, and David Silver. Meta-gradien reinforcement learning. In Neural Information Processing Systems, 2018
[34] Qi Yi, Rui Zhang, Jiaming Guo, Xing Hu, Zidong Du, Qi Guo, Yunji Chen, et al. Object- category aware reinforcement learning. Advances in Neural Information Processing Systems, 35:36453-36465, 2022.
12

---

[35] Qi Yi, Rui Zhang, Shaohui Peng, Jiaming Guo, Yunkai Gao, Kaizhao Yuan, Ruizhi Chen, Siming Lan, Xing Hu, Zidong Du, al. Online prototype alignment for few-shot policy transfer. arXiv preprini arXiv:2300.07307, 2023.
[36] Qi Yi, Rui Zhang, Shaohui Peng, Jiaming Guo, Xing Hu, Zidong Du, Qi Guo, Ruizhi Chen, Ling Li, and Yunji Chen. Learning controllable elements oriented representations for reinforcement learning. Neurocomputiing, page 126455, 2023.
[37] Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. In International Conference on Machine Learning, pages 25747-25759. PMLR, 2022.
[38] Jin Zhang, Jianhao Wang. Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjie Zhang. Metacure: Meta reinforcement learning with impowerment-driver exploration. In International Conference on Machine Learning. pages 13600-12610. PMLR, 2021.
[39] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning, Xrviv preprint arXiv:1910.088 08.08888,, 2019.
13

---

A Pseudo-Code
Algorithm 1 CSRO Meta-training
Input:
Offline Datasets D (er of a set of training tasks (M)Err, initialize learned policy ne, Q-function Q.. context encoder qs. and CLUB encoder gg. hyperparameter À
Parameter: 0, w, ob
1: while not done do
2: for step in training steps do
3:
Sample buffer D; D and context from buffer {(8), (g, (y.r); Di, history transitions h ~ D;.
4:
Compute each transition embedding q0(2((s, z (0(2(ss,)) and task representation z (z(c))
5:
Compute Ly (
6:
Update & to minimize (vD(
7:
Compute Leneoder (Ø) LmazM1 + XLminder(
8:
Update e to minimize Leneoder (1)
9:
Use history transitions h to compute Terstic(w), Tractor (0)
10:
Update 0,w to minimize (eritti(ww) Tractor (0)
end for
12: end while
Algorithm 2 CSRO Meta-testing
Input: A set of testing tasks {M,), Nen", learned policy ne, context encoder Js. random explore step
1: for each task M, do
2:
c= 0
3:
fort 0 1 do
4:
ift < f, then
5:
Agent samples random action a, to roll out (st, ae, re, 8/)
6:
else
7:
Compute posterior go(z/c). Agent use ne Ta(als, 2) roll out (st. az Te, si)
8:
9:
end if
10:
c = c=cU (st; Ot, Tt.
11:
end for
12:
Compute posterior 2 = qs (2(c).
13: Roll out policy Te(a|s, 2) for evaluation
14:
end for
B Proof of Proposition 1
For task Mj and its corresponding behavior policy TडS when exploration policy Te and behavior policy Ty are different, if and only if mutual information between task representations and policies T(2; m.) 0, the expected test return for both the same JAr, (mO, Te) JM, (mo: nT/))
Proposition 1.
Proof. The meta-policy gets the highest expected return when exploring policy and behavior policy are the same: Te T/. However, for any exploration policy Te, we hope Jas, (TT, Te) JM. (mo. mT)). and only if for any mes (=|c) is the same. Since is collected by M, and wes (=(c)c) can be written (z. is the same. So and ne are independent, that is, the
as (2) M;, Re), i.e. for any mutual information ((2; n 0.
C Environment Details
In this section, show details about the environments of our experiment.
14

---

Point-Robot: A problem of control point robot navigation in 2D space. The start position is fixed to (0, 0). The goal of each task located on unit semicircle centered on the start position. Each task needs to control the robot from the start position to the goal. The state space is IR?, comprising the XY position of the robot. The action space [-1,- -1)3, with each dimension corresponding to the moving distance in the XY direction. The reward function is defined as the negative distance from the goal.
Half- Cheetah- Vel: Control a Cheetah to move forward and achieve goal velocity. The target velocity sampled from [1.3]. The state space is IR comprising the position and velocity of the cheetah; the angle and angular velocity of each joint. The action space is 1]", with each dimension corresponding the torque of each joint. The reward function the absolute difference between the agent's s velocity and the target velocity plus the control cost.
Ant-Goal:
The Ant-Goal task consists of controlling task is located on a circle with radius 2 centered on (0, 0). The state space is IR229 comprising the position and velocity of the ant as well the angle and angular velocity of joints. The action space is [=,, with each dimension corresponding to the torque of each joint. The reward function is defined as the negative distance from the goal plus the control cost.
"ant" robot navigate. The goal of each
The Humanoid-DD task consists controlling a "humanoid" robot in the target direction. The target direction of each task is sampled from [0, 2 The state space is 18376 and the action space [-1, 11"7 The reward function is the dot between the velocity of the robot and the target direction plus the staying alive bonus and control cost.
Humanoid-Dir:
The Hopper- Kand-Params controls one-legged robot move forward. The source code is taken from the Bardd__aa___eeee envs repository. | The tasks are varied in body mass, body inertia, joint damping, and friction. Each parameter is the product of the default value and the coefficient sampled from [15 1.5%). The state space is R and the action space is [-1, 1J3. The reward function forward velocity plus the staying alive bonus and control cost.
Hopper-Rande
BalkkerRaand-aatte The Walker-Rant Rand--arand Paramss controls a bi bi-pedal robot to move forward, also from the rand param envs repository. Each parameter obtained in the same way as Hopper- Saaper-Rand-Prramte and the reward function is the same as Hopper-Rane Rand- -ardmms....... The state space is IR and the action space [-1,1]
D Offline Data Collections
For each task, we sample 40 environments from environment distribution. Out of these, 30 environ- ments are designated as training environments, while the remaining 10 environments serve test environments. We employ SAC [10] to train an agent on each training environment and save the policy different training steps. To create offline datasets, we generate 50 trajectories using each policy from every environment. Table presents the hyperparameters employed during the collection of offline datasets.
Table 3: Hyperparameters used in offline datasets collection.
Hyperparameters
Point Point-Robol Robot
Half-Cheetah-Vel
Ant- Dir
Humanoid-Dir
Hopper-Rande Rand-Params
Maiker-Rand-Paane
Training steps
5000
le6
le6
leb
le6
le6
Initial steps
5e4
Eval frequency
200
5e4
5e4
5c4
5e4
5c4
Sampling episodes
50
Learning rate
le-4
le-4
le44
le-4
le-4
le-d
Batch size
1024
1024
1024
1024
1024
1024
E Experimental Setting
For each task, We use offline datasets collected different times to train. Details of using offline datasets in Table
We list other hyperparameters in the offline meta-training phase in Table 5
hhanss/gittubb.om///mmmss////// param_eaaaa envs.
15

---

Table 4: Details of using offline datasets: The "Checkpoints" column indicates the data collected by policies at different steps during the meta-training phase. The three numbers denote the starting steps, ending steps, and step spacing.
Eny Point-Roo Robot
Checkpoints
[2200, 4800, 200]
Val
[100000, 950000, 50000]
Ant-Goal
[100000, 950000, 50000]
Humanoid-Dir
[50000, 950000, 50000]
Saappe-Raadd-Paaaae
[50000, 950000, 50000]
Walker- -ee---annd-aannn
[50000, 950000, 50000]
Table 5: Hyperparameters used in offline meta-training.
Hyperaaaameters
Point-Robot
Half-Cheetab-Nel Vel
Ant-Dir
Humanoid-Dir
Hopper-Rand-Parame
Walker-Rand-
Reward
5
Latent dimension
20
20
20
20
40
40
BRAC
False
True
Truc
True
True
True
Batch size
256
256
256
256
256
256
Meta batch size
16
16
16
16
16
Embedding batch size
1024
100
512
256
256
256
Actor Learning rate
3c-4
3c-4
30-4
3044
30-4
30-4
Learning rate
3e.4
3e.4
3e-4
Encoder Learning rale
3e-4
3e-4
3e.4
3e.4
3e-4
3e
Maximum episade length
200
MinMi loss weightÄ
25
10
50
50
25
25
behavior regularization
50
50
30
50
Discount factor
0.9
0.99
0.99
0.99
0.99
0.99
F Additional Experiments
F.1
Comparison Offline Test Results
Offline testing is an ideal evaluation method where the context used is sampled from the pre-collected offline data in the test task, thereby disregarding the context shift problem. To assess performance in the offline test phase, we adopt the same approach as the training environment and collect offline datasets context on the testing environment.
We compare CSRO with other methods across all six environments and plot the mean and standard deviation curves of returns based on 8 random seeds in Figure [ In most environments, CSRO demonstrates competitive performance compared to other baselines. Experimental results demonstrate the effectiieness of our algorithm, even in the absence of addressing the context shift problem.
F.2
Ablation Offline and Online Test
Under the offline test scenario that ignores the context shift problem, the algorithm can achieve its highest performance. We compare the performance of the online testing method that uses the non-prior context collection strategy(Np) with offline testing. We conduct experiments on six environments and plot the mean and standard deviation curves of returns across 8 random seeds in Figure 7
In most environments, the performance of CSRO that uses Np is close to the offline test. There exists gap between Point-Robot and Ant-Goal environments due to the particularly severe and challenging context shift problem in these two environments. However, our approach still outperforms previous methods. The experimental results highlight the efficacy of our approach in addressing the context shift issue, albeit with some remaining challenges in these specific environments.
F.3 Visualize Contexts and Trajectory of Online Test
Lastly, we further study the non-prior context collection strategy. Figure showcases two different tasks in the Point-Robot environment. We illustrate the context gathered through the non-prior context collection strategy and the corresponding trajectory navigation. Observing the visualizations, we notice that the agent's perception of the task improves after random exploration. Subsequent explorations enhance the agent's understanding of the environment, enabling it to accurately navigate toward the goal.
16

---

Point Foint-Robot
Half-Chectah-vel
Ant-Goal
-440
Average Beturn
Average Return
-200
-100
-200
Average
C3R0
D300
Offiurt
Office .onno
-500
Officet
conno
-350
DOnRC
Steps(1e5) Humanoid-DD
Train Stepsíle5)
Steps(1e5)
Hopper-Rand-Paree Faraams
Wallter-Rand-Paammm
254
Average Return
Averoge
300
Average
250
000
200
Offear
300
005A Office
Offaar
30
Train Steps(le5)
Train Steps(le5)
Train Steps(le5)
Figure 6: Compared other OMRL methods, CSRO's offline testing averages returns unseen testing tasks
Point: Robot
Half-Chectah-Vel
Ant-Goal
Aucrage Return
Awer age (eetur
-100
Average Return
-190
-400
CSRO
CSSD
-500
CSRO
Train Steps(1e5)
Train Stepsile5)
Train Steps(1e5)
Humanoid-n Dir
Hopper-Rand- Rand-Params
WalkerrRand.. -Paaamss
252
Averoge Pgge Returm
Average Return
Average
20:0
CGRO
150
Tain Steps/1e5)
Train Steps(1e5)
Train Steps(1e5)
Figure 7: The average return of the offline test and the online test that uses the non-prior context collect strategy on unseen test tasks.
F.4 Additional Visualization
We have added visualizations of task representations for CSRO, FOCAL 20 and CORRO 37 in
the Point-Robot environment, as shown Figure g We can see that FOCAL is worse. In CORRO, similar tasks are closer, and many points of the same task are far apart.
17

---

Contexts Test Task
Trajectory of Test Task 1
geal
Contexts of Test Task 2
Trajectory of Test Task 2
goo
gosl
Figure 8: Visualization of contexts and trajectory after using the non-prior context collection strategy in the Foint-Robot environment.
CSRO
FOCAL
CORRO
1.0
1.0
1.0
0.5
0.5
0.0
0.0
MP9
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Figure 9: t-SNE visualization of the learned task representation space in Point-Robot Each point represents context sampled by the non-prior context collection strategy from tasks. The direction of the goal of the task is from 0 degrees to 180 degrees, from red to purple.
F.5 Additional Exploration Policy
discussed how current online meta-RL attempts to train exploration policy for collecting context, an approach that isn't suitable for addressing the context shift problem in offline meta-RL. Here, we made an attempt to train exploration policy in MetaCURE (38] using offline data, and the experimental results are presented Table 6
In Section
Observably, there is substantial drop in experimental performance when using MetaCURE's exploration policy. This is because, in an offline setting, the exploration policy cannot explore environment, and training the exploration policy with offline data confines it to the vicinity of that data. Consequently, during testing, the exploration policy becomes ineffective when encountering previously unobserved states.
18

---

Table 6: Training MetaCURE's exploration policy with offline datasets for the purpose of collecting context during testing.
Env
Point-Robot
Half-Cheetah-Vel
CSRO
-6.4 -13.2+-2.3
.48.4
CSRO+MetaCURE
-87.9+8.99
E.6 Compare with IDAQ
We compare our work with IDAQ 31. which is a recent study addressing the context shift issue. IDAQ the same training approach as FOCAL, but it tackles this problem by attempting to collect context during testing that closely matches the distribution seen during training. As depicted in Table 7 the experimental results reveal that the performance of CSRO comparable
that of IDAQ. CSRO exhibits slightly lower performance on the Point-Robot compared to IDAQ, but it significantly outperforms IDAQ on the
Vel and Walker-Rand Rand-Params tasks.
Table 7: CSRO and IDAQ compare performance on three environments.
Env
Point-Robot
Half-Cheetah-Ve Vel Walker-Rand Rande--Rad--aramms
CSRO
-6.4+0.8
.48.4+3.9
344.2+38.0
IDAQ
-5.2+0.11
-60.9
297.0=23.6
Furthermore, we also compared the performance of several methods in sparse settings, shown in Table [ can be observed that, except for IDAQ, all other methods yield poor results. For CSRO, with non-prior context collection strategy, it becomes challenging to obtain context with rewards and infer the environment, rendering ineffective.
Table 8: Compared the performance of several methods in Sparse-Point-Rooot.
Method
CSRO
FOCAL
OffPearl IDAQ
Sparse- 0.1
0.8L0.1 0.6.0.1 6.5.1.33
We conduct a detailed analysis of IDAQ's performance the Sparse-Point Point-Robot. First, we introduce the configuration of the Sparse-Point Robot- environment: The task's s goals are randomly distributed along semicircular arc with a radius of from the starting point. We select 30 goals from this distribution for training tasks and 10 goals for testing tasks. The agent only receives a non-zero reward when its distance to the goal is less than 0.2. This implies that, on average, there are four training goals within radius of 0.2 for each tested goal.
Next, we will analyze IDAQ in Dparse-Point-Robee IDAQ retains all the from the training tasks and maintains used z-set during the testing phase. During the prior exploration phase, IDAQ samples from the training tasks which has the maximum distance from the used z-set (the difference between the direction of movement of sampled and used z-set is the largest, and adds it to the used æ-set) to gather context. If the accumulated reward from the collected context surpasses that of all previously collected contexts, the current context is retained; otherwise, it is discarded. IDAQ samples the with the maximum distance a total of 10 times.
This signifies that the agent explores 10 directions of training goals which have large differences in direction, effectively covering substantial portion of the entire space. Considering that there are 4 training goals near each test goal, this implies that the sampling can capture the vicinity of the test goal, enabling the agent to attain a non-zero reward. By retaining only high-reward contexts, this indicates that after the prior exploration phase, the remaining contexts exhibit small context shift and offer non-zero rewards (containing environmental information). As a result, during the subsequent posterior sampling, IDAQ can update its belief and make improved inferences about the testing task.
Based on the aforementioned analysis, we can identify the key environmental factors under which IDAQ operates effectively in the Sparse-Point-Roose presence of training goals within a radius 0.2 from the testing goal is crucial, without this condition, the prior exploration process may encounter difficulty in collecting rewarding contexts. To validate this hypothesis, we modified the configuration of the Sparse-Point Point-Poitttttoe Robot environment, resulting in three distinct setups: Setting A reduces the radius
19

---

for obtaining rewards from 0.2 to 0.05; Setting B removes the four training goals near the testing goal. Setting C introduuess out-o--distribution (OOD) tasks, where training goals are located within the are of [0, r), while testing goals are situated in the interval 1T, x In these three setups, there are no training goals within the reward radius of the testing goals. We conducted performance testing of IDAQ under these three setups, as shown in Table above. Significant performance degradation can be observed, confirming our conjecture.
Table Comparing IDAQ's performance on the Sparse-Point-Robe under various settings.
Sparse- Sparse-Point-Ron Point-Robot Setting A Setting B Setting C
IDAQ
6.5
0.410.1 1.1
1.3+0.2
Hence, we can conclude that for sparse environments, IDAQ is suitable when both the training tasks and testing tasks follow independent and identically distributed (IID). In addition, it is essential for the reward range of testing tasks to encompass certain training tasks.
Overall, both CSRO and IDAQ aim to address the issue of context shift in offline meta RL. CSRO's goal is acquire a more essential context encoder and reduce the shift during prior exploration. IDAQ, on the other hand, aims to collect context during the testing phase that shares the same distribution as that of the training phase. Hence, both methods exhibit certain limitations. For instance, CSRO exhibits poor performance in sparse environments; IDAQ needs more exploration steps and has limitations its applicability to certain sparse environments.
20


Reviews:
[{"review_id": "gyFF9czBQH", "reviewer": "Reviewer_4mWP", "summary": "This manuscript proposes Context Shift Reduction (CSRO) for the offline meta reinforcement learning problem. It aims at solving the context shift problem with only offline datasets, and demonstrates superior empirical performance.", "questions": "1. Why is BRAC chosen as the offline backbone algorithm instead of CQL [1] or IQL [2]?\n2. If the reviewer understands correctly, in line 125 the context is defined as a subset from the offline dataset $\\{(s_j,a_j,r_j,s_j’)\\}_{j=1}^n$. Is there any specific the context need to defined in this way rather than defined in a general context space?\n\n[1] Kumar, Aviral, et al. \"Conservative q-learning for offline reinforcement learning.\" Advances in Neural Information Processing Systems 33 (2020): 1179-1191.\n[2] Kostrikov, Ilya, Ashvin Nair, and Sergey Levine. \"Offline reinforcement learning with implicit q-learning.\" arXiv preprint arXiv:2110.06169 (2021).", "limitations": "See weaknesses and questions.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper is easy to understand and the experiments look reasonable.", "weaknesses": "One major weakness the reviewer identifies is limited discussion of the disadvantages of prior methods / or major novelties of the proposed method. For example, there are several potential improvements the authors can take:\n1. Provide theoretical justifications for the proposed method, such as under what condition the method can outperforms other methods\n2. Provide intuition on why prior methods do not solve the context shift problem well enough\n3. What are the main benefits/novelties of the CSRO method in solving context shift problem\n\nOverall the reviewer thinks the empirical results look solid and promising, and the reviewer would like to adjust the rating if the author can adjust the writing to better present the proposed method. Minor presentation issues:\n1. There is a missing space in line 78, between “[6,28]” and “methods”.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for providing your comprehensive review. We greatly appreciate your insights and are glad to address each of your concerns in a detailed manner.\n\n> **Q1**: Provide theoretical justifications for the proposed method, such as under what condition the method can outperforms other methods\n\n**A1**: Denote context $c=\\\\{(s,a,r,s')\\\\}$ as an experience collected by explore policy $\\\\pi_{e}$, test task $M_{i}=(S,A,P_{i},\\\\rho, R_{i})\\\\sim p(M)$. The expected return of test task $M_{i}$ which is evaluated by a learned meta-policy $\\\\pi_{\\\\theta}(a|s,z)$ is\n $J_{M_{i}}(\\\\pi_{\\\\theta},\\\\pi_{e})=E_{s_{0}\\\\sim\\\\rho(s_{0}),z\\\\sim q_{\\\\phi}(\\\\cdot|c),a_{t}\\\\sim\\\\pi(\\\\cdot|s_{t},z), r_{t}\\\\sim R_{i}(\\\\cdot|s_{t},a_{t}),s_{t}'\\\\sim P(\\\\cdot|s_{t},a_{t})}[\\\\sum_{t=0}^{H-1}r_{t}]$. \n The meta-policy gets the highest expected return when explore policy $\\\\pi_{e}=\\\\mu_{i}$. However, for any exploration policy $\\\\pi_{e}$, we hope $J_{M_{i}}(\\\\pi_{\\\\theta},\\\\pi_{e})=J_{M_{i}}(\\\\pi_{\\\\theta},\\\\mu_{i})$, if and only if for any $\\\\pi_{e}$, $q_{\\\\phi}(z|c)$ is the same.\n\nSince $c$ is collected by $M_{i}$ and $\\\\pi_{e}$, $q_{\\\\phi}(z|c)$ can be written as $q_{\\\\phi}(z|M_{i},\\\\pi_{e})$, i.e. for any $\\\\pi_{e}$, $q_{\\\\phi}(z|M_{i},\\\\pi_{e})$ is the same. So $z$ and $\\\\pi_{e}$ are independent, that is, the mutual information $I(z;\\\\pi_{e})=0$.\n\nTherefore, we minimize the mutual information between the policy and $z$ to alleviate the context shift problem. The mutual information upper bound on minimizing mutual information has been proved in the paper \"Efficient off-policy meta-reinforcement learning via probabilistic context variables\". \n\nIn addition to the offline datasets, other methods also require additional information related to the environment, which does not solve the problem well. We focus on the case of using only offline datasets and alleviate this problem by training a more essential encoder and a more appropriate collection context strategy, so our solution is better than the mentioned method.\n\n\n\n> **Q2**: Provide intuition on why prior methods do not solve the context shift problem well enough\n\n**A2**: As for why the prior methods did not solve this problem, we briefly described it in the introduction. \n\nSpecifically, FOCAL does not consider this problem and directly uses the pre-collected context in the test environment to test evaluation. BOReL thinks that the problem comes from MDP ambiguity, by assuming that the reward function of all tasks is known and using the reward function to relabel different task data to solve it. SMAC thinks that the problem comes from the policy of collecting context during training is different from that of testing. It is solved by destroying the complete offline setting and performing some online training. The latter two methods require additional information beyond the use of offline datasets alone, and the solutions are not ideal, we hope to solve this problem with only offline datasets. We will revise the article later with a more detailed description.\n\n\n\n> **Q3**: What are the main benefits/novelties of the CSRO method in solving context shift problem\n\n**A3**: Unlike prior methods that used information beyond offline datasets, our proposed method uses only offline datasets. We propose that the cause of the problem is the influence of policy information on task embeddings, and it is not enough to maximize the mutual information between the environment and task embedding, and it is also necessary to minimize the mutual information between the policy and task embedding. We empirically find that a non-a prior context collection strategy consisting of a small number of random explorations can alleviate this problem.\n\n\n\n> **Q4**: There is a missing space in line 78, between \"[6,28]\" and \"methods\".\n\n**A4**: Thanks for your advice. We will fix it later.\n\n\n\n> **Q5**: Why is BRAC chosen as the offline backbone algorithm instead of CQL or IQL?\n\n**A5**: Because the baseline method FOCAL uses BRAC, for a fair comparison, we also use BRAC.\n\n\n\n> **Q6**: If the reviewer understands correctly, in line 125 the context is defined as a subset from the offline dataset $(s_{j},a_{j},r_{j},s_{j}')^{n}_{j=1}$. Is there any specific the context need to defined in this way rather than defined in a general context space?\n\n**A6**: Your understanding is correct, context is defined as a subset of offline datasets. This is because we only have collected offline datasets to train the encoder, so the context is defined as a subset of the offline datasets. The prior work is also defined in this way, such as \"Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning\".\n\n"}, {"review_id": "ysQJ2ckNg5", "reviewer": "Reviewer_mr2B", "summary": "The paper presents a new method called Context-Shift Robust Offline Meta-Reinforcement Learning (CSRO) to tackle the issue of context shift in offline meta-reinforcement learning. The main contributions of the paper lie in introducing max-min mutual information representation learning during meta-training to lessen the impact of behavior policy, and employing a non-prior context collection strategy during meta-testing to alleviate the impact of the exploration policy. The experimental results demonstrate that CSRO surpasses prior methods in effectively addressing context shift and enhancing performance in demanding domains with reward or dynamic function variations.", "questions": "* The authors could consider incorporating my insight on the mutual information objective into the main paper for a clearer explanation.\n* Are there alternative methods for non-prior context exploration that are more efficient and suitable for CSRO?\n* It would be helpful to include a figure in Figure 1 that illustrates the performance drop of prior works like FOCAL when using online exploration to acquire context.\n* In Figure 5, where CSRO, FOCAL, and CORRO appear to perform similarly, what metric was used to compare these three methods?\n* In addition to the HalfCheetah-Vel task, it would be beneficial to include more task representation visualization results in the Appendix.\n* It would be advantageous to include a comparison with the recent context correction work.", "limitations": "This paper lacks a discussion on its limitations. One limitation I identified is the inefficient random exploration strategy used during the meta-testing stage with non-prior context. ", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The paper is well-written and well-organized.\n* This paper addresses an important issue in offline meta RL, specifically the context shift problem that arises due to disparities between training and testing contexts.\n* The proposed method introduces a mutual information objective to reduce the reliance of the behavior policy on task representations utilizing FOCAL, and incorporates context-independent random exploration during the initial meta-testing stage.\n* Empirical evidence substantiates that the proposed method consistently outperforms other baselines in online test experiments. Additionally, the thorough ablation study validates the effectiveness of the individual components.\n", "weaknesses": "* Regarding the mutual information objective, an additional insight is that $(s,a)$ can be shared across different tasks, while the reward function $r$ plays a vital role in task inference. Equation (8) in the paper focuses the predictions more on the reward rather than solely on $(s,a)$. This insight holds particular significance for point and ant-goal tasks where state-action sharing is more prominent. However, for two Rand-Param tasks, the task can be inferred from the $(s,a)$ pairs, resulting in similar performance between CSRO and CSRO w/o minMI. If this insight holds true, I suggest the authors discuss it in the main paper.\n* Another concern is that the non-prior context exploration method directly employs random action exploration, which can be inefficient. Are there other more efficient non-prior context exploration methods that could be utilized for CSRO? The prior exploration method in off-policy meta RL [1] could provide inspiration in this regard.\n* To enhance comprehension, it would be beneficial to include a figure in Figure 1 that demonstrates the performance drop of prior works, such as FOCAL, when online exploration is employed to acquire context, as compared to offline context.\n* In Figure 5, it appears that CSRO, FOCAL, and CORRO exhibit similar performance. Could you clarify the metric used to compare these three methods? Additionally, to provide a comprehensive understanding, it would be beneficial to include more task representation visualization results in the Appendix, beyond just the HalfCheetah-Vel task.\n* Furthermore, it would be advantageous to include a comparison with the recent context correction in offline meta RL [2].\n\n[1] Zhang J, Wang J, Hu H, et al. Metacure: Meta reinforcement learning with empowerment-driven exploration[C]//International Conference on Machine Learning. PMLR, 2021: 12600-12610.\n\n[2] Wang J, Zhang J, Jiang H, et al. Offline Meta Reinforcement Learning with In-Distribution Online Adaptation[J]. arXiv preprint arXiv:2305.19529, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your advice on further improving this paper. We would like to discuss them one by one.\n\n> **Q1**: The authors could consider incorporating my insight on the mutual information objective into the main paper for a clearer explanation.\n\n**A1**: Thank you very much for your suggestion. We will modify the mutual information part in the future so that it can be explained better. First of all, ideally, the agent should pay more attention to the reward function of the environment where the reward changes, and the dynamic function of the environment where the dynamic changes, and only infer the task from there. However, our behavior policy and task are highly correlated. Both the environment of reward change and the environment of dynamic change can infer the task from $(s,a)$, but the policy and task during the test are not related, which will cause context shift problems, so we need Mutual information processing.\n\nRegarding the Rand-Param tasks you mentioned, because we use medium data, there is some interference with task inference. In addition, if the agent samples the same $(s,a)$ as the training environment when exploring the test environment, this will also cause interference, so it is still necessary to minimize the mutual information. In the ablation experiments in Figure 4, minMI improves the performance of both environments.\n\n\n\n> **Q2**: Are there alternative methods for non-prior context exploration that are more efficient and suitable for CSRO?\n\n**A2**:We add the MetaCURE method to CSRO and use the offline dataset to train the exploration policy. The experimental results are as follows:\n\n| Env           | Point-Robot | Half-Cheetah-Vel |\n| ------------- | ----------- | ---------------- |\n| CSRO          | -6.4        | -48.4            |\n| CSRO+MetaCURE | -13.2       | -87.9            |\n\nWe can see that experimenting with MetaCURE does not work well. Because in offline settings, there is no way to interact with the environment, and using offline datasets to train the exploration policy, the exploration policy is limited to the vicinity of the datasets, and the conservatism of offline is in conflict with exploration, so by training an exploration policy to solve this problem in offline settings is difficult. We will discuss in the paper why the way of training exploration strategies in online meta RL is difficult to use in offline meta RL, and cite this article.\n\n\n\n> **Q3**: It would be helpful to include a figure in Figure 1 that illustrates the performance drop of prior works like FOCAL when using online exploration to acquire context.\n\n**A3**: Thank you for your suggestion. The two test results of FOCAL and OffPearl are given below, and the obvious performance degradation can be seen.  We will add it to the paper later to enhance comprehension.\n\n| Env      | Point-Robot |        | Half-Cheetah-Vel |        |\n| -------- | ----------- | ------ | ---------------- | ------ |\n|          | offline     | online | offline          | online |\n| FOCAL    | -4.4        | -14.9  | -45.7            | -69.5  |\n| OffPearl | -5.1        | -17.8  | -123.0           | -162.8 |\n\n\n\n\n> **Q4**: In Figure 5, where CSRO, FOCAL, and CORRO appear to perform similarly, what metric was used to compare these three methods?\n\n**A4**: The metric we use to compare several methods is whether the task embeddings of the same task can be clustered together and whether different tasks can be distinguished. In Figure 5, although FOCAL and CORRO cluster the same tasks together, the degree of differentiation of different tasks is not as good as CSRO. Similar colors in Figure 5 represent similar tasks. We can see that similar tasks of FOCAL and CORRO are connected together.\n\n\n\n> **Q5**: In addition to the HalfCheetah-Vel task, it would be beneficial to include more task representation visualization results in the Appendix.\n\n**A5**: Thank you for your suggestion. We added the t-sne visualization of CSRO, FOCAL, and CORRO in the Point-Robot in the pdf of the global response. We can see that FOCAL is worse. In CORRO, similar tasks are closer, and many points of the same task are far apart. We will add more visualizations of the environment in the appendix.\n\n\n\n> **Q6**: It would be advantageous to include a comparison with the recent context correction work\n\n**A6**: We reproduced the results of IDAQ on our offline datasets, and we can see that CSRO and IDAQ are comparable. We will add this experiment later, citing this article.\n\n| Env  | Point-Robot | Half-Cheetah-vel | Walker-Rand-Params |\n| ---- | ----------- | ---------------- | ------------------ |\n| CSRO | -6.4        | -48.4            | 344.2              |\n| IDAQ | -5.2        | -60.9            | 297.0              |\n\n"}, {"review_id": "fgQRvVpsPF", "reviewer": "Reviewer_DbDu", "summary": "This paper studies the context shift problem of task representation learning in offline meta-reinforcement learning (OMRL). The proposed method, CSRO, optimizes a combination of FOCAL's objective and an adversarial objective, to maximize task information and minimize behavior policy information in task representations. Experiments in various MuJoCo tasks show that CSRO outperforms baseline methods and learns good representations.", "questions": "1. Since all the baselines are reimplemented with BRAC, why do the results for CORRO and BOReL are presented with horizontal lines rather than training curves in Figure 3?\n\n2. Is there any difference between FOCAL and the ablation CSRO w/o minMI & Np?\n\n3. Performance of ablation methods is close to CSRO in most experiments according to Figure 4. Does this mean datasets in Half-Cheetah, Humanoid and Hopper cannot reflect the context shift problem?", "limitations": "Limitations are not discussed in this paper. I hope the authors address the above issues to improve the paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The problem of context shift in in OMRL is significant. The proposed adversarial method for minimizing the information of behavior policies is novel and makes sense. \n\n2. In experiments, the baselines selected are representative. The test performance of the method is great. ", "weaknesses": "Some claims in the paper are not very accurate and can be improved:\n\nLine 175: Equation 5 is not equal to the mutual information. It should be explained as an approximation.\n\nEquation 6~8: The meaning of expectation over i and j should be explained.\n\nLine 225: In context collection, taking random actions can also cause context shift to the training distribution. Also, the context collection strategy does not appear to be an original contribution, since CORRO (section 5.6) also uses a random exploration policy to collect context.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad to answer your questions and would appreciate any further response.\n\n> **Q1**: Line 175: Equation 5 is not equal to the mutual information. It should be explained as an approximation.\n\n**A1**:Thank you for your suggestion. This is an approximation, we will modify it later to make it clear. \n\n\n\n> **Q2**: Equation 6~8: The meaning of expectation over i and j should be explained.\n\n**A2**: Here, $z_{i}$ represents  task embedding obtained by $(s_{i}, a_{i}, r_{i}, s_{i}')$ through the context encoder. $E_{j}[\\log p(z_{j}|(s_{i},a_{i}))]$ means fixing $i$ and calculating the mean of all $z$. $E_{i}[\\log p(z_{i}|(s_{i},a_{i}))\\cdots]$ means calculating the mean of each corresponding $z_{i}$ and $(s_{i}, a_{i})$. We will modify it later to make it clearer.\n\n\n\n> **Q3**: Line 225: In context collection, taking random actions can also cause context shift to the training distribution. Also, the context collection strategy does not appear to be an original contribution, since CORRO (section 5.6) also uses a random exploration policy to collect context.\n\n**A3**:There is also some context shift in the random strategy, and in Appendix F we show that this also brings about a slight performance gap. But it produces a smaller distribution shift than the common exploration strategy. The purpose of CORRO is to learn a more robust meta-policy, and then test performance under the context collected by different policies, including random policy. We do it differently, we use warmup data collection with a non-priori random strategy and continuously update the posterior distribution of the context to continue collecting, finding that this can alleviate the context shift.\n\n\n\n> **Q4**: Since all the baselines are reimplemented with BRAC, why do the results for CORRO and BOReL are presented with horizontal lines rather than training curves in Figure 3?\n\n**A4**: This is because CORRO and BOReL train the encoder first, and then train the policy. It is more appropriate to draw a horizontal line. The rest of the methods are to train the encoder and policy at the same time, so it is more appropriate to draw a curve. In this way, it is easier to compare fairly.\n\n\n\n> **Q5**:Is there any difference between FOCAL and the ablation CSRO w/o minMI & Np?\n\n**A5**: FOCAL and CSRO w/o & Np are the same.\n\n\n\n> **Q6**: Performance of ablation methods is close to CSRO in most experiments according to Figure 4. Does this mean datasets in Half-Cheetah, Humanoid and Hopper cannot reflect the context shift problem?\n\n**A6**: In fact, the gap is quite large. From Figure 4, we can see that CSRO has a significant improvement over the CSRO w/o minMI. As for the performance close to CSRO w/o Np, This is because the minMI part has basically solved this problem. In Appendix F, we gave the offline results, as you can see that the performance after using minMI is close to the highest offline performance, and NP has no performance room for improvement.\n\n"}, {"review_id": "JggVNVPrBe", "reviewer": "Reviewer_qvwn", "summary": "This paper proposes CSRO, an offline Meta-RL algorithm that deals with the distributional shift problem in offline meta-RL with online adaptation. CSRO addresses this problem by constraining the task encoding to only contain information about transition and reward, but not contain information about state-action distribution. CSRO also proposes to use random exploration at the start of exploration to further address the distribution mismatch problem. Experiment results show improved performance on MuJoCo task sets.", "questions": "1. Can the authors evaluate CSRO on Meta-World ML1? These tasks are more challenging and convincing.\n2. There is a recent work that also addresses the problem of offline meta-RL with online adaptation [1]. Although this work is contemporary to CSRO and I do not require the authors to compare these two algorithms empirically, I expect the authors to discuss pros and cons of CSRO compared to [1].\n3. Is random exploration a reasonable choice? Will it fail on more complex or sparse-reward tasks?\n\n\n\n\n\n\n[1] Offline Meta Reinforcement Learning with In-Distribution Online Adaptation. https://openreview.net/forum?id=dkYfm01yQp", "limitations": "I encourage the authors to add some discussion on CSRO in the paper. I think the random exploration is one important limitation.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The distributional shift problem is a fundamental problem in offline meta-RL with online adaptation.\n2. The information-theoretic regularization on task embeddings is novel and interesting.\n3. Presentation is clear, and the paper is easy to follow.", "weaknesses": "1. The evaluation tasks are a bit too simple. I expect the authors to evaluate on more complex task distributions like Meta-World ML1, which is more challenging and convincing.\n2. I am concerned about the efficiency of the proposed exploration method. Random exploration can be very ineffective and may struggle on hard tasks like Meta-World or sparse reward tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your detailed review. We are glad to discuss your concerns one by one. \n\n> **Q1**: Can the authors evaluate CSRO on Meta-World ML1? These tasks are more challenging and convincing.\n\n**A1**: We conduct experiments on meta-world ML1, and we can see that CSRO has higher performance than FOCAL, which proves that our method is effective.\n\n| Env   | Reach-v2 |\n| ----- | -------- |\n| CSRO  | 0.19     |\n| FOCAL | 0.10     |\n\n\n\n> **Q2**: There is a recent work that also addresses the problem of offline meta-RL with online adaptation [1]. Although this work is contemporary to CSRO and I do not require the authors to compare these two algorithms empirically, I expect the authors to discuss pros and cons of CSRO compared to [1].\n\n**A2**:We reproduced the performance of IDAQ on the three environments and found that CSRO and IDAQ are comparable.\n\n| Env  | Point-Robot | Half-Cheetah-vel | Walker-Rand-Params |\n| ---- | ----------- | ---------------- | ------------------ |\n| CSRO | -6.4        | -48.4            | 344.2              |\n| IDAQ | -5.2        | -60.9            | 297.0              |\n\nBecause the IDAQ collection context tends to choose a context with a higher reward, its performance in the expert datasets should be better. However, due to its approximate greedy iteration, if the number of adaptation steps is reduced, the performance will drop. We found that IDAQ on Half-Cheetah-Vel, adapting to 1000 steps drops to using 600 steps, and the performance will drop from -60.9 to -90.1, while CSRO uses 600 steps and performance is  -48.4.\n\n\n\n> **Q3**: Is random exploration a reasonable choice? Will it fail on more complex or sparse-reward tasks?\n\nA3: We described in Figure 1 and Section 4.3 that the common exploration method collects context, and the collected context will be related to the sampled initial $z_{0}$. Subsequent posterior $q_{\\phi}(z|c)$ are also affected by $z_{0}$, leading to wrong task inferences. Therefore, we propose that we need to eliminate this kind of erroneous prior, first, use a small amount of random data for initial context collection and then meta-policy $\\pi_{\\theta}(a|s,z)$ continues to explore and collect context. Since it is infeasible to train an exploration strategy like online meta RL using only offline data, our approach is reasonable.\n\nThe CORRO and OffPearl methods we compared were also not tested in sparse environments, and our method may not work in sparse environments.\n\n"}]

Author rebuttal:
Dear reviewer, PDF has our supplementary experimental diagram.

Meta-review: This paper studies an important problem in offline meta RL, and proposes a novel and very sensible approach. All the reviewers agreed to accept the paper.

I have a comment to the authors in respect to the comparison with BOREL. First, BOREL (and similarly VariBAD) are trained to be Bayes-optimal, while the proposed method is based on Thompson sampling (and also includes random exploration at test, which is even less optimal). This means that potentially, if context shift wasn’t a problem, BOREL could reach  much higher scores at test time (see the BOREL paper for a discussion and comparison to PEARL). This should be more clearly mentioned in the related work, as from the current writing it appears that BOREL and CSRO solve the same problem, which is not correct. Another difference is that CSRO requires knowing the task labels in the training data, while BOREL does not (yet BOREL requires either the reward functions, or to collect data in a very specific way). The two points above make the comparison with BOREL less meaningful, and the authors should clarify it in the text.
Decision: Accept (poster)