Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning
Authors: Hai Zhang
Boyuan Zheng
Tianying Ji
JinHang Liu
Anqi Guo
Junqiao Zhao
Lanqing Li
Venue: ICLR 2025 Poster (2025)
Keywords: offline meta reinforcement learning
performance improvement guarantee
task representation shift
Primary area: reinforcement learning
TL;DR: We give performance improvement guarantee of prior COMRL works and identify a new issue called task representation shift, theoretically demonstrating that reining in task representation shift properly can achieve monotonic performance improvements.

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$ and its latent representation $Z$ ($I(Z;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 observe that the condition for monotonic performance improvements does not consider the variation of the task representation. When these variations are considered, the previously established condition may no longer be sufficient to ensure monotonicity, thereby impairing the optimization process.
We name this issue \underline{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)$ 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.

Published as a conference paper at ICLR 2025
SCRUTINIZE WHAT WE IGNORE: REINING IN TASK REPRESENTATION SHIFT OF CONTEXT-BASSE OFFLINE META REINFORCEMENT LEARNING
Hai Zhang', Boyuan Zheng', Tianying Ji, Jinhang Liu', Anqi Guo!
Junqiao Zhao'; :Lanqing Li3
Tongji University, 2 Tsinghua University. ZZejjing Lab, 4 The Chinese University of HongKong
(zhanghei12138, zhengboyuan, Jinhangliu, thao @tong edu...e.nn http::mmmmmissssssssneeeeee edu.
(anqiguc098, lanqingli1999 @emaee.....om
ABSTRACT
Offline meta reinforcement learning (OMRL) has emerged promising approach for interaction avoidance and strong generalization performance by leveraging pre- collected data and meta-learning techniques. Previous context- -aseed approaches predominantly rely on the intuition that alternating optimization between the con- text 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 and its latent representation Z (1(Z; M)) while the policy adopts the standard offline reinforcement learning (RL) algorithms condition- ing on the learned task representation. Despite promising results, the theoretical justification of performance improvements for such intuition remains underex- plored. 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 observe that the condition for monotonic performance improvements does not consider the variation of the task representation. When these variations are con- sidered, the previously established condition may longer be sufficient to ensure monotonicity, thereby impairing the optimization process. We name this issue task representation shift and theoretically prove that the monotonic performance improvements can be guaranteed with appropriate context encoder updates. We different settings to rein in the task representation shift on three widely adopted training objectives concerning maximizing / [(2; M) across different data qualities. Empirical results show that reining in the task representation shift indeed im- prove performance, Our work opens up new avenue for OMRL, leading to better understanding between the task representation and performance improvements,
INTRODUCTION
has driven impressive advances in many complex decision-making problems in recent years (Silver et al.) 2018 Schrittiieser et al.) 2020 Zhang et af.] 20236) 2024, Ma et al. 20244ba a). 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 (Wul
al. 2019 Levine et al., 2020) 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 (Fatemi et al., [022 Tang et al.)
2022), robotics (Sinha et al. 2022 Kumar et al. 2022) and games (Schrittwieser et al. 2021
"Corresponding Author
'Code: https//gthhuh.co//httrry//1 12138 Aerere/aene

---

Published as a conference paper at ICLR 2025
Though demonstrating its superiority, offline RL holds notable drawback towards generalizing to the unknown (Ghosh et al. (202T).
As remedy, by combining the meta-learning techniques, OMRL (Li ct al. 2021 Xu et al. 2022) 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 (Li et al.
2020 Mitchell et al.
2020 2024) hold a popular paradigm that seeks optimal meta-policy conditioning on the context of Markov Decision Processes (MDPs). Specifically, these methods (Li et al. 2020 Gao et al. 2024 Li et al. 2024) propose to train context encoder via maximizing I [2; M eaan 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 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.
Canditin
nut Iust
Macinbe
Maxiniox
I(.M)
Updaae 2
KZ.M)
Update Z
Updatee
Replay Butler
Contert
Encester
Replay Butler
Cumborst
Encesler
offline RL algorithu
Update Policy
offline RL
Update Policy 2
algorithas
Taslos
Taslos
original
ours
Alternating update process:
@-
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, our training framework considers the previously ignored variation of task representation by introducing an extra condition to decide whether the context encoder should be updated.
Intriguingly, we formalize this training framework alternating two-stage optimization framework and then link it with the general RL objective of maximizing the expected return. To be specific, maximizing [(Z; M)) and adopting standard offline RL algorithms be interpreted as consistently raising the lower bound of the expected return conditioning on the optimal task representation distribution (See Section 4.1). This achieved by extending the return discrepancy scheme (Janner et al. 2019)) to the COMRL framework. Thus, this analysis provides 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 better approximation to the optimal task representation distribution. By explicitly modeling the variation of task representation, we conclude that it is critical part of monotonic performance improvement. 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 context encoder updates (See Section 4.2).
To show the impacts of this issue, we use different settings to rein in the task representation shift on three widely adopted objectives concerning maximizing I(Z: M), covering the upper bound, the lower bound, and the direct approximation. 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 better understanding between the task representation and performance improvements.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORKS
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 (Li et al. 2020; Yuan &
Lu 2022 Gao et al. 2024 Li et al. 2024) 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 perform meta-adaptation, According to the theoretical insights from (Li al. 2024), prior COMRL works (Li et al. 2020, Yuan & Lu 2022 Gao et al. 2024 Li et al.
2024) mainly focus on designing a context encoder learning algorithm to better approximate I(Z M) However, our work turns the focus to refining the condition of monotonic performance improvements based on a given context encoder learning algorithm. Centering 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 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 (Kakade
Langford 2002; Schulman et al. 2015 Ji et al. 2022 Zhang al., 2023a), return discrepancy (Luo
ct al. 2018 Janner et al. 2019), , and regret bounds (Osband & Van Roy 2014 Curi et al. 2020). In offline RL, CaL-based methods (Kumar et al. 2020 Yu et al. (022 also enjoy safe policy improvement guarantees. However, most works focus on the single-task setting, leaving the perfor- mance improvement guarantees in the context-based meta-RL settings largely unexplored. While
ContrBAR (Choshen & Tamar (2033) also benefits from the performance improvement guarantee, the theoretical insight focuses on the online setting. Additionally, it is tailored to one particular approximation of [Z: ; .M) as it makes assumptions specific to this approximation. In contrast, our work focuses on the general offline setting, addressing broader class of algorithms that maximize various bounds of 7 (Z; M).
3 PRELIMINARIES
3.1 PROBLEM STATEMENT
A task in RL is generally formalized a fully observed MDP (Puterman 2014). which is represented by a tuple (S. A, P. Po. P.Po, R,y) with state space S, action space A. transition function P(5'|s, a). reward function R(s, a), initial state distribution (o(s) and discount factor y E [0, 1]. For arbitrary
policy T (als), we denote the state-action distribution at timestep as ds,,(5,a) _ Pr(s: = 8, 0g
also Po: ~ т, St-1 P.Vt 2 0). The discounted state-action distribution of given i[ass) is E (s, a). The ultimate goal is to find an optimal policy x (a|s) to maximize the expected return (( a)))
denoted as d=(s,a) a) _ (1 -1-1 7) tt=
It is widely assumed that the task m in the meta RL setting is randomly sampled from the task distribution p(M) (Li et al. 2020 Yuan & Lu 2022 Gao et al. 2024 Li et al. 2024). We denote I,..., 1, N IJ. an offline dataset D, {(sL.f, ai.j, S )), j== is collected in advance, where K denotes the number of transitions. The learning algorithm required to train meta-policy metaa with access only to the given offline datasets. At test time, given an unseen task, the meta policy Tmeta performs task adaptation to get a task-specific
the number of training tasks as N. For each task 10,
policy and then will be evaluated in the environment.
3.2 CONTEXT-BASED OMRL
Previous COMRI endeavors employ a compact representation z to capture/quantify the variation over tasks Yuan & Lu 2022). In practice, they choose to train context encoder Z(- σ) to extract the task information (Li ct al. 2020 Dorfman et al. (222I). where x denotes the given context and ó denotes the parameters of the context encoder. Then, the policy is learned by conditioning on the task representation as (aass, (- Φ): 0). We assume there exists an optimal task representation distribution Z (-|x;; σ²) that meets the optimal expected return for any policy parameter @ determined by T (als, Z(- o0") ); 0). Hence, the objective of COMRL is formalized as:
max (0) = Em [E(e.a)dd (1..0 (9) [m (8, a)]]
3

---

Published as a conference paper at ICLR 2025
where Rm, a) denotes the ground-truth reward function of the task
Previous COMRL works (Li ct al. 2020 2021; Yuan Lu 2022; Gao et al. 2024 Li ct al. 2024) proven to optimize the context encoder by maximizing the approximate bounds of (Z; M), ), as shown in Theorem 3.1 They hold the intuition that using such kind of approach to optimize the optimize the policy
context encoder and adopting the standard offline RL algorithms to performance improvements. However, the theoretical justification of performance improvement for such intuition has been less explored. For later use, we denote Z(-|æ; dmmatua)) as the optimal solution for maximizing I (Z; M). Whether (-1 d)mutuaa) ) is equivalent to (-12;; (+;=)")) remains an open problem. To avoid confusion, we introduce an extra notation here.
lead
Theorem 3.1 ((Li et al. 2024)), Denote X,, and X₁ are the behavior-relatee (s, a) component and task-related (s compponett of the context X, with X (Xp, X2). We have:
I(Z; Xx)x)) S [(; M) - X()
(2)
where 1) = = -I;
I(Z; 2) CCORRO = x 3)
CCSRO (X - - Xx), and denotes equality up constant.
4 METHODS
We first introduce Lipschitz assumption, a widely used technique in both model-free (Song & Sun 2019 Ghosh et al., 2022) and model-based (Ji et al. 2022] Zhang et al.) (023a) RL frameworks.
Assumption 4.1. The policy function is L Lipschitz w.r.t some norm || in the sense that
(2(12;;)1) 2112;; (b2) E E |m(,|,, ((110; (-18, 7(- 02): S T, (2(,a2;;)) 7( ([; 02))
(3)
The proof's for all the following sections are provided in Appendix Section 8.2
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 (Janner et al. 2019) in the model-based RI. field, we can similarly build a tractable lower bound for J((
Definition 4.2 (Return discrepancy in COMRL). The return discrepancy in the COMRL setting can be defined as
(') - ((0) 2 (J( - (+))
(4)
where J(O) = Bmm (F (s..)) (011000))) as) is the expected return of the policy condi- tioning on the learned task representation Z(-1x;)).
From Definition 4.2 we know that if the estimation error (8) -(A) can be upper-bounded, the lower bound for (+)) can be established simultaneously. Then, Eq. (4) can be interpreted as unified training objective for both the context encoder and the policy. Specifically, we can alternate the optimization target to lift the lower bound for JO))), where the whole optimization process can be seen as a two-stage alternating framework with the first stage updating the context encoder to minimize (0) ((0) and the second stage updating the policy to maximize (() (0) (() as (0) 2 (0) - (0) ] Hence, the performance improvement guarantee can be achieved.
The following theorem induces tractable upper bound for (0) J()), bridging the gap be- tween 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 Rmoz- For an arbitrary policy parameter 0, when meeting Assumption 4.1 the return bound in COMRL can be formalized as:
2Rmaa Loz Em (1m-=((ZZ)z0))) - 2f muutual)] }| + |Z( (1: ÿ mmutuall 2(-12r; (²))))
(J)(0) J(0)) < (1- 7)²
(5)
4

