Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning

## Abstract

Abstract Offline meta reinforcement learning (OMRL) has emerged as a promising approach for interaction avoidance and strong generalization performance by leveraging pre-collected data and meta-learning techniques.
Previous context-based approaches predominantly rely on the intuition that alternating optimization between the context encoder and the policy can lead to performance improvements, as long as the context encoder follows the principle of maximizing the mutual information between the task variable M 𝑀 M italic_M and its latent representation Z 𝑍 Z italic_Z ( I ⁢ ( Z ; M ) 𝐼 𝑍 𝑀 I(Z;M) italic_I ( italic_Z ; italic_M ) ) while the policy adopts the standard offline reinforcement learning (RL) algorithms conditioning on the learned task representation.
Despite promising results, the theoretical justification of performance improvements for such intuition remains underexplored.
Inspired by the return discrepancy scheme in the model-based RL field, we find that the previous optimization framework can be linked with the general RL objective of maximizing the expected return, thereby explaining performance improvements.
Furthermore, after scrutinizing this optimization framework, we find it ignores the variation of the task representation in the alternating optimization process, which weakens the condition necessary for monotonic performance improvements, and may therefore violate the monotonicity.
We name this issue task representation shift and theoretically prove that the monotonic performance improvements can be guaranteed with appropriate context encoder updates.
We use different settings to rein in the task representation shift on three widely adopted training objectives concerning maximizing I ⁢ ( Z ; M ) 𝐼 𝑍 𝑀 I(Z;M) italic_I ( italic_Z ; italic_M ) across different data qualities.
Empirical results show that reining in the task representation shift can indeed improve performance.
Our work opens up a new avenue for OMRL, leading to a better understanding between the task representation and performance improvements.

## 1 Introduction

RL has driven impressive advances in many complex decision-making problems in recent years [silver2018general, schrittwieser2020mastering, zhang2023safe], primarily through online RL methods.
However, the extensive interactions required by online RL entail high costs and safety concerns, posing significant challenges for real-world applications.
Offline RL [wu2019behavior, levine2020offline] offers an appealing alternative by efficiently leveraging pre-collected data for policy learning, thereby circumventing the need for online interaction with the environment.
This advantage extends the application of RL, covering healthcare [fatemi2022semi, tang2022leveraging], robotics [sinha2022s4rl, pmlr-v164-kumar22a] and games [schrittwieser2021online, mathieu2023alphastar].
Though demonstrating its superiority, offline RL holds a notable drawback towards generalizing to the unknown [ghosh2021generalization].

As a remedy, by combining the meta-learning techniques, OMRL [li2020focal, mitchell2021offline, xu2022prompting] has emerged as an effective training scheme toward strong generalization performance and fast adaptation capability while maintaining the merits of offline RL.
Among the OMRL research, context-based OMRL (COMRL) algorithms [li2020focal, li2024towards] hold a popular paradigm that seeks optimal meta-policy conditioning on the context of Markov Decision Processes (MDPs).
Specifically, these methods [li2020focal, gao2024context, li2024towards] propose to train a context encoder via maximizing $I(Z;M)$ to learn the task representation from the collected context (e.g. trajectories) and train the downstream policy with standard offline RL algorithms conditioning on the task representation, as shown in Figure 1 left.
The context encoder and policy are trained in an alternating manner, with each being updated once per cycle.
Despite promising results, the theoretical justification of performance improvements for such intuition remains underexplored.

Figure: Figure 1: Our training framework compared to the previous training framework. They both adopt the alternating optimization framework to train the context encoder and the policy. However, ours considers the variation of task representation ignored by the previous training framework by imposing the extra condition to determine whether needs to update the context encoder.
Refer to caption: x1.png

Intriguingly, we formalize this training framework as an alternating two-stage optimization framework and then link it with the general RL objective of maximizing the expected return.
To be specific, maximizing $I(Z;M)$ and adopting standard offline RL algorithms can be interpreted as consistently raising the lower bound of the expected return conditioning on the optimal task representation distribution among the expectation of tasks (See Section 4.1).
This is achieved by extending the return discrepancy scheme [janner2019trust] to the COMRL framework.
Thus, this analysis provides a feasible explanation for the performance improvement guarantee.

More importantly, after scrutinizing this optimization framework, we find it ignores the impacts stemming from the variation of the task representation in the alternating process.
This would cause the optimization framework to incorrectly conclude that the monotonic performance improvement can be guaranteed with a better approximation to the optimal task representation distribution.
By explicitly modeling the variation of task representation, we prove that it is a critical part of monotonic performance improvement.
Without this, the monotonicity can be violated.
To highlight the characteristic, we name this issue task representation shift and theoretically prove that it is possible to achieve monotonic performance improvement with appropriate updates (See Section 4.2).

To show the impacts of this issue, we use different settings to rein in the task representation shift on contrastive objective adopted in [li2020focal, gao2024context, li2024towards] and reconstruction objective adopted in [zintgraf2019varibad, dorfman2021offline, li2024towards] across different data qualities. These two objectives are proven to maximize the approximate bounds of $I(Z;M)$ [li2024towards]. We also introduce the cross-entropy objective, which is a straightforward discrete approximation to $I(Z;M)$.
Empirical results show that reining in the task representation shift can indeed improve performance.
Our work opens up a new avenue for OMRL, leading to a better understanding between the task representation and performance improvements.

## 2 Related Works

Offline RL. Offline RL [fujimoto2019off] leverages a fully offline dataset to perform policy optimization and value function learning, avoiding explicit interaction with the environment.
To overcome the notorious distribution shift and the induced value overestimation issues, prior works adopt methods based on policy regularization [fujimoto2019off, kumar2019stabilizing, fujimoto2021minimalist], value penalty [wu2019behavior, kumar2020conservative] and sequence modeling [chen2021decision, janner2021offline].
While these methods have demonstrated promising performance on a single task, their generalization performance to the unknown remains limited [ghosh2021generalization].
Our work follows the offline RL setting, attempting to adopt meta-learning techniques to improve the generalization performance.

Meta RL. Meta RL has emerged as an enticing avenue for effectively adapting to unknown environments [finn2017model].
Unlike traditional gradient-based meta RL methods [finn2017model, mitchell2021offline], context-based meta RL methods demonstrate superior advantages in both sample efficiency and asymptotic performance [rakelly2019efficient, zintgraf2019varibad, li2024towards].
Our work delves into this promising domain to advance the capabilities of context-based meta RL further.

Context-based Offline Meta RL. As the marriage between context-based meta RL and offline RL, COMRL combines the merits of both sides.
Specifically, COMRL methods [li2020focal, yuan2022robust, gao2024context, li2024towards] leverage the offline dataset to train a context encoder to learn robust task representations, and then pass the representation to the policy and value function as input.
At test time, COMRL methods leverage the generalization ability of the learned context encoder to perform meta-adaptation.
Acknowledging the theoretical insights from [li2024towards], prior COMRL works [li2020focal, yuan2022robust, gao2024context, li2024towards] mainly focus on how to design a context encoder learning algorithm to better approximate $I(Z;M)$.
However, our work moves the focus on how to refine the condition of monotonic performance improvements based on a given context encoder learning algorithm.
Center around this motivation, we identify a new issue called task representation shift, which is ignored by the previous COMRL endeavors.

Performance Improvement Guarantee.  Ensuring performance improvement is a key concern in both online and offline reinforcement learning settings. In online RL, performance improvement guarantees are often established through methods such as performance difference bounds [schulman2015trust, kakade2002approximately, ji2022update, zhang2023how], return discrepancy [janner2019trust, luo2018algorithmic], and regret bounds [osband2014model, curi2020efficient].
In offline RL, CQL-based methods [kumar2020conservative, yu2021combo] also enjoy safe policy improvement guarantees.
However, most works focus on the single-task setting, leaving the performance improvement guarantees in the COMRL settings largely unexplored.
While ContrBAR [choshen2023contrabar] also benefits from the performance improvement guarantee in the OMRL setting, it is specifically served for its own framework.
In contrast, our work addresses a broader class of algorithms that maximize $I(Z;M)$, including but not limited to ContrBAR.

## 3 Preliminaries

### 3.1 Problem Statement

