Title: Context Shift Reduction for Offline Meta-Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: offline meta-reinforcement learning

offline reinforcement learning

meta-reinforcement learning
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

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 C ontext S hift R eduction for O MRL ( 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.

1 Introduction

Meta-Reinforcement 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-RL methods exhibit swift adaptation to new tasks with a limited number of trials. Context-based meta-RL methods [27, 39, 7] have emerged as the most popular 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 set of transitions) and subsequently condition the policy on the task representation to execute appropriate actions.

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 [17, 32, 14] presents a 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-reinforcement 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 1, we compared the test results of FOCAL [20] and OffPearl [27] in 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, but they do not fully address the context shift problem. Some approaches [20, 21, 19] utilize precollected 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, [4] assumes that the reward functions of all tasks are known, [26] requires the collection of a substantial volume of online data post offline training.

In this paper, we introduce a novel approach called C ontext S hift R eduction for O MRL ( 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 1 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 a 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, 4, 26] conditions on an initial task representation. The collected context corresponds to the policy of the initial task representation, which significantly differs from the behavior policy, as shown in Figure 1 middle. This context leads to incorrect task inference and control, as depicted in Figure 1 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 a 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 MuJoCo 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

Env

Point-Robot

Half-Cheetah-Vel

context A

context B

context A

context B

FOCAL

-4.4

±

0.1

-14.9

±

1.1

-45.7

±

2.7

-69.5

±

9.6

OffPearl

-5.1

±

0.1

-17.8

±

1.5

-123.0

±

11.5

-162.8

±

28.8

OfflineDatasetsofMeta-Training

ContextsofMeta-Test

TrajectoriesofMeta-Test

test 0

testo

test1

test 2

goal

test1

test2

goal

policy parameter initialization that requires only a 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, 11] address this problem as a PODMP. These methods employ a 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].

