SlowGuess/NeurIPS_2023_OCR huggingface dataset:
By combining the above two mutual information, we can get the total loss of context encoder as:
$$
L _ {\text {e n c o d e r}} (\phi) = L _ {\max M I} + \lambda L _ {\min M I}, \tag {10}
$$
where  $\lambda$  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_{\phi}(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_{\phi}(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_{\phi}(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_{\phi}(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  $\pi_{\theta}$  and behavior policy  $\pi_{\beta}$  by punishing the distribution shift of learned policy  $\pi_{\theta}$ . The optimization objectives under the actor-critic framework become:
$$
L _ {c i r t i c} (\omega) = \mathbb {E} _ {(s, a, r, s ^ {\prime}) \sim D, z \sim q _ {\phi} (z | c), a ^ {\prime} \sim \pi_ {\theta} (\cdot | s ^ {\prime}, z)} [ (Q _ {\omega} (s, a, z) - r - \gamma Q _ {\omega} ^ {t a r g e t} (s ^ {\prime}, a ^ {\prime}, z)) ^ {2} ]. \tag {11}
$$
$$
L _ {\text {a c t o r}} (\theta) = - \mathbb {E} _ {(s, a, r, s ^ {\prime}) \sim D, z \sim q _ {\phi} (z | c), a ^ {\prime \prime} \sim \pi_ {\theta} (\cdot | s, z)} [ Q _ {\omega} (s, a ^ {\prime \prime}, z) - \alpha D (\pi_ {\theta} (\cdot | s, z), \pi_ {\beta} (\cdot | s, z)) ]. \tag {12}
$$
where  $D$  is the divergence estimate of learned meta-policy  $\pi_{\theta}(\cdot |s,z)$  and behavior policy  $\pi_{\beta}(\cdot |s,z)$ , and we use the KL divergence.
We summarized our meta-training and meta-test method in the pseudo-code of Appendix A.
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.
**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.
Figure 3: The average return of the online test phase on unseen test tasks versus other OMRL methods. Point-Robot, Half-Cheetah-Vel, Ant-Goal, and Humanoid-Dir are all reward functions changing environments. Hopper-Rand-Params and Walker-Rand-Params are dynamic functions changing environments. The shaded region shows a standard deviation across 8 seeds.
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 CVAE [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.
Table 2: Comparing CSRO with other baselines in the online test phase without and with the non-prior context collection strategy(Np) in Point-Robot, Half-Cheetah-Vel, and Walker-Rand-Params environments.
Env	Point-Robot	Half-Cheetah-Vel	Walker-Rand-Params
w/ Np	w/o Np	w/ Np	w/o Np	w/ Np	w/o Np
CSRO	-6.4±0.8	-9.2±0.6	-48.4±3.9	-68.5±13.9	344.2±38.0	319.7±38.4
FOCAL	-11.8±1.6	-14.9±1.1	-60.9±5.7	-69.5±9.6	253.3±42.7	247.5±29.4
OffPearl	-17.0±1.6	-17.8±1.5	-133.7±18.9	-162.8±28.8	284.5±30.9	262.0±24.5
CORRO	-7.8±1.9	-10.5±3.0	-65.6±9.3	-92.1±23.2	312.5±46.6	275.2±73.9
BORL	-21.6±3.9	-23.2±5.8	-90.1±28.3	-56.1±10.7	260.6±40.2	245.8±32.9
Figure 4: Ablation study in six environments, to compare CSRO with methods that without CLUB minimize mutual information (minMI) and non-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-priori 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.
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 1 to 3 are mapped to rainbow colors, from red to 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-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.
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.
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.
[35] 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.
[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. Neurocomputing, 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 empowerment-driven exploration. In International Conference on Machine Learning, pages 12600-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. arXiv preprint arXiv:1910.08348, 2019.
A Pseudo-Code
Algorithm 1 CSRO Meta-training
Input: Offline Datasets  $D = \{D_{i}\}_{i = 1}^{N_{env}}$  of a set of training tasks  $\{M_i\}_{i = 1}^{N_{env}}$  , initialize learned policy   
 $\pi_{\theta}$  Q-function  $Q_{\omega}$  , context encoder  $q_{\phi}$  , and CLUB encoder  $q_{\psi}$  , hyperparameter  $\lambda$    
Parameter:  $\theta ,\omega ,\phi ,\psi$    
1: while not done do   
2: for step in training steps do   
3: Sample buffer  $D_{i}\sim D$  and context from buffer  $c = \{(s_j,a_j,r_j,s_j')\} \sim D_i$  , history transitions  $h\sim D_i$    
4: Compute each transition embedding  $z = q_{\phi}(z|(s,a,r,s'))$ $z = q_{\psi}(z|(s,a))$  and task representation  $z = q_{\phi}(z|c)$    
5: Compute  $L_{VD}(\psi)$    
6: Update  $\psi$  to minimize  $L_{VD}(\psi)$    
7: Compute  $L_{encoder}(\phi) = L_{maxMI}(\phi) + \lambda L_{minMI}(\phi)$    
8: Update  $\phi$  to minimize  $L_{encoder}(\phi)$    
9: Use history transitions  $h$  to compute  $L_{critic}(\omega)$ $L_{actor}(\theta)$    
10: Update  $\theta ,\omega$  to minimize  $L_{critic}(\omega)$ $L_{actor}(\theta)$    
11: end for   
12: end while
Algorithm 2 CSRO Meta-testing
Input: A set of testing tasks  $\{M_i\}_{i=1}^{N_{env}}$ , learned policy  $\pi_\theta$ , context encoder  $q_\phi$ , random explore step  $t_r$   
1: for each task  $M_i$  do  
2:  $c = \{\}$   
3: for  $t = 0, \dots, T - 1$  do  
4: if  $t < t_r$  then  
5: Agent samples a random action  $a_t$  to roll out  $(s_t, a_t, r_t, s_t')$   
6: else  
7: Compute posterior  $z = q_\phi(z|c)$ .  
8: Agent use  $\pi_\theta(a|s, z)$  roll out  $(s_t, a_t, r_t, s_t')$   
9: end if  
10:  $c = c \cup (s_t, a_t, r_t, s_t')$   
11: end for  
12: Compute posterior  $z = q_\phi(z|c)$ .  
13: Roll out policy  $\pi_\theta(a|s, z)$  for evaluation  
14: end for
Proof. The meta-policy gets the highest expected return when exploring policy and behavior policy are the same:  $\pi_e = \pi_\beta^i$ . However, for any exploration policy  $\pi_e$ , we hope  $J_{M_i}(\pi_\theta, \pi_e) = J_{M_i}(\pi_\theta, \pi_\beta^i)$ , if and only if for any  $\pi_e$ ,  $q_\phi(z|c)$  is the same. Since  $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, p_e)$  is the same. So  $z$  and  $\pi_e$  are independent, that is, the mutual information  $I(z; \pi_e) = 0$ .
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  $\mathbb{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  $\mathbb{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  $\mathbb{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\pi]$ . The state space is  $\mathbb{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. 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  $\mathbb{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  $\mathbb{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.
Table 3: Hyperparameters used in offline datasets collection.
Hyperparameters	Point-Robot	Half-Cheetah-Vel	Ant-Dir	Humanoid-Dir	Hopper-Rand-Params	Walker-Rand-Params
Training steps	5000	1e6	1e6	1e6	1e6	1e6
Initial steps	2e3	5e4	5e4	5e4	5e4	5e4
Eval frequency	200	5e4	5e4	5e4	5e4	5e4
Sampling episodes	50	50	50	50	50	50
Learning rate	1e-4	1e-4	1e-4	1e-4	1e-4	1e-4
Batch size	1024	1024	1024	1024	1024	1024
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.
2https://github.com/dennisl88/random-param_envs.
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.
Env	Checkpoints
Point-Robot	[2200, 4800, 200]
Half-Cheetah-Val	[100000, 950000, 50000]
Ant-Goal	[100000, 950000, 50000]
Humanoid-Dir	[50000, 950000, 50000]
Hopper-Rand-Params	[50000, 950000, 50000]
Walker-Rand-Params	[50000, 950000, 50000]
Table 5: Hyperparameters used in offline meta-training.
Hyperparameters	Point-Robot	Half-Cheetah-Vel	Ant-Dir	Humanoid-Dir	Hopper-Rand-Params	Walker-Rand-Params
Reward scale	100	5	5	5	5	5
Latent dimension	20	20	20	20	40	40
Use BRAC	False	True	True	True	True	True
Batch size	256	256	256	256	256	256
Meta batch size	16	16	10	16	16	16
Embedding batch size	1024	100	512	256	256	256
Actor Learning rate	3e-4	3e-4	3e-4	3e-4	3e-4	3e-4
Critic Learning rate	3e-4	3e-4	3e-4	3e-4	3e-4	3e-4
Encoder Learning rate	3e-4	3e-4	3e-4	3e-4	3e-4	3e-4
Maximum episode length	20	200	200	200	200	200
MinMI loss weightλ	25	10	50	50	25	25
behavior regularization	50	50	50	50	50	50
Discount factor	0.9	0.99	0.99	0.99	0.99	0.99
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.
Figure 6: Compared to other OMRL methods, CSRO's offline testing averages returns on unseen testing tasks
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