A task in RL is generally formalized as a fully observed MDP [puterman2014markov], which is represented by a tuple $(\mathcal{S},\mathcal{A},P,\rho_{0},R,\gamma)$ with state space $\mathcal{S}$, action space $\mathcal{A}$, transition function $P(s^{\prime}|s,a)$, reward function $R(s,a)$, initial state distribution $\rho_{0}(s)$ and discount factor $\gamma\in[0,1]$.
For arbitrary policy $\pi(a|s)$, we denote the state-action distribution at timestep $t$ as $d_{\pi,t}(s,a)\triangleq\text{Pr}(s_{t}=s,a_{t}=a|s_{0}\sim\rho_{0},a_{t}\sim
\pi,s_{t+1}\sim P,\forall t\geq 0)$.
The discounted state-action distribution of given $\pi(a|s)$ is denoted as $d_{\pi}(s,a)\triangleq(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}d_{\pi,t}(s,a)$.
The ultimate goal is to find an optimal policy $\pi(a|s)$ to maximize the expected return $\mathbb{E}_{(s,a)\sim d_{\pi}}[R(s,a)]$.

Given a distribution of tasks $p(M)$, every task $m$ sampled from $p(M)$ is an MDP sharing the same state space and action space but differing from transition dynamics and reward functions.
It is widely assumed that the tasks in the meta RL setting are randomly sampled from $p(M)$ [li2020focal, yuan2022robust, gao2024context, li2024towards].
We denote the number of training tasks as $N$.
For each task $i\in[0,1,...,N-1]$, an offline dataset $D_{i}=\{(s_{i,j},a_{i,j},s^{\prime}_{i,j},r_{i,j})\}_{j=1}^{K}$ is collected in advance, where $K$ denotes the number of transitions. The learning algorithm is required to train a meta-policy $\pi_{\text{meta}}$ with access only to the given offline datasets. At test time, given an unseen task, the meta policy $\pi_{\text{meta}}$ performs task adaptation to get a task-specific policy and then will be evaluated in the environment.

### 3.2 Context-based OMRL

Previous COMRL endeavors hold that the variation over different tasks can be represented by a compact representation $z$ [yuan2022robust].
In practice, they choose to train a context encoder $Z(\cdot|x)$ to extract the task information [li2020focal, dorfman2021offline], where $x$ denotes the given context. Then, the policy is learned by conditioning on the task representation as $\pi(a|s,Z(\cdot|x);\theta)$.
Naturally, there exists an optimal task representation distribution $Z^{*}(\cdot|x)$ that meets the optimal expected return for any policy parameter $\theta$ is determined by $\pi(a|s,Z^{*}(\cdot|x);\theta)$. Hence, the objective of COMRL is formalized as:

$$ $\displaystyle\max_{\theta}J^{*}(\theta)=\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d _{\pi(\cdot|s,Z^{*}(\cdot|x);\theta)}}[R_{m}(s,a)]]$ (1) $$

where $R_{m}(s,a)$ denotes the ground-truth reward function of the task $m$.

Previous COMRL works [li2020focal, yuan2022robust, gao2024context, li2024towards] are proven to optimize the context encoder by maximizing the approximate bounds of $I(Z;M)$, as shown in Theorem 3.1.
They hold the intuition that using such kind of approach to optimize the context encoder and adopting the standard offline RL algorithms to optimize the policy can lead to performance improvements.
However, the theoretical justification of performance improvement for such intuition has been less explored.
For later use, we denote $Z_{mutual}(\cdot|x)$ as the optimal solution for maximizing $I(Z;M)$.
Whether $Z_{mutual}(\cdot|x)$ is equivalent to $Z^{*}(\cdot|x)$ is still an open problem.
Hence, we introduce an extra notation here to avoid confusion.

###### Theorem 3.1 ( [ li2024towards ] ) .

Denote $X_{b}$ and $X_{t}$ are the behavior-related $(s,a)$-component and task-related $(s^{\prime},r)$-component of the context $X$, with $X=(X_{b},X_{t})$. We have:

$$ $\displaystyle I(Z;X_{t}|X_{b})\leq I(Z;M)\leq I(Z;X)$ (2) $$

where 1) $\mathcal{L}_{FOCAL}\equiv-I(Z;X)=-I(Z;X_{t}|X_{b})-I(Z;X_{b})$; 2) $\mathcal{L}_{CORRO}\equiv-I(Z;X_{t}|X_{b})$; 3) $\mathcal{L}_{CSRO}\geq(\lambda-1)I(Z;X)-\lambda I(Z;X_{t}|X_{b})$

###### Theorem 3.1 ( [ li2024towards ] ) .

## 4 Methods

We first introduce a Lipschitz assumption, a widely used technique in both model-free [song2019efficient, ghosh2022provably] and model-based [ji2022update, zhang2023how] RL frameworks.

###### Assumption 4.1 .

The policy function is $L_{z}-$ Lipschitz w.r.t some norm $||\cdot||$ in the sense that

$$ $\displaystyle\forall Z_{1}(\cdot|x),Z_{2}(\cdot|x)\in\mathbb{Z},\,\,|\pi(\cdot |s,Z_{1}(\cdot|x);\theta)-\pi(\cdot|s,Z_{2}(\cdot|x);\theta)|\leq L_{z}\cdot|Z _{1}(\cdot|x)-Z_{2}(\cdot|x)|$ (3) $$

###### Assumption 4.1 .

### 4.1 A Performance Improvement Perspective Towards Prior Works

Our goal is to arrive at an optimization objective for COMRL to get the performance improvement guarantee.
Motivated by the return discrepancy scheme [janner2019trust] in the model-based RL field, we can similarly build a tractable lower bound for $J^{*}(\theta)$.

###### Definition 4.2 (Return discrepancy in COMRL) .

The return discrepancy in the COMRL setting can be defined as

$$ $\displaystyle J^{*}(\theta)-J(\theta)\geq-|J(\theta)-J^{*}(\theta)|$ (4) $$

where $J(\theta)$ is the expected return of the policy $\pi$ among the expectation of tasks conditioning on the learned task embedding $Z(\cdot|x)$ as $J(\theta)=\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z(\cdot|x);
\theta)}}[R_{m}(s,a)]]$.

From Definition 4.2, we know that if the estimation error $|J^{*}(\theta)-J(\theta)|$ can be upper-bounded, the lower bound for $J^{*}(\theta)$ can be established simultaneously.
Then, Eq. (4) can be interpreted as a unified training objective for both the context encoder and the policy.
Specifically, we can alternate the optimization target to lift the lower bound for $J^{*}(\theta)$, where the holistic optimization process can be seen as a two-stage alternating framework with the first stage updating the context encoder to minimize $|J^{*}(\theta)-J(\theta)|$ and the second stage updating the policy to maximize $J(\theta)-|J^{*}(\theta)-J(\theta)|$ as $J^{*}(\theta)\geq J(\theta)-|J^{*}(\theta)-J(\theta)|$.
Hence, the performance improvement guarantee can be achieved.

The following theorem induces a tractable upper bound for $|J^{*}(\theta)-J(\theta)|$, connecting the bridge between the intuition of previous works and the theoretical justification concerning the performance improvement guarantee.

###### Theorem 4.3 (Return bound in COMRL) .

Assume the reward function is upper-bounded by $R_{max}$. For an arbitrary policy parameter $\theta$, when meeting Assumption 4.1, the return bound in COMRL can be formalized as:

$$ $\displaystyle|J^{*}(\theta)-J(\theta)|\leq\frac{2R_{max}L_{z}}{(1-\gamma)^{2}} \mathbb{E}_{m,x}(|Z(\cdot|x)-Z_{mutual}(\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}( \cdot|x)|)$ (5) $$

*Proof:*
See Section 8.2 in the Appendix. ∎

By directly unrolling the Theorem 4.3, we can establish the lower bound of $J^{*}(\theta)$ as:

$$ $\displaystyle J^{*}(\theta)\geq J(\theta)-\frac{2R_{max}L_{z}}{(1-\gamma)^{2}} \mathbb{E}_{m,x}(|Z(\cdot|x)-Z_{mutual}(\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}( \cdot|x)|)$ (6) $$

Note that $|Z_{mutual}(\cdot|x)-Z^{*}(\cdot|x)|$ is a constant w.r.t the learned task representation $Z(\cdot|x)$.
Hence, this term can be ignored when optimizing $Z(\cdot|x)$. Admittedly, there may exist a smarter algorithm to further reduce the gap of $|Z_{mutual}(\cdot|x)-Z^{*}(\cdot|x)|$, but this is not our main focus and we leave this to future work.

Theorem 4.3 and Eq. (6) indicate the following two principles for the context encoder and the policy. For the context encoder learning, we should minimize the return bound, namely minimizing $\mathbb{E}_{m,x}|Z(\cdot|x)-Z_{mutual}(\cdot|x)|$. For the policy learning, we should lift the lower bound of $J^{*}(\theta)$, namely maximizing $J(\theta)-\frac{2R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}(|Z(\cdot|x)-Z_{
mutual}(\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}(\cdot|x)|)\equiv J(\theta)$, as the task representation related terms are constants for policy optimization. To maximize $J(\theta)$, we can adopt standard offline RL algorithms [wu2019behavior, fujimoto2021minimalist, kumar2020conservative] similar to [yang2022unified].