Offline 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 environment during the test phase. Offline RL methods [15, 14, 12, 8] use pre-collected data generated by behavior 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, 8, 13] 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-learning on pre-collected offline datasets. Existing OMRL methods [20, 18, 4, 37, 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 is impractical and ignores the context shift problem. In contrast, online test methods [4, 26] use the context collected by the agent in the test task. While these online test methods can mitigate the context shift problem, they require additional information beyond just offline datasets, such as 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.

Mutual Information Optimization. In order to learn more about essential representation, mutual information optimization is often employed to enhance the correlation between the representation

Meta-TrainingPhase

Meta-TestPhase

qs(zlc)

LmaxMI

Explore

(s,a,r,s')}→

q(z(s,a,r,s'))

Q(s,a,z)

1

Leric

TNP

→gs(zlc)→π(als,z)

n

OfflineDatasets

y

→q,(z/(s,a))

π(als,z)

↑

LVD

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. [1, 3] 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.

3 Preliminaries

Problem Formulation. In RL, we formalize each task as a Markov Decision Process (MDP), defined by M = ( S, A, P, ρ 0 , R, γ ) with state space S , action space A , transition function P ( s ′ | s, a ) , reward function R ( s, a ) , initial state distribution ρ 0 ( s ) and discount factor γ ∈ (0 , 1) . The distribution of policy is π ( a | s ) . We aim to maximize the expected cumulative rewards J ( π ) = E τ [ ∑ ∞ t =0 γ t R ( s t , a t )] , where τ is trajectory following s 0 ∼ ρ 0 , a t ∼ π ( a t | s t ) , s t +1 ∼ P ( s t +1 | s t , a t ) . We define the value function and action-value function as:

In offline meta-reinforcement learning (OMRL), we assume each task sampled from a prior task distribution, M i = ( S, A, P i , ρ 0 , R i , γ ) ∼ p ( M ) . Different tasks have the same state space and action space, but the reward function or transition dynamic is different due to different parameters, So the task distribution can be defined as p ( R,P ) . The offline datasets D i = { ( s i,j , a i,j , r i,j , s ′ i,j ) } N size j =1 of task M i is composed of the state, action, reward, and next-state tuples collected by the taskdependent behavior policy π i β ( a | s ) . Since the parameters of different tasks are unknown, for each task, context-based OMRL samples a mini-batch of historical data from the offline datasets as the context c = { ( s j , a j , r j , s ′ j ) } n j =1 . The context encoder transforms the context c into a task representation z = q ϕ ( z | c ) . Policy, value function, and action-value function are all conditioned on z . We use offline data to train the meta policy π θ ( a | s, z ) with parameters θ , and will not interact with the environment during the training. OMRL aims to maximize the expected cumulative rewards in the test task:

4 Method

4.1 Description of Context Shift Problem

OMRL uses pre-collected offline data from task-dependent behavior policy to maximize the 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 i for

collecting data is π i β ( a | s ) , and the offline data is D i = { ( s i,j , a i,j , r i,j , s i,j ) ′ } N size j =1 . The distribution of offline data is denoted by

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 to 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 1 left, the task goal is to 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-test phase, the common exploration policy is conditioned on an initial task representation z 0 sampled from a prior p ( z ) to collect context. The initial task representation z 0 contains a specific direction, so the collected context will move to this specific direction which is different from behavior policy, as shown in Figure 1 middle. Affected by this context, the agent will go along the specific direction instead of going to the task goal, as shown in Figure 1 right.

To address the context shift problem, we propose that CSRO reduce the effect of policy in the context in both the 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 q ϕ to distinguish tasks. We first sample the context c i = { ( s i,j , a i,j , r i,j , s ′ i,j ) } N c j =1 from the offline datasets D i of task M i . And then the context encoder extracts each transition embedding from each transition of the context: z i,j = q ϕ ( z | ( s i,j , a i,j , r i,j , s ′ i,j )) . Finally, we use the mean of transition embedding z i = E j [ z i,j ] as the task representation of task M i . 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 is a 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 z and the behavior policy π β .

Since the task representation needs to represent the task information, we need to maximize the mutual information between task representation and the task. We use distance metric learning in [20] to obtain mutual information approximation with task:

̸

where c i is context sampled from task M i , y i is task label, z i = q ϕ ( z | c i ) 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.

Definition 1. In the test task M i = ( S, A, P i , ρ, R i ) ∼ p ( M ) , context c as a context collected by exploration policy π e . The expected return of test task M i which is evaluated by a learned meta-policy π θ ( a | s, z ) is:

Proposition 1. For task M i and its corresponding behavior policy π i β , when exploration policy π e and behavior policy π i β are different, if and only if mutual information between task representations and policies I ( z ; π e ) = 0 , the expected test return for both is the same J M i ( π θ , π e ) = J M i ( π θ , π i β ) .

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 representations and policies π β to 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, a, r, s ′ ) is jointly determined by task and behavior policy p ( a, r, s ′ | s ) = π β ( a | s ) p ( s ′ | s, a ) r ( s, a ) , we use ( s, a ) tuple to represent behavior policy. Following [3], we introduce using CLUB to minimize mutual information between transition embedding z and ( s, a ) . Mutual information upper bound estimation of CLUB is defined as:

However the conditional distribution p ( z | ( s, a )) is unavailable on our tasks, therefore we use a variational distribution q ψ ( z | ( s, a )) to approximate p ( z | ( s, a )) by minimizing Log-likelihood:

We minimize mutual information between transition embedding z and ( s, a ) by minimize mutual information upper bound estimation:

which z i represents task embedding obtained by ( s i , a i , r i , s ′ i ) through the context encoder. Here, the variational distribution estimator q ψ and context encoder q ϕ are used for adversarial training. If each transition embedding z obtained by the context encoder q ϕ ( z | ( s, a, r, s ′ )) contains more ( s, a ) information, the q ψ ( z | ( s, a )) has higher accuracy, but mutual information upper bound L minMI ( ϕ ) will be higher. If the transition embedding z does not contain ( s, a ) information, the ability to predict z i through ( s i , a i ) is the same as that from other ( s j , a j ) , L minMI ( ϕ ) 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, we can get the total loss of context encoder as:

where λ is a hyperparameter to adjust the weight of the two mutual information. Based on 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 c and update the posterior distribution q ϕ ( z | c ) according to context c . The current exploration strategy of OMRL [27, 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 z 0 sampled from the prior distribution p ( z ) to explore the new task. After that, the agent updates posterior q ϕ ( z | c ) by collected context c and samples new z . Then the agent continues to explore the task to collect context conditioned on z . The steps of sampling z and collecting context are performed iteratively. Finally, the agent executes task inference according to all collected contexts.

Because there is 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 c will correlate with the sampled initial z 0 . In this way, the posterior q ϕ ( z | c ) will be affected by z 0 , 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 [7, 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, so by training an exploration policy to solve this problem in offline settings is difficult.

We propose a non-prior context collection strategy to eliminate this problem. Our strategy is not conditioned on sampled initial z 0 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 to the posterior q ϕ ( z | c ) and updates the posterior at 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 z 0 .

4.4 Algorithm Summary

Finally, in order to avoid the problem of overestimation of Q caused by the distribution shift of the actions in offline learning, we use BRAC [32] to constrain the similarity between learned policy π θ and behavior policy π β by punishing the distribution shift of learned policy π θ . The optimization objectives under the actor-critic framework become:

where D is the divergence estimate of learned meta-policy π θ ( ·| s, z ) and behavior policy π β ( ·| s, z ) , and we use the KL 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 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 https: //github.com/MoreanP/CSRO.git .

5.1 Experimental Settings

Environments. We evaluate our method on the Point-Robot and MuJoCo [29] that are often used as the offline meta-RL benchmarks:(1) Point-Robot is a 2D navigation task of a point in continuous space, and the goal is located on a 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 to control the humanoid agent to move in the direction provided by the task. (5) Hopper-Rand-Params is to control the velocity of the agent as fast as possible. (6) Walker-Rand-Params 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 in different environments have different physical parameters, including mass, inertia, stamping, and friction.

Point-Robot

Half-Cheetah-Vel

5.0

-50

375

7.5

-100

-400

Average Return

-10.0

Average Return

Average Return

12.5

150

425

-15.0

200

450

-17.5

475

-20.0

CSRO

250

500

CSRO

FOCAL

FOCAL

FOCAL

22.5

OffPearl

CORRO

0OE-

OffPearl

CORRO

-525

OffPearl

CORRO

25.0

BOReL

-350

BOReL

550

BOReL

1

2

E

4

5

6

1

2

3

4

5

6

7

8

6

2

4

8

10

Train Steps(1e5)

Train Steps(1e5)

Train Steps(1e5)

Humanoid-Dir

Hopper-Rand-Params

Walker-Rand-Params

400

700

350

350

Average Return

600

Return

OOE

Average Return

300

500

AverageR

250

250

200

400

CSRO

150

CSRO

200

CSRO

FOCAL

FOCAL

FOCAL

300

OffPearl

CORRO

100

OffPearl

CORRO

150

OffPearl

CORRO

BOReL

50

BOReL

BOReL

2

4

6

8

2

3

4

5

6

7

8

3

4

5

6

8

9

Train Steps(1e5)

Train Steps(1e5)

Train Steps(1e5)

Offline Data Collections. For each environment, we 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 at different training times as behavior policy. We use each policy to 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 to 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 CV AE [28] or adds noise to produce negative samples, and uses InfoNCE [22] to train robust context encoder. (4) BOReL [4] builds on VariBAD [39]. For a fair comparison, we use a variant of BOReL that does not utilize oracle reward functions, as introduced in the original paper [4].

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 are available in the Appendix.

Point-Robot

Half-Cheetah-Vel

50

375 -

Average Return

-8

Average Return

-100

Average Return

400

-10-

425

-12

150

450

-141

200

475-

-16

500

CSRO

-250

CSRO

CSRO

-18

CSROw/o Np

CSOR w/o minMI

CSOR w/o minMI

CSRO w/o Np

525

CSOR w/o minMI

CSRO w/o Np

20

1

2

Train Steps(1e5)

3

CSRO w/o minMI& Np

4

5

6

-300

1

2

E

Train Steps(1e5)

4

5

CSRO w/o minMI & Np

6

7

8

6

-550

2

Train Steps(1e5)

4

6

CSRO w/o minMI & Np

8

10

Humanoid-Dir

Hopper-Rand-Params

Walker-Rand-Params

400

700

350

350

Average Return

600

Average Return

300

Return

000

250

500

200

Average

250-

400

150

200

300

CSOR w/o minMI

CSRO

CSRO w/o Np

CSRO w/o minMI & Np

100

50

CSRO w/o Np

CSROw/ominMI&Np

CSOR w/o minMI

CSRO

150

CSRO w/o Np

CSRO w/o minMI & Np

CSRO

CSOR w/o minMI

1

4

5

6

6

1

2

3

5

6

2

4

5

6

7

8

Train Steps(1e5)

Train Steps(1e5)

Train Steps(1e5)

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 we proposed.

We plot the mean and standard deviation curves of returns across 8 random seeds in Figure 3. In the six environments, CRSO outperforms all the baselines, especially in Point-Robot 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-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, we notice that if both components are not used, the average return of the online test is 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-Robot and Ant-Goal, the behavior policy of different tasks moves in different directions, which makes the influence of context shift more serious.

In the online test phase, We also compare the average return of CSRO with other baselines without and with Np. As shown in Table 2, in the online 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 is lower than CSRO, because others do not fully reduce the influence of behavior policy on context encoder.

CSRO

FOCAL

OffPearl

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

0.0

0.5

1.0

0.0

0.5

1.0

0.0

0.5

1.0

CORRO

BOReL

1.0

1.0

0.5

0.5

0.0L

0.0

0.0

0.5

1.0

0.0

0.5

1.0

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 HalfCheetah-Vel, we use the non-prior context collection strategy to sample 100 contexts to visualize task embedding. In Figure 5, 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.

6 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 on 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 on a 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(YSBR-029), Youth Innovation Promotion Association CAS and Xplore Prize.

References

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410 , 2016.

Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint arXiv:1801.04062 , 2018.

Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. Club: A contrastive log-ratio upper bound of mutual information. In International conference on machine learning , pages 1779-1788. PMLR, 2020.

Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learningidentifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems , 34:4607-4618, 2021.

Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. Rl2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning , pages 1126-1135. PMLR, 2017.

Haotian Fu, Hongyao Tang, Jianye Hao, Chen Chen, Xidong Feng, Dong Li, and Wulong Liu. Towards effective context for meta-reinforcement learning: an approach based on contrastive learning. ArXiv , abs/2009.13891, 2020.

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.

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:2107.12216 , 2021.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

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.02729, 2021.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Aviral Kumar, Justin Fu, G. Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Neural Information Processing Systems , 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:11791191, 2020.

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning , pages 45-73. Springer, 2012.

Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning , pages 5757-5766. PMLR, 2020.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

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.

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:2102.10774 , 2021.

Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.01112 , 2020.

Sen Lin, Jialin Wan, Tengyu Xu, Yingbin Liang, and Junshan Zhang. Model-based offline meta-reinforcement learning with regularization. arXiv preprint arXiv:2202.02929 , 2022.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748 , 2018.

Shaohui Peng, Xing Hu, Qi Yi, Rui Zhang, Jiaming Guo, Di Huang, Zikang Tian, Ruizhi Chen, Zidong Du, Qi Guo, et al. Self-driven grounding: Large language model agents with automatical language-aligned skill learning. arXiv preprint arXiv:2309.01352 , 2023.

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-conditioned tasks. arXiv preprint arXiv:2303.05069 , 2023.

Shaohui Peng, Xing Hu, Rui Zhang, Ke Tang, Jiaming Guo, Qi Yi, Ruizhi Chen, Xishan Zhang, Zidong Du, Ling Li, et al. Causality-driven hierarchical structure discovery for reinforcement learning. arXiv preprint arXiv:2210.06964 , 2022.

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.

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.

Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research , 9(11), 2008.

Jianhao Wang, Jin Zhang, Haozhe Jiang, Junyu Zhang, Liwei Wang, and Chongjie Zhang. Offline meta reinforcement learning with in-distribution online adaptation. arXiv preprint arXiv:2305.19529 , 2023.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Zhongwen Xu, H. V. Hasselt, and David Silver. Meta-gradient reinforcement learning. In Neural Information Processing Systems , 2018.

Qi Yi, Rui Zhang, Jiaming Guo, Xing Hu, Zidong Du, Qi Guo, Yunji Chen, et al. Objectcategory aware reinforcement learning. Advances in Neural Information Processing Systems , 35:36453-36465, 2022.

Qi Yi, Rui Zhang, Shaohui Peng, Jiaming Guo, Yunkai Gao, Kaizhao Yuan, Ruizhi Chen, Siming Lan, Xing Hu, Zidong Du, et al. Online prototype alignment for few-shot policy transfer. arXiv preprint arXiv:2306.07307 , 2023.

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. Neurocomputing , page 126455, 2023.

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.

Jin Zhang, Jianhao Wang, Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjie Zhang. Metacure: Meta reinforcement learning with empowerment-driven exploration. In International Conference on Machine Learning , pages 12600-12610. PMLR, 2021.

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 metalearning. arXiv preprint arXiv:1910.08348 , 2019.

A Pseudo-Code

Algorithm 1 CSRO Meta-training Input : Offline Datasets D = { D i } N env i =1 of a set of training tasks { M i } N env i =1 , initialize learned policy π θ , Q-function Q ω , context encoder q ϕ , and CLUB encoder q ψ , hyperparameter λ Parameter : θ, ω, ϕ, ψ 1: while not done do 2: for step in training steps do 3: Sample buffer D i ∼ D and context from buffer c = { ( s j , a j , r j , s ′ j ) } ∼ D i , history transitions h ∼ D i . 4: Compute each transition embedding z = q ϕ ( z | ( s, a, r, s ′ )) , z = q ψ ( z | ( s, a )) and task representation z = q ϕ ( z | c ) 5: Compute L V D ( ψ ) 6: Update ψ to minimize L V D ( ψ ) 7: Compute L encoder ( ϕ ) = L maxMI ( ϕ ) + λL minMI ( ϕ ) 8: Update ϕ to minimize L encoder ( ϕ ) 9: Use history transitions h to compute L critic ( ω ) , L actor ( θ ) 10: Update θ , ω to minimize L critic ( ω ) , L actor ( θ ) 11: end for 12: end while

Algorithm 2

B Proof of Proposition 1

Proposition 1. For task M i and its corresponding behavior policy π i β , when exploration policy π e and behavior policy π i β are different, if and only if mutual information between task representations and policies I ( z ; π e ) = 0 , the expected test return for both is the same J M i ( π θ , π e ) = J M i ( π θ , π i β ) .

Proof. The meta-policy gets the highest expected return when exploring policy and behavior policy are the same: π e = π i β . However, for any exploration policy π e , we hope J M i ( π θ , π e ) = J M i ( π θ , π i β ) , if and only if for any π e , q ϕ ( z | c ) is the same. Since c is collected by M i and π e , q ϕ ( z | c ) can be written as q ϕ ( z | M i , π e ) , i.e. for any π e , q ϕ ( z | M i , p e ) is the same. So z and π e are independent, that is, the mutual information I ( z ; π e ) = 0 .

C Environment Details

In this section, we show details about the environments of our experiment.

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 is located on a 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 R 2 , comprising the XY position of the robot. The action space is [ -1 , -1] 2 , 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 is sampled from [1 , 3] . The state space is R 20 , comprising the position and velocity of the cheetah; the angle and angular velocity of each joint. The action space is [ -1 , 1] 6 , with each dimension corresponding to the torque of each joint. The reward function is the absolute difference between the agent's velocity and the target velocity plus the control cost.

Ant-Goal: The Ant-Goal task consists of controlling an "ant" robot to navigate. The goal of each task is located on a circle with radius 2 centered on (0 , 0) . The state space is R 29 , comprising the position and velocity of the ant as well as the angle and angular velocity of 8 joints. The action space is [ -1 , 1] 8 , 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.

Humanoid-Dir: The Humanoid-Dir task consists of controlling a "humanoid" robot in the target direction. The target direction of each task is sampled from [0 , 2 π ] . The state space is R 376 and the action space is [ -1 , 1] 17 . The reward function is the dot between the velocity of the robot and the target direction plus the staying alive bonus and control cost.

Hopper-Rand-Params: The Hopper-Rand-Params controls a one-legged robot to move forward. The source code is taken from the rand_param_envs repository. 2 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 [1 . 5 -3 , 1 . 5 3 ] . The state space is R 11 and the action space is [ -1 , 1] 3 . The reward function is forward velocity plus the staying alive bonus and control cost.

Walker-Rand-Paras: The Walker-Rand-Params controls a bi-pedal robot to move forward, also from the rand_param_envs repository. Each parameter is obtained in the same way as Hopper-Rand-Params and the reward function is the same as Hopper-Rand-Params. The state space is R 17 and the action space is [ -1 , 1] 6 .

D Offline Data Collections

For each task, we sample 40 environments from environment distribution. Out of these, 30 environments are designated as training environments, while the remaining 10 environments serve as test environments. We employ SAC [10] to train an agent on each training environment and save the policy at different training steps. To create offline datasets, we generate 50 trajectories using each policy from every environment. Table 3 presents the hyperparameters employed during the collection of offline datasets.

E Experimental Setting

For each task, We use offline datasets collected at different times to train. Details of using offline datasets in Table 4:

We list other hyperparameters in the offline meta-training phase in Table 5.

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 our performance in the offline test phase, we adopt the same approach as the training environment and collect offline datasets as 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 6. In most environments, CSRO demonstrates competitive performance compared to other baselines. Experimental results demonstrate the effectiveness 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 a 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 8 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.

Point-Robot

Half-Cheetah-Vel

-50

-300

Average Return

Average Return

-100

Average Return

-8

150

350

-10-

-200

400-

-12

250

-450

14

CSRO

CSRO

CSRO

FOCAL

00E-

FOCAL

-500

FOCAL

OffPearl

OffPearl

OffPearl

-16

CORRO

OsE-

CORRO

CORRO

1

2

Train Steps(1e5)

3

4

5

6

1

2

Train Steps(1e5)

3

4

5

6

7

8

6

2

Train Steps(1e5)

4

6

8

10

Humanoid-Dir

Hopper-Rand-Params

Walker-Rand-Params

700

350

350

Average Return

600

Average Return

000

Average Return

300

250

500

250

200

400-

150

200

CSRO

CSRO

CSRO

FOCAL

OffPearl

100

FOCAL

OffPearl

150

FOCAL

OffPearl

CORRO

CORRO

CORRO

50

1

2

Train Steps(1e5)

3

4

5

6

7

8

9

1

2

Train Steps(1e5)

3

4

5

6

7

8

2

Train Steps(1e5)

3

4

5

8

6

Point-Robot

Half-Cheetah-Vel

-300

50

6

350

Average Return

Average Return

-100

Average Return

-8

-10-

150

400

-12-

200

450

14

-16

CSRO

250

CSRO

500

CSRO

CSROoffline test

CSRO offline test

CSRO offline test

2

3

4

5

6

2

3

4

5

6

7

8

4

6

8

10

Train Steps(1e5)

Train Steps(1e5)

Train Steps(1e5)

Humanoid-Dir

Hopper-Rand-Params

Walker-Rand-Params

400

700

350

350

Average Return

600

Average Return

Average Return

OOE

300

500

250

250

400

200

200

00E

150

CSRO

CSRO

CSRO

CSRO offline test

100

CSRO offline test

150

CSRO offline test

2

3

4

6

7

8

6

1

2

3

5

6

8

1

2

3

4

6

7

8

Train Steps(1e5)

Train Steps(1e5)

Train Steps(1e5)

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 in Figure 9. We can see that FOCAL is worse. In CORRO, similar tasks are closer, and many points of the same task are far apart.

Contexts of Test Task 1

Trajectory of Test Task 1

goal

goal

Contexts of Test Task 2

Trajectory of Test Task 2

1eob

Meob

CSRO

FOCAL

CORRO

1.0

1.0

1.0

0.5

0.5

0.5

0.0

0.0

0.0

0.0

0.5

1.0

0.0

0.5

1.0

0.0

0.5

1.0

F.5 Additional Exploration Policy

In Section 4.3, we discussed how current online meta-RL attempts to train an 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 an exploration policy in MetaCURE [38] using offline data, and the experimental results are presented in Table 6.

Observably, there is a substantial drop in experimental performance when using MetaCURE's exploration policy. This is because, in an offline setting, the exploration policy cannot explore the 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.

F.6 Compare with IDAQ

We compare our work with IDAQ [31], which is a recent study addressing the context shift issue. IDAQ uses 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 is comparable to that of IDAQ. CSRO exhibits slightly lower performance on the Point-Robot compared to IDAQ, but it significantly outperforms IDAQ on the Half-Cheetah-Vel and Walker-Rand-Params tasks.

Furthermore, we also compared the performance of several methods in sparse settings, as shown in Table 8. It can be observed that, except for IDAQ, all other methods yield poor results. For CSRO, with a non-prior context collection strategy, it becomes challenging to obtain context with rewards and infer the environment, rendering it ineffective.

We conduct a detailed analysis of IDAQ's performance in the Sparse-Point-Robot. First, we introduce the configuration of the Sparse-Point-Robot environment: The task's goals are randomly distributed along a semicircular arc with a radius of 1 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 a radius of 0.2 for each tested goal.

Next, we will analyze IDAQ in Sparse-Point-Robot: IDAQ retains all the z from the training tasks and maintains a used z -set during the testing phase. During the prior exploration phase, IDAQ samples a z from the training tasks which has the maximum distance from the used z -set (the difference between the direction of movement of sampled z and used z -set is the largest, and adds it to the used z -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 z 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 a 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 a 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-Robot: presence of training goals within a radius of 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-Robot environment, resulting in three distinct setups: Setting A reduces the radius

for obtaining rewards from 0.2 to 0.05; Setting B removes the four training goals near the testing goal. Setting C introduces out-of-distribution (OOD) tasks, where training goals are located within the arc of [0 , 3 4 π ) , while testing goals are situated in the interval [ 3 4 π, π ] . 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 9 above. Significant performance degradation can be observed, confirming our conjecture.

Table 9: Comparing IDAQ's performance on the Sparse-Point-Robot under various settings.

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 to acquire a more essential context encoder and reduce the shift during prior exploration. IDAQ, on the other hand, aims to collect a 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 in its applicability to certain sparse environments.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nI encourage the authors to add some discussion on CSRO in the paper. I think the random exploration is one important limitation.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nSome 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\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nLimitations are not discussed in this paper. I hope the authors address the above issues to improve the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe paper is easy to understand and the experiments look reasonable.\n\nWEAKNESSES:\nOne 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\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nSee weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}