---

Published as a conference paper at ICLR 2025
By directly unrolling the Theorem 4.3 we can establish the lower bound of J (6) as:
2 2Rmar I,z
(0) 2 ((0)
E (l [x; Φ) Z t ) + 1Z((|22; o mutuel)) Z( Z-112; x)))
(1-
Note that (-|a; m [(|la; 66 is a constant w.r.t the learned task representation Z(- a Φ). Hence, this term can be ignored when optimizing Z( Z(|22; Φ). Admittedly, there may exist smarter algorithm to further reduce the gap of [Z xx; @mutual) (1|22; o²)], 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, should minimize the return bound, namely minimizing Em (-|aa; Φ) Z(-12; @ )|- For the policy learning. we should lift the lower bound of
(O), namely maximizing (0) 2Rmea
(1-7)² Em, (1Z(-|2;)) (-ka;; () 11 [ d)mutual)- 7(1a2; ) = J(0), the task representation related terms are constants for policy optimization.
To maximize (Ø). we can adopt standard offline RL algorithms (Wu et al. 2019 Kumar et al. 2020
Fujimoto & Gu 2021) similar to (Yang et al. 2022).
Recall that previous COMRL works adopt alternatingly training the context encoder by maximizing
[(Z; M). which can be approximately seen minimizing Em. l (-|x;; Ø) Z(-la;; ) ) and training the policy by conditioning on the learned task representation well as applying the standard offline RI. algorithms. Therefore, we argue that our theoretical analyses provide an explanation for the performance improvement guarantee of previous COMRL works.
However, this optimization framework only considers the discrepancy between Z(-|2; Φ) and Z (-|2;; (e @ ), without considering the variation of (-1|a; Φ). In the next section, we will show that this may violate the monotonicity the performance improvements.
4.2 MONOTONIC PERFORMANCE IMPROVEMENT CONCERNING TASK REPRESENTATION SHIFT
In this section, we aim to demonstrate that as the previous optimization framework in Section 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 the previous framework.
Corollary 4.4 (Monotonic performance improvement condition for previous COMRL works). When the condition for monotonic performance improvement of previous COMRL
meeting Assumption 4.1 works is:
ARmar Lus Bm, 12 -'())
(1Z()12;) Bm. Z(-12; (')l
- (1-
Z(-lx; ②) should close to ((1r,; d¹) such that the lower bound small enough for finding policy to achieve monotonic performance improvement. Next, we will introduce the performance difference bound framework to model variation of task representation.
As shown in Corollary 4.4
Definition 4.5 (Performance difference bound COMRL). For an alternating update process, denote (0:) Emm.u((s.a)-daci 2e-0e1 [Rm (s, a)) the expected return of the policy
(-1s8, 71-xx; (b);: 01) before update of the context encoder and the policy. Similarly, denote (0)) = Em, (s.a) ( (a RRm (2,0)) as the expected return of the policy n(-1,, Z &2) ); 02) after update of the context encoder and the policy. The performance difference bound in the COMRL setting can be defined as
J²(02) - (+)))) C
(8)
when C is non-negative, the algorithm allows a monotonic performance improvement.
According to Definition 4.5 we need to find a positive C' to improve the performance monotonically. To achieve this, we can derive the lower bound of the performance difference.
5

---

Published as a conference paper at ICLR 2025
Theorem 4.6 (Lower bound of performance difference in COMRL). Assume the reward function is upper-bounded by Rmoz- When meeting Assumption 4.1 the lower bound of the performance difference in COMRL can be formalized as:
2Rmax La E x (2] Z(122;))) [Z( (2) Z(-la; #1)1
(2(02) - J' (01) 2 6/
(1-7)²
(9)
For achieving the monotonic performance improvement guarantee, we further let the right-hand side of Eq. (9) be positive. Then, we would have the following condition to achieve monotonicity.
2Rmaa La 62 2}- Em (-|2(+
7(+|2;;
x x²)l [+(+| 92) - 7-110; Φ1)1] 2 0
(10)
Compared to the condition in Corollary 4.4 Eq. 10 presents an additional term |Z(-|=; (2)
Z1|12; Φ1) to achieve monotonic performance improvement, which corresponds exactly the previ- ous 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(-|2: ) but also the magnitude of the update. For example, if we assume that ((|aa; o1)) trained from scratch and the update process from Z [-I|a;; Φ1) to ( (-Jar; (22) brings 7(1a2; 22) close to 7-|12; o²)). then with the condition in Corollary [.4 the monotonic performance improvement can be easily achieved with small (/2- However, for the condition Eq. (10), small c|2 may cause the violation of monotonicity [Z Φ1) Z(-læ; 62)l remains large.
To highlight the impacts of [= (2) (=|a;; +1 Φ1),. 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 eî2 on the condition of appropriate encoder updates, shown in Theorem 4.10
Assumption 4.7. The impacts of task representation shift are upper-bounded by 3 and less than the policy improvement 6ī2 with certain coefficient.
Assumption 4.8. The space of the task representation is discrete and limited.
Assumption 4.9. There exists an for any given (-|ax; 61), |Z(- [x; 62) Z(- 02)1 < 8. where 2(|2;; 42) denotes the context encoder updated by maximizing I [Z; M) with the data size b randomly sampled from the training dataset based on (-Iar; Φ1) and 7(-|x: (22) denotes the context encoder updated by fitting the empirical distribution on i.i.d samples from (1-1r; Φ²) based on Z(-1x; Φ1 B
Theorem 4.10 (Monotonic performance improvement guarantee on training process). Denote K as (1-2)²
4Rmoo L., <1 - 1,3 and Z as the cardinality of the task representation space. Given that the context encoder has already been trained by maximizing T [Z; M) some extent. When meeting Assumption
4.1 4.7 4.8and|43 with probability greater than 1 &, 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:
2 -2
k:
(1
2log
+ va)²
(11)
N²
E
Here, & E [0. is a constant.
Theorem 4.10 shows the connection between the needed update data size k and the performance improvement brought only by the policy update <ī2. If the calculated is larger than the given batch size, the context encoder should avoid this update and wait for the accumulation of eiz. As si increases, decreases. When 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 right. right. Compared the original alternating framework, Theorem 44.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 where the way to adjust the update of the context encoder colored by red.
Forr details justify these assumptions, please refer Appendix 8.3

---

Published as a conference paper at ICLR 2025
Algorithm 1 General Algorithmic Framework Towards Reining In The Task Representation Shift
Input:
Offline training datasets X, initialized policy Te, context encoder Zg. given task representa- tion shift threshold / and given training batch size for context encoder Nb.
for iter in alternating iterations do
2: // Update the context encoder
3: Estimate the conditions (e.g.use Eq. to approximate k)
4: if Conditions for updating the context encoder are met (e.g. k <= Nb.) then
5: while Accumulated task representation shift is less than 3 do
6: Sample context from X
7: Obtain task representations from Zg with inputting the sampled context
8: Compute Lencnder concerning maximizing [[Z; M)
9:
Update & to minimize Lencoder
end while
11:
end if
// Update the policy
13: Detach the task representations
Sample training data from X
15:
Compute pplicc via standard offline RL algorithms
Update I to minimize Spolicy
17: end for
4.3 PRACTICAL IMPLEMENTATION
As outlined in Algorithm] 1 the way to rein in the task representation shift can be seen aspects, namely 1) when to update the context encoder determined by A and 2) how many times to update the context encoder determined by To cover these two aspects, introduce and Nace- Here, Mr denotes that the context encoder needs to be updated every n updates of the policy, and Nase times. Our settings include 1) Nr = 2, Nase = 1, 2) N₂ = 3, Nace 1, N₂ = 1, Nace = 2, and 4) N₂ 1, Nace 3. All settings as well the original setting, namely Ng Nace 1, are performed for 8 different random seeds.
parameters Ng
denotes when the context encoder needs be updated, it is updated
With respect to the specific context encoder learning algorithms concerning maximizing I(Z; M), we choose three representative algorithms that have been widely adopted in previous COMRL endeavors. The details of these algorithms are shown in Appendix Section
8.4
Contrastive-base is applied in (Li et al., 2020 Gao et al. be the upper bound TZZ M).
2024
2024) and it is proven to Li et al.
Reconstructionn-bsse is applied in Ziingraft et al. 2019 Dortman et al.
2021 Li et al. 2024) and
is proven to be the lower bound of [(; M), which is equivalent CORRO
Yuan Lu 2022
under the offline setting.
Cross-entropy-bbete based is proposed our work and is a direct approximation w.r. I(Z; M).
With respect to the policy learning algorithm, maintain the consistency of previous COMRL works, we directly adopt BRAC (Wu al. 2019) to train the policy.
EXPERIMENT
Our experiments are conducted to show that the previous optimization framework that ignores the impacts of task representation shift 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 Todorov et al. 2012) and MetaWorke (Yu et al. 2020) benchmarks to evaluate
the algorithms. Following the protocol of UNICORN (Li et al. 2024). randomly sample 20
7

---

Published as a conference paper at ICLR 2025
training tasks and 20 testing tasks from the task distribution. We train the SAC (Haarojj al. 2018) agent from scratch for each task and use the collected replay buffer as the offline dataset. During the meta-testing phase, randomly samples a trajectory from the dataset as the context to obtain the task representation and then passes it to the meta-policy complete the testing in the true environments.
adopt fully offline setting in few-shot manner like (Li et al. 2020) that
5.2
MAIN RESULTS
Roach
Buston Prace
Puth
3rgial
4. May
61 2.8
1.
Figure Testing returns of different settings to rein in the task representation shift. Solid curves refer to the mean performance of trials over 8 random seeds, and the shaded areas characterize the standard deviation of these trials.
Figure 2 and Table 3 illustrate the performance of our different settings to rein in the task repre- sentation shift on 20 testing tasks. We find that these specific objectives exhibit clear trends across all benchmarks that the best setting showeases statistically significant performance improvements compared to the worst setting, thereby emphasizing the indispensability of considering the task repre- sentation shift issue. We also notice that the previous training framework, namely Ng = 1, = 1, fails to deliver the best asymptotic performance across all benchmarks. Though N; = 1, Nace = can be seen as an implicit way to rein in the task representation shift, the lack of explicit control still hinders the capability of these base algorithms to reach their full potential. Furthermore, as shown in Table we observe that settings with N₂ N achieve better performance more frequently than those with N₂ = 1, Nace > w.r.t the original setting N 1, Nace = 1. Based on recommend prioritizing the adjustment of N₂ when tuning parameters. This only tends to yield better performance but also offers the advantage of reducing training time. In contrast, increasing Nace would increase training time due the need for multiple context encoder updates within a single alternating step.
this observation,
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 and Table 3 demonstrate the performance of our different settings to rein the task representation shift Ant-Dir. We find that the contrastive-base algorithm fails the random dataset. Except for this case, the others show consistent results that reining in the task representation shift can 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 appealing direction and we leave this to future work.
8

---

Published as a conference paper at ICLR 2025
Expert
100
300
106
28
Figure 3: Testing returns of different settings to rein in the task representation shift different data qualities in 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.
DISCUSSION
This section seeks to raise some interesting issues derived from our work. We hope these issues can further enable pondering the significance of the task representation shift.
unoss-nntrup
recorotructio
origena)
Figure 4: Testing returns of the pre-training scheme against training from scratch in Ant-Dir. Solid curves refer the mean performance of trials over random seeds, and the shaded areas characterize the standard deviation of these trials.
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 be trained from scratch. Hence, a question is raised naturally. If 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 in 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 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 xy pretraan)) as the task representation distribution after pre-training. When meeting Assumption
4.1 the monotonic performance improvement condition for pre-training scheme is:
4Rmaz Em. [-<x;) pretrrin Z o o') 0
(12)
7)-
While maximizing I(Z; M)) can help Z( o (pretrain) to approach Z(-|x;; gmutual) the sub-optimal
gap between 7(-|x;; (matuad) and 7(-12; (²) still persists. Furthermore, given the inherent approxima-

---

Published as a conference paper at ICLR 2025
tion error between Z- lx; ( and ((1 @ the discrepancy between Z( d and (-112; Φ*) maybe large. 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 pre-training algorithm better approximate Z(-lx; Notice that there exists the decomposi- tion (- @Pretrasn) - 2 (-1=; Φ²)) < [2(]2z; ) Z [z; d ) + [- e ) ((|2;; æ I.)]. Since the cross erosspentropy-ioo entropy-bassee based objective the direct approximation of [Z; M), holding better performance under the pre- 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; (22) (- [æ; Φ1)1 in the alternating process. Nevertheless, how to achieve better performance under the pre-training scheme is also an interesting problem.
6.2 CAN THE VISUALIZATION OF TASK REPRESENTATION BE STRONGLY RELIED UPON TO IMPLY THE ASYMPTOTIC PERFORMANCE?
Previous COMRL endeavors Lii al. 2020
Yuan & Lu 2022 Gao et al. 2024) mostly apply I-SNE on the learned task representation to show the differentiation of the tasks. They hold the be- lief that better performance corresponds to better differentiation of the tasks. However, based on our analysis, the visualization result only de- pends on the task representation at convergence and ignores the task representation shift during the whole optimization process. According to the illustration shown in Figure(3 5 less desirable differentiation results also lead to better per- formance. Hence, the visualization results may represent the true task distribution but cannot sufficiently imply the final performance.
X00
incontract Rarel, patamaaren225
Figure 5: The 2D projection of the learned task representation space in Ant-Dir. Points are uni- formly sampled from the evaluation datasets. Tasks of given goals from 6 are mapped to rainbow colors, ranging from purple to red.
7 CONCLUSION & LIMITATION
Finding that the performance improvement guarantee 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 ignores the variation of the task representation, which may impair performance improvement. 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. 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, direction that merits further research is to design policy learning algorithm to achieve better policy improvements conditioning on the optimal task representation distribution. We leave these interesting problems future work.
ACKNOWLEDGMEENTS
The work is supported by the National Key Research and Development Program of China (No. (220YF00711022) by the Young Scientists Fund of National Natural Science Foundation of China (No. 62406295), and by the "Pioneer" and "Leading Goose" R&D Program of Zhejiang (2024SSYS0007).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Era Choshen and Aviv Tamar. Contrabar: Contrastive bayes-adaptive deep rl. In International Conference on Machine Learning, pp. 6005-6027. PMLR, 2023.
Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model-based reinforcement learning through optimistic policy search and planning. Advances in Neural Information Processing Systems, 33:14156-14170, 2020.
Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta reinforcement learning- identifiability challenges and effective data collection strategies. Advances in Neural Information Processing Systems, 14:4607-46188 2021.
Mehdi Fatemi, Mary Wu, Jeremy Petch, Walter Nelson, Stuart J Connolly, Alexander Benz, Anthony Carnicelli, and Marzyeh Ghassemi. Seri-markov offline reinforcement learning for healthcare. In Conference on Health, Inference. and Learning, pp. 119-137. PMLR, 2022.
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. Advances in neural information processing systems, 34:20132-20145, 2021.
Yunkai Gao, Rui Zhang, Jiaming Guo, Fan Wu, Qi Yi, Shaohui Peng. Siming Lan, Ruizhi Chen, Zidong Du, Xing Hu, al. Context shift reduction for offline meta-reinnforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Arob Ghosh, Xingyu Zhou, and Ness Shroff. Provably efficient model-free constrained rl with linear function approximation. Advances in Neural Information Processing Systems, 35:13303-13315
2022.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability, Advances in neural information processing systems, 14:25502-25555 2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy
maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning. pp. 1861-1870. PMLR, 2018.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Tianying Ji, Yu Luo, Fuchun Sun, Mingxuan Jing, Fengxiang He, and Wenbing Huang. When to update your model: Constrained model-basee based reinforcement learning. Advances in Neural Information Processing Systems, 35:23150-2311 2022.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning, pp. 267-274, 2002.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances Neural Information Processing Systems, 33:1 119911991 2020.
Aviral Kumar, Anikait Singh, Stephen Tian, Chelsea Finn, and Sergey Levine. A workflow for offline model-free robotic reinforcement learning. In Aleksandra Faust, David Hsu, and Ger- hard Neumann (eds.), Proceedings of the 5th Conference on Robot Learning, volume 164 of Proceedings of Machine Learning Research, pp. 417 428. PMLR, 08-11 Nov 2022. URL https::/proceeiingg .ml www.s/ppress.vv... 64
ppprees/vv.44/kuurr22eee. www.kkkaaa2aa. html
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01663 2020.
Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reiinforeement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.0//11 2020.
Lanqing Li, Yuanhao Huang, Mingzhe Chen, Siteng Luo, Dijun Luo, and Junzhou Huang. Provably improved context-based offline meta-rl with attention and contrastive learning. arXiv preprint arXiv:2100 10774, 2021.
11

---

Published as a conference paper at ICLR 2025
Lanqing Li, Hai Zhang, Xinyu Zhang, Shatong Zhu, Yang YU, Junqiao Zhao, and Pheng-Ame Heng. Towards an information theoretic framework context- based offline renfrreeiforre learning. In The Thirty-eighan Annual Conference Neural Information Processing Systems, 2024. URL
https: www.reeeeeeeeeeeeeeeeeeeee net. wwtttt/o/r/rrttteeeeee 1= 9mx
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint (rXiv:/807.038550 2018.
Haozhe Ma, Kuankuan Sima, Thanh Vinh Vo, Di Fu, and Tze- Yun Leong. Reward shaping for reinforcement learning with an assistant reward agent. In Forty-first International Conference Machine Learning, 2024a. URL htt.tttss///ppeeeeeiiieee hotp:////pppeeeeeeeeeeee/orrrrn 3=a
Haozhe Ma. Thanh Vinh Vo, and Tze-Yur Leong. Mixed-inttiative bayesian sub-goal optimization in hierarchical reinforcement learning. In Proceedings of the 23rd international conference on autonomous agents and multiagent systems, pp. 1328- 1336, 2024b.
Eric Mitchell, Rafael Rafailov, Xue Bin Peng. Sergey Levine, and Chelsea Finn. Offlin meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pp. 7780-7791. PMLR, 2021.
Ian Oshand and Benjamin Van Roy. Model-based reinforcement learning and the eluder dimension. Advances in Neural Information Processing Systems, 27, 2014.
Long Peng, Yang Cao, Yuejin Sun, and Yang Wang. Lightweight adaptive feature de-drifting for compressed image classification. IEEE Transactions on Multimedia, 2024.
Long Peng, Wenbo Li, Renjing Pei, Jingjing Ren, Jiaqi Xu, Yang Wang, Yang Cao, and Zheng- Jun Zha. Towards realistic data generation for real-world super-resolution. The Thirteenth International Conference on Learning Representations, 2025.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Julian Schritwieser, Ioannis Antonoglou. Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604 609, 2020.
Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with learned model. Advances in Neural Information Processing Systems, 14:27580-2759) 2021.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-18977 PMLR, 2015.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144 2018.
Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4r1: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning, pp. 907-917. PMLR, 2022.
Zhao Song and Wen Sun. Efficient model-free reinforcement learning in metric spaces, arXiv preprint arXiv:/905.00 190500045, 2019.
Shengpu Tang, Maggie Makar, Michael Sjoding, Finale Doshi- Velez, and Jenna Wiens. Leveraging factored action spaces for efficient offline reinforcement learning in healthcare. Advances in Neural Information Processing Systems, 35:34272-34286, 2022.
12

---

Published as a conference paper at ICLR 2025
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033 IEEE, 2012.
Tsachy Weissman, Erik Ordentlich, Gadiel Seroussi, Sergio Verdu, and Marcelo Weinberger. Inequalities for the deviation of the empirical distribution. Hewlett-Puckard Labs, Tech. Rep, pp. 125, 2003.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: 1911.:1991.111666 2019.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learning, pp. 24631-24643 PMLR, 2022.
Shentao Yang, Shujian Zhang, Yihao Feng, and Mingyuan Zhou. unified framework for alternating offline model training and policy learning. Advances in Neural Information Processing Systems, 15:17216-17233 2022.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100. PMLR, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 34:28954-28967 2021.
Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcemmnt learning via contrastive learning. In International Conference on Machine Learning, pp. 25747-25759.
PMLR, 2022.
Di Zhang, Bowen Lv. Hai Zhang, Feifan Yang, Junqiao Zhao, Hang Yu, Chang Huang, Hongtu Zhou, Chen Ye, and changjun jiang. Focus on what matters: Separated models for risual-based RL generalization. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL WtR.httss //tp:///eeee...vee.eeeeee net./ fo d=w::l
Hai Zhang, Hang Yu, Junqiao Zhao, Di Zhang, Chang Huang, Hongtu Zhou, Xiao Zhang, and Chen Ye. How fine-tune the model: Unified model shift and model bias policy optimization. In Teiitty-eeenthh Conference on Neural Information Processing Systems, 2023a. URL https
//openreview, net.
Xiao Zhang, Hai Zhang, Hongtu Zhou, Chang Huang, Di Zhang, Chen Ye, and Junqiao Zhao. Safe reinforcement learning with dead-ends avoidance and recovery. IEEE Robotics and Automation Letters, 9()):9114998 2023b.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes adaptive deep rl via meta- neta-learning. arXiv preprint arX1v:9910.08348, 2019.
13

---

Published as a conference paper at ICLR 2025
8 APPENDIX
8.1 USEFUL LEMMAS
Lemma 8.1 (Return bound. Zhang 2023a)). Let Rmaa denote the bound of the reward function. ex denote max, DTV (T (2)) and EM, denote ((s (Dr((PM,) llPM,)). For two arbitrary policies 1,, E II, the expected return under two arbitrary models M1, M2 E M can be bounded as,
Ma
V M2
VMl S 2Bmaz
(1 71-
Ex
+
(1 (1-7) 71- MM
7
(13)
Proof:
MM
Mu,
(Pt, (8, - Th
-Pe,bs, (s. ())
t=0
8,a
Rmaz
PiAss (8, Pi,,r, a)
(14)
1=
2Bman
72 Drv (8, a)ll, (s,a))
t=0
According to Theorem 2 (Return Bound) in (Zhang al., 2023a).
x
Σ
DDD((I (8, (1 (s, ( (1 7)²
ET
+
(1 - 7)2
(1- "
7 Ma2
(15)
We bring this result back and can draw the conclusion.
L
Lemma 8.2 (Inequility for L, deviation of the empirical distribution. (Weisssaa et al. 2003)). Let P be a probability distribution the set A (1,-..,a) a). For sequence of samples 21, Fm P. let P be the empirical probability distribution on A defined by P(j) m E I(z, = j). The Li-dviiation of the true distribution P and the empirical distribution P over A from m independent identically samples is bounded by.
Pr([P P 2 12 50 (2) (eexp(mme"
(16)
Proof:
According the inequality for Lı deviation of the empirical distribution (Weissman et al. 2003). they
conclude that Pr(PP- P1 2 < (2) 2) exp((mó /4), where Φ(p) 1-2p log 1-2 and FP Max min(P
- P(A)). The paper (Weissman al. 2003) point out that 0 < < 1/2, VP. Based on this, we can derive the function @(p) and find that @(p) is monotonically
decreasing on (0, 1/2). Therefore, we can get dpp) 2 ((1/2) 2. Then, we bring this result back
and can conclude that Pr(PP- Ph 2 e) < (2|4) 2) exp( mme2 m
8.2 MISSING PROOFS
Theorem 8.3 (Return bound in COMRL). Assume the reward function is upper-bounded by Rmaz- For an arbitrary policy parameter when meeting Assumption 4.1 the return bound in COMRL can be formalized as:
(J) J(0) < 2Rmaa Loz Bm
(1 7)2 Bm,x(1Z(+frr (1Z(+1z;)) - 2(:12;; muutual }1 + IZ( (2: ÿ mmutuall 2(-12r; 1b" (²)))
(17)
14

---

Published as a conference paper at ICLR 2025
Proof:
[J(0) J
= (a) 1o (01
RRm (8, []]] E (1 RRm (s, (0,,)]]
(18)
Em (s (dd te (11
Rmm (s, a)]
( (1ls, 2(1m0)))) Rm (s, a)ll
(19)
[Em, (9 wdde 2(
RRm (s, a)]
(s)a) (10b2( (m
Rm (s, (s,a)] E (1 (6 190
[Rim a)]l
(20)
& a (5 (0. 21110))))))
(s, a)] (( (( (1102 [m,
(s) rdd 1,1,,, 20
Bm a)] a) 2(11)) Bn (s, a]]l
(21)
+ V s,
V
(22)
Rmar
< Em,al 7)2
((,|x; mu :0) max D (TV Z(- 0)|= (-1,,
2Rmax max Drv
+ (1 (1-
(T( (T Z(- muutwal) ); 0) [T|| 2-| x));
(23)
2Rmar bz Em,.. (Z c; Φ) Z-|12; @ matual)) ) + | 2: muttual) Z(|la; (*)))
(24)
7)2 (1-7)²
Eq.
(3 is the result of directly applying Lemma 8.1 and Eq. 22 is the result of directly applying
Assumption 4.1
Corollary 8.4 (Monotonic performance improvement condition for previous COMRL works). When the condition for monotonic performance improvement of previous COMRL
meeting Assumption 4.7
works is:
4Rmaz Liz E Bm (1Z(=(z;) -(11x;
i (((2) ( 2 (1 7)
(25)
Proof:
(02) ((0)) ((02) J (02) + (²2))) (0)) (²,), ((2))
(26)
Denote
2Rmaa
as K. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get: (1-7)²
J(B2) - J" (02) + (B2) J (01) .J(1)
(27)
2 Em, 1 max
((12=; 0): (-12;; (")=;"");;;0))0) 02)) + E
(28)
max Dry ((1 Z(12;; (1):
(29)
2 1 Z(1 Z(+l=;)))] ci2l
(30)
To get the monotonic performance improvement, need Eq. 30 to be larger than 0. Hence, can get:
4Rmarbz Wm, [2(()2;)) Z(-1x; 6*)l 2
f
(31)
(1 71-7)²
Theorem 8.5 (Lower bound of performance difference COMRL), Assume the reward function is upper-bounded by Rmax- When meeting Assumption 4.1 the lower bound of the performance difference in COMRL can be formalized as:
2Rman Loz
(²((2) - (0,) 2 E (1 71 Em 1m.22222((42:00)) Z(-l=;0") (Z 1Z Z(-1a; +1)1
(32)
Proof: We can introduce the following decomposition (0) - (0)
(33)
Em, Ble.a) a. (21 RRm (8, a)] - Em, F ( ( Rmm (8, a))
(34)
= Em,rr (F (B(s,a)ddd (1120))))
Rmm (8, a)) - B( (11 R (m, (s,a)]
(35)
15

---

Published as a conference paper at ICLR 2025
Denote (1 2,,00000 (s, a)]] as G2 then can get the following derivation. Here, only consider the value within the brackets.
(s, a)] (a (Rmm (8,a))
(36)
22
(37)
G tz Gz - Gz + Gz - -% + G 5%
(38)
Denote
2llma (1-7)²
as k. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get:
G Gz + Gz 22 0% 22 + 2%
G
(39)
2 max D m 02) ( ;));0)) ci2
max (
( 02);
max Dr (1)))
(40)
2 -xL, Z- [z]
- -kLglZ(11x; [z; -1,
2(11=;11)
(41)
Simplifying Eq.
411 we can get:
2Rmax Lg
(2(02) - (0,)
12 (1 71-
Bm (2|2(122; Z(-12; 1Z(11;022) Z(11=;111
(42)
Theorem 8.6 (Monotonic performance improvement guarantee on training process). Denote K as (1-2)²²
ARmaz i12 3 and |7| as the cardinality of the task representation space. 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 probability greater than 1 - &, we can get the monotonic performance improvement guarantee by updating the context encoder via maximizing I IZ: M) from at least extra k samples, where:
2 2 + va)²
k
2log
(43)
R2
Here, E [0, 1] is a constant.
Proof:
Let Z1 have already been trained by maximizing (Z: to some extent. Given that get Zz, we need to train Z, by extra & samples from the training dataset by maximizing I(Z; M) to get the monotonic performance improvement guarantee.
We begin with the following inequality:
IZ(- [a; Φ2) Z(-la; o')l < [Z(-|x; 62) Z(-12: 02) + (2)))) Ō2)
Z(-12;;))
(44)
where (-:aa; 52) denotes the context encoder updated by fitting the empirical distribution on i.i.d
samples from (-| x²) based on ((11r; Φ1).
According to Assumption 4.9
we have:
12(-|;; - 2(-12; S
(45)
Now, we move our focus on solving [Z(-|x; )2) - Z(-1a;; Φ²){-
By applying Lemma 8.2 the I.1 deviation of the empirical distribution 2(-aa; 22) and true 7-1x;;*)
over [Z] is bounded by:
-
exp (-ke?
(46)
Pr((Z(((=; - ( <= 2 (2 2) exp (-ke? k /2)
(47)
16

---

k Published as a conference paper at ICLR 2025
Then for a fixed x, with probability greater than - we have:
2
-
12(-)a; (11la;; (()1);)))
log
(48)
V
k
5
Hence, we can get:
2
21 2
a
I((12;))) - Z(-1=;0")15 S
V
log
&
+
(49)
Let
(1-2)² 4ftmaa La E/
3, and recall that denotes
4flmaa Loz C12
[1-7)²
213, then we can get the & as:
212|
2
a
)
log
f1 2
(50)
6
k
4 PRmax 11z
1
2/
2 log
+
va)
(51)
&
Since we need the least number of samples that update the context encoder, we take k
き v log 21
E + Va)².
Corollary 8.7 (Monotonic performance improvement condition for pre-training scheme). Denote
2 20|r; ar; spretrain) the task representation distribution after pre-training. When meeting Assumption
4.1 the monotonic performance improvement condition for pre-training scheme is:
4Rmaz L2 Em (1Z(a/;;; spretrain)- - (-12;; 0 d')1]
f1 (1 7)2
(52)
Proof: Denote Jeretrain (01) as the expected return of the policy r--|s, Z (x; geretrain); 01) be- fore update of the policy. Similarly, denote Juretrain (82) as the expected return of the policy < Z-112; [x; garetrain); f2) after update of the policy.
jpretrain (02) j (01)
(53)
Jaretrain (02) J (02) (O2) (01) + J+) jeretrain (0:)
(54)
Denote 2lmaa (1-7)² as K. Taking Lemma 8.1 and Assumption 4.1 into the above formulation, we can get:
jpretrain (02) J (62) + (0,) - (0, + (0, - jeretrain (8,)
(55)
2 Em..l -K max D2 (T(-)s, Z(-1x;; ppretrain) 02) 02)11
+ +
(56)
max Drv (T ZZ- ( (pretrain): ) (001)1 [ [T]
0")=;;");;011
(57)
2 E Em, 2 [- opretrain) - (-|x;; x-) + <12l
(58)
To get the monotonic performance improvement, need Eq. 588 to be larger than 0. Hence, can get:
4R 4Rmaa Lu
c/2
Em Il spretrain) Z(-12;; d')l] |] 2
(59)
(1 - 7)3
8.3 JUSTIFICATION OF ASSUMPTIONS
Assumption 4.7
As our aim is to rein in the task representation shift, setting threshold to bound the task representation shift is natural. Nevertheless, how to set this threshold smartly or automatically adjust this threshold would need further research and we leave this to future work. The task representation shift less
17

---

Published as a conference paper at ICLR 2025
than the policy improvement <i2 with certain coefficient is to help us ensure in Theorem 4.10 is larger than O. Since we can update the policy consistently (e.g. & determines when to update the context encoder, but the policy is updated normally), ei2 can accumulate gradually. Therefore, this assumption is reasonable.
Assumption 4.8
Since the task representation is obtained through sampling, whether during the guidance of the downstream policy the training of the context encoder by maximizing [Z; M), it is reasonable to assume that the space of the task representation is discrete and limited. We think this assumption is general it can cover various cases, e.g. sampling, generated from the deterministic network, or discretization.
Assumption 4.9
Ant classifler
Ant Dr contrestive
Olt: reconstruction
320
Valliar-Prram dassiffirr
Walker-Paam contrastive
ValiianParamm. reconstruction
Figure 6: The numerical experiments Ant-Dir and Walker-Parru Paraam.
To justify this assumption in practice, we design a numerical experiment.
Firstly, notice that the sub-optimality gap between 7 [=|x;; (fmustua) and 7(-| o0") is constant. We can introduce notation Z--læ; )²) that denotes the context encoder updated by fitting the empirical distribution on i.i.d samples from Z- |x; gimutual) based on 2-|lx; Similar to Eq. (48), we can get the upper bound of the discrepancy between Z(-lx;; )2) and 7(-|x;; 62):
[7 (2)) 7(-1x;;0 x2)1
(60)
< (= 22) Z Z--10; | ( mutuaal) ) + [[-12x; muttual) (111x;0))
- 2(-12; 22)
(61)
constant
constant2
(- maluel)) - Z(-12; 0²) +
(62)
sub-optimality gap
constant3
constant4
(63)
Secondly, we use the pre-trained context encoder to be Z(-12; montual)) use the initialized context encoder to be (-|æ;; Φ1) and simulate two update processes: 1) use the specific objective w.r.t
maximizing 1 [Z; M) update Z(|12; Φ1) to get ((|xx; Φ2); 2) use task representations randomly
sampled from the pre-trained context encoder to update Z [z; Φ1) to get (1- 62). We set different
numbers of samples to complete these two update processes.
To express the discrepancy, we compute nn.MS1. loss of the task representation randomly sampled from Z(-|z; d2) and (-|2x; 62). As shown in Figure 6, it does fit an inversely proportional trend as the training size increases. Therefore, can practically denote the discrepancy between 7(-12; (2)
18

---

Published as a conference paper at ICLR 2025
and Z(-/2; @2) as:
constant5
(1(=))2)) - Z(112; 62)1
(64)
By combining the above formulations, we can get:
( r; 02) - Z(|1a;; Ø2)1
(65)
$ (=<a|=; )22) (-|12; 22)1 (1122;)2) 2(112; 62)
(66)
constant6
constant4
(67)
Since the training sample is finite, the discrepancy would not converge to 0. Hence, with an appropriate a, Assumption 4.9 can be grounded in practice.
Furthermore, even if we assume the discrepancy would converge to 0 and take the constant into Assumption 4.9 it only needs to modify the formulation of in Theorem 4.10
Proof: If we take the constant e into Assumption
4.9 then it becomes:
((11))22) - 2-1 Ø2)1 < Z(+ xz; 62) Z( 62) + (-| )2) - 2(11=; = V
a
+
(68)
Therefore, we need to modify Eq. 499 as:
2
2 -
a
[Z(12;;)2) Z(-12;)
V
log
V
+
(69)
Let 6 = (1-3)²
TRmus Le 2i 3, and update the notation of k. as TRmas bo 17
(1-g)²
3/3 - c, then we can get
the 1: as:
2
21 -
+ (1 7)
1
log
ci2
2
(70)
5
4Rmaz Lig
21 -
k
2log
+ va)²
(71)
R²
Since we need the least number of samples that update the context encoder, we take k a / log 2121 + va)².
To further enhance the theoretical rigor, we can even set the bound of Assumption 4.9 as a constant since the bound hetween 7(-Ja: Φ22) and 7-1|2; Φ22) will be reduced consistently with larger data sizes. Setting this constant only needs the accumulation cī2 to become larger, which also acceptable under our theoretical framework. For specific algorithms, e.g. contrastive, reconstruction.... if there exist theoretical guarantees for the convergence bound, it is nice to derive a more precise calculation method for k.
Proof: If we set the bound in Assumption 4.9 to constant ax then Eq.
40
becomes:
2 -
IZ(-|z 62) Z(-12; x") <
log g
+ Va
(72)
Let (1->)²
L.maa f2 3, and update the notation of Allmaa
(1~3)²
91 3
Va, then
get the l as:
2
2 -2
(1 71- 1
log
+ Va <
4Rman Loz Ej - -8
(73)
K
{
2
2
2 2
log
(74)
22
19

---

Published as a conference paper at ICLR 2025
21 6
Since we need the least number of samples that update the context encoder, we take 222 log
8.4 IMPLEMENTATION DETAILS
In this section, we report the details of all the objectives in our main paper.
Contrastive-based (Li et 2020). The contrastive-baser algorithm uses the task representation to compute distance metric learning loss. We adopt the open-source code of FOCAL
Revonstruction-basee (Li et 2024). The reconstruction- based algorithm passes the state, action, and task representation through the decoder to reconstruct the next state and the reward signal. We use the open-source code of UNICORN and adopt UNICORN-O as the reconstruction based objective training framework.
and
Cross-entropy-based. The cross-entropy- algorithm a straightforward discrete approximation to I(; .). 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.
According to Theorem 3.1 the previous methods are optimizing the approximate bounds of I(Z; M). To better approximate T[; M), choose face this term directly. Notice that [[;; M) (M) H ZM) For the first part, it is a constant w.r.t the variable Z. Hence, it can be ignored. For the second part, have:
H(M = EE Σ> p(m|2) log [(m|2))
(75)
The ideal condition is task representation Z can uniquely identify the task M. Hence, minimizing I (M IZ) is equivalent to maximizing Eg [log p(m[==], as we need p(mlzz to approach 1. As maxi- mizing log p(m|z) can be instantiated the cross-entropy-loss, we claim that cross-entropy-loss is direct approximation towards optimizing [(Z
Built upon the code base of FOCAL, cross-entropy- basedd only replaces the distance metric learning loss with cross-entropy loss.
To maintain fairness, we make sure all benchmark-irrelevant parameters are consistent, shown in Table [
Table I: Benchmark-irellvaan parameters setting in the training process.
training tasks
20
testing tasks
20
taok training batch size
16
rl batch size
256
context size
trajectory
actor-network size
[256, 256]
critic-network
[256.
task encoder network size
[64, 64]
learning
3e-4
For each environment, we make the benchmark-relevan parameters the same for each algorithm, as shown in Table 2
Table 3 reports the performance in our experimental setting. The results are averaged by 8 random seeds with each seed averaged by the last 5 evaluation performances. Furthermore, to demonstrate whether the performance improvement w.r.t the worst case statistically significant, we conduct a paired t-test for the experimental results. The statistically significant results are highlighted in red while others are colored in blue.
20

---

Published as a conference paper at ICLR 2025
Table 2: Benchmark-rellvant parameters setting in the training process.
Configurations
Reach
Ant-Dir
Button-Press
Dial-Tum
Walker-Param
Push
datasel size
3e5
9c4
3e5
3e5
4.5e5
3e5
task representation dimension
Table 3: The performance in experiments section. Each result is averaged by 8 random seeds.
Bawksark
Algerilans
2 L.N. -5
A- L.N. -2
N,11,% 1, - Eorggeal
-,
N.-B S.
50
952-55--5 26
200035 M0GA3
IT
Rach
Oeicerin
200022 E 22501 10 1685
27N7A1 E HISH OF HEEH
2500.00 14-TT
AN E 16146 6 BEERG
Macerraten:
Denoew
15
130 I 70007
11N TRETTIINN IF 1390
1951
150/222013 1p irss
1Y00A BETS (p
327.12
Serve Pirss
15
179 TET TTTTT
1:1:1 DIL
PIEE i 0200
(:35
Park
Unthare
THEEE
HNLL 10-M or 11:
11EM 10- 1000
13811
WakasPars
Orleans
W01.0
15.09 12 11SS
EISTEE 1530
4054 1937 1010
20130 SSON 1330
20L3 10,20
21EAN
04
(o
to
2014 11310
180 11100
205.25 1974 0033330
Maceeraatin
175.7 2038
25231 1671 (a (100))
21075 31.8 ( 000
1 6 50050
247 25.00
10.00
77.47
Ast Ne.Dionandei
Receerstion
$1.55 1TB ( (ret
1761 1010 1F 5/900
1010 i camm
5:71 ip boa0
136 1737
Chasestaay
o
3040 to
18372 io
115.99 S4.1
200 1930 m100)
1915 1080
I7S16'S NCU 2,02
198 b 1rmn
o mram
10777
39:31
ip
(y
Eape
Gemeiir
2
26.4 0850
W0L.1 23,00 (200)
260 à 35.00 DaED
B17 23.35 0.3000
1215
8.5
ADDITIONAL. RESULTS
8.5.1
PERFORMANCE CONTROL.
Classifier
Contrastive
Reconstruction
260
Performance Control
N₂ Mase
No M,
original
No = 2.
Ne Ma =2
Figure 7: The experimental results of using the evaluation performance to guide the learning of the context encoder on Ant-Dir, Each result of Performance Control is averaged by random seeds.
Here, we give a potential way to achieve better performance. shown in Theorem 4.10 & has a close connection with eig- When ei is sufficient, A would become less than the given batch size, then the context encoder can be updated. The straightforward way is to use the performance evaluated through online interaction with the environment as the guideline. Specifically, performance of the learned policy conditioning on the pre-trained task representation to approximate J*))), thereby providing way to approximate e12. To resolve performance fluctuations, we calculate ii through ei2 + 6ī + max (J* (02) - J" (0)), 0). According to Theorem 4.10 we simply set a hyper-parameter to calculate 1: as l ( +1c-9' During this experiment, we a 0.07. As for 3, we simply set the context encoder to be updated once when an update is required.
the evaluation
As shown in Figure 7 we observe that all three algorithms exhibit improved asymptotic performance on Ant-Dir. This suggests a potential way to reduce parameter sensitivity by leveraging performance to guide the training of the context encoder.
21

---

Published as a conference paper at ICLR 2025
Nevertheless, the current method is not practically feasible. Calculating online performance evalua- tions across 20 training tasks for each training step would require up to a month on an NVIDIA 3090 GPU. To speed up the process, we randomly select only 3 training tasks for calculating, reducing the training time to approximately 3-6 days, depending on the algorithm. Note that the performance on learning curves is still averaged by 20 testing tasks. Exploring methods guide the training of the context encoder through more efficient performance estimation techniques could be a promising direction for future work. We also envision our work being generalized into other areas Peng et al. (002)) Zhang et al. 2024 Peng et al. (2025) that use the (context) encoder.
8.5.2
EVALUATIONN OF CSRO
Since our goal is to validate that reining in the task representation shift also makes a difference for CSRO Gao et al. (2024), do not carefully tune its hyper-parameters. As shown in Figure 8 and Figure 9 the conclusions still hold in CSRO.
WalkerPararr
Balifure
Button Press
Pysh
crggral
-1
Figure 8: Testing returns of different settings to rein in the task representation shift of CSRO.
Ausdon
20
800
ongenl
6gg 2
Figure 9: Testing returns of different settings to rein in the task representation shift on different data qualities in Ant-Dir CSRO.
22


Reviews:
[{"review_id": "IaNQTVWl92", "reviewer": "Reviewer_FVd6", "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.", "questions": "Can you elaborate more on how different batch sizes induce different task level contexts?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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\".\n2. Theoretical justification for the claims are well presented and it outlined a complete algorithmic framework. \n3. It presents rigorous experimental validation using several environments from two different benchmarks with 3 types of objective functions and 3 types of data sets.", "weaknesses": "1. The task variation used in the experiments could be improved by including more distinct tasks. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9i7Yps0SCk", "reviewer": "Reviewer_SUPu", "summary": "This paper addresses the overlooked issue of task representation shift in offline meta-reinforcement learning (OMRL), which can prevent consistent performance improvement. \n\nThe 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.\n\nOverall, 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.", "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?\n2. 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.\n3. The motivation for using the cross-entropy-based algorithm is somewhat unclear. Could you explain why it replaces the distance metric learning loss?\n4. From the experimental results, the improvement seems marginal. Could you analyze the reasons behind this?\n5. 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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-structured and clear.\n2. The authors identified a unique challenge existing in offline meta-RL, task representation shift, which is highly novel.\n3. 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.\n4. The paper concludes with some interesting discussions, which have the potential to motivate future research.", "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.\n2. The experiments are limited and need improvement; it would be beneficial to verify the impact of task representation shift in more diverse testing scenarios.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yQWU89tZu0", "reviewer": "Reviewer_FSsk", "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.", "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. \n2. The authors introduce a large number of assumptions in Section 4. Has the reasonableness of these assumptions been justified? \n3. When choosing the Ant environment for experiments under Mujoco, was it randomly selected, or was it only the environment that showed the best results? \n4. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.Provides a new perspective for learning better task representations in COMRL.  \n2.The paper includes a detailed mathematical derivation process.  \n3.The authors provide their code implementation in the supplementary materials.", "weaknesses": "Experimental Section:\n1. 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. \n2. The authors only consider four environments, while in CSRO and UNICORN, each algorithm considers six environments under Mujoco.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yhMfAclpKf", "reviewer": "Reviewer_jUhN", "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.\n\nThe paper claims two contributions:\n\n(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.\n\n(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.", "questions": "Update: I am happy with most of the authors' answers and recommending accepting the paper. I raised the rating to a strong 6.\n\nOld questions:\n\n1) Statistical significance analysis should be run in all experiments or claims on experimental results changed significantly, 2) Please answer all technical questions and comments.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n\nFor 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.\n\nFor claim (2), mathematical proofs under the made assumptions appear correct.\n\nThe paper is overall well written and understandable.", "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.\n\nThe 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.\n\nvol(Z) needs to be explained.\n\nOther technical details, discussed below, need to be described in more detail.\n\nDETAILS:\n\nIn 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.\nStatistical significance testing is needed to draw conclusions about the experimental results. This applies to all the results.\n\nIn \"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.\n\nIn 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)?\n\nOn Line 281, the paper refers to Eq. (32) which has not been introduced yet.\n\nNotation: 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}}.\n\nThe 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.\n\n\nRELATED WORK:\n\nIn \"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.\n\n\nLANGUAGE/PRESENTATION:\n\nIn \"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.\n\nThe 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.\n\nOn 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?\n\n\"we need to find the positive C to improve the performance monotonically.\" ->\n\"we need to find a positive C to improve the performance monotonically.\"\n\n\"Center around this motivation\" -> \"Centering around this motivation\"\n\n\"an policy learning algorithm\" -> \"a policy learning algorithm\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)