Paper: Task Adaptation from Skills: Information Geometry, Disentanglement, and New Objectives for Unsupervised Reinforcement Learning (ICLR-2024)

Reviewer comment (Presentation):
All experiments are provided in the Appendix. It seems that providing the main experiments and analyses in the main text can help the readers to better understand this work.

Author response:
## Q2&Q4: About experimental results in the main text and proof sketch

Thank you for your valuable suggestion on enhancing our paper's presentation. We will certainly take your feedback into account.

Reviewer comment (Theory):
What will happen if we change the W-distance to other distance metrics in PWSEP(i)? In my opinion, it seems that the proof of Thm 3.5 holds for any distance metric (owns Non-negativity, Symmetry, and Triangular Inequalities). Is it right? Or are there some special properties of W-distance necessary for proving Thm 3.5?

Author response:
## Q3: What will happen if we change the W-distance to other distance metrics in PWSEP(i)?

Thm 3.5 holds for any distance metric that owns "Non-negativity, Symmetry, and Triangular Inequalities" and strict convexity. 

We briefly discussed this topic in appendix F.4 after the proof of theorem 3.5. Although total variation distance and Hellinger distance are also true distance metrics, they're less used in RL compared to KL divergence and Wasserstein distance, resulting in limited research on their efficient approximations for state distributions. In addition, we haven't found valid literature to support the strict convexity of total variation distance and Heilinger distance. Additionally, Wasserstein distance has the potential to take advantage of the choice of transportation cost to provide smooth and stable measurements with meaningful information as discussed in appendix G.4

Reviewer comment (Theory):
It's better to provide some proof sketch of theorems in the main text, like Thm 3.5.

Author response:
## Q2&Q4: About experimental results in the main text and proof sketch

Thank you for your valuable suggestion on enhancing our paper's presentation. We will certainly take your feedback into account.

Reviewer comment (Writing):
Some definitions of words are not defined well (diversity and separability).

Author response:
First of all, LSEPIN is defined as $\min_z I(S;\mathbf{1}_{z})$, its mean is not $I(S,Z)$ (details in "formal difference" of appendix B.2).

As we have mentioned at the beginning of Section 3:
*"Separability means the discriminability between states inferred by different skills."* 
In inequality (61) of appendix E, we show that increasing $I(S;\mathbf{1}_{z})$ promotes the KL divergence between $p(S|z)$ and $p(S|Z\neq z)$(average state distribution of non-$z$ skills). Also discussed in appendix E, higher KL divergence between $p(S|z)$ and $p(S|Z\neq z)$ means skill $z$ is more distinctive and has less overlap with other skills, so this means better separability.

For a set of skills, they are diverse because all skills in the set are distinctive and have little overlap with each other (they are "far" from each other in terms of KL divergence). The standard MISL objective $I(S;Z)$ can not guarantee the seperability of each skill, as discussed in Section 3 after the list of informal results, so MISL without LSEPIN can not guarantee diversity.

Reviewer comment (Novelty):
The findings from the theoretical derivation are not surprising. WAC and adaptation cost.

Author response:
Following the above clarification of the WAC definition, we hope it is clear now that our analysis is dealing with the fundamental question for URL: "How the learned skills can be used for downstream task adaptation, and what properties of the learned skills are desired for better downstream task adaption?"

> The findings ... not surprising.

The findings may appear "not surprising", given that promoting diversity and separability of learned skills has been an intuitive heuristic in prior practical algorithms [2][3]. However, our unique contribution lies in offering a theoretical justification for this heuristic.

Reviewer comment (Theory):
We might incorrectly understand the contribution of this work. Why is the theoretical derivation of Theorem 3.1 important for a task adaptation?

Author response:
Theorem 3.1 and the theoretical results in Section 3.2 shed light on what kind of skills are favored for downstream task adaptation and how to quantitatively measure them.

Reviewer comment (Theory):
What can we additionally learn from Theorem 3.1 or how can we use the theorem for other task adaptation works? For my understanding, the relationship is: increasing LSEPIN results in lower WAC.

Author response:
Theorem 3.1 and the theoretical results in Section 3.2 shed light on what kind of skills are favored for downstream task adaptation and how to quantitatively measure them.

Reviewer comment (Evaluation):
In LSEPIN, the metric measures the least skill code. I guess an abundant skill code may hurt the calculation. Isn't it problematic? How could you ensure that all the codes are meaningful in the computation of LSEPIN?

Author response:
It is important to ensure every learned skill is distinctive and separable from others. High $I(S;\mathbf{1}_{z})$ means that this skill covers a specific region of the state space that is less covered by other skills. If it is not separable from other skills, its state coverage could have a huge overlap with other states, so deleting this skill would make no difference for exploration or downstream task adaptation.


In [1], they implement many existing algorithms with a low number of skills, eg. only 4 skills for SMM. If one of its skills exhibits a low $I(S;\mathbf{1}_{z})$, it may have a huge overlap with other skills, leading to a situation where 25% of the skills are underutilized, offering no meaningful contribution to exploration and downstream task adaptation.

