Task Adaptation from Skills: Information Geometry, Disentanglement, and New Objectives for Unsupervised Reinforcement Learning

Thanks for the fast reply!

We hope the figure provides the intuition that WSEP can learn the vertices that are not with maximum "distances".

Here is a simple MDP that produces the feasible state distribution polytope in the figure, and the quantitative results below shows why in this case, WSEP learns more vertices than MISL.

An MDP example that produces the polytope in the figure

An MDP with $|\mathcal{S}|=3$ for states and $|\mathcal{A}|=3$ for actions. The transition matrix for action $a_1$ is:

$\begin{pmatrix} 0.2 & 0.7 & 0.1\\\\ 0.2 & 0.7 & 0.1\\\\ 0.2 & 0.7 & 0.1 \end{pmatrix}$

The transition matrix for action $a_2$ is:

$\begin{pmatrix} 0.2& 0.1& 0.7\\\\ 0.2& 0.1& 0.7 \\\\ 0.2& 0.1& 0.7 \end{pmatrix}$

The transition matrix for action $a_3$ is:

$\begin{pmatrix} 0.4 & 0.3 & 0.3\\\\ 0.4 & 0.3 & 0.3 \\\\ 0.4 & 0.3 & 0.3 \end{pmatrix}$

Because the transition probabilities only depend on actions, state distribution $p(S)$ is determined by distribution $p(A)$: $p(S) = p(a_1)[0.2,0.7,0.1]+p(a_2)[0.2,0.1,0.7]+p(a_3)[0.4,0.3,0.3]$

We can see $p(A)$, in this case, is the convex coefficient of three probabilities, so the feasible state distribution is in the convex polytope $\mathcal{C}$ with these three probabilities being the vertices of $\mathcal{C}$.

This MDP accords with the figure, where we have labeled the vertices with $z_1,z_2,z_3$,

$p(S|z_1)= [0.2, 0.7, 0.1]$

$p(S|z_2)= [0.2, 0.1, 0.7]$

$p(S|z_1)= [0.4, 0.3, 0.3]$

Quantitative comparision of I(S;Z) and WSEP

MISL only learns skills at $z_1,z_2$:

For this MDP, the unique center of the "circle" with "maximum radius" (We showed that it's uniquely determined by the MDP in appendix D) would be $[0.2,0.4,0.4]$.

MISL would learn only $z_1,z_2$ with $p(z_1)=p(z_2)=0.5,p(z_3)=0$ to have the average state distribution $p(S)=\sum_z(S|z)$ at the unique center $[0.2,0.4,0.4]$.

By this solution, $I(S;Z)$ is maximized by $I(S;Z)=\sum_z p(z)D_{KL}(p(S|z)\|p(S))\approx 0.253$.

Putting weight on the other vertex $z_3$ would lower $I(S; Z)$. For example, if $p(z_1)=p(z_2)=0.45,p(z_3)=0.1$, the average $p(S)$ would be $[0.22, 0.39, 0.39]$ which is not the center of the maximum "circle" any more, and their $I(S; Z)\approx 0.237$.

WSEP can learns skills at all vertices $z_1,z_2,z_3$:

For WSEP, when the skill number is set to $|Z|=3$, when two skills are at $z_1,z_2$ like MISL, putting the last skill at $z_1$ or $z_2$ would have lower WSEP than putting the last skill at $z_3$, due to the triangle inequality.

Suppose the transportation cost between every two states is 1, then

$W(p(S|z_1),p(S|z_2))=0.6$

$W(p(S|z_1),p(S|z_3))=W(p(S|z_2),p(S|z_3))=0.4$

WSEP for $2$ skills at $z_1$ and $1$ skill at $z_2$ would be $2.4$,

WSEP for $2$ skills at $z_2$ and $1$ skill at $z_1$ would also be $2.4$.

WSEP for $3$ skills at $z_1,z_2,z_3$ respectively would be $2.8$, and this is the solution for maximizing WSEP for this case (By convexity, moving any of the three vertices would lower its distance to the others).

Therefore, WSEP would favor learning all $3$ vertices instead of only $2$ skills at $z_1,z_2$ like MISL. In this case, MISL can not discover $z_3$ that is not with maximum "distances" but WSEP can.

Thanks again and we hope this addressed the question!

We sincerely thank you for your thoughtful review and recognition of our work. We are delighted to address your questions. Please see our response below:

Q1: About whether "WSEP can find more skills or even all skills"

A1: We consider that, by 'skills' you are referring to those at the vertices.

By "it is possible for WSEP to discover all vertices of the feasible polytope" we mean that WSEP is not restricted by maximum "distances", so it might discover the vertices that can not be discovered by MISL in certain cases. For example, in the situation shown by the figure from this anonymous link, MISL can only learn $z_1,z_2$. Even when the skill number is set to be $|Z|=3$, the third skill learned by MISL will be duplicated at $z_1$ or $z_2$ but not lie at $z_3$. Instead, WSEP can discover all 3 vertices in this case.

However, as mentioned in Section 3.3.4 and appendix F.6, unlike PWSEP, optimizing WSEP is not guaranteed to always discover all vertices. Among the vertices that can not be discovered by WSEP, there can also be ones with maximum "distance" that can be discovered by MISL. This motivates us to look into PWSEP, which is guaranteed to always discover all vertices with enough skills.

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

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.

Thanks again for your time and attention!

We thank the reviewer for your valuable feedback and your willingness to engage in discussion. We would like to address your concerns and answer your questions. Please see the following for our response.

1. Regarding the concerns related to the "WAC" cost and Theorem 3.1

About:

• [Overall] Although the authors conducted several .. the authors study the worst-case adaptation.
• WAC is measured with the worst ... the importance of $max_p$ part in WAC is not persuasive.
• What can we .. For my understanding, the relationship: increasing LSEPIN results in lower WAC.
• The findings from the theoretical derivation are not surprising. WAC and adaptation cost.*

1.1 Clarification of the definition of WAC

To address these concerns, we would like to first clarify the definition of the WAC cost and mitigate potential misunderstandings. The WAC is defined as: $\text{WAC} = \max_r \min_{z\in \mathcal{Z}^*}D_{KL}({p(S|z)}\|{p^r})$ $p^r$ is the optimal state distribution for downstream task $r$, and $\mathcal{Z}^*$ is the set of learned skills.

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$.

1.2 Practicality and meaningness

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.

1.3 About the contribution of our findings

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.

2. Regarding concerns related to "diversity and separability"

About:

• Why measuring $min_z I(S;z)$ can be used to measure diversity and separability? ..., but the $min_z I(S;z)$ is just a mutual information for a single code.
• Some Definitions of words are not defined well (diversity and separability).

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.

3. Regarding why focus on single skill code

In LSEPIN, the metric measures the least skill code. ... How could you ensure that all the codes are meaningful in the computation of LSEPIN?

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.

Thanks again for your valuable feedback on our paper. Our rebuttal addresses the concerns, especially the ones related to the WAC cost definition. By WAC, we actually consider the practical adaptation procedure you mentioned, which is to adapt from the 'closest' skill. We hope this clarifies our main contribution.

We wonder whether you have any remaining concerns, we are looking forward to addressing any additional concerns you may have.

Dear Authors,

Thank you for the kind response.

Although the paper still exhibits weaknesses in its experimental support and the organization of the flow, the responses provided have clarified the major concerns. Therefore, I would raise my score to 'weak accept (6).

Sincerely,

Reviewer j7aV

We appreciate your time and attention. We thank you for the review and comments and please see our response to your questions below.

Q1: 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?

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 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. 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.

Q2: Is there a measure for adaptation other than Worst-case Adaptation Cost?

A2: This question is inspiring and we have been also considering this question recently. One idea is to consider adapting from a convex combination of learned skills instead of from one of the learned skills because a good convex combination of learned skills can be "closer" to the optimal state distribution and practically feasible to obtain.

For example in Figure 2 of our main paper, a convex combination of $p(S|z_1),p(S| z_5)$ could be closer to $p^r$ than $p(S|z_1)$. In practice, the convex combined state distribution $p_\lambda(S)=\sum_z \lambda_z p(S|z)$ can be obtained by sampling $z$ from the distribution $p_\lambda(z)=\frac{\lambda_z}{\sum_{z'}\lambda_{z'}}$.

For the practical adaptation procedure, we can first find the $\lambda$ resulting in a $p_\lambda(S)$ with the best accumulated reward. Because we can not directly update the parameters for $p_\lambda(S)$, we could use a new parametric model to perform offline RL with data collected by $p_\lambda(S)$ and relabeled by downstream task reward function $r$. Through this approach, we concurrently distill knowledge from pretraining and execute adaptation for the downstream task.

Theoretical analysis of this adaptation procedure can be an idea for future works.

Q3: Is there any advantage to the KL divergence over the Wasserstein distance? theoretically or computationally?

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.

Q4: What are the limitations of the PWSEP algorithms (specifically the WSEP objective)?

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.

Thanks again for the constructive feedback!

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