Recall that previous COMRL works adopt alternatively training the context encoder by maximizing $I(Z;M)$, which can be approximately seen as minimizing $\mathbb{E}_{m,x}|Z(\cdot|x)-Z_{mutual}(\cdot|x)|$ and training the policy by conditioning on the learned task representation as well as applying the standard offline RL algorithms.
Therefore, we argue that previous COMRL works can be interpreted by this framework, thus providing an explanation concerning performance improvement guarantee.

However, this optimization framework only considers the discrepancy between $Z(\cdot|x)$ and $Z_{mutual}(\cdot|x)$, without considering the variation of $Z(\cdot|x)$.
In the next section, we will show that this may violate the monotonicity of the performance improvements.

###### Definition 4.2 (Return discrepancy in COMRL) .

###### Theorem 4.3 (Return bound in COMRL) .

### 4.2 Monotonic Performance Improvement Concerning Task Representation Shift

In this section, we aim to demonstrate that as the previous optimization framework in Section 4.1 doesn’t model the variation of task representation explicitly, this framework would provide an insufficient condition for monotonic performance improvement.
We first show conditions for monotonic performance improvement of this framework.

###### Corollary 4.4 (Monotonic performance improvement condition for previous COMRL works) .

When meeting Assumption 4.1, the condition for monotonic performance improvement of previous COMRL works is:

$$ $\displaystyle\epsilon^{*}_{12}\triangleq J^{*}(\theta_{2})-J^{*}(\theta_{1}) \geq\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}(|Z(\cdot|x)-Z^{*}( \cdot|x)|)$ (7) $$

*Proof:*
See Section 8.2 in the Appendix.
∎

As shown in Corollary 4.4, the monotonic performance improvement can be guaranteed with only better approximation to $Z^{*}(\cdot|x)$.
Next, we will introduce the performance difference bound framework to model variation of task representation.

###### Definition 4.5 (Performance difference bound in COMRL) .

Denote $J^{1}(\theta_{1})$ as the expected return of the policy $\pi(\cdot|s,Z_{1}(\cdot|x);\theta_{1})$ among the expectation of tasks before update of the context encoder and the policy.
Similarly, denote $J^{2}(\theta_{2})$ as the expected return of the policy $\pi(\cdot|s,Z_{2}(\cdot|x);\theta_{2})$ among the expectation of tasks after update of the context encoder and the policy.
The performance difference bound in the COMRL setting can be defined as

$$ $\displaystyle J^{2}(\theta_{2})-J^{1}(\theta_{1})\geq C$ (8) $$

when $C$ is non-negative, the algorithm allows a monotonic performance improvement.

According to Definition 4.5, we need to find the positive $C$ to improve the performance monotonically.
To achieve this, we can derive the lower bound of the performance difference.

###### Theorem 4.6 (Lower bound of performance difference in COMRL) .

Assume the reward function is upper-bounded by $R_{max}$. When meeting Assumption 4.1, the lower bound of the performance difference in COMRL can be formalized as:

$$ $\displaystyle J^{2}(\theta_{2})-J^{1}(\theta_{1})\geq\epsilon^{*}_{12}-\frac{2 R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}[2|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)| +|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|]$ (9) $$

*Proof:*
See Section 8.2 in the Appendix. ∎

For the need of the monotonic performance improvement guarantee, we further let the right-hand side of Eq. (32) be positive.
Then, we would have the following condition to achieve monotonicity.

$$ $\displaystyle\epsilon^{*}_{12}-\frac{2R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_ {m,x}[2|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)|+|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|]\geq 0$ (10) $$

Compared to the condition in Corollary 4.4, Eq. 10 presents an additional term $|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|$ to achieve monotonic performance improvement, which corresponds exactly to the previous ignored impacts stemming from the variation of the task representation.
To achieve monotonic performance improvements, the optimization of task representation should consider not only the approximation to $Z^{*}(\cdot|x)$, but also the magnitude of the update.
For example, if we assume that $Z_{1}(\cdot|x)$ is trained from scratch and the update process from $Z_{1}(\cdot|x)$ to $Z_{2}(\cdot|x)$ brings $Z_{2}(\cdot|x)$ close to $Z^{*}(\cdot|x)$, then with the condition in Corollary 4.4, the monotonic performance improvement can be easily achieved with small $\epsilon^{*}_{12}$.
However, for the condition in Eq. (10), small $\epsilon^{*}_{12}$ may cause the violation of monotonicity as $|Z_{1}(\cdot|x)-Z_{2}(\cdot|x)|$ remains large.

To highlight the impacts of $|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|$, we name this issue Task Representation Shift.
Under some mild assumptions, we can conclude that it is possible to achieve monotonic performance improvement with sufficient policy improvement $\epsilon^{*}_{12}$ on the condition of appropriate encoder updates, as shown in Theorem 4.10.

###### Assumption 4.7 .

The impacts of task representation shift are upper-bounded by $\beta$ and less than the policy improvement $\epsilon^{*}_{12}$ with a certain coefficient.

###### Assumption 4.8 .

The space of the task representation is discrete and limited.

###### Assumption 4.9 .

There exists an $\alpha$ for any given $Z_{1}(\cdot|x)$, $|Z_{2}(\cdot|x)-\tilde{Z}_{2}(\cdot|x)|^{2}\leq\frac{\alpha}{b}$, where $Z_{2}(\cdot|x)$ denotes the context encoder updated by maximizing $I(Z;M)$ with the data size $b$ randomly sampled from the training dataset based on $Z_{1}(\cdot|x)$ and $\tilde{Z}_{2}(\cdot|x)$ denotes the context encoder updated by fitting the empirical distribution on $b$ i.i.d samples from $Z^{*}(\cdot|x)$ based on $Z_{1}(\cdot|x)$ (^1^11We set discrepancy to be subject to an inverse proportional function is because as the training samples increase, $Z_{2}$ and $\tilde{Z}_{2}$ should become similar.).

###### Theorem 4.10 (Monotonic performance improvement guarantee on training process) .

Denote $\kappa$ as $\frac{(1-\gamma)^{2}}{4R_{max}L_{z}}\epsilon^{*}_{12}-\frac{1}{2}\beta$ and $vol(\mathbb{Z})$ as the volume of the task representation coverage simplex. Given that the context encoder has already been trained by maximizing $I(Z;M)$ to some extent. When meeting Assumption 4.1,
4.7, 4.8 and 4.9, with a probability greater than $1-\xi$, we can get the monotonic performance improvement guarantee by updating the context encoder via maximizing $I(Z;M)$ from at least extra $k$ samples, where:

$$ $\displaystyle k=\frac{1}{\kappa^{2}}(\sqrt{2\log\frac{2^{vol(\mathbb{Z})}-2}{ \xi}}+\sqrt{\alpha})^{2}$ (11) $$

Here, $\xi\in[0,1]$ is a constant.

*Proof:*
See Section 8.2 in the Appendix. ∎

Theorem 4.10 shows the connection between the needed update data size $k$ and the performance improvement brought only by the policy update $\epsilon^{*}_{12}$.
If the calculated $k$ is larger than the given batch size, the context encoder should avoid this update and wait for the accumulation of $\epsilon^{*}_{12}$.
As $\epsilon^{*}_{12}$ increases, $k$ decreases.
When $k$ is smaller than the given batch size, the monotonic performance improvement can be achieved by updating the context encoder.
Notice that this updating process can be performed multiple times, as long as satisfying Assumption 4.7.
The updating process is visualized in Figure 1 right.
Compared to the original alternating framework, Theorem 4.10 unveils that the core part for better performance improvement is to adjust the update of the context encoder to rein in the task representation shift, thereby showcasing the advantage over the previous works.
The general algorithmic framework is shown in Algorithm 1, where the way to adjust the update of the context encoder is colored by red.

Figure: Algorithm 1 General Algorithmic Framework Towards Reining In The Task Representation Shift

###### Corollary 4.4 (Monotonic performance improvement condition for previous COMRL works) .

###### Definition 4.5 (Performance difference bound in COMRL) .

###### Theorem 4.6 (Lower bound of performance difference in COMRL) .

###### Assumption 4.7 .

###### Assumption 4.8 .

###### Assumption 4.9 .

###### Theorem 4.10 (Monotonic performance improvement guarantee on training process) .