Reviewer comment (Evaluation):
WAC is measured with the worst state distribution. How could this be meaningful and practical? That is, why do we need to measure the worst case adaptation? To the best of my knowledge, skill adaptation is applied to the state distributions which are similar to the state distribution for the target skill. Therefore, the importance of $\max_p$ part in WAC is not persuasive.

Author response:
WAC is indeed considering adaptation from **closest** skills in the set of learned skills because of $\min_{z\in \mathcal{Z}^*}D_{KL}({p(S|z)}\|{p^r})$. The $\max_r$ in WAC means to choose a downstream task $r$ such that its optimal state distribution $p^r$ is the far from all learned skills, even for the skill $z^*=argmin_{z\in \mathcal{Z}^*}D_{KL}({p(S|z)}\|{p^r})$ that is **closest** to $p^r$. 

In practice, although we can adapt from the closest skill after knowing the downstream task $r$, we do not have prior knowledge about $r$ during unsupervised pre-training. Therefore, we can only prepare for the worst downstream task (the one far from even the closest learned skill).

Reviewer comment (Theory):
Could you please additionally describe the necessity of symmetry and triangle inequality?

Author response:
This topic is discussed in Appendix B.3. We showed an example of what could happen if we replace the Wasserstein distance with KL divergence in appendix F.7. The example in appendix F.7 shows that if we replace the Wasserstein distance in WSEP with KL divergence, due to the lack of triangle inequality, maximizing this metric could result in skills too close together, harming diversity.

Symmetry is important when comparing distances between different pairs of points, for example, we have pair $a,b$ and pair $b,c$ and a non-symmetric measure $d$. Without symmetry, $d(a,b)\neq d(b,a)$ and $d(c,d)\neq d(d,c)$. Then it's possible that there exists a situation where $d(a,b)>d(c,d)$ but $d(b,a)<d(d,c)$. From the $d$ measurement you cannot know whether pair $a,b$ is more separable than pair $c,d$.

Besides, state distributions of two different skills could share different domains, which is a problem for a well-defined (not infinity) KL divergence.

Reviewer comment (Evaluation):
What is the limitation of WSEP?

Author response:
- Theoretically, as discussed in Section 3.3.4 and appendix F.6, although WSEP also discovers vertices, it is not guaranteed to discover all $|\mathcal{V}|$ vertices with only $|\mathcal{V}|$ skills. This motivates us to propose the PWSEP algorithm to solve the vertex discovery problem using only $|\mathcal{V}|$ skills.

- Practically, the implementation of Wasserstein distance depends on the choice of transportation cost. For continuous control, it is common that the transportation cost is chosen as the L2 norm. However, the L2 norm between two states might not accurately measure how difficult it is to travel from one state to the other, because there might be obstacles between them. 

As discussed in appendix G.4, an idea for future work is to learn state representations so that the L2 norm in the representation space can reflect the actual difficulty of traveling from one state to the other.

Reviewer comment (Evaluation):
Although the authors conducted several theoretical derivations and analyses, the choice of metrics and the relationship between them have little meaning. Mostly because the task adaptation might assume close state distribution $p$ for a skill $p(s|z)$. However, the authors study the worst-case adaptation. Additionally, the organization of the paper is not well constructed and hard to follow.

Author response:
WAC is indeed considering adaptation from **closest** skills in the set of learned skills because of $\min_{z\in \mathcal{Z}^*}D_{KL}({p(S|z)}\|{p^r})$. The $\max_r$ in WAC means to choose a downstream task $r$ such that its optimal state distribution $p^r$ is the far from all learned skills, even for the skill $z^*=argmin_{z\in \mathcal{Z}^*}D_{KL}({p(S|z)}\|{p^r})$ that is **closest** to $p^r$. 

In practice, although we can adapt from the closest skill after knowing the downstream task $r$, we do not have prior knowledge about $r$ during unsupervised pre-training. Therefore, we can only prepare for the worst downstream task (the one far from even the closest learned skill). 

Theorem 3.1 and the theoretical results in Section 3.2 shed light on what kind of skills are favored for downstream task adaptation and how to quantitatively measure them. 

The logic flow of our paper is like this:
1. We try to answer the fundamental question of URL: "How the learned skills can be used for downstream task adaptation and what properties of the learned skills are desired for better downstream task adaption?”
2. We found that LSEPIN captures the desired properties necessary for preparing skills for downstream task adaptation.
3. We found that LSEPIN and MISL are essentially optimizing KL divergences and MISL with LSEPIN also has limitations like the one mentioned in remark 3.2.1, this inspires us to investigate whether we can overcome the limitations of MISL by optimizing true metrics between state distributions.
4. We found that WSEP and PWSEP with true metric Wasserstein distance can overcome the limitation and can discover potentially optimal skills that can not be discovered by MISL.

Thank you for pointing out the readability and organization issues that can be further enhanced. We could replace some sentences with tables or diagrams.

