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

Dear Reviewer jUhN,

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. We hope our explanation below would clarify your concerns:

$**Q1.Statistical significance analysis**$

Please refer to the general response Q1.

$**Q2.The details w.r.t the Assumptions**$

Please refer to the general response Q2.

$**Q3.Vol(Z) and coverage simplex**$

We thank the reviewer for this valuable feedback! “Coverage simplex”[1] represents a geometric structure that encompasses all possible task representation Z, each task representation can be seen as a point within this simplex.

We use the volume vol(Z) of this simplex to quantify the span of task representations. Here, we use vol(Z) in place of |A| as we are treating vol(Z) as an analog to the cardinality. According to our Assumption 4.8, the task representation space is discrete and limited, hence, using the cardinality of the task representation space to replace |A| is natural.

To avoid confusion, we use |Z| to replace the notation of vol(Z), and clarify the meaning of |Z| in our Theorem 4.10. This term is treated as a constant within our theoretical framework and does not affect the algorithm. Nevertheless, if you have suggestions for a better notation, we would greatly appreciate your feedback.

[1] An Information Geometry of Statistical Manifold Learning

$**Q4.Definition 4.5: a more detailed definition**$

“Among the expectation of tasks” denotes the outermost expectation over tasks, namely $\mathbb{E}_m$. To avoid potential confusion, we remove the expression w.r.t “Among the expectation of tasks” from our paper and use precise mathematical formulation instead. Additionally, we update the statement of Definition 4.5.

$**Q5.Eq. (32) has not been introduced yet**$

We thank the reviewer for pointing out this typo and we update our submission accordingly.

$**Q6.Notation: Using multiple characters to denote quantities**$

We thank the reviewer for this nice advice and we update our statement in Section 4.3 and the algorithmic box accordingly.

$**Q7.The notation is slightly confusing. Now$\theta$denotes policy parameters but it is good to add a symbol to denote the context encoder.**$

We thank the reviewer for this feedback. We update the notation of the task representation as $Z(\cdot|x;\phi)$, where $\phi$ is the parameter of the context encoder, and accordingly update the whole paper. We also update the notation in Algorithm 1 from $q_\phi$ to $Z_\phi$ to avoid confusion.

$**Q8.Related works**$

We add more details of the ContraBAR in our related works to highlight two points: 1) the theoretical insight of ContraBAR lies in the online setting while ours lies in the offline setting, and 2) ContraBAR optimizes a specific approximate bound of $I(Z;M)$, namely $I(z;s_{t+1},r_t|a)$, and set an assumption on this approximation while we face a large class of algorithms that optimize various approximate bounds of $I(Z;M)$.

To further highlight the core claim of our paper, we move the full related work into appendix and leave the COMRL and performance improvement guarantee parts.

$**Q9.Weaken the condition necessary for monotonic performance improvement...**$

We thank the reviewer for pointing out this misleading expression. Here, we use “weaken” to point out that the previous theoretical framework does not account for the variation of the task representation, making the condition for monotonic performance improvement insufficient.

In fact, $**the task representation varies**$. Hence, the previous framework can not truly capture the condition for monotonic performance improvement (weak).

To avoid this confusion, we rephrase our sentence as: “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.”

$**Q10.However, ours considers the variation of task representation ignored by the previous training framework,...**$

We rephrase this sentence: 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.

$**Q11.As shown in Corollary 4.4, the monotonic performance improvement can be guaranteed with only better approximation.**$

We thank the reviewer for your advice! We rephrase this sentence: $Z(\cdot|x;\phi)$ should be close to $Z(\cdot|x;\phi^*)$ such that the lower bound is small enough for finding a policy to achieve monotonic performance improvement.

$**Q12.Other Typos**$

We thank the reviewer for pointing out these typos, and we update our paper to “find a positive C”, “centering around”, “a policy learning algorithm”.

Dear Reviewer jUhN,

We would like to express our sincere gratitude to you for reviewing our paper and providing valuable feedback. We have tried our best to respond to your concerns/questions and hope your concerns can be potentially addressed.