### 4.3 Practical Implementation

To show the impacts of the task representation shift issue, we use different settings to adjust the context encoder.
We denote the given batch size as $bs$ and the accumulation steps of task representation shift as $acc$.
Our settings include 1) $k=2\times bs,acc=1$, 2) $k=3\times bs,acc=1$, 3) $k=1\times bs,acc=2$, and 4) $k=1\times bs,acc=3$.
All settings as well as the original setting, namely $k=1\times bs,acc=1$ are performed for 8 different random seeds.

With respect to the specific context encoder learning algorithms concerning maximizing $I(Z;M)$, we choose 3 representative algorithms that are also widely adopted in previous COMRL endeavors.
The reconstruction-based algorithm, passing the state, action, and task representation through the decoder to reconstruct the next state and the reward signal, is applied in [zintgraf2019varibad, dorfman2021offline, li2024towards] and is proven to be the lower bound of $I(Z;M)$.
The contrastive-based algorithm, using the task representation to compute distance metric learning loss, is applied in [li2020focal, gao2024context, li2024towards] and is proven to be the upper bound of $I(Z;M)$ as well as contain the lower bound of $I(Z;M)$.
The cross-entropy-based algorithm is a straightforward discrete approximation to $I(Z;M)$ and is proposed by our work.
Specifically, after extracting from the context encoder, the task representation would be passed through a fully connected layer to get the probabilities for each task.
Then, the context encoder is trained by the supervision of the task label and back-propagation of the cross-entropy loss.
The implementation details of these algorithms are shown in Appendix Section 8.3.

With respect to the policy learning algorithm, to maintain the consistency of previous COMRL works, we directly adopt BRAC [wu2019behavior] to train the policy.

## 5 Experiment

Our experiments are conducted to show that the previous optimization framework that ignores the impacts of task representation shift is not sufficient.
We hope to illustrate the potential of reining in the task representation shift, laying the foundation for further research towards better performance improvement of COMRL.

### 5.1 Environments Settings

We adopt MuJoCo [todorov2012mujoco] and MetaWorld [yu2020meta] benchmarks to evaluate the algorithms.
Following the protocol of UNICORN  [li2024towards], we randomly sample 20 training tasks and 20 testing tasks from the task distribution.
We train the SAC [haarnoja2018soft] agent from scratch for each task and use the collected replay buffer as the offline dataset.
During the meta-testing phase, we adopt a fully offline setting in a few-shot manner like [li2020focal] that randomly samples a trajectory from the dataset as the context to obtain the task representation and then passes it to the meta-policy to complete the testing in the true environments.

### 5.2 MAIN RESULTS

Figure: Figure 2: Testing returns of different settings to rein in the task representation shift against the original setting concerning cross-entropy, contrastive, and reconstruction based objectives for maximizing $I(Z;M)$. Solid curves refer to the mean performance of trials over 8 random seeds, and the shaded areas characterize the standard deviation of these trials.
Refer to caption: x2.png

Figure 2 illustrates the learning curves of our different settings to rein in the task representation shift against the original setting on 20 testing tasks for each benchmark concerning 3 widely adopted objectives for maximizing $I(Z;M)$.
We find that these specific objectives hold unanimous performance demonstration on all benchmarks that the previous training framework, namely $k=1\times bs,acc=1$, cannot provide the best asymptotic performance.
These objectives focus solely on the optimization target of the task representation, namely the optimal solution of $I(Z;M)$.
While they are guaranteed to produce reasonable results, as shown in Section 4.1, ignoring the impacts of task representation shift still hinders the capability of these algorithms to reach their full potential.
As shown by Figure 2, even with minor changes to the algorithms, the performance improvements are substantial.

Figure: Figure 3: Testing returns of different manners to rein in the task representation shift against the original manner concerning cross-entropy, contrastive, and reconstruction based objectives for maximizing $I(Z;M)$ on different data qualities of Ant-Dir. Solid curves refer to the mean performance of trials over 8 random seeds, and the shaded areas characterize the standard deviation of these trials.
Refer to caption: x3.png

### 5.3 Can The Results Show Consistency Across Different Data Qualities?

To make our claim more universal, we conduct the experiment on different data qualities.
We collect three types of datasets with size being equal to the dataset used in Section 5.2.
Depending on the quality of the behavior policies, we denote these datasets as random, medium, and expert respectively.

Figure 3 demonstrates the learning curves of our different settings to rein in the task representation shift against the original setting concerning 3 widely adopted objectives for maximizing $I(Z;M)$ on Ant-Dir.
We find that the contrastive-based algorithm fails on the random dataset.
Except for this case, the others show consistent results that reining in the task representation shift may lead to better performance improvements.

Based on these impressive results, we believe that developing a smarter algorithm to control the task representation shift automatically is an appealing direction and we leave this to future work.

## 6 Discussion

This section seeks to raise some interesting issues derived from our work.
We hope these issues can further enable pondering on the significance of the task representation shift.

Figure: Figure 4: Testing returns of the pre-training scheme against training from scratch of cross-entropy, contrastive, and reconstruction based objectives on Ant-Dir. Solid curves refer to the mean performance of trials over 8 random seeds, and the shaded areas characterize the standard deviation of these trials.
Refer to caption: x4.png

### 6.1 Can The Pretraining Scheme Be Adopted To Achieve Better Performance Improvement?

According to our analysis, the task representation shift issue only happens in the case that the context encoder needs to be trained from scratch.
Hence, a question is raised naturally.
If we train the context encoder in advance and use this pre-trained context encoder to train the policy directly, can we achieve better performance improvement?

To answer this question, we conduct the experiment on Ant-Dir and still use the cross-entropy, contrastive, and reconstruction based objectives to pre-train the context encoder.
As shown in Figure 4, there is a significant performance gap between pre-training and training from scratch.
To better theoretically interpret this, we introduce the following Corollary.

###### Corollary 6.1 (Monotonic performance improvement condition for pre-training scheme) .

Denote $Z_{pretrain}(\cdot|x)$ as the task representation distribution after pre-training.
When meeting Assumption 4.1, the monotonic performance improvement condition for pre-training scheme is:

$$ $\displaystyle\epsilon^{*}_{12}-\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_ {m,x}[|Z_{pretrain}(\cdot|x)-Z^{*}(\cdot|x)|]\geq 0$ (12) $$

*Proof:*
See Section 8.2 in the Appendix. ∎

Due to the implementation error and approximation error, $|Z_{pretrain}(\cdot|x)-Z^{*}(\cdot|x)|$ can not equal to 0.
Hence, according to Corollary 6.1, pre-training cannot actually achieve monotonic performance improvements.
As the task representation cannot be changed under the pre-training scheme, compared to training from scratch, this scheme loses some degrees of freedom to improve performance.

Additionally, this scheme can be intuitively regarded as degrading to the problem that how to design a pre-training algorithm to better approximate $Z^{*}(\cdot|x)$.
Notice that there exists the decomposition $|Z_{pretrain}(\cdot|x)-Z^{*}(\cdot|x)|\leq|Z_{pretrain}(\cdot|x)-Z_{mutual}(
\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}(\cdot|x)|$.
Since the cross-entropy based objective is the direct approximation of $I(Z;M)$, holding better performance under the pre-training scheme is also in line with expectation.

This analysis further enhances the importance of considering the impacts of task representation shift $|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|$ in the alternating process.
Nevertheless, how to achieve better performance under the pre-training scheme is also an interesting problem.

Figure: Figure 5: The 2D projection of the learned task representation space in Ant-Dir. Points are uniformly sampled from the evaluation datasets. Tasks of given goals from 0 to 6 are mapped to rainbow colors, ranging from purple to red.
Refer to caption: x5.png

###### Corollary 6.1 (Monotonic performance improvement condition for pre-training scheme) .

### 6.2 Can The Visualization Of Task Representation Be Strongly Relied Upon to Imply The Asymptotic Performance?

Previous COMRL endeavors [li2020focal, yuan2022robust, gao2024context] mostly apply t-SNE on the learned task representation to show the differentiation of the tasks.
They hold the belief that better performance corresponds to better differentiation of the tasks.
However, based on our analysis, the visualization result only depends on the task representation at convergence and ignores the task representation variation during the whole optimization process.
According to the illustration shown in Figure 5, less desirable differentiation results can also lead to better performance.
Hence, it is insufficient to imply performance based on such evaluation principles.
Nevertheless, we recognize that the visualization result can be seen as an auxiliary metric to assist in determining the task representation.

## 7 Conclusion & Limitation

Finding that the performance improvement guarantee of the COMRL training framework remains under-explored, we constructively provide a theoretical framework to link the training scheme in COMRL with the general RL objective of maximizing the expected return.
After scrutinizing this framework, we further find it ignores the variation of the task representation, which weakens the condition necessary for monotonic performance improvements.
Based on this finding, we propose a new issue called task representation shift, refine the condition for monotonicity, and prove that monotonic performance improvements can be achieved with appropriate context encoder updates.
With minor changes to the previous training framework for reining in the task representation shift, the performance improvements can be substantial.
We prospect that deeply exploring the role of task representation shift can make a profound difference in the COMRL setting.
Naturally, one limitation is that there may exist more advanced algorithms to control the task representation shift.
Additionally, our work focuses on the representation part for monotonic performance improvement, leaving the policy learning part alone.
Hence, one direction that merits further research is to design a policy learning algorithm to achieve better policy improvements conditioning on the optimal task representation distribution.
We leave these interesting problems to future work.

## 8 Appendix

### 8.1 Useful Lemmas

###### Lemma 8.1 (Return bound. [ zhang2023how ] ) .

Let $R_{max}$ denote the bound of the reward function, $\epsilon_{\pi}$ denote $\max_{s}D_{TV}(\pi_{1}||\pi_{2})$ and $\epsilon^{M_{2}}_{M_{1}}$ denote $\mathbb{E}_{(s,a)\sim d^{\pi_{1}}_{M_{1}}}[D_{TV}(p_{M_{1}}||p_{M_{2}})]$. For two arbitrary policies $\pi_{1},\pi_{2}\in\Pi$, the expected return under two arbitrary models $M_{1},M_{2}\in\mathcal{M}$ can be bounded as,

$$ $\displaystyle|V^{\pi_{2}}_{M_{2}}-V^{\pi_{1}}_{M_{1}}|\leq 2R_{max}(\frac{ \epsilon_{\pi}}{(1-\gamma)^{2}}+\frac{\gamma}{(1-\gamma)^{2}}\epsilon^{M_{2}}_ {M_{1}})$ (13) $$

*Proof:*

$$ $\begin{split}|V^{\pi_{2}}_{M_{2}}-V^{\pi_{1}}_{M_{1}}|&=|\sum^{\infty}\limits_ {t=0}\gamma^{t}\sum\limits_{s,a}(p^{\pi_{2}}_{t,M_{2}}(s,a)-p^{\pi_{1}}_{t,M_{ 1}}(s,a))r(s,a)|\\ &\leq R_{max}\sum^{\infty}\limits_{t=0}\gamma^{t}\sum\limits_{s,a}|p^{\pi_{2}} _{t,M_{2}}(s,a)-p^{\pi_{1}}_{t,M_{1}}(s,a)|\\ &=2R_{max}\sum^{\infty}\limits_{t=0}\gamma^{t}D_{TV}(p^{\pi_{1}}_{t,M_{1}}(s,a )||p^{\pi_{2}}_{t,M_{2}}(s,a))\\ \end{split}$ (14) $$

According to Theorem 2 (Return Bound) in [zhang2023how],

$$ $\displaystyle\sum^{\infty}\limits_{t=0}\gamma^{t}D_{TV}(p^{\pi_{1}}_{t,M_{1}}( s,a)||p^{\pi_{2}}_{t,M_{2}}(s,a))\leq(\frac{\epsilon_{\pi}}{(1-\gamma)^{2}}+ \frac{\gamma}{(1-\gamma)^{2}}\epsilon^{M_{2}}_{M_{1}})$ (15) $$

.

We bring this result back and can draw the conclusion.

∎

###### Lemma 8.2 (Inequility for $L_{1}$ deviation of the empirical distribution. [ weissman2003inequalities ] ) .

Let $P$ be a probability distribution on the set $\mathcal{A}=\{1,...,a\}$. For a sequence of samples $x_{1},...,x_{m}\sim P$, let $\hat{P}$ be the empirical probability distribution on $\mathcal{A}$ defined by $\hat{P}(j)=\frac{1}{m}\sum_{i=1}^{m}\mathbb{I}(x_{i}=j)$. The $L_{1}$-deviation of the true distribution $P$ and the empirical distribution $\hat{P}$ over $\hat{A}$ from $m$ independent identically samples is bounded by,

$$ $\displaystyle Pr(|P-\hat{P}|\geq\epsilon)\leq(2^{|\mathcal{A}|}-2)\exp{(-m \epsilon^{2}/2)}$ (16) $$

*Proof:*

According to the inequality for $L_{1}$ deviation of the empirical distribution [weissman2003inequalities], they conclude that $Pr(|P-\hat{P}|_{1}\geq\epsilon)\leq(2^{|\mathcal{A}|}-2)\exp(-m\phi(\pi_{P})
\epsilon^{2}/4)$, where $\phi(p)=\frac{1}{1-2p}\log\frac{1-p}{p}$ and $\pi_{P}=\max_{A\in\mathcal{A}}\min(P(A),1-P(A))$.
The paper [weissman2003inequalities] point out that $0\leq\pi_{P}\leq 1/2,\forall P$.
Based on this, we can derive the function $\phi(p)$ and find that $\phi(p)$ is monotonically decreasing on $(0,1/2)$.
Therefore, we can get $\phi(p)\geq\phi(1/2)=2$.
Then, we bring this result back and can conclude that $Pr(|P-\hat{P}|_{1}\geq\epsilon)\leq(2^{|\mathcal{A}|}-2)\exp(-m\epsilon^{2}/2)$.

∎

###### Lemma 8.1 (Return bound. [ zhang2023how ] ) .

###### Lemma 8.2 (Inequility for L 1 subscript 𝐿 1 L_{1} italic_L start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT deviation of the empirical distribution. [ weissman2003inequalities ] ) .

### 8.2 Missing Proofs

###### Theorem 8.3 (Return bound in COMRL) .

Assume the reward function is upper-bounded by $R_{max}$. For an arbitrary policy parameter $\theta$, when meeting Assumption 4.1, the return bound in COMRL can be formalized as:

$$ $\displaystyle|J^{*}(\theta)-J(\theta)|\leq\frac{2R_{max}L_{z}}{(1-\gamma)^{2}} \mathbb{E}_{m,x}(|Z(\cdot|x)-Z_{mutual}(\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}( \cdot|x)|)$ (17) $$

*Proof:*

$$ $\displaystyle|J(\theta)-J^{*}(\theta)|$ $\displaystyle=|\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z(\cdot|x );\theta)}}[R_{m}(s,a)]]-\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s ,Z^{*}(\cdot|x);\theta)}}[R_{m}(s,a)]]|$ (18) $\displaystyle=|\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z(\cdot|x );\theta)}}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z^{*}(\cdot|x); \theta)}}[R_{m}(s,a)]]|$ (19) $\displaystyle=|\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z(\cdot|x );\theta)}}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{mutual}(\cdot| x);\theta)}}[R_{m}(s,a)]$ $\displaystyle+\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{mutual}(\cdot|x);\theta) }}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z^{*}(\cdot|x);\theta)}}[R _{m}(s,a)]]|$ (20) $\displaystyle\leq\mathbb{E}_{m,x}[|\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z( \cdot|x);\theta)}}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{mutual} (\cdot|x);\theta)}}[R_{m}(s,a)]|$ $\displaystyle+|\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{mutual}(\cdot|x);\theta )}}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z^{*}(\cdot|x);\theta)}}[ R_{m}(s,a)]|]$ (21) $\displaystyle=\mathbb{E}_{m,x}[|V^{\pi(\cdot|s,Z(\cdot|x);\theta)}_{m}-V^{\pi( \cdot|s,Z_{mutual}(\cdot|x);\theta)}_{m}|+|V^{\pi(\cdot|s,Z_{mutual}(\cdot|x); \theta)}_{m}-V^{\pi(\cdot|s,Z^{*}(\cdot|x);\theta)}_{m}|]$ (22) $\displaystyle\leq\mathbb{E}_{m,x}[\frac{2R_{max}}{(1-\gamma)^{2}}\max_{s}D_{TV }(\pi(\cdot|s,Z(\cdot|x);\theta)||\pi(\cdot|s,Z_{mutual}(\cdot|x);\theta))$ $\displaystyle+\frac{2R_{max}}{(1-\gamma)^{2}}\max_{s}D_{TV}(\pi(\cdot|s,Z_{ mutual}(\cdot|x);\theta)||\pi(\cdot|s,Z^{*}(\cdot|x);\theta))]$ (23) $\displaystyle\leq\frac{2R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}(|Z(\cdot |x)-Z_{mutual}(\cdot|x)|+|Z_{mutual}(\cdot|x)-Z^{*}(\cdot|x)|)$ (24) $$