The findings may appear "not surprising", given that promoting diversity and separability of learned skills has been an intuitive heuristic in prior practical algorithms [2][3]. However, our unique contribution lies in offering a theoretical justification for this heuristic.

Reviewer comment (Theory):
Is it possible to rigorously prove that adding a loss that promotes uniform p(Z_input) to objective (1) does not promote any diversity? Or promotes less diversity than the LSEPIN loss in all cases?

Author response:
**A1:** It might be hard to derive a quantitative bound to show adding a loss like $H(Z_{input})$ to objective (1) promotes less than LSEPIN in all cases, because the diversity of state distributions $p(S|z)$ depend on not only $z_{input}$ but also parameter $\theta$ and the function class for policy $\pi_\theta$. However, we can illustrate why such loss is not guaranteed to promote diversity like LSEPIN by the following example where maximizing $I(S; Z)$ with uniform $p(Z_{input})$ results in limited diversity:

For a case of with $|S|=3$, the feasible polytope $\mathcal{C}$ allows a maximum "circle" centered at $[0.2,0.4,0.4]$ with a maximum "radius" of $0.253$ as this fig in [the anonymous link](https://anonymous.4open.science/r/iclrreb2024-62C5/) shows. There are 5 vertices $v_{1:5}$ of $\mathcal{C}$, of which $v_1, v_2, v_4, v_5$ lie on the maximum "circle". Vertex $v_1$ and vertex $v_4$ are $[0.2,0.7,0.1]$ and $[0.2,0.1,0.7]$.

Because $\theta$ is shared by all skills and $z=\\{\theta,Z_{input}\\}$, we can assume $p(Z)=p(Z_{input})$

When there are 4 skills to learn, uniform $p(Z_{input})$ should have $p(z_{input})=p(z)=0.25$ for all $z$. In order to achieve both uniform $p(Z)$ and maximization of $I(S; Z)$, the optimal skill set $\\{z_1,z_2,z_3,z_4\\}$ should be the one containing two skills $z_1,z_2$ at $v_1$ and the other two skills $z_3,z_4$ at $v_4$, as shown in the [figure](https://anonymous.4open.science/r/iclrreb2024-62C5/reb_fig.png). Only with this skill set, the uniform average of skills $p(S)=\sum_z p(z)p(S|z)$ with $p(z)=0.25$ could be the center of the maximum "circle" $[0.2,0.4,0.4]$.

We can see from this example that although skill set $\\{z_1,z_2,z_3,z_4\\}$ maximizes both $I(S; Z)$ and $H(Z_{input})$, there are two pairs of skills being not separable thus resulting in limited diversity.  


In appendix E, we have shown higher $I(S;1_z)$ explicitly increases $D_{KL}(p(S|Z\neq z) \| p(S|z))$ thus promoting $p(S|z)$ to be separable from $p(S|Z\neq z)$ (average state distribution of skills other than $z$). Therefore, LSEPIN promotes diverse skills explicitly. Compared to skills only at $v_1$ and $v_4$, MISL with LSEPIN would favor skills at $v_1, v_2, v_4, v_5$ respectively, and for $p(S)$ to be the center of maximum "circle" $[0.2,0.4,0.4]$, distribution $p(Z)$ is not necessarily unform.

Reviewer comment (Theory):
Is there any advantage to the KL divergence over the Wasserstein distance? theoretically or computationally?

Author response:
**A3:** Computationally, for non-parametric estimation with collected data, KL divergence may be preferred over Wasserstein distance. This is attributed to the fact that KL divergence can be directly estimated using samples, whereas the estimation of Wasserstein distance requires solving a linear program, introducing additional computational complexity. Theoretically, KL divergence is always strictly convex, while the strict convexity of Wasserstein distance depends on the choice of transportation cost as mentioned in our proofs in appendix F.

Reviewer comment (Theory):
What are the limitations of the PWSEP algorithms (specifically the WSEP objective)?

Author response:
**A4:** As mentioned in Section 3.3.4, Although lemma 3.3 shows that maximizing WSEP also discovers vertices, the discovered vertices could be duplicated (shown in appendix F.6), so it might not be able to discover all $|V|$ vertices with only $|V|$ skills. This motivates us to propose PWSEP algorithm that iteratively optimizes a projected Wasserstein distance to make sure every new iteration learns a new skill.

Another limitation of WSEP is discussed in appendix F.5, showing that although a higher WSEP lowers an upper bound of the adaptation cost as our theoretical analysis shows, because of the gap between the upper bound and actual adaptation cost, sometimes high WSEP could not result in lower adaptation cost. 

As for PWSEP, despite its favorable theoretical property of discovering all vertices, in practice, each iteration could only learn a locally optimal skill, as mentioned in the empirical results of appendix H.

Reviewer comment (Presentation):
The empirical results in Appendix H should ideally be a part of the main paper, since including them makes the submission more well-rounded and gives empirical validation for the results presented in Section 3.

Author response:
We sincerely appreciate your time and attention and thank you for your valuable feedback! We will consider your suggestion for the presentation.