SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yunkai Gao $^{1,2,3}$  Rui Zhang $^{2}$  Jiaming Guo $^{2}$  Fan Wu $^{2,3,4,5}$  Qi Yi $^{1,2,3}$  Shaohui Peng $^{5}$  Siming Lan $^{1,2,3}$  Ruizhi Chen $^{5}$  Zidong Du $^{2,6}$  Xing Hu $^{2,6}$  Qi Guo $^{2}$  Ling Li $^{4,5}$  Yunji Chen $^{2,4*}$
<sup>1</sup> University of Science and Technology of China, USTC, Hefei, China
$^{2}$  State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences, Beijing, China
3 Cambricon Technologies
<sup>4</sup> University of Chinese Academy of Sciences, UCAS, Beijing, China
<sup>5</sup> Intelligent Software Research Center, Institute of Software, CAS, Beijing, China
$^{6}$  Shanghai Innovation Center for Processor Technologies, SHIC, Shanghai, China
{gyk314, yiqi, lansm}@mail.ustc.edu.cn
{wufan2020, pengshaohui, ruizhi, liling}@iscas.ac.cn
{zhangrui, guojiaming, duzidong, huxing, guoqi, cyj}@ict.ac.cn
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.
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 pre-collected offline datasets from the test tasks to acquire task representation, disregarding the context shift problem. However, this approach is impractical as obtaining pre-collected offline datasets for new and unseen tasks is extremely challenging. Other methods attempt to mitigate the context shift problem but require additional information beyond the use of offline datasets alone. For instance, [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 Context Shift Reduction for OMRL (CSRO) to address the context shift problem using only offline datasets. The key insight of CSRO is to minimize the impact of the policy in context during both the meta-training and meta-test phases. Firstly, the context is collected from the task-dependent behavior policy, as depicted in Figure 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-priori 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.
Meta Reinforcement Learning. Meta-RL leverages training task sets sampled from the task distribution to learn and generalize to new tasks. Gradient-based meta-RL [6, 33] methods seek
Table 1: Context shift leads to testing performance decline. Context A refers to the use of offline data collected by behavior policy as the context during testing, with no context shift problem. Context B, on the other hand, involves using a trained meta-policy to explore and collect context based on an initial task representation, resulting in a context shift problem.
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
Figure 1: Offline meta-RL on the Point-Robot: one of the tasks where the task objective is to control the agent to reach the goal on the semicircle. Left: During meta-training, the offline dataset of this task is collected by task-dependent behavior policy and comprised solely of trajectories leading towards the goal. Middle: The agent was trained using the OffPearl [27] method. During the meta-test phase, the agent utilizes the current exploration policy of OMRL [27, 4, 26] and employs the learned meta-policy, which is conditioned on an initial task representation sampled from a prior distribution, to collect contexts. Right: The agent navigates out trajectory based on the context depicted in the middle subfigure, the same color represents the corresponding context and trajectory.
policy parameter initialization that requires only 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
Figure 2: CSRO framework. Left: In meta-training, the context encoder  $q_{\phi}$  extracts each transition embedding from  $(s,a,r,s')$  transition tuple sample from offline datasets, the CLUB variational distribution estimator  $q_{\psi}$  is used to approximate  $p(z|(s,a))$ . We use the mean of transition embedding as task representation which conditions both the critic  $Q_{\omega}$  and the actor  $\pi_{\theta}$ . Right: In test phase, agent uses  $\pi_{Np}$  to explore new tasks.
and other properties. The greater the mutual information, the greater the correlation between the two variables. [2, 22] maximizes the mutual information between variables by maximizing the estimated mutual information lower bound. [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.
Problem Formulation. In RL, we formalize each task as a Markov Decision Process (MDP), defined by  $M = (S, A, P, \rho_0, R, \gamma)$  with state space  $S$ , action space  $A$ , transition function  $P(s'|s, a)$ , reward function  $R(s, a)$ , initial state distribution  $\rho_0(s)$  and discount factor  $\gamma \in (0, 1)$ . The distribution of policy is  $\pi(a|s)$ . We aim to maximize the expected cumulative rewards  $J(\pi) = E_{\tau}[\sum_{t=0}^{\infty} \gamma^t R(s_t, a_t)]$ , where  $\tau$  is trajectory following  $s_0 \sim \rho_0$ ,  $a_t \sim \pi(a_t | s_t)$ ,  $s_{t+1} \sim P(s_{t+1} | s_t, a_t)$ . We define the value function and action-value function as:
$$
V _ {\pi} \left(s _ {t}\right) = \mathbb {E} _ {a _ {t} \sim \pi , s _ {t + 1} \sim P \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right)} \left[ \sum_ {l = 0} ^ {\infty} \gamma^ {l} R \left(s _ {t + l}, a _ {t + l}\right) \right] \tag {1}
$$
$$
Q _ {\pi} \left(s _ {t}, a _ {t}\right) = R \left(s _ {t}, a _ {t}\right) + \gamma \mathbb {E} _ {s _ {t + 1} \sim P \left(s _ {t + 1} \mid s _ {t}, a _ {t}\right)} \left[ V _ {\pi} \left(s _ {t + 1}\right) \right] \tag {2}
$$
In offline meta-reinforcement learning (OMRL), we assume each task sampled from a prior task distribution,  $M_{i} = (S,A,P_{i},\rho_{0},R_{i},\gamma) \sim 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}^{\prime})\}_{j = 1}^{N_{size}}$  of task  $M_{i}$  is composed of the state, action, reward, and next-state tuples collected by the task-dependent behavior policy  $\pi_{\beta}^{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')\}_{j = 1}^n$ . The context encoder transforms the context  $c$  into a task representation  $z = q_{\phi}(z|c)$ . Policy, value function, and action-value function are all conditioned on  $z$ . We use offline data to train the meta policy  $\pi_{\theta}(a|s,z)$  with parameters  $\theta$ , and will not interact with the environment during the training. OMRL aims to maximize the expected cumulative rewards in the test task:
$$
J (\pi_ {\theta}) = \mathbb {E} _ {M \sim p (M)} [ J _ {M} (\pi_ {\theta}) ]. \tag {3}
$$
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  $\pi_{\beta}^{i}(a|s)$ , and the offline data is  $D_{i} = \{(s_{i,j},a_{i,j},r_{i,j},s_{i,j})'\}_{j=1}^{N_{size}}$ . The distribution of offline data is denoted by
$$
p \left(D _ {i}\right) = \rho \left(s _ {0}\right) \prod_ {j = 0} ^ {N _ {\text {s i z e}}} \pi_ {\beta} ^ {i} \left(a _ {j} \mid s _ {j}\right) R \left(s _ {j}, a _ {j}\right) P \left(s _ {j + 1} \mid s _ {j}, a _ {j}\right). \tag {4}
$$
In the meta-test phase, the data is online collected by the exploration policy.
The different sources of offline datasets used in meta-training and online data used in the meta-test lead to the context shift problem: distribution discrepancy between context from behavior policy and context collected from exploration policy. As shown in Eq 4, the distribution of collected offline datasets is jointly determined by task and behavior policy, which means the context used to train the context encoder contains some characteristics specific 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_{\phi}$  to distinguish tasks. We first sample the context  $c_{i} = \{(s_{i,j},a_{i,j},r_{i,j},s_{i,j}^{\prime})\}_{j = 1}^{N_c}$  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_{\phi}(z|(s_{i,j},a_{i,j},r_{i,j},s_{i,j}^{\prime}))$ . Finally, we use the mean of transition embedding  $z_{i} = \mathbb{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  $\pi_{\beta}$ .
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:
$$
L _ {m a x M I} (\phi) = 1 \left\{y _ {i} = y _ {j} \right\} \| z _ {i} - z _ {j} \| _ {2} ^ {2} + 1 \left\{y _ {i} \neq y _ {j} \right\} \frac {\beta}{\left\| z _ {i} - z _ {j} \right\| _ {2} ^ {n} + \epsilon}, \tag {5}
$$
where  $c_{i}$  is context sampled from task  $M_{i}$ ,  $y_{i}$  is task label,  $z_{i} = q_{\phi}(z|c_{i})$  is a task representation. 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}, \rho, R_{i}) \sim p(M)$ , context  $c$  as a context collected by exploration policy  $\pi_{e}$ . The expected return of test task  $M_{i}$  which is evaluated by a learned meta-policy  $\pi_{\theta}(a|s, z)$  is:
$$
J _ {M _ {i}} \left(\pi_ {\theta}, \pi_ {e}\right) = E _ {s _ {0} \sim \rho \left(s _ {0}\right), z \sim q _ {\phi} (\cdot | c), a _ {t} \sim \pi_ {\theta} (\cdot | s _ {t}, z), r _ {t} \sim R _ {i} \left(s _ {t}, a _ {t}\right), s _ {t} ^ {\prime} \sim P _ {i} (\cdot | s _ {t}, a _ {t})} \left[ \sum_ {t = 0} ^ {H - 1} r _ {t} \right]. \tag {6}
$$
Proposition 1. For task  $M_{i}$  and its corresponding behavior policy  $\pi_{\beta}^{i}$ , when exploration policy  $\pi_{e}$  and behavior policy  $\pi_{\beta}^{i}$  are different, if and only if mutual information between task representations and policies  $I(z; \pi_{e}) = 0$ , the expected test return for both is the same  $J_{M_{i}}(\pi_{\theta}, \pi_{e}) = J_{M_{i}}(\pi_{\theta}, \pi_{\beta}^{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  $\pi_{\beta}$  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) = \pi_{\beta}(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:
$$
I _ {C L U B} (z, (s, a)) = \mathbb {E} _ {i} [ \log p (z _ {i} | (s _ {i}, a _ {i})) - \mathbb {E} _ {j} [ \log p (z _ {j} | (s _ {i}, a _ {i})) ] ]. \tag {7}
$$
However the conditional distribution  $p(z|(s,a))$  is unavailable on our tasks, therefore we use a variational distribution  $q_{\psi}(z|(s,a))$  to approximate  $p(z|(s,a))$  by minimizing Log-likelihood:
$$
L _ {V D} (\psi) = - \mathbb {E} _ {M \sim p (M)} \mathbb {E} _ {i} [ \log q _ {\psi} \left(z _ {i} \mid \left(s _ {i}, a _ {i}\right)\right) ]. \tag {8}
$$
We minimize mutual information between transition embedding  $z$  and  $(s, a)$  by minimizing mutual information upper bound estimation:
$$
L _ {\min  M I} (\phi) = \mathbb {E} _ {M \sim p (M)} \mathbb {E} _ {i} \left[ \log q _ {\psi} \left(z _ {i} \mid \left(s _ {i}, a _ {i}\right)\right) - \mathbb {E} _ {j} \left[ \log q _ {\psi} \left(z _ {j} \mid \left(s _ {i}, a _ {i}\right)\right) \right] \right], \tag {9}
$$
which  $z_{i}$  represents task embedding obtained by  $(s_{i}, a_{i}, r_{i}, s_{i}^{\prime})$  through the context encoder. Here, the variational distribution estimator  $q_{\psi}$  and context encoder  $q_{\phi}$  are used for adversarial training. If each transition embedding  $z$  obtained by the context encoder  $q_{\phi}(z|(s,a,r,s'))$  contains more  $(s,a)$  information, the  $q_{\psi}(z|(s,a))$  has higher accuracy, but mutual information upper bound  $L_{minMI}(\phi)$  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}(\phi)$  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.