Notably, given that we are approaching the deadline for the discussion phase, please feel free to reach out if you have any additional questions or require further clarification.

Best,

All authors

Dear Reviewer FSsk,

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:

$**Q1. The description of the task representation shift**$

We thank the reviewer for the feedback! There might be some misunderstandings regarding our method. Our method does not aim to learn a better task representation but focuses on how to adjust the learning process of the task representation based on previous COMRL methods.

According to our theory, previous works suggest that the condition for monotonic performance improvements is shown in Eq. (7). However, $**they ignore the fact that the task representation also varies during the optimization process**$. If we take this variation into consideration, the condition for monotonic performance improvement becomes Eq. (10). This theoretical insight highlights the importance of reining in the task representation variation, specifically $|Z_2-Z_1|$. Without such control, it is more likely the condition necessary for performance improvement would be violated.

For example, if we assume that $Z_1$ is trained from scratch and the update process from $Z_1$ to $Z_2$ brings $Z_2$ close to $Z^*$, then with the condition in Eq. (7), 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_2-Z_1|$ remains large.

Based on this, our algorithm framework points out that to get better performance improvement, we need to consider two aspects to rein in the task representation shift: 1) When to update the context encoder and 2) How many times to update the context encoder when it needs to be updated.

$**Q2.The experimental results are obtained within the improved FOCAL framework, lacking experiments on important baselines**$

In our experiment, $**we are not limited to be within the FOCAL**$, which is the upper bound of $I(Z;M)$.

We cover the reconstruction, which is the lower bound of $I(Z;M)$ and $**equivalent to the objective of CORRO in the offline setting**$ (please see theoretical details in UNICORN [1]). We also cover the cross-entropy, which is the direct approximation towards $I(Z;M)$ (please refer to Appendix 8.4). $**These two algorithmic backbones have nothing to do with FOCAL at the algorithmic implementation level**$.

To avoid potential confusion, we update the statement in Section 4.3 and the last paragraph of the Introduction to make it clearer.

We appreciate your suggestions to expand the baselines and we include CSRO as an addition since it linear interpolates between the upper-bound and the lower-bound. We believe it indeed helps us to make our results more convincing. Thank you again for your insights and encouragement to strengthen our work.

$**Q3.The authors only consider 4 environments**$

Please refer to the general response Q1.

$**Q4.Large number of assumptions**$

Please refer to the general response Q2.

$**Q5.When choosing the Ant environment for experiments under Mujoco, was it randomly selected?**$

Based on the performance reported in UNICORN [1], we observe that, among the five MuJoCo environments, $**only Ant-Dir and Walker-Param demonstrate a noticeable performance differentiation**$. Therefore, in our work, we randomly select Ant-Dir for MuJoCo benchmark. For the other three environments, we choose more complex settings from MetaWorld to better validate our theory. Additionally, $**we include Walker-Param as a newly added environment**$, where our experimental results also show considerable performance.

$**Q6. How are the hyper-parameters for different environments? Is it possible to analyze the results of different hyper-parameters**$

Please refer to the general response Q4.

[1] Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning

Dear Reviewer FSsk,

We would like to express our sincere gratitude to you for reviewing our paper and providing valuable feedback. We have tried our best to respond to your concerns/questions and hope your concerns can be potentially addressed.

Notably, given that we are approaching the deadline for the discussion phase, please feel free to reach out if you have any additional questions or require further clarification.

Best,

All authors

Dear Reviewer SUPu,

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:

$**Q1. The algorithm box about$k$and$acc$and the confusing setup in Section 4.3.**$

We sincerely thank the reviewer for pointing out the confusion in Section 4.3. The purpose of Section 4.3 is to explain the ways we use to rein in the task representation shift. According to our algorithmic box, the way to control the task representation shift can be seen as two aspects. 1) When to update the context encoder and 2) How many times to update the context encoder.