Eq. (23) is the result of directly applying Lemma 8.1 and Eq. (24) is the result of directly applying Assumption 4.1.
∎

###### Corollary 8.4 (Monotonic performance improvement condition for previous COMRL works) .

When meeting Assumption 4.1, the condition for monotonic performance improvement of previous COMRL works is:

$$ $\displaystyle\epsilon^{*}_{12}\triangleq J^{*}(\theta_{2})-J^{*}(\theta_{1}) \geq\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}(|Z(\cdot|x)-Z^{*}( \cdot|x)|)$ (25) $$

*Proof:*

$$ $\displaystyle J(\theta_{2})-J(\theta_{1})=J(\theta_{2})-J^{*}(\theta_{2})+J^{* }(\theta_{2})-J^{*}(\theta_{1})+J^{*}(\theta_{1})-J(\theta_{1})$ (26) $$

Denote $\frac{2R_{max}}{(1-\gamma)^{2}}$ as $\kappa$. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get:

$$ $\displaystyle J(\theta_{2})-J^{*}(\theta_{2})+J^{*}(\theta_{2})-J^{*}(\theta_{ 1})+J^{*}(\theta_{1})-J(\theta_{1})$ (27) $\displaystyle\geq\mathbb{E}_{m,x}[-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z(\cdot|x) ;\theta_{2})||\pi(\cdot|s,Z^{*}(\cdot|x);\theta_{2}))]+\epsilon^{*}_{12}$ (28) $\displaystyle-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z(\cdot|x);\theta_{1})||\pi( \cdot|s,Z^{*}(\cdot|x);\theta_{1}))]$ (29) $\displaystyle\geq\mathbb{E}_{m,x}[-2\kappa L_{z}|Z(\cdot|x)-Z^{*}(\cdot|x)|+ \epsilon^{*}_{12}]$ (30) $$

To get the monotonic performance improvement, we need Eq. (30) to be larger than 0. Hence, we can get:

$$ $\displaystyle\epsilon^{*}_{12}-\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_ {m,x}[|Z(\cdot|x)-Z^{*}(\cdot|x)|]\geq 0$ (31) $$

∎

###### Theorem 8.5 (Lower bound of performance difference in COMRL) .

Assume the reward function is upper-bounded by $R_{max}$. When meeting Assumption 4.1, the lower bound of the performance difference in COMRL can be formalized as:

$$ $\displaystyle J^{2}(\theta_{2})-J^{1}(\theta_{1})\geq\epsilon^{*}_{12}-\frac{2 R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}[2|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)| +|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|]$ (32) $$

*Proof:*
We can introduce the following decomposition

$$ $\displaystyle J^{2}(\theta_{2})-J^{1}(\theta_{1})$ (33) $\displaystyle=\mathbb{E}_{m,x}\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{2}(\cdot |x);\theta_{2})}}[R_{m}(s,a)]-\mathbb{E}_{m,x}\mathbb{E}_{(s,a)\sim d_{\pi( \cdot|s,Z_{1}(\cdot|x);\theta_{1})}}[R_{m}(s,a)]$ (34) $\displaystyle=\mathbb{E}_{m,x}[\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{2}( \cdot|x);\theta_{2})}}[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{1}( \cdot|x);\theta_{1})}}[R_{m}(s,a)]]$ (35) $$

Denote $\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{i}(\cdot|x);\theta_{j})}}[R_{m}(s,a)]]$ as $G^{\theta_{j}}_{Z_{i}}$, then we can get the following derivation. Here, we only consider the value within the brackets.

$$ $\displaystyle=\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{2}(\cdot|x);\theta_{2})} }[R_{m}(s,a)]-\mathbb{E}_{(s,a)\sim d_{\pi(\cdot|s,Z_{1}(\cdot|x);\theta_{1})} }[R_{m}(s,a)]$ (36) $\displaystyle=G^{\theta_{2}}_{Z_{2}}-G^{\theta_{1}}_{Z_{1}}$ (37) $\displaystyle=G^{\theta_{2}}_{Z_{2}}-G^{\theta_{2}}_{Z^{*}}+G^{\theta_{2}}_{Z^ {*}}-G^{\theta_{1}}_{Z^{*}}+G^{\theta_{1}}_{Z^{*}}-G^{\theta_{1}}_{Z_{2}}+G^{ \theta_{1}}_{Z_{2}}-G^{\theta_{1}}_{Z_{1}}$ (38) $$

Denote $\frac{2R_{max}}{(1-\gamma)^{2}}$ as $\kappa$. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get:

$$ $\displaystyle G^{\theta_{2}}_{Z_{2}}-G^{\theta_{2}}_{Z^{*}}+G^{\theta_{2}}_{Z^ {*}}-G^{\theta_{1}}_{Z^{*}}+G^{\theta_{1}}_{Z^{*}}-G^{\theta_{1}}_{Z_{2}}+G^{ \theta_{1}}_{Z_{2}}-G^{\theta_{1}}_{Z_{1}}$ (39) $\displaystyle\geq-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z_{2}(\cdot|x);\theta_{2})| |\pi(\cdot|s,Z^{*}(\cdot|x);\theta_{2}))+\epsilon^{*}_{12}$ $\displaystyle-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z^{*}(\cdot|x);\theta_{1})||\pi (\cdot|s,Z_{2}(\cdot|x);\theta_{1}))$ $\displaystyle-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z_{2}(\cdot|x);\theta_{1})||\pi (\cdot|s,Z_{1}(\cdot|x);\theta_{1}))$ (40) $\displaystyle\geq-\kappa L_{z}|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)|+\epsilon^{*}_{12}$ $\displaystyle-\kappa L_{z}|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)|-\kappa L_{z}|Z_{2}( \cdot|x)-Z_{1}(\cdot|x)|$ (41) $$

Simplifying Eq. (41) we can get:

$$ $\displaystyle J^{2}(\theta_{2})-J^{1}(\theta_{1})\geq\epsilon^{*}_{12}-\frac{2 R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_{m,x}[2|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)| +|Z_{2}(\cdot|x)-Z_{1}(\cdot|x)|]$ (42) $$

∎

###### Theorem 8.6 (Monotonic performance improvement guarantee on training process) .

Denote $\kappa$ as $\frac{(1-\gamma)^{2}}{4R_{max}L_{z}}\epsilon^{*}_{12}-\frac{1}{2}\beta$ and $vol(\mathbb{Z})$ as the volume of the task representation coverage simplex. Given that the context encoder has already been trained by maximizing $I(Z;M)$ to some extent. When meeting Assumption 4.1,
4.7, 4.8 and 4.9, with a probability greater than $1-\xi$, we can get the monotonic performance improvement guarantee by updating the context encoder via maximizing $I(Z;M)$ from at least extra $k$ samples, where:

$$ $\displaystyle k=\frac{1}{\kappa^{2}}(\sqrt{2\log\frac{2^{vol(\mathbb{Z})}-2}{ \xi}}+\sqrt{\alpha})^{2}$ (43) $$

Here, $\xi\in[0,1]$ is a constant.

*Proof:*

Let $Z_{1}$ have already been trained by maximizing $I(Z;M)$ to some extent. Given that to get $Z_{2}$, we need to train $Z_{1}$ by extra $k$ samples from the training dataset by maximizing $I(Z;M)$ to get the monotonic performance improvement guarantee.

We begin with the following inequality:

$$ $\displaystyle|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)|\leq|Z_{2}(\cdot|x)-\tilde{Z}_{2}( \cdot|x)|+|\tilde{Z}_{2}(\cdot|x)-Z^{*}(\cdot|x)|$ (44) $$

where $\tilde{Z}_{2}(\cdot|x)$ denotes the context encoder updated by fitting the empirical distribution on $k$ i.i.d samples from $Z^{*}(\cdot|x)$ based on $Z_{1}(\cdot|x)$.

According to Assumption 4.9, we have:

$$ $\displaystyle|Z_{2}(\cdot|x)-\tilde{Z}_{2}(\cdot|x)|\leq\sqrt{\frac{\alpha}{k}}$ (45) $$

Now, we move our focus on solving $|\tilde{Z}_{2}(\cdot|x)-Z^{*}(\cdot|x)|$.

By applying Lemma 8.2, the $L_{1}$ deviation of the empirical distribution $\tilde{Z}_{2}(\cdot|x)$ and true $Z^{*}(\cdot|x)$ over $vol(\mathbb{Z})$ is bounded by:

$$ $\displaystyle Pr(|\tilde{Z}_{2}(\cdot|x)-Z^{*}(\cdot|x)|\geq\epsilon)\leq(2^{ vol(\mathcal{\mathbb{Z}})}-2)\exp{(-k\epsilon^{2}/2)}$ (46) $\displaystyle Pr(|\tilde{Z}_{2}(\cdot|x)-Z^{*}(\cdot|x)|<=\epsilon)\geq 1-(2^{ vol(\mathcal{\mathbb{Z}})}-2)\exp{(-k\epsilon^{2}/2)}$ (47) $$

Then for a fixed $x$, with probability greater than $1-\xi$, we have:

$$ $\displaystyle|\tilde{Z}_{2}(\cdot|x)-Z^{*}(\cdot|x)|\leq\sqrt{\frac{2}{k}\log \frac{2^{vol(\mathbb{Z})}-2}{\xi}}$ (48) $$

Hence, we can get:

$$ $\displaystyle|Z_{2}(\cdot|x)-Z^{*}(\cdot|x)|\leq\sqrt{\frac{2}{k}\log\frac{2^{ vol(\mathbb{Z})}-2}{\xi}}+\sqrt{\frac{\alpha}{k}}$ (49) $$

Let $\epsilon=\frac{(1-\gamma)^{2}}{4R_{max}L_{z}}\epsilon^{*}_{12}-\frac{1}{2}\beta$, and recall that $\kappa$ denotes $\frac{(1-\gamma)^{2}}{4R_{max}L_{z}}\epsilon^{*}_{12}-\frac{1}{2}\beta$, then we can get the $k$ as:

$$ $\displaystyle\sqrt{\frac{2}{k}\log\frac{2^{vol(\mathbb{Z})}-2}{\xi}}+\sqrt{ \frac{\alpha}{k}}\leq\frac{(1-\gamma)^{2}}{4R_{max}L_{z}}\epsilon^{*}_{12}- \frac{1}{2}\beta$ (50) $\displaystyle k\geq\frac{1}{\kappa^{2}}(\sqrt{2\log\frac{2^{vol(\mathbb{Z})}-2 }{\xi}}+\sqrt{\alpha})^{2}$ (51) $$

Since we need the least number of samples that update the context encoder, we take $k=\frac{1}{\kappa^{2}}(\sqrt{2\log\frac{2^{vol(\mathbb{Z})}-2}{\xi}}+\sqrt{
\alpha})^{2}$.
∎

###### Corollary 8.7 (Monotonic performance improvement condition for pre-training scheme) .

Denote $Z_{pretrain}(\cdot|x)$ as the task representation distribution after pre-training.
When meeting Assumption 4.1, the monotonic performance improvement condition for pre-training scheme is:

$$ $\displaystyle\epsilon^{*}_{12}-\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_ {m,x}[|Z_{pretrain}(\cdot|x)-Z^{*}(\cdot|x)|]\geq 0$ (52) $$

*Proof:*
Denote $J^{pretrain}(\theta_{1})$ as the expected return of the policy $\pi(\cdot|s,Z_{pretrain}(\cdot|x);\theta_{1})$ among the expectation of tasks before update of the policy.
Similarly, denote $J^{pretrain}(\theta_{2})$ as the expected return of the policy $\pi(\cdot|s,Z_{pretrain}(\cdot|x);\theta_{2})$ among the expectation of tasks after update of the policy.

$$ $\displaystyle J^{pretrain}(\theta_{2})-J^{pretrain}(\theta_{1})$ (53) $\displaystyle=J^{pretrain}(\theta_{2})-J^{*}(\theta_{2})+J^{*}(\theta_{2})-J^{ *}(\theta_{1})+J^{*}(\theta_{1})-J^{pretrain}(\theta_{1})$ (54) $$

Denote $\frac{2R_{max}}{(1-\gamma)^{2}}$ as $\kappa$. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get:

$$ $\displaystyle J^{pretrain}(\theta_{2})-J^{*}(\theta_{2})+J^{*}(\theta_{2})-J^{ *}(\theta_{1})+J^{*}(\theta_{1})-J^{pretrain}(\theta_{1})$ (55) $\displaystyle\geq\mathbb{E}_{m,x}[-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z_{ pretrain}(\cdot|x);\theta_{2})||\pi(\cdot|s,Z^{*}(\cdot|x);\theta_{2}))]+ \epsilon^{*}_{12}$ (56) $\displaystyle-\kappa\max_{s}D_{TV}(\pi(\cdot|s,Z_{pretrain}(\cdot|x);\theta_{1 })||\pi(\cdot|s,Z^{*}(\cdot|x);\theta_{1}))]$ (57) $\displaystyle\geq\mathbb{E}_{m,x}[-2\kappa L_{z}|Z_{pretrain}(\cdot|x)-Z^{*}( \cdot|x)|+\epsilon^{*}_{12}]$ (58) $$

To get the monotonic performance improvement, we need Eq. (58) to be larger than 0. Hence, we can get:

$$ $\displaystyle\epsilon^{*}_{12}-\frac{4R_{max}L_{z}}{(1-\gamma)^{2}}\mathbb{E}_ {m,x}[|Z_{pretrain}(\cdot|x)-Z^{*}(\cdot|x)|]\geq 0$ (59) $$

∎

###### Theorem 8.3 (Return bound in COMRL) .

###### Corollary 8.4 (Monotonic performance improvement condition for previous COMRL works) .

###### Theorem 8.5 (Lower bound of performance difference in COMRL) .

###### Theorem 8.6 (Monotonic performance improvement guarantee on training process) .

###### Corollary 8.7 (Monotonic performance improvement condition for pre-training scheme) .

### 8.3 Implementation Details

In this section, we report the implementation details of all the objectives in our main paper.

- •
Contrastive-based [li2020focal]. We adopt the open-source code of FOCAL(^2^22https://github.com/FOCAL-ICLR/FOCAL-ICLR/).
- •
Reconstruction-based [li2024towards]. We obtain the source code of UNICORN from the authors and adopt UNICORN-0 as the reconstruction-based objective training framework.
- •
Cross-entropy-based. Built upon the code base of FOCAL, cross-entropy-based only replaces the distance metric learning loss with cross-entropy loss.

To maintain fairness, we make sure all benchmark-irrelevant parameters are consistent, as shown in Table 1.

**Table 1: Benchmark-irrelevant parameters setting in the training process.**
| training tasks | 20 |
| --- | --- |
| testing tasks | 20 |
| task training batch size | 16 |
| rl batch size | 256 |
| context size | 1 trajectory |
| actor-network size | [256, 256] |
| critic-network size | [256, 256] |
| task encoder network size | [64, 64] |
| learning rate | 3e-4 |

**Table 2: Benchmark-relevant parameters setting in the training process.**
| Configurations | Reach | Ant-Dir | Button-Press | Dial-Turn |
| --- | --- | --- | --- | --- |
| dataset size | 3e5 | 9e4 | 3e5 | 3e5 |
| task representation dimension | 5 | 5 | 5 | 5 |

For each environment, we make the benchmark-relevant parameters the same for each algorithm, as shown in Table 2.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Can you elaborate more on how different batch sizes induce different task level contexts?
soundness: 3
strengths: 1. This work considers a nuanced issue in traditional OMRL optimization framework. It critically investigates the components and show that variation in task representation is fundamental to monotonic performance improvement. It introduces the phenomena called "task representation shift".
2. Theoretical justification for the claims are well presented and it outlined a complete algorithmic framework. 
3. It presents rigorous experimental validation using several environments from two different benchmarks with 3 types of objective functions and 3 types of data sets.
summary: This work attempts to explain the performance improvement of the Offline Meta RL (ORML) optimization framework. It identifies that the variation of task representation learned through the optimization process if often ignored. Such issue violates the condition of monotonic performance improvements. Thus, addressing task representation shift with carefully designed encoder updates is necessary. Experimental results across two different benchmarks with 3 different types of training objective and data qualities have been presented.
weaknesses: 1. The task variation used in the experiments could be improved by including more distinct tasks. 
2. Lack of results with more statistically significant metrics such as interquartile mean with confidence interval. Such comparison would help the reader as the general standard deviation highly overlaps.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. The setup in Section 4.3 is a bit confusing. Could the authors clarify what "accumulation steps of task representation shift" refers to? Also, does setting $k = 2 \times bs$ refer to the initial value for training?
2. It would be better to explain in the appendix how the cross-entropy-based loss can replace the loss in FOCAL, preferably by providing the corresponding expression.
3. The motivation for using the cross-entropy-based algorithm is somewhat unclear. Could you explain why it replaces the distance metric learning loss?
4. From the experimental results, the improvement seems marginal. Could you analyze the reasons behind this?
5. The authors state in section 6.2, "We recognize that the visualization result can be seen as an auxiliary metric to assist in determining the task representation." Why is the visualization result insufficient to fully represent the true task representation? The visualized convergence results being imperfect yet leading to better outcomes seem counterintuitive.
soundness: 3
strengths: 1. The paper is well-structured and clear.
2. The authors identified a unique challenge existing in offline meta-RL, task representation shift, which is highly novel.
3. The proof provided is detailed and logically rigorous, highlighting a flaw overlooked by previous work: it ignores the variation of task representation in the optimization process. I believe this is the most significant contribution of this paper.
4. The paper concludes with some interesting discussions, which have the potential to motivate future research.
summary: This paper addresses the overlooked issue of task representation shift in offline meta-reinforcement learning (OMRL), which can prevent consistent performance improvement. 

The authors provide theoretical insights and practical strategies to control this shift, demonstrating that such control enables more stable and higher asymptotic performance. They propose adjustments in training, such as tuning batch sizes and accumulation steps, to manage task representation effectively. Experimental results across benchmarks validate that these adjustments lead to performance gains, highlighting the importance of considering task representation dynamics in OMRL.

Overall, this is a very interesting paper with great potential to inspire future work. I would be willing to increase the score if the authors could address my following concerns.
weaknesses: 1. The algorithm box does not clearly explain how $k$ and $acc$ are utilized. Adding a brief explanation in the red-highlighted part of the algorithm box about how these are calculated would make the algorithm more understandable.
2. The experiments are limited and need improvement; it would be beneficial to verify the impact of task representation shift in more diverse testing scenarios.
3. The results largely depend on the settings of hyper-parameters, such as $k$ and $acc$, which seem to be unstable. The paper lacks an analysis of the experimental effects caused by adjusting these two parameters.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. The authors' description of the task representation shift is unclear. Can it be understood as the phenomenon where the task representation deteriorates after updating the encoder? A specific example would be helpful for better understanding this point. 
2. The authors introduce a large number of assumptions in Section 4. Has the reasonableness of these assumptions been justified? 
3. When choosing the Ant environment for experiments under Mujoco, was it randomly selected, or was it only the environment that showed the best results? 
4. How are the hyperparameters selected for different environments? It seems from the results in Figure 2 that the experimental outcomes are quite sensitive to the choice of k/acc. Is it possible to analyze the results of different hyperparameters in relation to the characteristics of the environments themselves?
soundness: 2
strengths: 1.Provides a new perspective for learning better task representations in COMRL.  
2.The paper includes a detailed mathematical derivation process.  
3.The authors provide their code implementation in the supplementary materials.
summary: This paper explains the reasons for the effectiveness of previous context-based offline meta reinforcement learning methods and introduces a task representation shift problem. The authors further demonstrate both theoretically and empirically that improving the update process of the context encoder for this problem can significantly enhance the performance of the original COMRL algorithm.
weaknesses: Experimental Section:
1. The experimental results are obtained only within the improved FOCAL framework, lacking experiments on important baselines such as CORRO [1], CSRO [2]. I believe conducting experiments in just one framework is insufficient to verify effectiveness. It would be more convincing if significant results could be demonstrated across multiple baselines. 
2. The authors only consider four environments, while in CSRO and UNICORN, each algorithm considers six environments under Mujoco.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Update: I am happy with most of the authors' answers and recommending accepting the paper. I raised the rating to a strong 6.

Old questions:

1) Statistical significance analysis should be run in all experiments or claims on experimental results changed significantly, 2) Please answer all technical questions and comments.
soundness: 3
strengths: According to my knowledge the theoretical analysis of task representation shift in context based offline meta reinforcement learning is original and may be of interest to researchers working in this domain and potentially others where context encoders are used.

For claim (1), the theoretical analysis is a solid contribution in case the authors can clarify the questions/comments further below. The insight that task representation shift influences the requirements for monotonic performance improvements is valuable.

For claim (2), mathematical proofs under the made assumptions appear correct.

The paper is overall well written and understandable.
summary: The paper focuses on task representation shift in context encoder updates when the context encoder is updated in offline meta reinforcement learning (OMRL). In particular, the paper analyzes the relationship between performance improvement in policy updates and the amount of data used for context encoder updates.

The paper claims two contributions:

(1) Identifying task representation shift as a major issue in context based offline meta reinforcement learning. The task representation shift makes it harder to satisfy the condition necessary for monotonic performance improvements. For the "major issue" claim, the paper claims to provide empirical evidence.

(2) Proof of monotonic performance improvement taking into account task representation shift and number of samples used to improve the context encoder. This information can be used in practice, for example, such that an algorithm updates the context encoder only when the improvement due to policy updates is sufficiently large such that the number of samples used for the context encoder updates is sufficient.
weaknesses: For claim (1), regarding the claim that task representation shift is a "major issue", the experimental results do not at the moment provide evidence for this as claimed in the paper. Statistical significance analysis of the results should provide further information.

The Assumptions 4.7, 4.8, and 4.9 need to be motivated in more detail. In particular, in Assumption 4.9, assuming that fitting error decreases inversely proportional to the number of samples is questionable.

vol(Z) needs to be explained.

Other technical details, discussed below, need to be described in more detail.

DETAILS:

In more detail, the paper says "As shown by Figure 2, even with minor changes to the algorithms, the performance improvements are substantial." but that the performance improvements are "substantial" is not at all obvious but rather a result of random chance.
Statistical significance testing is needed to draw conclusions about the experimental results. This applies to all the results.

In "vol(Z) as the volume of the task representation coverage simplex.", please define "coverage simplex". What exactly is it? Furthermore, on lines 860 - 863, Lemma 8.2 is used. However, this Lemma applies in the case of discrete values and uses the number of values '|A|' in '2^|A|' but here '|A|' is replaced with vol(Z) resulting in 2^vol(Z). Why can this be done? This needs detailed explanation.

In Definition 4.5, I do not understand what "among the expectation of tasks before update of the context encoder and the policy" and "among the expectation of tasks after update of the context encoder and the policy." mean. Maybe the word "among" is here confusing. Can you provide a more explicit detailed definition? Just provide the equations for J^1(\theta_1) and J^2(\theta_2)?

On Line 281, the paper refers to Eq. (32) which has not been introduced yet.

Notation: using multiple characters to denote quantities such as "bs", "acc" etc. is not a good way. If this kind of textual description for variables is desired, one way is to use something like N_{\text{acc}}.

The notation is slightly confusing. Now \theta denotes policy parameters but it would be good to add also a symbol to denote the context encoder parameters to distinguish clearly in the equations what is being optimized.


RELATED WORK:

In "While ContrBAR (Choshen & Tamar, 2023) also benefits from the performance improvement guarantee in the OMRL setting, it is specifically served for its own framework", what does "it is specifically served for its own framework" mean? Since the proofs for performance improvement guarantee are one of the claimed contributions in this paper it is important to describe this in sufficient detail and discuss the differences.


LANGUAGE/PRESENTATION:

In "weakens the condition necessary for monotonic performance improvements, and may therefore violate the monotonicity.", rephrasing may be needed. "weakens" is slightly misleading since the sufficient conditions are actually stricter, not weaker, that is, when taking the task distribution shift into account larger policy improvements are needed to satisty monotonicity according to the analysis in this paper.

The sentence "However, ours considers the variation of task representation ignored by the previous training framework by imposing the extra condition to determine whether needs to update the context encoder." is missing some words.

On line 261: "As shown in Corollary 4.4, the monotonic performance improvement can be guaranteed with only better approximation to Z^*(·|x).": I maybe understand the intention of the text here but would be good to describe this more explicitly, that is, that Z(·|x) should be close to Z^*(·|x) such that the lower bound is small enough for finding a policy that improves on the old policy?

"we need to find the positive C to improve the performance monotonically." ->
"we need to find a positive C to improve the performance monotonically."

"Center around this motivation" -> "Centering around this motivation"

"an policy learning algorithm" -> "a policy learning algorithm"

Meta-review: This work investigates task representation shift in offline meta reinforcement learning, providing theoretical insights and practical strategies to control this shift, such as tuning batch sizes and accumulation steps.  Experiments are well designed and executed, across two different benchmarks with three different types of training objective and data qualities, validating the proposed methods. The reviewer believes that this is an important but overlooked issue by the meta RL community and this paper is a solid work in identifying this issue and proposing the first working solutions. All reviewers agree that this is a good work and should be accepted and I agree with the reviewers' evaluation and believe this work could open up an interesting future research direction.

Decision: Accept (Poster)