Thus, to cover these two aspects, we introduce two parameters $k$ and $acc$. Here, $k=n×bs$ denotes that the context encoder needs to be updated every $n$ updates of the policy (according to our theory, this means $\epsilon^*_{12}$ accumulates $n$ times to meet the condition for updating the context encoder) and $acc=n$ denotes that when the context encoder needs to be updated, it is updated $n$ times. Please note that this parameter $k$ is distinct from $k$ used in our theoretical framework.

To clarify, we update the statement in Section 4.3, the algorithmic box and update the notations in the experiment correspondingly.

$**Q2. The experiments are limited and need improvement**$

Please refer to our general response Q1.

$**Q3. The results largely depend on the settings of hyper-parameters and lacks of analysis of the experimental effects.**$

Please refer to our general response Q4.

$**Q4. The cross-entropy-based loss**$

We add a description w.r.t the cross-entropy-based objective in the Appendix. Please refer to Appendix 8.4 for more details.

$**Q5. The improvements seem marginal**$

According to our paired t-test results in Table 3, there always exist instances where the p-value is less than 0.05, indicating statistical significance. This suggests that even simple modifications can have a meaningful impact.

Also, as shown in Appendix 8.6, if we use the evaluation performance to guide the learning of the context encoder (where the calculation of $k$ is determined by our theoretical framework), the performance improvement can get further enhancement. Hence, we also acknowledge the need for smarter algorithms to achieve stronger performance improvements (which has been stated in our limitation).

Additionally, the final performance may also be restricted by the offline dataset. For example, on the Ant-Dir dataset, the average return is only 21, which may limit the extent of achievable improvements.

$**Q6. Concerns in Section 6.2**$

This is a good question. We would like to state that while visualization results produced by t-SNE can provide valuable insights into the quality of the learned task representation, they may not fully capture the performance of the downstream policy.

Our theoretical analysis suggests that achieving better performance improvement involves not only optimizing Z towards the desired target $**but also effectively reining in the task representation shift**$. If the task representation shift is not properly adjusted, it may impede the attainment of the performance. We believe this highlights a crucial aspect of our work: the importance of both optimizing the task representation and reining in the task representation shift.

Dear Reviewer SUPu,

We thank the reviewer for the feedback! Our response to your concerns/questions:

$**Q1. Determining some key parameters through heuristics seems simple and lead to only marginal performance improvements.**$

We agree with the reviewer that we adopt a simple strategy to show the potential of reining in task representation shift (which we have stated in our limitation) since $**we do not claim to solve this issue by providing a strong algorithm, but stand for proposing this issue**$.

We kindly remind the reviewer to focus on the pre-training scheme, $**which is the case where task representation shift is completely ignored**$. When comparing the pre-training with the best case of reining in the task representation shift, the performance improvement is no longer marginal, e.g. the mean of pre-training cross-entropy is near 200 while the mean of the best condition in our heuristic setting is near 300 (50% improvement); the mean of pre-training contrastive is near 175 while the mean of the best is near 265 (50% improvement),...,

The reason for less performance gains against the original is that $**$N_{k**=1, N_{acc}=1 can indeed be seen as a way to rein in the task representation shift.}

Nevertheless, we also invite the reviewer to see Table 3 in Appendix 8.4, the results of the paired t-test can also support that the performance improvement against the original setting has statistical significance, thereby $**achieving our goal to be a starting point for future research**$.

And, as we use the offline dataset provided by UNICORN, according to the performance reported in UNICORN, the SOTA algorithms $**after carefully sweeping the parameters**$ namely UNICORN and CSRO, achieve the best performance 276 on Ant-Dir, 407 on Walker-Param, 2774 on Reach, and the original FOCAL only achieves 217 on Ant-Dir, 308 on Walker-Param, 2423 on Reach. By using the simple heuristic strategy, $**some previously weak baselines can beat the SOTA**$, e.g. cross-entropy can achieve 291 on Ant-Dir, FOCAL can achieve 450 on Walker-Param, FOCAL can achieve 2802 on Reach. $**From the view of simplifying the algorithm and reducing the effort for parameter-tuning, the task representation shift is also valuable.**$

Since our main focus is to propose the task representation shift issue, our contributions in theory are $**three folds**$. Note that $**the baseline works like FOCAL, CORRO, and CSRO do not hold such strong theoretical contributions**$:

(1) Provide a performance improvement guarantee for previous COMRL methods, $**which is the first in the offline setting**$.

(2) Consider the variation of task representation explicitly and refine the condition for monotonic performance improvement guarantee.

(3) Give theoretical proof of how we can achieve monotonic performance improvement guarantee.

Additionally, the experiments conducted in previous works have three main components: baseline comparison, ablation and visualization.

In our work, we $**integrate the baseline comparison and ablations**$ (if we view the heuristic strategy as the hyper-parameter) into Section 5.1, since our framework builds upon these previous algorithms. To extend the applicability of our theory, we also add an experiment on different-quality datasets, which is shown in Section 5.2. We also add discussions $**concerning visualization of the task representation**$ in our Discussion.

Hence, compared to other baseline works, our experimental design is reasonable and can validate our claims.

$**Q2. The concerns for question 6**$

We can use another example to explain this further. The policy has converged at 300K training steps. Hence, we use the example that pre-training cross-entropy at 300K v.s. the reconstruction for $N_{k}=3,N_{acc}=1$ at 300K. Though the cross-entropy demonstrates better visualization results, the final performance for pre-training cross-entropy is less than the reconstruction for $N_{k}=3,N_{acc}=1$, which reins in the task representation shift. Hence, $**ignoring task representation shift would cause an effect that “less desirable differentiation results lead to better performance”**$. Based on this effect, we aim to claim that using the visualization results to $**imply the final performance**$ is unreliable. We agree with the reviewer that the reason behind this (even the performance converges) may be the sub-optimal learned policy as ignoring the task representation shift would impede policy learning.

If we go further, and both algorithms being compared take task representation shift into account, we also think it is not certain that better differentiation would lead to better performance. We encourage the reviewer to see the visualization results of UNICORN[1] in its Appendix C.3. Although FOCAL demonstrates better visualization results, its final performance still falls behind UNICORN. We speculate that excessive differentiation may lead to a failure to capture the similarities between tasks. Nevertheless, for this point, we are open to further discussions.

In general, the visualization results $**may**$ represent the true task distribution, however, it cannot sufficiently imply the final performance.

To avoid confusion, we update our statement in Section 6.2 "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." to "Hence, the visualization results may represent the true task distribution but cannot sufficiently imply the final performance." and also update the example in Figure 5.

[1] Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning.

We hope our response could potentially address your concerns/questions. If you need any further elaboration, please feel free to reach out!

Dear Reviewer FVd6,

We sincerely thank the reviewer for the constructive comments and the appreciation of our work. Our response to your concerns/questions:

$**Q1. The task variation used in the experiments could be improved by including more distinct tasks. / Lack of results with more statistically significant metrics**$

Please refer to the general response Q1.

$**Q2. Can you elaborate more on how different batch sizes induce different task level contexts?**$

Thank you for the question. We’re not entirely clear on what you mean by “different task-level context”. Nevertheless, regarding the different batch sizes, we do not need to tune the training batch size of the context encoder. In our approach, we follow the settings in UNICORN [1], where we randomly select a trajectory and use all transitions from that trajectory as our batch to train the context encoder. This keeps our batch size fixed for each update.

According to our proposed theory, consistently updating the policy allows for increasing $\epsilon^*_{12}$, which in turn decreases $k$. Therefore, we only update the context encoder when $k$ is less than our given batch size, eliminating the need to adjust the context encoder's training batch size. Notably, if the training batch size is too small, it may insufficiently guide the downstream policy learning, but our update mechanism can effectively avoid this issue. We hope this can potentially address your question. If that’s not the case, could you kindly clarify so we can provide more detailed information?

[1] Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning