Task Adaptation from Skills: Information Geometry, Disentanglement, and New Objectives for Unsupervised Reinforcement Learning
Authors: Yucheng Yang
Tianyi Zhou
Qiang He
Lei Han
Mykola Pechenizkiy
Meng Fang
Venue: ICLR 2024 spotlight (2024)
Keywords: unsupervised skill learning
reward-free RL
downstream task adaptation
wasserstein distance
theoretical analysis
Primary area: reinforcement learning

Abstract: Unsupervised reinforcement learning (URL) aims to learn general skills for unseen downstream tasks. Mutual Information Skill Learning (MISL) addresses URL by maximizing the mutual information between states and skills but lacks sufficient theoretical analysis, e.g., how well its learned skills can initialize a downstream task's policy. Our new theoretical analysis shows that the diversity and separatability of learned skills are fundamentally critical to downstream task adaptation but MISL does not necessarily guarantee them. To improve MISL, we propose a novel disentanglement metric LSEPIN and build an information-geometric connection between LSEPIN and downstream task adaptation cost. For better geometric properties, we investigate a new strategy that replaces the KL divergence in information geometry with Wasserstein distance. We extend the geometric analysis to it, which leads to a novel skill-learning objective WSEP. It is theoretically justified to be helpful to task adaptation and it is capable of discovering more initial policies for downstream tasks than MISL. We further propose a Wasserstein distance-based algorithm PWSEP can theoretically discover all potentially optimal initial policies.

Published as a conference paper at ICLR 2024
TASK ADAPTATION FROM SKILLS: INFORMATION GE- OMETRY, DISENTANGLEMENT, AND NEW OBJECTIVES FOR UNSUPERVISED REINFORCEMENT LEARNING
Yucheng Yang", Tianyi Zhou', Qiang He", Lei Han', Mykola Pechenizkiy?, Meng Fang
Findhoven University of Technology, University of Maryland, College Park, RRuh University Bochum, "Tencent Robotics X, University of Liverpool (y.yang, m. m.pechenizkiy) @ttu.nl, tianyi tianyi@umd.edu, Meng. liverpool.ac.uk
ABSTRACT
Unsupervised reinforcement learning (URL. aims to learn general skills for unseen downstream tasks. Mutual Information Skill Learning (MISL. addresses URL by maximizing the mutual information between states and skills but lacks sufficient theoretical analysis, e.g., how well its learned skills can initialize a downstream task's ' policy. Our new theoretical analysis in this paper shows that the diversity and separability of learned skills are fundamentally critical downstream task adaptation but MISL does not necessarily guarantee these properties. To complement MISL, we propose a novel disentanglement metric LSEPIN. Moreover, we build an cntemaation-gemmettn task adaptation cost. For better geometric properties, we investigate a new strategy that replaces the KL divergence in information geometry with Wasserstein distance.
connection between LSEPIN and downstream
We extend the geometric analysis to it, which leads to a novel skill-leaming objective WSEP It is theoretically justified to be helpful to downstream task adaptation and it is capable of discovering more initial policies for downstream tasks than MISL. We finally propose another Wasserstein distance-based algorithm PWSEP that can theoretically discover all optimal initial policies.
INTRODUCTION
Reinforcement learning (RL) has drawn growing attention by its success in autonomous control (Ki- umarsi et al., 2017), Go (Silver et al., 2016) and video games (Mnih et al., 2013; Vinyals et al., 2019). However, primary limitation of the current RL is its high sample complexity. Inspired by the successful inctrain-finetume paradigm in other deep learning fields like natural language processing (Radford et al., 2019; Devlin et al., 2019) and computer vision (Henaff, 2020; He et al., 2020), there has been growing work studying the pretraining of RL. RL agent receives no task-related reward during pretraining and learns by its intrinsic motivations (Oudeyer & Kaplan, 2009). Some of these intrinsic motivations can help the agent to learn representations of the observations (Schwarzer al., 2021) and some learn the dynamics model (Ha & Schmidhuber, 2018: Sekar al., 2020). In this work, we focus on Unsupervised RL (URL) that learns set of skills without external reward and the learned skills are expected to be quickly adapted unseen downstream tasks.
A common approach for skill discovery of URL is Mutual Information Skill Learning (MISL) (Eysen- bach al., 2022) that maximizes the mutual information between state and skill latent (Eysenbach et al., 2019; Florensa et al., 2017; Hansen et al., 2020; Liu & Abbeel, 2021b). The intuition is that by maximizing this mutual information the choice of skills can effectively affect where the states distributed that these skills could be potentially used for downstream tasks. There are more algorithms using objectives modified on this mutual information. For example, Lee et al. (2019); Liu & Abbeel (2021b) added additional terms for better exploration, and Sharma et al. (2020); Park et al. (2022a) focus on modified input structure to prepare the agent for specific kinds downstream tasks.
Despite the popularity of MISL, there has been little theoretical analysis of how well the MISL- learned skills can be applied as downstream task initializations. Previous work Eysenbach et al. (2022) has tried to analyze MISL but they consider an impractical downstream task adaptation procedure

---

Published as a conference paper at ICLR 2024
that uses the average state distribution of all learned skills initialization instead of directly using the learned skills. Therefore, it is still unclear how well the MISL-learned skills can be applied as downstream task initializations.
In this work, we theoretically analyze the connection between the properties of learned skills and their downstream task performance. Our results show that the diversity and separability of learned skills are fundamentally critical to downstream task adaptation. Separability, or the distinctiveness of skill distributions, is key for diverse skills. Without it, even large number of skills may cover only limited range resulting in limited diversity. The importance of diversity is empirically demonstrated in previous works (Eysenbach al., 2019; Kim al., 2021: He et al., 2022; Laskin et al., 2022). Our results also show that MISL alone does not necessarily guarantee these properties. To complement MISI, we propose a novel disentanglement metric that able measure the diversity and separability of learned skills. Our theoretical analysis relates the disentanglement metric to downstream task adaptation.
In particular, we introduce novel disentanglemen metric "Least SEParability and INformativeness (LSEPIN)". which is directly related to the task adaptation cost from learned skills and comple- mentary to the widely adopted mutual information objective of MISL. LSEPIN captures both the informativeness, diversity, and separability of the learned skills, which are critical to downstream tasks and can be used to design better URL objectives. We relate LSEPIN to Worst-case Adaptation Cost (WAC), which measures the largest possible distance between a downstream task's optimal feasible state distribution and its closest learned skill's state distribution. Our results show increasing LSEPIN could potentially result in lower WAC.
In addition, we show that optimizing MISL and LSEPIN are essentially maximizing distances measured by KL divergences between state distributions. However, a well-known issue is that KL divergence not true metric, i.e., is not symmetric and does not satisfy the triangle inequality. This motivates us to investigate whether an alternative choice of distance can overcome the limitations of MISL. Wasserstein distance is a symmetric metric satisfying the triangle inequality and has been feasibly applied for deep learning implementations (Arjovsky et al., 2017; Dadashi et al., 2020), so we investigate new strategy that replaces the KI. divergence MISL with Wasserstein distance and exploits its better geometric properties for theoretical analysis. This leads to new skill learning objectives for URL and our results show that the objective built upon Wasserstein distance, "Wasserstein SEParatibility (WSEP)", is able to discover more potentially optimal skills than MISL. Furthermore, we propose and analyze an unsupervised skill-learning algorithm "Projected SEP" (PWSEP) that has the favored theoretical property discover all potentially optimal skills and is able to solve the open question of "vertex discovery" from Eysenbach et al. (2022).
Analysis of LSEPIN is complement to prior work extend the theoretical analysis of MISL to practical downstream task adaptation, while the analysis of WSEP and PWSEP opens up a new unsupervised skill learning approach. Our results also answer the fundamental question of URL about what properties of the learned skills lead to better downstream task adaptation and what metrics can measure these properties.
Our main contributions can be summarized in the following:
1. We theoretically study novel but practical task adaptation cost (i.e., WAC) for MISL, which measures how well the MISL-learned skills can be applied as downstream task initializations.
2. We propose a novel disentanglement metric (i.c., LSEPIN) that captures both the informa- tiveness and separability of skills. LSEPIN is theoretically related to WAC and can be used to develop URL objectives.
3. We propose a new URL formulation based on Wasserstein distance and extend the above theoretical analysis to it, resulting in novel URL objectives for skill learning. Besides also promoting separability, they could discover more skills than existing MISL that are potentially optimal for downstream tasks.
Although our contribution is mainly theoretical, in appendices and I we show the feasibility of practical algorithm design with proposed metrics and empirical examples to validate our results. summary of our proposed metrics and algorithm is in appendix A and frequently asked questions answered in appendix B.
2

---

Published as a conference paper at ICLR 2024
2 PRELIMINARIES
without external rewards with dis- crete states S and actions A, dynamics P(s++11, 8t,0s) 0z), initial state distribution po(s)), and discount factor 7 [0, 1]. A policy T(a]s) has its discounted state occupancy measure as (*()) - (1 - 7) Et=n 7 Pr (s), where (T) the probability that policy и visits state time t. There can be downstream tasks that define extrinsic reward as a state-dependent function r(s), where uction-dependent reward functions can be handled by modifying the state include the previous action. The cumulative reward of the corresponding downstream task is Ep (pp((s)(r) (s)[r(s)]-
We consider infinite-horizon MDPs M =
We formulate the problem of unsupervised skill discovery as learning skill-conditioed policy π (a,]();,2) 2) where 2 E 2 represents the latent skill and 2 is a discrete set. (-)) and I(::-) denote entropy and mutual information, respectively. ((,+) denotes Wasserstein distance. We use upper- case letters for random variables and lower-case letters for samples, eg. p(S).
2.1 MUTUAL INFORMATION SKILL LEARNING
Unsupervised skill learning algorithms aim to learn a policy =AASS Z) conditioned on a latent skill Their optimization objective is usually the mutual information [(S; Z) and they differ the prior or approximation of this objective (Gregor et al., 2017; Eysenbach et al., 2019; Achiam et al., 2018; Hansen et al., 2020).
In practical algorithms, the policy is generally denoted as To( Zimpus) with parameters 0 and conditioned skill latent Pinput ~ p( Zippor) Let y (S) Pinpur) denote the state distribution of policy. The practical objective of MISL could be:
max T(S; Zippu) Zu ( (p (mppl)) pT= 0.p(Zupe)
(I)
Policy parameters θ and the latent variable Ziputt can be composed into single representation, (8, Zmput).. then To (AIS, Zimpui) = T S, 2). We call representation "skill" in the following paper. Then, MISL is learned by finding an optimal p( ()) that solves
max / (S; 2) (P (p(S)=) || ((8)) 2(2)
where p(S) Ex( [I())))], is the average state distribution of discovered skills.
2.2 INFORMATION GEOMETRY OF MISL
1901
(a)
(b)
Figure 1: Visualized examples: (a) C is the feasible state distribution set, and the blue simplex the probability simplex for state distribution with [S] 3. (b) MISL discovers skills the vertices (21, 22, 23} on the "circle" with maximum "radius" centered in their average state distribution p(S).
Prior work Eysenbach et al. (2022) shows that the set C of state distributions feasible under the dy- namics the MDP constitutes a convex polytope lying on probability simplex of state distributions, where every point in the polytope C is represented by skill latent and its state distribution is p(S(z). For any downstream task defined by a reward function r : S - R. because of the linearity of Ep(s) (rs)) and convexity of C, the state distribution that maximizes the cumulative reward Ep(e) [r()]
3

---

Published as a conference paper at ICLR 2024
lies at one of the vertices of C (Boyd & Vandenberghe, 2014). Equation (2) shows that MISL learns skill distribution 2 ~ p( Z) to put weight on skills that have maximum KL divergence to the average state distribution. It can be considered as finding skills that lie on the unique (uniqueness proved in appendix E) "circle" with maximum "radius" inside the polytope C, thus the discovered skills lie at the vertices of polytope C, as shown in Lemma 6.5 of Eysenbach al. (2022) by Theorem 13.11 of Cover & Thomas (2006). So the skills discovered by MISL are optimal for some downstream tasks, An intuitive example of the skills discovered by MISL shown in fig. 1b.
3 THEORETICAL RESULTS
Although MISL discovers some vertices that are potentially optimal for certain downstream tasks, when the downstream task favors target state distributions at the undiscovered vertices, which often happens in practice that the learned skills are not optimal for downstream tasks, there exists a "distance" from discovered vertices to the target vertex, and the "distance" from the initial skill for adaptation to the target state distribution can be considered as the adaptation cost. The prior work only analyzes the adaptation cost from the average state distribution of skills p(S) = E, [[]]] to the target state distribution. Because most practical MISI algorithms initialize the adaptation procedure from one of the learned skills (Lee et al., 2019; Eysenbach et al., 2019; Liu & Abbeel, 202 Laskin et al., 2021) instead of the average PS)). the prior analysis provides little insight on why these practical algorithms work. The fundamental question for unsupervised skill learning remains unanswered: 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?
We have answered this question with theoretical analysis in this section, empirical validation of the theories is in appendix I. Our informal results are follows:
1. In order to have a low adaptation cost when initializing from one of the learned skills, the learned skills need to be diverse and separate from each other. Separability means the discriminability between states inferred by different skills.
2. MISL alone does not necessarily guarantee diversity and scparability. We propose a disentanglement metric LSEPIN to complement MISI. for diverse and separa- ble skills.
plsks PUS
Figure 2: Example of concyclic ver- tices:
3. MISL discovers limited vertices, we propose WSEP metric based on Wasserstein distance that can promote diversity and separability as well as discover more vertices than MISL. One Wasserstein distance- based algorithm PWSEP can even discover all vertices.
21, 22, 23, 24, 25 can all be vertices optimal for MISL, p(S) can be identified by [S] = 3 of them, so eq. (2) can be solved by any of these vertices. The blue arrow is to adapt from learned skills to target distribution for downstream task (our setting) and the brown ar- row is to adapt from the average state distribution (setting in Eysen- bach et al. (2022))
The first point is intuitive that the diverse and separable skills are likely to cover more potentially useful skills, as shown by empirical results in Eysenbach et al. (2019); Park et al. (2022b); Laskin et al. (2022). The second point claims MISL alone does not guarantee diversity and separability, and this can be seen from the example in fig. 2. In this case, there are two sets of (S) = 3 skills Z, [21, =4, )5} and Z, : { 23, #5} both on the maximum "circle" solving MISL. Because 22 and 23 have close state distributions. skills of Z,, are less diverse and less separable. There can be more than |S] vertices on the maximum "circle" in the case of fig. 2 because, unlike prior work Eysenbach et al. (2022), we do not take into account the "non-concyclic" assumption that limits the number of vertices on the same "circle" to be (S]. Our proposed disentanglement metric LSEPIN would favor Z, over Zn, and theoretical analysis of LSEPIN is conducted in section 3.2 to show its relation to downstream task adaptation cost. The downstream task procedure we consider is initialized from one of the learned tasks, for the case in fig. 2, when the target state distribution is p", consider adapting from the skill in Z, that is closest to p (blue arrow), which is 21, while the prior work adapts from pSS) (brown arrow).
4

---

Published as a conference paper at ICLR 2024
The advantages of Wasserstein distance are that it is true metric satisfying symmetry and triangle in- equality. We can use it to measure distances that can not be measured by KL divergences. Optimizing these distances also promotes diversity and separability as well as results in better vertex discovery, even capable of discovering all vertices and solving the open question of "vertex discovery" from Eysenbach et al. (2022). Details about Wasserstein distance skill learning are shown in section 3.3. A summary of all proposed metrics and algorithm is in appendix A.
3.1 How TO MEASURE DIVERSITY AND SEPARATABILITY OF LEARNED SKILLS
Many previous MISL algorithms (Eysenbach et al., 2019; Gregor al., 2017; Sharma al., 2020)
emphasized the importance of diversity and tried to promote diversity by using uniform p( (imptt) for
eq. (I). However, uniform p( (Zinpu)) for objective eq. (1) does not ensure diverse for p(Z) in eq. (2) since 2 = (0, Zinput) also depends on the learned parameter 0. We show an example in appendix D when maximizing (S; 2) with uniform p( impu)) results in inseparable skills. Empirical discussions in Park et al. (2022b); Laskin et al. (2022) also mentioned that the learned skills of these MISL methods often lack enough diversity and separability. Furthermore, as mentioned previously by the example in fig. 2, even when [(S; 2) in eq. (2) is maximized, the learned skills could still lack diversity and separability of the skills. To complement MISL, we propose a novel metric to explicitly measure the diversity and separability of learned skills.
We consider / ((;;1) (1, is the binary indicator function of Z = z) to measure the informativeness and separability of an individual skill z. In the context of unsupervised skill learning, informativeness should refer to the information shared between a skill and its inferred states. As mentioned, separabil- ity means the states inferred by different skills should be discriminable. We analyze the minimum of [(;; 12) over learned skills. We name it Least SEParability and INformativeness (LSEPIN)
LSEPIN min / (S;12).
[[;; 1.) is related to how much states inferred by skill and states not inferred by z are discriminable from each other, so it covers not only informativeness but also separability of skills, In the context of representation learning, the metrics capturing informativeness and separability are called the disentanglement metries (Do & Tran, 201 9b; Kim et al., 2021), so we also call LSEPIN as a disentanglement metric for unsupervised skill learning. More details about the difference between disentanglement for representation learning and disentanglement for our skill learning setting are in appendix F.
3.2 HOW DISENTANGLEMEN AFFECTS DOWNSTREAM TASK ADAPTATION
We provide a theoretical justification for the proposed disentan- glement metrie, showing that it can be complement of / (S; Z) to evaluate how well the URL agent is prepared for downstream tasks by the following theorems.
Definition 3.1 (Worst-case Adaptation Cost). Worst-case Adap- tation Cost (WAC) is defined as
WAC max min (kI(p(()=)) I"), r =E2'
(4)
where p" is the optimal feasible state marginal distribution for the downstream task defined by r, and 2" is the set of learned skills.
p(5)
The following theoretical results show how the LSEPIN metric is related to the WAC in definition 3.1.
pts 2))
pa
00512
the MISL objective [(S; 2) is maximized by solutions
Theorem 3.1. When learned skill sets -,,i=1,2, with N < [S] skills (N skills have p(2) > 0) sharing the same skill are all MISI solutions, The skill set with the higher T(S; 12) will have higher p(z) and lower adaptation cost for all r= in the RR,, where R, is the set of downstream tasks always satisfying Vi, Vr Ry z arg max DKL (x(z))))) II ). And the maximum of this adaptation cost has the following formulation: C2(r) - D==))
Figure 3: Example of two MISL
solutions: Zi {21, 22, 23 ) and Z {21, 252 2/ Z has higher
[[S
IC_ max rER:
(5)
1 1- p(2))
5

---

Published as a conference paper at ICLR 2024
where
C2(r) ((8 Drr. (p() ll)", D2(r) = [I
(6)
(7)
Theorem 3.1 provides a correlation between our proposed metric [(S; 1:) the adaptation cost IC; from skill in 2 {2] that is closest the downstream task optimal distribution and its detailed proof is in appendix C.1. To better understand the claim of this theorem, we can look at the intuitive example shown in fig. 3. In this case |S| 3. When MISL is maximized by 3 skills, the skill combinations MISL solutions could be 2 (21, 22, 23) and 2 {=1, 22 2/)- 2 has higher T(S; 1,,) than 2j. By theorem 3.1, solution 2g should have lower cost to adapt to the optimal distribution p'= of the downstream task T21-
Corollary 3.1.1. When the MISL objective 1(8,2) maximized by N & |S| skills, WAC is bounded of solution Z' by
- ((r)-p(=)D&(() 1 ((2)
WAC S max IC2 max max =EZ rER:
(8)
=EZ.
WAC the worst-case adaptation cost defined definition 3.1, C. and D, are as defined in eqs. (6) and (7). R2 here needs to satisfy Vr E R2,, arg max DKL (p((22)) p").
Corollary 3.1.1 provides an upper bound for WAC. The proof is deferred to appendix C.2. The results in Corollary 3.1.1 and theorem 3.1 considered situations when MISL is solved and [(S; Z) is maximized, we also discussed how 7 (S; 1, and LSEPIN affects learned skills and adaptation cost when T(S; Z) is not maximized in appendix C.4.
By theorem 3.1 we know that higher ((S; 1:) implies lower IC,, but how much IC, associated with an individual skill contribute to the overall WAC can not be known in prior and it depends on specific C2 and D.. Moreover, specific C₂ and D2 depend on the "shape" the undiscovered parts of C and can not be known before the discovery of all vertices. Therefore, in practice, like existing work (Durugkar et al., 2021; He et al., 2022) treating the desired properties of each skill equally in practical algorithms,, we could treat every IC, equally. We have the following theorem showing under which assumptions we can treat every IC, equally for WAC.
Theorem 3.2. When 1. the optimal state distribution for the downstream task is far from S) and 2. The state space is large, iee. |S] is large. IC, of all learned skills can be considered equally contribute to WAC
Both assumptions for this theorem are practical and can commonly happen in complex and high- dimensional environments. When every 1C2, is treated equally for WAC, optimizing LSEPIN could lead to lower WAC. It is formally analyzed and proven in appendix C.3.
In summary, we have provided theoretical insight on how (S; 1.) affects downstream task adaptation and how optimizing LSEPIN could lower WAC under practical assumptions. We do not assume "non-concyclic" vertices and we consider the practical approach of directly adapting from learned skills instead of the average state distribution. LSEPIN is complement to the mutual information objective T(S; Z). Compared to T(S; 2), it provides a better metrie to evaluate the effectiveness of learned MISL skills for potential downstream tasks. Our results have shown the diversity and separability of the Iearned skills measured by I(S; 1g) and LSEPIN are desired for better downstream task adaptation.
Remark 3.2.1. One limitation with MISI even with LSEPIN is that even without the limitation of the number of skills to have p(z) > 0, it still can not discover vertices such that
II p(S)) (KL()(Sl))
Eax [DKL(p(S12) P(S)) p(z)
Vertex pTs in fig. belongs to such vertices.
3.3 SKILL LEARNING WITH WASSERSTEIN DISTANCE
In this subsection, we analyze strategy that replaces the KL divergence in information geometry with Wasserstein distance for better geometric properties to overcome the limitation of MISL shown in remark 3.2.1.

---

Published as a conference paper at ICLR 2024
Maximizing [(S; Z) and LSEPIN essentially maximizing distances measured by KL divergences between points in a polytope. KL divergence is not symmetric and does not satisfy the triangle inequality, KL divergences between points of the polytope could be incomparable when two KI. divergences don't share a same point. We study the strategy that replaces the KI. divergence in MISL with Wasserstein distance since Wasserstein distance true metric. Then we conduct further theoretical analysis to exploit its better geometrical properties such as symmetry and triangle inequality.
In this section, we will introduce the learning objectives as well as evaluation metrics for Wasserstein Distance Skill Learning (WDSL), analyze what kind of skills these objectives can learn, where the learned skills lie in the polytope, and how these learned skills contribute to downstream task adaptation. Theoretically, the favored property of WDSL is that discovers more vertices in C that are potentially optimal for downstream tasks than MISI, and one WDSL algorithm can discover all vertices.
3.3.1 OBJECTIVES FOR WASSERSTEIN DISTANCE SKILL LEARNING
First of all, we can trivially replace the KI. divergences in the MISI. objective eq. (2) with Wasserstein distance and obtain a basic WDSL objective
max E P(=) (v p(z)
We name it Average Wasserstein skill learning Distance (AWD), similar to the MISL objective in eq. (2), this objective also learns skills that lie on a hyper ball with a maximum radius. Because this objective is not main proposition and also suffers from the limitation of remark 3.2.1, we put the analysis of this objective in appendix G.I.
We mainly analyze this objective for WDSL:
WSEP
Σ Σ ZEE (2E)
(10)
where Z is the set of skills with p(:) > We call this objective Wasserstein SEParatibility (WSEP), it can be considered as disentanglement for WDSL as it measures the Wasserstein distance between learned skills. Recall that separability for MISL defined as how discriminable the state is, Wasser- stein distances between skills not only represent discriminability but also express the distance between trajectories when there are no overlappings.
3.3.2 GEOMETRY OF LEARNED SKILLS
As mentioned hefore in section 2.2, the skills that potentially optimal for downstream tasks lie the vertices of the polytope of feasible state distributions. By the following lemma, we show that optimizing WSEP will push the learned skills to the vertices of the polytope.
Lemma 3.3. When WSEP is maximized, all learned skills with p(z) > 0 must lie at the vertices of the polytope.
Proof of this lemma is in appendix G.2.
The previous theoretical results of disentanglement metric LSEPIN depend on the maximization of [(S; 2). so as mentioned in remark 3.2.1, it still only discover vertices with maximum "distances" the average distribution p(S). However, WSEP does not depend on the maximization of other objectives, e.g., eq. (9), so there is no distance restriction on the vertices discovered by WSEP. Therefore, is possible for WSEP to discover all vertices of the feasible polytope C, thus discovering all optimal skills for potential downstream tasks. For example, in environment with polytope shown fig. 1b, MISL only discovers 3 vertices on the "circle" with maximum "radius" while WSEP is able to discover all 5 vertices.
Remark 3.3.1. When there is no limitation the number of skills with positive probability, Maxi- mizing WSEP could discover more vertices than MISL in some cases, and even potentially discover all vertices, shown in the example appendix G.6.
7

---

Published as a conference paper at ICLR 2024
3.3.3 How WSEP AFFECTS DOWNSTREAM TASK ADAPTATION
Then, we propose a theorem about how the WSEP metrie is related to downstream task adaptation when there is a limitation on the quantity of learned skills.
Definition 3.2. Mean Adaptation Cost (MAC): mean of the Wasserstein distances between the undiscovered vertices and the learned skills closest to them.
MAC = 2
min Σ
W(p(S)2)),
(11)
EEZ-
ZEVIZZ
V is the set of all skills that have their conditional state distribution at vertices of the MDP's feasible state distribution polytope, and zt is all learned skills with p(z) 0.
Theorem 3.4. When WSEP is maximized by 2 skills, the MAC can be upper-bounded:
MAC - E=EZ- Li (12g) WJWE
(12)
MAC EvEz by - (127) 1)WSEP (V] 2+!|
(13)
where
13 Σ W W(r(Slu), VEV
(14)
Ly VEV
W(p($lo),p(51o))
Theorem 3.4 shows the relation between WSEP and the upper bounds of adaptation cost MAC the practical setting, where the number of skills to be learned is limited. The proof is in appendix G.3.
In stationary MDP, the polytope fixed, the edge lengths Lý and Ly are constant. Larger WSEP seems to tighten the bounds, but different WSEP also means a different 2" set of learned skills, thus a different Σ=ez. T, Therefore, increasing WSEP only tightens the bound in eq. (13) but not necessarily the bound in eq. (12).
Remark 3.4.1. More distance is not always good: WSEP as disentanglement metric promotes the distances between learned skills and These bounds of MAC show that maximizing WSEP can indeed help with downstream task adaptation, but this does not mean that learned skills with more WSEP will always result in lower MAC. An illustrative example is shown in appendix G.5, where more distant skills with higher WSEP do not have lower adaptation costs
Remark 3.4.2. If we replace the Wasserstein distances in WSEP with KL. divergences, we get symmetric formulation of K K.SEP = Ex,ez Ezze2.04) (KL(p(S)zz)) p(S)z;))- is symmetric, but it does not promote diversity and separability because KL divergence does not satisfy the triangle inequality. More details are analyzed appendix G.8
WSEP does not suffer from the limitation of remark 3.2.1 because it does not try to find skills on maximum "circle". Although WSEP can potentially discover more vertices than MISL, we find that it may not be able to discover all vertices of the feasible state distribution polytope C in appendix G.7.
3.3.4 SOLVING THE VERTEX DISCOVERY PROBLEM
The following theorem shows a learning procedure based on Wasserstance distance capable of discovering all vertices of feasible state distribution polytope C.
Theorem 3.5. When V is the set of all vertices of the feasible state distribution polytope C, all | vertices can be discovered by [V] iterations of maximizing
PWSEP(1) min W (p(S)z)).
Σ XP(S12))). 2,E2,
(15)
where Z, is the set of skills discovered from iteration 0 to i 1 and is the skill being learned at ith iteration. X is a convex coeffcient of dimension 1 that every element X 2 0, Vj E (0, 1,.,, i 1}
and {j X= =
8

---

Published as a conference paper at ICLR 2024
In the initial iteration when 2, 0. PWSEP(O) can be (p((S)2o) p( (([))))))) with 2rand to be randomly initialized skill.
(i) can be considered as a projection to the convex hull of Z,, so we call it Projected WSEP and this learning procedure the PWSEP algorithm. It can discover all [V] vertices with only [V] skills. Although lemma 3.3 shows that maximizing WSEP also discovers vertices, the discovered vertices could be duplicated (shown in appendix G.7). Maximizing projected distance PWSEP(i) could ensure the vertex learned at each new iteration is not discovered before. Proof and more analysis of the vertex discovery problem can found in appendix G.4.
RELATED WORK
MISL is widely implemented and has been the backbone of many URL algorithms (Achiam et al., 2018; Florensa et al., 2017; Hansen al., 2020). Prior work Eysenbach et al. (2022) tried to provide theoretical justification for the empirical prevalence of MISL from information geometric (Amari & Nagaoka, 2000), but their analysis mainly considered an unpractical downstream task adaptation procedure. Works like Eysenbach al. (2019); Park et al. (2022b; 2023); He et al. (2022); Laskin et al. (2022) showed the empirical advantages of favored properties such as diversity and separability of learned skills. Our theoretically justified these properties and showed they benefit practical adaptation.
In Kim et al. (2021) the concept of disentanglement was mentioned. They used the SEPIN k and WSEPIN metries from representation learning (Do & Tran, 2019b) promote the informativeness and separability between different dimensions of the skill latent. However, properties of latent representations could be ensured by optimization only in the representation space, so they do not explicitly regulate the state distributions of learned skills like our proposed LSEPIN and WSEP do. Appendix F discussed more details.
Recent practical unsupervised skill learning algorithms (He et al., 2022; Durugkar et al., 2021) maximize a lower bound of WSEP, so our analysis on WSEP provides theoretical insight on why these Wasserstein distance-based unsupervised skill learning algorithms work empirically. Their empirical results showed the feasibility and usefulness of skill discovery with Wasserstein distance.
Successor feature (SF) method SFOLS (Alegre et al., 2022) can also discover all vertices but learns an over-complete set of skills, which our PWSEP algorithm efficiently avoids. In appendix G.4.2, the difference between the SF setting and our skill learning setting is discussed in detail, as well as the comparison theoretical properties between our proposed PWSEP and SFOLS. Other methods like Hansen et al. (2020); Liu & Abbeel (2021b) combined MISL with SF for URL, and they shown to accelerate downstream task adaptation. Since they are MISL methods adapting from one of the learned skills, our theoretical results also apply to them.
5 CONCLUSION
We investigated the geometry of task adaptation from skills learned by unsupervised reinforcement learning. We proposed disentanglement metric LSEPIN for mutual information skill learning to capture the diversity and separability of learned skills, which are critical to task adaptation. Unlike the prior analysis, we are able build a theoretical connection between the metric and the cost of downstream task adaptation. We further proposed a novel strategy that replaces KL divergence with Wasserstein distance and extended the geometric analysis and algorithm PWSEP for unsupervised skill learning. Our theoretical result shows why they should work, what could be done, and what limitations they have. Specifically, we found that optimizing the proposed WSEP objective can discover more optimal policies for potential downstream tasks than previous methods maximizing the mutual information objective T(S; 2). Moreover, the proposed PWSEP algorithm based on Wasserstein distance can theoretically discover all optimal policies for potential downstream tasks.
which leads to novel objective WSEP
Our theoretical results could inspire new algorithms using LSEPIN Wasserstein distance for unsupervised skill learning. For Wasserstein distance, the choice of transport cost is important, which may require strong prior knowledge. Our future work will develop practical algorithms that learn deep representations such that common transport costs such as L2 distance in the representation space can accurately reflect the difficulty of traveling from one state to the other.
9

---

Published as a conference paper at ICLR 2024
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint arXiv: 1807. 10299, 2018.
Lucas Nunes Alegre, Anale L. C. Bazzan, and Bruno C. da Silva. Optimistic linear support and successor features basis for optimal policy transfer. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 394-413. PMLR, 2022. URL https: pesttp::/prroeeee /ptpceedicc...... mlr. www.prress/veeee / wttpv////aaaaeeeeeee
html.
Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry. 2000.
Brandon Amos and J. Zico Kolter. Optnet: Differentiable optimization as layer in neural net- works. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 136-145. PMLR, 2017. URL http:/ pttp::///ccceedngsss www..pree./a....m 20070//mmm7aa...m/mm html.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214 223. PMLR, 2017.
Stephen P. Boyd and Lieven Vandenberghe. Convex Optimizatiion. Cambridge University Press, 2014. ISBN 978-0-521-833788 doi: 10. 111017/C8097780511884444 URL https:/ /web. stanford, edu/ 7Eboyd/ cvxbook/.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https://openreviee net/forum? 1d=#113JaR5ym.
Victor Campos, Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giró-i-Nieto, and Jordi Torres. Explore, discover and learn: Unsupervised discovery of state-covering skills. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1317-1327. PMLR, 2020. URL http: http///pooceeen.eeeeee/ee/c ..r.m.m..... mir. wwwppprsssilll
campos20a html.
Eric Carlen and Wilfrid Gangbo. Constrained steepest ascent in the 2-wassrsteein metric. Annals of Mathematics, 157:807-888 846, 05 2003. doi: 10. 2..0007/mmaass.200..557....
Thomas M. Cover and Joy A. Thomas. Elements of Information Theory (Wiley Series in Telecommu- nications and Signal Processing). Wiley-Interscienee, USA, 2006. ISBN 0471241954.
Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning, arXiv preprint urXiv:2006.04677 2020.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 2019.
Kien Do and Truyen Tran. Theory and evaluation metrics for learning disentangled representations, 2019a. URL https: ////aaivvv.//./////v//.//////aav org/ abs 190../99... 09961.
Kien Do and Truyen Tran. Theory and evaluation metries for learning disentangled representations. arXiv preprint wwXpv:190.0099611 2019b.
Ishan Durugkar, Steven Hansen, Stephen Spencer, and Volodymyr Mnih. Wasserstein distance maximizing intrinsic control. CoRR, 10.21110..33, 2021. URL https://acxiv.org abs/2110. 15331.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without reward function. In ICLR, 2019.
10

---

Published as a conference paper at ICLR 2024
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. The information geometry of unsupervised reinforcement learning. The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https:: Settps://openreeee wpwpneeeevvveeeeeeeeeeeeeeeeeeee ew net/ forum? wtwpummmddww0U000vc
Carlos Florensa, Yan Duan, and Pieter Abbeel. Stochastic neural networks for hierarchical rein- forcement learning. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https:/ httpps://ppenreveww net wtw.ee/wwweeeooommeeeeeeeeee
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11): 131444413,,14, 2020.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. In International Conference on Learning Representations, 2017.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint urXiv: / 2018.
Nikolaus Hansen and Andreas Ostermeier. Completely derandomized sclfaddattoooo in evolution strategies. Evol. Comput., 9()):199-195, 2001. doi: 10.1 10.1162/1053365501750999999 URL https: //dii.///ooi.. wwoi..ogg//00 100 1ww.11122111166601 11111111111111111111111111111111 9.3399988.1
Steven Hansen, Will Dabney, André Barreto, David Warde- Farley, Tom Van de Wiele, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. International Conference on Learning Representations, 2020.
Kaiming He, Haoqi Fan, Yuxin Wu. Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9733 2020.
Shuncheng He, Yuhang Jiang, Hongchang Zhang, Jianzhun Shao, and Xiangyang Ji. Wasserstein unsupervised reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 36, pp. 2884-6892 2022.
Olivier Henaff. Data-efficient image recognition with contrastive predictive coding. International conference on machine learning, pp. 4182-4192. PMLR, 2020.
J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science Engineering, 9()): 90-95, 2007. doi: 10. 109/MCSE.2000..55
Jaekyeom Kim, Seohong Park, and Gunhee Kim. Unsupervised skill discovery with bottleneck option learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 5572-5582. PMLR, 2021. URL http: //proceedings... mll
wtt.:mmw......eeeeeeeeeeeeeeeeee www.s/ressviii /vtsss//vv///// http//11////n// kim21j html.
Bahare Kiumarsi, Kyriakos G Vamvoudakis, Hamidreza Modares, and Frank L Lewis. Optimal and autonomous control using reinforcement learning: A survey. IEEE transactions on neural networks and learning systems, 29(6):2042-2022 2017.
Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced wasserstein distances. Advances in neural information processing systems, 32, 2019.
Yehuda Koren. On spectral graph drawing. In Tandy J. Warnow and Binhai Zhu (eds.), Computing and Combinatorics, 9th Annual International Conference, COCOON 2003, Big Sky, MT, USA, July 25-28, 2003, Proceedings, volume 2697 of Lecture Notes in Computer Science, pp. 496- 508. Springer, 2003. doi: 10. 1100773--50-45555 _5, URL ht : //ddi.org//..100000 3-540-455 45071 1 500
11

---

Published as a conference paper at ICLR 2024
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Cather- ine Cang. Lerrel Pinto, and Pieter Abbeel. URLB: unsupervised reinforcement learning benchmark. In Joaquin Vanschoren and Sai-Kit Yeung (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks
1, NeurlPS Datasets and Benchmarks 2021, December 2021, virtual, 2021. URL
neurips httppss.ccppppppeeee http://eeeeeaaceeeeeesca./aee//p cc/pappr/2021/ hass:/999288ffecd00o222222622222
Abstract round2. html.
Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel. CIC: contrastive intrinsic control for unsupervised skill discovery. CoRR, abs/2202.00161,, 2022. 2022, URL
https:/ /aap:i//aaivvv//iii org/abs/////ooooooooo/ggst/..arr
abs 2202. 00161.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric P. Xing, Sergey Levine, and Ruslan Salakhut- dinov. Efficient exploration via state marginal matching. CoRR, abss/196.052211 2019.
Hao Liu and Picter Abbeel. Behavior from the void: Unsupervised active pre-training. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 18459-18473, 2021a. URL https: httpsts///prrceeeeee hash/ 9tt//99ff3333 3d15 53344ffffe (0774764001aaaa
http:////ccccceeeeeeeoeoe neurips. cc/paper/2021/ Abstract.htm
Hao Liu and Pieter Abbeel. APS: active pretraining with successor features. In International Conference on Machine Learning, 2021b.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Pierre-Yves Oudeyer and Frederic Kaplan. What is intrinsic motivation? a typologgy of computational approaches. Frontiers in neurorobotics, 1:6, 2009.
Seohong Park, Jongwook Choi, Jaekyeom Kim, Honglak Lee, and Gunhee Kim. Lipschitz-constrained unsupervised skill discovery. ArXiv, abs/2202.00914, 2022a.
Seohong Park, Jongwook Choi, Jaekyeom Kim, Honglak Lee, and Gunhee Kim. Lipschitz-conssrraieed unsupervised skill discovery, 2022b. URL htt.t::tt///xiiiie /arxiv..vvvvvv org/ /trg//bs/// /2tp..222..o2222 00914.
Scohong Park, Kimin Lee, Youngwoon Lee, and Pieter Abbeel. Controllabilit--wware unsupervised skill discovery. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning. ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 27225-27245 PMLR, 2023. URL https: /pwoceeedoo..............s / ml mlr. wtt..ng.s.....ss................ v202 2012/park2l pa..pa2k2.. 23h. html.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found. Trends Mach. Learn.. 11(5-6):355-600 2019, doi: 10.1561/2200000073.. URL 2200000073.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Mark Rowland, Jiri Hron, Yunhao Tang, Krzysztof Choromanski, Tamas Sarlos, and Adrian Weller. Orthogonal estimation of wasserstein distances. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 186-195. PMLR, 2019.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58663):944 2015.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self- predicreeie representations. In International Conference Learning Representations, 2021.
12

---

Published as a conference paper at ICLR 2024
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supe world models. In International Conference on Machine Learning, pp. -592. PMLR, 2020.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. In International Conference on Learning Representations, 2020.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):4444889 2016.
Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk. Nearest neighbor estimates of entropy. American Journal of Mathematical and Management Sciences, 23 (3-4):301-321, 2003a. doi: 00.1000019963222200..00000 1080/01966300 2003 10737616.
URL https: /ttp///diiio/ee/ org/10.
Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk. Nearest neighbor estimates of entropy. American Journal of Mathematical and Management Sciences, 23 (3-4)1301-321, 2003b.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-basee control.
In International Conference on Intelligent Robots and Systems, 2012.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-354 2019.
Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in RL: learning representations with efficient approximations. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL net / /tt///orrmmii////// 2id hett/Eorrmiid---1110000. HJ1NpoA5YO.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference Machine Learning, 2021.
13

---

Published as a conference paper at ICLR 2024
A SUMMARY OF PROPOSED METRICS AND ALGORITHM
Our results are mainly theoretical, showing how properties measured by the following metrics are related to downstream task adaptation. The results can inspire new algorithms with these metries.
We proposed (SS 1:) to measure the separability of an individual skill 2, more about separability is discussed in appendix F. It linked to the adaption cost IC2 that initializes from the closest skill other than skill (in the set 2 {=}) to the downstream task optimal distribution.
We proposed LSEPIN min; I(S; 1-) to measure the overall separability and diversity of learned skills. It is linked to the adaption cost WAC defined in eq. (4) that initializes from the closest skill to the downstream task optimal distribution in the learned skill set Z.
We proposed WSEP = Σx,ez E1,e2 (p((S==), p(S ((5]z))) to measure the overall separability and diversity of learned skills. is linked to the adaption cost MAC defined in eq. (1I) that initializes from the closest skill to the downstream task optimal distribution in the learned skill set 2. WSEP discover more skills that are potentially optimal for downstream tasks.
We proposed the PWSEP algorithm. Theoretically, can discover all skills potentially optimal for downstream tasks.
Table I: Properties of proposed metrics
[S: 1.)
LSEPIN
WSEP
Measurement target
An individual skill
Set of learned skills KL divergence
Set of learned skills Wasserstein distance
"Distance" metric
KL divergence
Related adaptation cost
IC²
WAC Yes
MAC
Practical objective
No
Yes
Table displays the summarised properties of the proposed metrics.
B FREQUENTLY ASKED QUESTIONS
B. WHAT'S THE OVERLAPS AND DIFFERENCES BETWEEN OUR WORK AND EYSENBACH ET AL. (2022)?
Their analysis of how MISL affects downstream task adaptation is limited to the adaptation procedure of initializing from the average state distribution p(S) = Egez [[]]]=]] of learned skill set Z. Because most practical MISL algorithms initialize the adaptation procedure from one of the learned skills (Lee et al., 2019; Eysenbach et al., 2019; Hansen et al., 2020; Liu & Abbeel, 2021b; Laskin al., 2021) instead of the average p(S), their analysis provides little insight on why these practical algorithms work. Unlike Eysenbach et al. (2022), our work analyzed the practical and popular adaptation procedure of initializing from the learned skills. Our theoretical results not only provide insight on why the above- mentionee practical algorithms work but also answer the fundamental question for unsupervised skill learning: 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? Moreover, our proposed metrics and algorithm are novel and our analysis for the Wasserstein distance offers an innovative perspective on unsupervised skill learning.
Differences:
The only overlap between our work and the prior work (Eysenbach et al., 2022) that we both analyze the geometry of the state distributions of the skills learned by unsupervised skill learning methods.
Overlap:
B.2 WHAT EXACTLY DOES SEPARABILITY MEAN AND HOW IT IS RELATED TO DIVERSITY?
Separability of a skill means the discriminability between states inferred by different skills. For example, high DKL (p(S12) || p(S # 2)) high W (p(S)²), p(S|Z * 2)) could mean skill is more "distant" from other skills. We show in appendix F that increasing I(S: 1.) leads to higher (K1(((()))) p(S Z * z)), thus promoting separability.
14

---

Published as a conference paper at ICLR 2024
Separability is the prerequisite for diversity. Without separability, with large number of skills, they could be close to each other and only cover a small region of the feasible state distribution polytope, skills need to be distinctive and separable first before they can be diverse.
B.3 WHAT'S THE DIFFERENCE BETWEEN I(S; 12) AND 2)?
High-level Intuitive difference:
The binary indicator 1; is random variable that takes value in {0. 1} with the distribution p(1.), which only describes skill 2, while Z is random variable with the distribution p(Z), which describes all skills. [(S; 1z) enables measure how well every single skill is learned terms of informative- ness and separability, so we can evaluate the minimum min, T(S; 12) by LSEPIN and the median by MSEPIN(L719): While the default Mutual Information Skill Learning (MISL) objective Z(S; Z) only measures the overall informativeness of all skills.
Formal difference:
[[; 12) can be decomposed as
I(S; p(=)Dxr + p(z 2))x((((SIZ #)]1S))
(16)
while 1(S; can be decomposed
I[S E Dx (p(S=)]]]]]
(17)
The distribution p(S 2 + 2) #''"p(2''>1 (())=²)) in the composition of [(S;1,) but not in
the composition of T(S; Z) Notice that the second term in the for decomposition of T(S; 1,):
(p + =)][s) + F (expo
(18)
So these are two different formulations.
Ad Advantages of T(S; 12):
1. 1_) can be used to evaluate every single skill.
2. Besides informativeness, [(S; 1.) explicitly encourages separability, more details about separability are discussed in appendix F.
B.4 WHY NOT REPLACE THE WASSERSTEIN DISTANCE IN WSEP OR PWSEP WITH KI. DIVERGENCE?
Because KL divergence does not satisfy the triangle inequality, maximizing the replaced WSEP objective could jeopardize the diversity of learned skills as shown in appendix G.8. Appendix G.4 also discussed why KL divergence can not replace the Wasserstein distance in PWSEP.
B.5 IS IT PRACTICAL TO TREAT THE DISENTANGLEMENT OF EVERY SKILL EQUALLY?
Yes, theoretically, as the proof of theorem 3.2 in appendix C.3 shows, the contribution of IC2 to WAC be treated equally for each skill under assumptions:
1. The downstream task favored state distribution is far from the average state distribution p(S).
2. The state space is large.
Both could commonly happen in nigh-dimensional practical environments. When the contribution of IC, to WAC can be treated equally for each skill 2, the disentanglement metric I(S; 1,) for each skill be considered equally important.
Practically, existing unsupervised skill learning algorithms such as He et al. (2022); Durugkar et al. (2021) treat the desired properties of every skill equally.
15

---

Published as a conference paper at ICLR 2024
B.6 CAN ONE INFER THE SHAPE OF C OR PERFORM MAP ESTIMATION TO GET THE EXACT WEIGHT OF EACH 2C2 FOR WAC?
It is unnecessary and could be infeasible. As mentioned, is practical to treat IC2 of each skill equally important for WAC. By Maximum A Posteriori (MAP) estimation, one still needs a prior of the polytope shape, then update the belief of the shape by the samples of state distributions instead of just state samples. One needs to know which kind of state distribution is feasible for this particular MDP instead of just exploring the state space.
Undiscovered vertices are not within the convex hull of discovered skills, so you can not combine the discovered skills get an undiscovered state distribution. You will have zero samples the KL divergence between learned skills and the undiscovered state distributions, on which the MAP update depends.
It might be possible for the specific setting where you can infer the possible state distributions by the explored states. However, the purpose this estimation is just to obtain a weight about treating some skills more important than others. which can be too expensive for practical algorithm design, and bias or lag in estimation could make it worse than just treating every skill equally.
DOES (S; 1.) AND LSEPIN AFFECT DOWNSTREAM TASK PERFORMANCE WHEN ((;; Z) B.7
IS NOT MAXIMIZED?
Yes, empirically, the importance of diversity and separability of skills are mentioned in many prior works (Eysenbach et al., 2019; Laskin al., 2022; He al., 2022; Park et al., 2022b), disentanglement metrics like [(S; 1.) and LSEPIN can explicitly measure these properties. Theoretically, it is discussed in appendix C.4.
C PROOF FOR THEORETICAL RESULTS IN SECTION 3.2
C.1 PROOF FOR THEOREM 3.1
Theorem 3.1. When learned skill sets i;,i=1
= 1,2, with N S [S] skills (N skills have p(2) > 0) sharing the same skill are all MISL solutions, The skill set with the higher T(S; 12) will have higher p(z) and lower adaptation cost for all vz in the set Rz, where R.z is the set of downstream tasks always satisfying Vi, Vr E R,, z = arg eax DKL(P(S122) II p²*). And the maximum of this adaptation cost has the following formulation:
C2(r) - pp)))(())==(1)
IC₂ max rER;
(5)
1 1- ((p()
where
C+(r) I(S: DRL (p(S) p))
(6)
D2(r)= = DRL (p(()())=) r").
Proof. First, is not necessary for the skill sets to be all MISL solutions. The proof of this theorem only needs the skill sets to satisfy these necessary conditions:
1. All Z share the same average state distribution p(S) = Exrez, (x(S2z))
2. DKL. (p((=²))) || ((S)) is constant for all in all Z,
3. All Z, share the same skill z
In this proof, skill in set 2,.
to denote the skill shared among the skill sets and o' e Z; denote general
A MISL solution is skill distribution p(Z) with N skills having positive probabilities. Different solutions have the same p(S) = Ev'z,,(S) Because for MISL solutions, every learned skill z' E Z, should satisfy p(S]=') arg maxpec DRI.(P Il p(S)). It like the skill-conditione state distributions of MISL. solutions are all on the "largest circle" that is centered in p(S), with the
16

---

Published as a conference paper at ICLR 2024
"radius" being maXpee (P || p(S)) (see Section 6.2 of Eysenbach al. (2022)). Therefore, MISL solutions sharing the same 2 satisfy the necessary conditions.
The proof of Lemma 6.3 in Eysenbach et al. (2022) has shown that average state marginal p(S) is
a vector in [S] -iimensiona space with S - degrees of freedom. MISL can recover at most |S]
unique skills because every learned skill z' should satisfy p(S]=') arg max,ec (p (p p(S)).
thus more than [8] unique skills put more than |S] distance constraints p(5), making it overly specified and ill-defined Using a fixed number of skills common in practice, so we use s [S] skills for the MISL problem.
Because we do not consider the assumption in Eysenbach et al. (2022) restricting at most [S] skills on the maximum "circle" solve MISL, the could be more than |8] skills the maximum "circle", so N can be |S] and different Z can all have [S] skills.
The sketch of the proof goes like the following:
1. The sum of KI. divergences from learned skills to the optimal target distribution for a downstream task is constant.
2. The upper bound of IC₂ depends this constant and p(z), and the upper bound decreases monotonically with higher p(z).
3. Increasing I(S; 1,) results in higher p(z) thus lower upper bound of IC2.
We that for solution 2, the weighted sum of the KL divergence between the optimal vertex of considered downstream task r E R2 and the MISL learned skills is:
Σ (p p')) z'EZ,
(19)
((())))))))))))))) flog p(S|z') p"
Σ z'EZ,
Σ Σ P(S | p(S
(20)
p
z'EZ,
-H((SZZ) + H(p(S),p'))
(21)
+ H(S) H
(22)
+ p'),
(23)
where H(p(S),p") denotes the cross-entropy of the p" relative p(S).
Because I(S; Z) = Ez'ez, (PKL( p(S)) and p(S) are constant under the necessary conditions and 2, the sum only depends on E R. and can consider
Σ P()) I1) (r), Vi FEEZ
(24)
Because the minimum is less than the mean, we have
(1 p(2))
min
DRL( 11'))
2'2i''
Σ (k(()))( 11))) Vi
(25)
2'E
then
min DKI.(((S12)) [1 C-(r) - DK) (p(S)2) p²) Vi #'2
(26)
1 p(z)
We have assumed = = arg max, DKI.(p(S)) p²), so arg min,. DKA. (p((=))))=)) p²) is chosen from z' E 2,, 2' # 2,80
min (KL(((S)2)) || pp) < C2(r) p(=)
Vi
(27)
1-p(2)
17

---

Published as a conference paper at ICLR 2024
and can bound min.' DKL (p((22) || p²) by eq. (27) and obtain the formulation of eq. (5).
For any given downstream task r E Rg, different MISL solutions 2, have the same C2(r) and (2(r) but different p(z), so next show how higher [[S; 1.) affects p(z) and then results in tighter upper bound of ming (RL(p(S)22) p").
Because
([S:1)) =p) (p()) p()))
(28)
+ (1 - (x( & 2) || (()))
and
P(S * p(S) -(())((()))
(29)
1 --p(z)
for the same skill 2, the value of /(S; 1g) could only be changed by p(z). 01(8:1.) 0p(z)
is non-negative,
a higher [(S; 1,) would indicate a higher p(z).
OI();1)) (1 - p(=)) 0 (p p(S)
0p(2)
8p(z)
(30)
+DKL (p((22 || p(S)) DKL + 2) p(S))
By the necessary condition and by convexity KL divergence, because pl ZZ + 2) is a convex combination of skills not equal to we have
DKL p(S) 2 (kL(p(S)Z l P(S))
(31)
a1(S:1.) is non-negative if then 0p(z)
(pri( # 2) p(5)) 20.
(32)
0p(z)
By eq. (29), we have:
p' p'(S)Z + = -
(33)
1-)((2)
where p'(=)= = =p==)++++++++ A, and a small positive value that keeps the changed (Z) satisfying the necessary conditions described at the beginning of the proof. Then
2 + =) =>p ( + +1 -)p(S)
(34)
1 - p²()) p(=)
(35)
1 - p(2) p())
So 1. By convexity of KL divergence, we have
Dkl(p(SIZ # P(S))
(36)
=DkL(Dd (S|Z # #) + (1- A)p(S) ((S)
(37)
( +)) ||)) - (1
(38)
<Dii (SIZ # =) (())
(39)
and
@DkL(p(S) # x) p())
(40)
0p(z)
DKL(v'(S1 2 + 2) p(S))
lim 4
A DRL(p(S1Z 4 x) P(S)
(41)
A
20
(42)
18

---

Published as a conference paper at ICLR 2024
aT(S): Therefore, 0p(z)
is non-negative, higher I(S;1;) would indicate higher p(2).
With this result, we go back eq. (27). Because DKL(P(512) p*) is required to be max DrL (KI(p((S)') P"), vi, Vr E R2, we have
DKL(p(S|z) || p¹) D2 2 E (1') Vi, R₂ (43)
2'E2,
and because of eq. (43)
C. p())) I-p(z)
0p(z) 00, E R₂
(44)
Now we have proved that a higher 1(5;1g) indicates a higher p(z), thus a tighter upper bound in eq. (27) for all R2, including the worst-case defined in eq. (5). The theorem is then proved.
C.2 PROOF OF COROLLARY 3.1.1
Corollary 3.1.1. When the MISL objective I(S, Z) maximized by N S |S] skills, WAC is bounded a solution& 2" by
C2(r) (p)))
WAC S max IC2 = max max ZEZ PER.
(8)
zEZ
1 p()(2)
WAC the worst-cuse adaptation cost defined definition 3.1, C2 and D2 are as defined in eqs. (6) and (7), R2 here needs satisfy Vr E R2, arg max DKL. (p((zz))))) v²).
Proof. The formulation can be directly obtained from eq. (27) and definition 3.1.
The remaining problem here is whether Uzrez- R2 contains all possible downstream tasks.
If there exists downstream task not in UKez- R2, then this should satisfy that arg = DKL (p(S)z') || p'*) is not in Z", which is contradictory. Therefore Urez Ra contains all possible downstream tasks and this corollary proved.
C.3 DIRECT CORRELATION BETWEEN WAC AND LSEPIN
By theorem 1 we know that increasing (S; 1.) tightens the upper bound of the adaptation cost IC, associated with an individual skill z. However, how much IC, an individual 2 contributes to the overall WAC depends the specific C₂ and D, of each skill 2, and C₂ and D2 cannot be known in prior. Therefore, / (S; 12) of certain skills could be more important than others for WAC, which can only be known when already discovered all vertices of the polytope C.
Without prior knowledge, we could only treat every IC= equally and assume that the cost associated with individual skills would have the same importance on the WAC. The following results provide assumptions that are sufficient to treat the cost of individual skills equally and directly associate LSSPIN with WAC under these assumptions, which require:
The optimal state distribution for the downstream task to be far away from p(S) (eq. (46))
The state space is large (eq. (45))
Both requirements are common in practice.
Lemma C.1. When the MISL objective [(S; Z) is maximized by N [S] skills, the solution set of skills is Z", assume that Vz E Z",
Dm min DKL. (p((=))))) II p²)
(45)
FER,
C = (KI((p(S) || p²), Vr R.
(46)
where C and Dm are two constants, R. the set of reward functions that satisfy Vr E R., arg max (xa((((S)z)) p²*). Then, higher LSEPIN results in lower WAC.
19

---

Published as a conference paper at ICLR 2024
Proof. We denote
D2(r) DKL(O 11)")
(47)
C2(r) = (KL('))LL(((SEE))
11)) lp))
(48)
First, by eqs. (23) and (46), we have
C2(r) -1 +(((s)))pp)
(49)
HH() -((()Z) + (KL( ll pr)
(50)
+[( (KL(((S))))) =[( +C= C.
(51)
(52)
Om is constant because for MISL solutions, [[S; Z) is maximized and constant.
Because 0 S p(z) < 1, we have, 04 v p(=) - 1:
(2(r) Dx arg rR,,
Cm -(==)D(rr) arg min D=(r)
arg max rER,
(53)
1 (2)(z>
1 1-p(2)
rER.
Combine eqs. (8), (45) and (53), we have
Cm pmz)
WAC S max .E2
(54)
1- p(2)
Because C. and Dm is constant across different 2, and Dm 2 C. by eqs. (45), (46) and (48), increasing mi p(z) decreases right hand side of eq. (54). By eqs. (28) and (29) and non- negativity of @1(S:1a] increasing LSEPIN could lead to higher min p(z) thus lower WAC. M 8p(z)
Theorem 3.2. When I. the optimal state distribution for the downstream task is far from p(S) and 2. The state space large, i.e. (S| is large. IC, of all learned skills can be considered equally contribute to WAC.
Proof. Lemma C.I shows the direct correlation between WAC and LSEPIN under assumptions in eq. (46) and eq. (45).
C' in Eq. 46 could assume the feasible optimal state distribution for downstream task rg to have maximum "distance" from the average distribution p(S). It is common in practice that the optimal state distributions for downstream tasks are "far" from the average state distribution p(S) MISL solutions. Eq. 45 assumes constant Dm this can apply to practical situations when [S] is large.
By definition of R3 in theorem 3.1, we can see
min DKL (x((s())))l = max max D (LL( (m) rERr
(55)
rER.
let Z' denote the set containing learned skills that maximizes max f maX (kr((((5=)'ll=")
and z'' + z. By eq. (48), we have
Cm
Dm(2) E p(2) + Elez)pl2)
(56)
Because eq. (46) assumed y" to be on a "maximum circle" centered in p(S). p" should have |S| - degrees of freedom, where r" solves min,eR, Dx (p()))] p²). Including Including z, there should be |S]
skills that satisfies DKL (p(ss'))))
(m) determine p" and Dm which together have |S] 1 degrees of freedom. So we have Iz' = |S| 2 excluding and Cm
Cm
(=) < p(z)+ E p(2))
(57)
EETEzzZZ(()) P(2")
Without loss of generality, we consider situations when 1 = 188, when 2 Z' (z) contains only one skill. We denote it z", then:
Cm
(m < 1 1- p(=")
(58)
20

---

Published as a conference paper at ICLR 2024
Because optimizing LSEPIN increases p(z) for all 2 the larger [8] is, the tighter the bound in eq. (58) is. And because Dm (z) 2 Cm by definition, when |S] is large, all (2) can considered constant close to C...
Therefore, combining the above analysis on the assumptions of eqs. (45) and (46) and lemma C.I, we proved the claim of theorem 3.2.
C.4 How 1(); AFFECTS THE LEARNED SKILLS WHEN MISL NOT OPTIMIZED
As mentioned by the proof of theorem 3.1 in appendix C.I, , the claim that higher 7(S; 1,) leads to lower _C2 applies also to any skills sets satisfying the necessary conditions:
1. All 2; share a common skill z and the same average state distribution p(S) =
Ez [p()
2. (pL(2())2 || p(S)) is constant for all a' in all Z,
In practice, there could be a gap between optimal maxp(2) Ep(z) DKL suboptimal d,n, max; Z DKL (pS)2) || Ezez, (p(S) [p(S]= s)))))) where 2, the skill sets that be learned. This gap could be a result of limitation in the skill number [ or of an optimization error.
( (p(S]z) p(S) and learned
In order to maximize 2) = Ep(z) DKL. (p(()) P(S)) The learned pl Z) would be positive
only for the skills that have (KL(p(S12) || p(S)) = dm- For two skill sets Z1, Z2 maximizing T(S; 7) the same suboptimal value dmn they satisfy the necessary condition 2 since every skill has (p((2(S)2)) p(S)) = dm-
Exploration is important when I(5; Z) is suboptimal empirically shown in Laskin al. (2021), If 2,, Z2 satisfy the necessary condition and share the same p(S), they would have the same exploration of the state space.
Therefore, our theoretical results in theorems 3.1 and 3.2 also apply to practical situations when MISI is not optimally solved, and they showed that when two sets have the same exploration while maximizing [ (S; Z) to the same degree, the set with better LSEPIN is preferred for downstream task adaptation.
D EXAMPLE WHEN MISL WITH UNIFORM ((Z) IS NOT GUARANTEED TO PROMOTE DIVERSITY
(5522
p(5iz
p(S
Figure Example MISL with uniform p(Z) does not promote diversity
For a case of with |S| = 3. the feasible polytope allows a maximum "circle" centered at [0.2, 0.4, 0.4] with maximum "radius" (I(S, Z)) of 0.253 this fig in shows. There are 5 ver- tices U1.5 of C, of which v1 V2 14, Uy lie on the maximum "circle". Vertex e and vertex 14 are [0.2, 0.7, 0.1] and [0.2, 0.1,
21

---

Published as a conference paper at ICLR 2024
Because 8 is shared by all skills and {0, Zimpue), we can assume p(Z) p Zmpput)
When there are 4 skills to learn, Uniform pl Zinput) should have p(zixput) = p(z) = 0.25 for all 2. For both uniform p(2) and maximization of I [S;; Z), the optimal skill set {21, 22, 23, 21 should be the one containing two skills 21, 22 at 112 and the other two skills 23, 24 at Da, as shown in the figure. Only with this skill set, the uniform average of skills p(5) E. ((z)p(S)s) could be the center of the maximum "circle" [0.2, 0.4, 0.4].
We see from this example that although skill set {21, 22, 23. 24) maximizes both [(S; Z) and H (Zimpux). ), there are two pairs of skills being not separable thus resulting in limited diversity
In appendix F, show higher ((; 1.) explicitly increases DKL (p(S) * =)([p(S|=)) thus pro- moting pl () } to be separable from p(5 2 + 2) (average state distribution of skills other than 2). Therefore, LSEPIN promotes diverse skills explicitly. Comparing skills only at 0j and v4, MISL with LSEPIN would favor skills at 11 U2, U4, respectively. In order to keep p(S) to be the center of maximum "circle" [0.2, 0.4, 0.4], p(Z) of skills U1, U2 04. V/ not necessarily uniform.
E UNIQUENESS OF THE "CIRCLE" LEARNED FROM MISL
Maximizing [(; Z) leads to unique average state distribution PSS), and it can be proven by the following:
1. By Lemma 6.5 of Eysenbach al. (2022), MISL objective is equivalent to milp(()) max, (pt(((S)z) p(S)).
2. Because of the strict convexity of KL divergence, max; (p((2) || p(S) is also strictly convex.
3. p(S) is constrained in a convex polytope C.
Therefore it is convex optimization problem with a strictly convex objective function. The solution is unique. With unique p(S) be the "center", the "circle" of maximum "radius" (pt(p(()=)) ll p(S) is also unique.
F DIFFERENCE BETWEEN DISENTANGLEMENT FOR REPRESENTATION LEARNING AND SKILL LEARNING SETTING
Kim et al. (2021) adopted metrics SEPIN @k and WSEPIN (Do & Tran, 2019a) from representation learning to measure the disentanglement of the skill latent. They focused on the informativeness and separability (independencies) between different dimensions of the representation.
SEPIN@K is the top average of I(S Z Zys) while WSEPIN is the average of I(S, Z Zy. weighed on (S, Z:), where Z, denotes ith dimension of the skill latent.
However, by [SS Z Z411 they care about how to efficiently represent the skill with the latent representation rather than promoting the separability and diversity of the trajectories induced by different skills. It can be maximized by only optimization in the latent space regardless of the state distributions in the state space. Therefore, this metric does not directly encourage diversity and separability of state distributions of skills thus not contributing to adaptation costs.
If we directly adopt T(S, Zz Z for our setting, Z should be one-hot vector. We have
[(S, Z,IZ,)) (S))) H Z,, 2zz),
Because Z is one-hot vector, Z; is completely dependent Z i.c. knowing Zzi also means knowing Z, Therefore, HS) Z+:)= = (S) (+:) and I(S, Z, (2x)) )) is constantly zero and we can not directly the disentanglement metrics composed of T(S, Z,(Z+)) to explicitly measure the diversity and separability of learned skills.
Instead, we propose use I(S;1, to measure the informativeness. and separability of each individual skill 2, where 1, is the binary indicator for Z = 2. It is a mutual information metric intrinsically measuring informativeness. We will show that it also measures separability.
22

---

Published as a conference paper at ICLR 2024
In the URL skill learning setting, separability between one skill and others entails that the states inferred by this skill should share little overlap with the states inferred by others. Therefore, means that the states should be certain to be inferred by one single skill, thus meaning low (11][S).
Because (S;; 15) = H(S) - H(]]1,), for situations with the same state and skill distributions, the with larger [(5; 1,) would be lower in 15) since (S] ( 1,) H(1,) HS)),
lower H 1,) means also low in (1gg5) thus better separabilityy.
Moreover, larger [(S;1.) 1.) explicity encourages larger DKr. (p(())(SS))))) + 2) p(S[2)).
By the proof of eq. (40) in Appendix B.I, we have shown when [S; Z) is fixed a higher (S; 1g)
results in higher DKL (p + z) p(S)).
By convexity of KL divergence:
DRL(p(S) # (S)) S (x)
(59)
- (o |l p(S) 2))
- p(=))kr.(v((112 (p)) [|)
||
p(S)=) +0
(60)
(61)
So increasing / (S; 1g) also increases the lower bound of (pL(p * 2) p(S]z)), thus explicitly encouraging skills to be separated from each other. Therefore, we compose metrics from [(S; 1,) to measure the disentanglement of learned skills.
G PROOF FOR SKILL LEARNING WITH WASSERSTEIN DISTANCES
When using KL divergence, solving eq. (2) learns skills that lie at the vertices of the polytope. This actually because of the strict convexity KL divergence not mentioned by prior work (Eysenbach et al., 2022), which leads to
(AI, + (1 A)p2 |1 p(S))
(62)
(PI |1 pS)) (1 DLL (P2 p(S)),
where 0 < A < 1 and pl * p2. Therefore, the skills at the vertices of the convex polytope always have higher KL divergence from the average state distribution than other skills within the polytope or on the faces of the polytope, because those skills are just convex combinations of the skills at the vertices.
In general, Wasserstein distance (-.pp) is commonly considered to be convex (Santambrogio, 2015; Peyré & Cuturi, 2019). For the 2-Wasserstein distance, its square ( p) considered strictly convex by (2.12) of Carlen & Gangbo (2003). In practice, for image observations, similar to Liu & Abbeel (202 that apply L2 norm-based particle based-entropy (Singh et al., 2003b) in the representation space, can also apply the 2-Wasserstein distance-based metrics to the representations of observations. So is practical to consider strictly convex Wasserstein distances for analysis.
G.1 ANALYSIS ON THE BASIC WDSL OBJECTIVE
We can trivially replace the KL divergences in the MISL objecttve eq. (2) with Wasserstein distance and obtain basic WDSL objective
max E (p)) [W (pp()))))p(5)))
(63)
p(*)
Similar to Lemma 6.5 of Eysenbach et al. (2022), optimizing eq. (9) is equivalent to solving
min max W (p(S)=),p(S)).
(64)
P(S) p(=)
We name it Average Wasserstein skill learning Distance (AWD). Therefore, when the adaptation cost for downstream tasks is defined as the Wasserstein distance between the state distributions of the initialization policy and the optimal policy since maximizing AWD learns an average state distribution that is close worst case policy, this average state distribution will be worst-case robust initialization.
We have the following Lemma:
23

---

Published as a conference paper at ICLR 2024
Lemma G.1. When AWD is maximized, all learned skills with p(z) 0 must lie at the vertices of the polytope.
Proof. If is vertex of the polytope, it can be represented by strictly convex combination of some vertices, and for any skill not at a vertex,
((((s))((S1)) < (
& max ((p(5), s'EV=
(6.5)
Z'EV:
where V. is the smallest set of vertices such that & is within the convex hull of them, namely p(S]z) Ex'ev, A='p(S)=') 0 =').
X for all a' E V/ and Σv'yy A=' = The first inequality by strict convexity, and the second is because the maximum should be not less than the mean.
We can see when p[S]2)
one of the vertices, there must be a vertex with higher distance to p(S). Therefore, just like MISL, WDSL with AWD also discovers skills at the vertices.
Although optimizing AWD discovers vertices like MISL, it also suffers from the limitation of re- mark 3.2.1 because it also only discovers vertices the maximum "circle" with maximum Wasserstein distance to the average.
G.2 PROOF FOR LEMMA 3.3
Lemma 3.3. When WSEP is maximized, all learned skills with p(z) > 0 must lie at the vertices of the polytope.
Proof. Similar the proof lemma G.1, for any skill 2,, when p(S]) is not the vertices, it can be represented by strictly convex combination of some vertices, so for all z; E 2,
Σ
W(1(S)))).
< Σ A
Σ ((((s))),
(=E))(=))
''EV., aJEZ((#)
(66)
max ''EV
Σ
W(p(S)<).
zJEZ((z))
where Vz, is the smallest set of vertices such that is within the convex hull of them, namely p(S]z,) = Evev, )x 0 < X < 1 for all x'' E V2, and Σvev. Ag' = 1.
When 2() is not at of the vertices, there must be vertex with higher (x,)2((=) W ((((S)=) for all z; E Z. Therefore, WSEP is not maximized when there is still skill not at the vertices.
G.3 PROOF FOR THEOREM 3.4
Theorem 3.4. When WSEP is maximized by 1Z"| skills, the MAC can be upper-bounded:
MAC < E3Ez- L2 (1Z*) 1)WSEP
(12)
V 2-
MAC < ExER. by (12*) 1)WSEP
(13)
where
13 Σ W(p(S)o).P(S1a)
VEV
(14)
Ly = max ev
W W(p($1o),<(S1b')
NEV
Proof. By definition,
MAC < Evevi Essz W p((S)))((((()))
(67)
V
2
24

---

Published as a conference paper at ICLR 2024
And
Σ M D(((5)))(((13)
2
Σ 1i ΣΣ ((b(Skk). SE2
(68)
ZEZ
Σ Li -SEPP
ZEZ
where the inequality in eq. (67) is by convexity.
Then the first inequality eq. (12) is proved by combining eq. (67) and eq. (68). Because by definition, for all 2 E Z. Ly S Ly, the second inequality eq. (13) also holds.
G.4 PROOF FOR THEOREM 3.5 AND THE VERTEX DISCOVERY PROBLEM
G.4.1 PROOF AND ALGORITHM FOR THEOREM 3.5
Theorem 3.5. When V is the set of all vertices of the feasible state distribution polytope C, all [|] vertices can be discovered by [V] iterations of maximizing
PWSEP(1): min W (p(S)zs),
Σ Xp(S1=))). S,
(15)
where 2, is the set of skills discovered from iteration 0 and zi is the skill being learned at ith iteration. X is convex coeffcient of dimension i - that every element X 2 0, E (0, 1,.,, and (j
=1
In the initial iteration when 2; = 0, PWSEP(0) can be ( Po). ((())rand)) with Zrand to be a randomly initialized skill.
Proof. We first see that for the initial iteration, by eq. (65), maximizing W (pp(S)20), ((S) (rand)) learns a 2g at one of the vertices.
Then in later iterations when 2, is not empty, we need to prove that maximizing PWSEP()) can learn a new vertex that was not discovered before.
For every skill 2/ not at the vertices, it can be represented by strict convex combination of p(S] = Every An (Sp()) where V=g is the smallest set of vertices such that z; within the convex hull of them, 0 < X < for all v E Ve and Σvevy X 1.
We have
min W (p(S)
Σ ND(S1zz)) min W
Σ {pp(S(o),
Σ (p((S1z))
2jEZ,
VEV.
(69)
S w
Σ
Σ X.p(S1z))
"EVi
'EV o4
where
A arg W (p(Slv),
Σ Np(S1zz))
(70)
Z,EZ,
Furthermore, we can see
W Σ P.P(S)
Σ
x Σ x(p(S1=)) < Σ
W. (p(S)c),
Σ Mo(S1z))) (71)
NEV 4
VEV. EE
VEV
= Σ
min W (p(S)o), Σ \p(S1zz) X
A
NEV #/
2,
(72)
25

---

Published as a conference paper at ICLR 2024
Now we have shown, for every skill 2f not at the vertices, when the maximum vertex is unique,
min W (p(s)=), Σ ND((11)) <
Σ
X min W (p(Slo),
Σ Nn(51z))
(73)
8,EZ,
vev.
s,E2,
max min W (p(S)o), V
Σ xp((1zz))
(74)
2JEZ.
Equation (74) strict when the vertex to maximize the projected distance unique. For extreme cases where face of the convex hull of Z, and a face of the polytope C are "parallel", there could be multiple vertices that maximize PWSEP()) and make it possible for all points on the "parallel" face of C. also to be optimal. However, this can be easily circumvented in practice by adding the skill being learned 2 into the Z, set for a temporary set of 2 and optimizing the projected distance based on 2 instead of 2, once 2/ reaches the optimal "parallel" face of C. By doing so, 2 contains a point z\ on the "parallel" face of C, and o will have O projected distance to the convex hull of the new temporary set 2/, so the points the "parallel" face of C will not be all optimal for 2/. Then, have excluded an extreme case of non-verter maxima and we can continue to update 2! to a vertex without reinitializing a new policy.
Therefore, maximizing PWSEP(i) in every iteration learns skill that lies at a vertex. The remaining part is to prove that the vertex discovered after every new iteration will be new and different from the vertices discovered earlier.
If the newly discovered vertex 2, is not new and has been previously discovered, it should be in Z, and PWSEP (i) would be O.
undiscovered new vertex v at iteration i should not be in the convex hull of 2,: If the convex hull of Z; contains v, the vertex u should be either one of the points in set 2, or a strict convex combination of points in Z, Since the polytope C is convex, its vertices cannot be strict convex combination of other points in C, the vertex can only be in the set of Z; as already discovered vertex if it is in the convex hull of Z,
Therefore, when the set 2, does not contain 0t, on should not be in the convex hull of 2,, resulting
min W ((S).)), Σ Xx((11z)) A
(75)
2,E2,
Then, if there still exists undiscovered new vertices, optimal PWSEP(i) would be larger than 0, and optimizing PWSEP(i) finds a new vertex at each new iteration.
In conclusion, PWSEP discovers a new vertex at each new iteration, so it can discover all | vertices with | iterations of policy learning.
Other statistical distances like total variation distance and Heilinger distance are also true metrics satisfying symmetry, and triangle inequality, but they are not strictly convex. KL divergence is strictly convex, and the information projection p* arg minpe DKi((p || q) has unique solution for the convex P set, so it seems that the Wasserstein distance projection in the PWSEP algorithm can be replaced by information projection. however, when the new skill z; does not cover all the states in the support of distributions in Z,, which could happen a lot in practice. this projected KL divergence would be infinite and no longer convex in the union of their supports. For the moment projection, which is the reverse information projection p* arg min pep Dx p), it would also be infinite whenever the new skill 2; discovers new states that are not the support of distributions in 2,.
minpgp
Notice that, the sum of the projected Wasserstein distances
SPWD = Σ min W (p(Skmm)
E (²p(S(zn))
(76)
S.E2
EZ
is not an evaluation metric for skill learning. This means that a higher value of SPWD does not mean
better skills for downstream task adaptation. An extreme example is that for [S] = 3 and the polytope C is a Chiliagon with [V] = 1000 vertices. Then the optimal skill set covering all 1000 vertices have SPWD near zero For another skill set with 999 skills at the same vertex and 1 skill at another
26

---

Published as a conference paper at ICLR 2024
vertex with maximum distance from the prior one, then it is possible that this set could have greater SPWD than the optimal skill set and this skill set is clearly not optimal since it only covers 2 vertices.
Therefore, unlike [S; Z), LSEPIN, and WSEP, SPWD is not an evaluation metric and can not be used as an optimization objective for unsupervised skill learning. To discover all vertices, the projected Wasserstein distance needs to be optimized iteratively the PWSEP algorithm.
Algorithm 1 PWSEP algorithm
Initialize a random point Trand, a real value v >> 0
TT. + arg max pTp(Sl))).pFrang
Initial iteration
II, < {
while v>0 do i ++i+1
Ts, - arg maxx,, PWSEP(i)
b (15)
11 U
V f- maxx, PWSEP(1)
end while
Algorithm is the algorithm described in theorem 3.5, its practical implementation discussed in appendix H.3 and its empirical performance is discussed in appendix I.
G.4.2 ABOUT VERTEX DISCOVERY AND SUCCESSOR FEATURE METHOD
To the best of our knowledge, the first method to solve the vertex discovery problem with a finite number of skills is the successor feature-based method SFOLS (Alegre et al., 2022).
SFOLS is designed to find a CCS set defined in (2.) from Appendix.A.5 in their paper. When written in our notation, it would be
CCS [Ew, s.t.p" = 2 p w=
(77)
where pT is the occupancy measure of skill is equivalent to the reward function of states This definition of CCS also applies to any skill that lies at the edges or faces of the polytope, not necessarily at the vertices; therefore, SFOLS might need to learn skill number that is much larger than the vertex number [V] before discovering all vertices. Even for their more rigorous definition in F, it still
their (6), where they constrained the skills be in the non-dominated multi-objective does not exclude skills at the edges or faces, because vertices do not Pareto dominate the points at one of its edges, i.c. there always exist tasks on which the vertices have fewer accumulative rewards than points at their edges.
Therefore, SFOLS might need a number of skills |Z] greater than the number of vertices | to discover all vertices and solve the vertex discovery problem, while PWSEP is guaranteed to discover all |V] vertices with only [V] skills.
To contain only vertices, set modified from CCS should be:
(p J.W. | 1, its elements are linearly independent and Vu E W,p" 2 w
(78)
vertex should be optimal for any task that has optimal solutions on its adjacent edges or faces. so a vertex should be optimal for at least [S] - linearly independent This could inspire future successor feature algorithms to learn more general policies that lie the vertices instead of faces of polytope C.
Moreover, there are fundamental differences between our setting of skill discovery and the setting of successor feature methods. For unsupervised skill learning with MISL or Wasserstein distance, the latent representation for skills is only used for indexing. its exact value does not affect the state distributions of the learned skills; While for the successor feature setting, the value of the weight directly determines its associated state distribution, which is the reason why optimization in the weight space is required for the successor feature methods.
27

---

Published as a conference paper at ICLR 2024
That is say, for an ideal unsupervised skill learning algorithm like PWSEP, it should be able to discover V vertices with a skill space containing only [V] points. This is impossible for successor feature learning that needs a weight (used as skills) space RI (Appendix.A.5 of Alegre et al. (2022)) containing infinite points.
However, successor feature methods can be considered directly solving eq. (2) instead of eq. (1). Be- cause each weight determines optimal parameter ((w) for its optimal policy To(w)- Equation (1) for successor feature setting becomes:
(S; W) By(w) (iL (S/w, 8) p(S.))) (,P(W)
By(w) (Rr))/RI((E (p (ww) || (pw))p(SS ((w))))
(79)
By(w)
p(S))
Therefore, methods that combine successor features with MISL such as Hansen al. (2020); Liu Abbeel (2021b) are the practical algorithms closer the setting (optimizing eq. (2)) of our theoretical analysis.
G.S EXAMPLE WHEN HIGHER WSEP RESULTS IN HIGHIER MAC
In fig. 5. It shows polytope with 4 vertices in 3-dimensional space. Suppose in a situation where two skills are learned, and the learned skills vertices and This situation would be optimal for maximizing WSEP, its MAC is the mean of the length of edge ed and edge cb, and its worst-case adaptation cost the length of edge situation where two skills lic at vertices and d, it has lower WSEP, but its MAC is the mean of the length of edge ed and edge db, which is lower than the former situation. Its worst-case adaptation cost is the length of edge od, which is also lower. this case, LÝ < Lÿ, so from skill combination and to skill combination and d, although WSEP decreased, Σ,ez Lý SEP also decreased due to the change Lý, resulting in tighter upper bound in eq. (12).
Then for another
Figure 5: polytope in a 3-dimensional space, there are 4 vertices.
For a practical example, this situation can be an environment with a state space of one-dimensional integers. If you have two skills, for maximum WSEP, they would put their probability mass close to -00 and oc which could be infinitely far away from any sampled downstream task optimal distribution.
G.6 AN MDP EXAMPLE WHEN WSEP DISCOVERS MORE VERTICES THAN MISL
Consider an MDP with S = 3 for states and |A| 3 for actions.
The transition matrix for action a1 is:
0.2
0.7
0.1
0.2
0.7
0.1
0.2
0.7
0.1
The transition matrix for action d2 is:
0.2
0.1
0.7
0.2
0.1
0.7
0.2
0.1
0.7
The transition matrix for action ag is:
0.4
0.3
0.3
0.4
0.3
0.4
0.3
0.3
Because the transition probabilities only depend on actions, state distribution p(S) is determined by distribution p( p(A):
28

---

Published as a conference paper at ICLR 2024
0.3,
(80)
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 C with these three probabilities being the vertices C.
The polytope C produced by this MDP can be shown the figure below:
P(S)
Figure 6: Example when WSEP discovers more vertices than MISL
We label the vertices with 01, 12, Us as shown in fig. 6,
p())e = (0.2,0.7,00 0.1 u2) [0.2,
(81)
p(S)ca) = |0.4,0.3,0.3
G.6.1 MISL ONLY LEARNS SKILLS AT
Firstly, want to show that the p(S) maximizing I) is uniquely determined by the MDP in appendix E. For this specific MDP, the optimal p(S) = [0.2, 0.4, 0.4). prove p(S) is the solution, start from the equivalent MISI objective mentioned in appendix E:
min max (p((2) P(S)
(82)
((S)
By strict convexity, for p(S), the solution x' of max, DKL (p(S)=)) p(S)), must lie at one of the vertices, as shown by eq. (62), also by strict convexity, the value of max; DKL (P(S)2) || p(S) for all p(S) * [0.2, 0.4, 0.4] should be strictly larger than its first order Taylor expansion at [0.2, 0.4, 0.4]. The gradients of the KL. divergences w.r.t. p(S) [0.2, 0.4, 0.4] are:
91 = [-,,-7/11 7/4, - 1/4]
92 [1 -14,-1/4,,7/44
(83)
= -344,
We represent any other p(S) with p(S) A [0.2, 0.4, 0.4], where A can be represented by [-a a a. bj (p(S) should be within the probability simplex). By eq. (83), we get the gradients of the KL divergences w.r.t. [a, to be:
G = [-0.75, 0.75]
G2 [0.75, -0.75
(84)
G3 [1.25, 1.25]
Therefore, we can see when (a, b) * [0,0]:
max (pL((p(S)))) p(S))
(85)
max(0.253 0.75a 0.75b, 0.253 + 0.75a 0.75b, 0.105 1.25a 1.25b)
29

---

Published as a conference paper at ICLR 2024
where 0.253 is KI divergence from in 02 to (0.2, and 0.105 is Kl divergence from Uy to [0.2, 0.4, 0.4].
We denote: f(a b) ma f2(a,, = max(0..23 - 0.75a + 0.75b, 0.253 + 0.75a 0.75b, 0.105 + 1.25a + 1.25b)
(86)
We can see that changing a, will either increase f1 or fz. because when (a, b] # [0,0]:
max max(Dk(pp(Sl)) II p(S)) f(a, b) 2 max mar(fila, b), f2(a, b))
(87)
We can see any p(5) that is not 0.2, 0.4] ((a, b] * [0.0] would result in a larger max; DKL (p(5)2) ) p(S)). so [0.2, 0.4, 0.4] is the only optimal p(S) for maximizing [(S; Z).
MISL would learn only p( (S1122 v1 22 with 2(2)) (22) 0.5 to have the average 0.4,
state distribution pSS) (S(2) the unique center 00.2,0.4,00 0.4].
By this solution, [[5; 2) is maximized by (() - Σ. p(z) Dx (1))))p(55)) ~ 0.253.
Putting weight the other vertex 23 would lower T(S; Z). For example, if p(z)) p(22) 0.45, p(23) = 0.1, p(S = 03, 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 [(5; 2) 0.237.
G.6.2 WSEP CAN LEARN SKILLS AT ALL VERTICES UJ, U2, 03
or WSEP, when the skill number is set to |Z| 3. putting skills at 2 and 1 skill at would have lower WSEP than putting 3 skills at 21, 22, 23 respectively, because of the triangle inequality, suppose the transportation cost between every two states is 1, then
W(v(S)=)), 0.6
W(p(S\=)), = ((((( = 0.4
WSEP for skills at 2 and skill 22 would be 2.4, while WSEP for skills at 21, 22, 23 respectively would be 2.8.
Therefore, unlike MISL, WSEP would favor learning vertices instead of only 2.
G.7 WSEP CAN NOT ALWAYS DISCOVER ALL VERTICES
For |S| = 4 case, suppose the polytope C of feasible state distributions is 3-dimensional and there are 4 vertices {1i, U2, V3, }. and
Σ W(p(SW)),
<<
Σ
((((S))).P(SIo))).
(88)
DE (E
then if en is very close to u2, comparing to a point v' on the surface of v1, and very close to U1 it possible that
oE (2. p((
Σ
Σ
((1(S))).
(E(Uims)
(89)
Σ
W(p(ST)),
Σ
W(P(S))).((Iu)))
22
(E(U2)0))
The skill set of (U1, U2, Va, o') result in higher WSEP than {ui, 02 U3 v4), so maximizing WSEP can not always discover all vertices.
G.8 ABOUT WSEP PORMULATION WITH KL DIVERGENCES
Defined in remark 3.4.2
KLSEP
Σ
Σ
DEL (p(12s))) ((S(=;))
#,Z 2,E2,11j
30

---

Published as a conference paper at ICLR 2024
sum of
DkL(p(S1z) II p(S1z))) + (xL(p(S)z)) p(S)z))
Because KL divergence does not satisfy the triangle inequality, it possible that
(KI(p(S|z:) || p(S]z))) + DKr. (p((zz))) p(S)z)))
> >KLL(p(S1e:) | (())=)) + DKL =x) | p(5)=))
(90)
+ (p(( 24) =x)) DKL (p()) ((S(z))).
Then learning another skill close to 2, 2j could result in higher sum of these KL divergences than learning a new skill z that is far from z and 2j. resulting in poor diversity and separability.
H PRACTICAL ALGORITHM DESIGN
Although our main contributions are theoretical analyses, provide some methods here for approxi- mation the proposed metrics and show how they can be used for practical algorithm design.
Table 2 shows an overview of methods to approximate the metrics.
Table 2: Approximation of metrics
Metric
Non- Non-parametric
Parametric
I(S: Z)
PBE (Singh al., 2003a)
const. po (2|s)
LSEPIN
PBE
const. + (1pp(1.|50
WSEP
PWD (Rowland al., 2019) SWD (Kolouri al., 2019)
E[D (2. (D)
H.1 PRACTICAL ALGORITHM TO INCORPORATE LSEPIN
For practical algorithms, LSEPIN could play an important role since it measures the informativeness and separability of individual skills and it encourages every learned skill to be distinctive and poten- tially useful. is totally feasible to approximate LSEPIN with Paticle-Basee Entropy (PBE) (Singh al., 2003a) (details in appendix H.5) a parametric discriminator that takes state and skill input (s 2). This discriminator can be learned by:
max Ez>p(Z)
(1 - Do(s, 2)) [s log (s. 2)]]
(91)
Then LSEPIN can be estimated by:
min 1(5;1,) H(1,) H(1 #E2
(92)
constant + E.~p(51=) log (1,
(93)
constant + E.~p(()=)) log (ø z)]
(94)
Practical algorithms solve eq. (1) instead of eq. (2), so the desired state distribution of skills nd (mpu)) is determined by the both p( impur) and learned 0. In order to let every input skill Zinput have equal steps of roll-out, the input skill distribution pp (iiinpt)) is often set uniform prior (Eysenbach et al., 2019; Sharma et al., 2020). Therefore, ((1,) can considered as constant. Then the intrinsic reward to increase LSEPIN can be log ( 2) for states inferred by skill 2, and log(1 (s, 2)) for states not inferred by skill 2.
H.2 PRACTICAL ALGORITHM FOR WSEP
How to efficiently compute the Wasserstein Distances between states of different skills is fundamental for WSEP In He et al. (2022). approaches like Sliced Wasserstein Distances (SWD) (Kolouri et al., 2019) and Projected Wasserstein distance (PWD) (Rowland et al., 2019) have been demonstrated useful for unsupervised skill discovery. Here propose alternative approach inspire future
31

---

Published as a conference paper at ICLR 2024
algorithm design. We can parametrize test function that also takes state and skill as input Dz 2). and this test function is learned by:
min Eg By Bal Eapp(s (s, =)] E Ee [ =)]
(95)
+ XI (o.
1|2 1)²]
where the last is a Lagrange term that restricts the test function to be 1-Lipschitz with respect to and =)),:) is the distribution of interpolated samples between states inferred by skill and states not inferred by skill z. Test function Do learned by this loss can be used to estimate the dual form Peyré & Cuturi (2019) of the I-Wasserstein distance between states inferred not inferred by any skill 2. Then, pertaining skill by the intrinsic reward +2 (s,ss2)
2) Dø(s, z')] could maximize WSEP.
We that both eq. (91) and eq. (95) share similarities with the discriminator loss in generative adversarial networks (GANs) (Goodfellow et al., 2020) and Wasserstein GANs (Arjovsky et al., 2017). Therefore, it encouraging better separability between states inferred by different skills thus resulting in more distinctive skills.
H.3 PRACTICAL IMPLEMENTATION FOR PWSEP
The key to implementing algorithm for PWSEP is how efficiently compute the projection
PWSEP(1) = min W (1
Σ 2y
(96)
2,E2,
We can see that this problem is convex for A. The derivative OPWSEP(:) 02 sensitivity analysis of linear programming min W (DTh, (81=:)) (S)=)) Es,ez, pp (S1zz) However, with larger batch size B of (s, 2) samples, it could be infeasible to obtain this derivative from the implicit function theorem and KKT conditions like the approach from OptNet (Amos & Kolter, 2017), because each derivative calculation requires solving inverse or pseudo-inverse of a B-dimensional matrix. We have compared it with the derivative-free method CMA-ES (Hansen & Ostermeier, 2001), finding that although CMA-ES ES computes the Wasserstein distance multiple times for update of X, the total time cost of CMA-ES is much more efficient. Besides, because we have separated parameter 0g for each skill policy, and p T (S]=z) of other skills remain unchanged when updating 0, the optimal À after an 0; update would be close to the optimal À before the update. Therefore, the computation for À to be optimal could be large the beginning then it gets small.
can be approximated by
H.4 CHOICE OF TRANSPORTATION COST FOR WASSERSTEIN DISTANCE
Wasserstein distance provides stable and smooth measurements when the transportation cost can be defined meaningfully. For example, in environments with physical dynamics, the transportation cost between states could be defined as the physical distance, in certain discrete environments, it can be defined as the Manhattan distance. These are not only more intuitive but also provide smooth and stable measurements with meaningful information on how difficult it is for an agent to reach from one state to the other. However, the choice of them requires prior knowledge of the environment. For completely unsupervised transport cost design, our idea for future work is to learn representation so that the basic distances like L2 in the representation space can capture how difficult to travel from state to another, which could be related previous work about graph Laplacian representation (Wu et al., 2019). Their method uses spectral graph drawing (Koren, 2003) to learn a state representation so that dynamically consecutive states have small 1.2 distance each other in the representation space.
H.5 PRACTICAL CON-PARAMETRIC ESTIMATION OF LSEPIN
LSEPIN can be approximated by particle-based entropy (Singh al., 2003a; Liu & Abbeel, 2021b):
T(S: 1.) Apu (S) Hsp(S)1.)
(97)
32

---

Published as a conference paper at ICLR 2024
Where Hyp(-) is the particle- basele entropy that can be estimated as:
Hpp (S)
Σ log c
Σ
||s
(98)
(N))
where (-) denotes the ls nearest neighbors around particle, is constant for numerical stability (commonly fixed to 1).
The second term of Equation 97 is the state entropy conditioned on knowing whether skill equals and it defined as
H(S HSS = H(SIZ #).
(99)
The first term of Equation 99 can be estimated by sampling from the states generated with skill 2. The second term is a little more tricky for parametric methods that try to approximate log p(s)2), 2), but for particle- pertcle-based entropy, can still be conveniently estimated by sampling states from all skills not equal to 2.
I EMPIRICAL VALIDATION
We provide some empirical examples to validate the proposed theoretical results. For more intuition, we consider the skill trajectories in a visualized maze environment with continuous state and action spaces from Campos et al. (2020) and a higher dimensional Mujoco Ant environment (Todorov et al., 2012). Maze navigation and high-dimensional control are major test grounds for previous MISL methods (Eysenbach et al., 2019; Kim et al., 2021; Park al., 2022b).
1.1 CONTINUOUS MAZE EN VIRONMENT
To get an intuitive understanding of what kind of skills are encouraged by our proposed metrics, we first look at the trajectories in the maze environment.
(a) MISL
(b) MISL+PLSEPIN (c) parametric WSEP (d) MISL+ALLEPIN
(c) PWSEP
Figure 7: Trajectory examples: Trajectories of different skills in different colors. (a) is learned by typical MISL with state entropy for exploration (Lee et al., 2019; Liu & Abbeel, 2021b), while (b)(e)d) are proposed in appendix H our paper. (b) is learned by MISL with state entropy and parametric LSEPIN proposed in appendix H.1. (c) is learned by parametric WSEP proposed in appendix H.2. (d) is learned by MISI, with state entropy and non-parametrie LSEPIN estimated by pasde entropy (Singh et al., 2003a; Liu & Abbeel, 2021a). (e) is learned by PWSEP, the projection in eq. (15) is found by the CMA-ES algorithm (Hansen & Ostermeier, 2001) as described in appendix H.3.
The corresponding metrics for the agents in fig. 7 are listed in table 3. For stable estimation, the metrics are estimated by non-pare con-parmettric approaches using large number samples. Metrics related (I Z) and LSEPIN) are estimated by particle-basee entropy eq. (98) with n 4000, 2000. WSEP is estimated by sliced Wasserstein distance (Kolouri et al., 2019) using 4000 samples of (s, =). The magnitude of particle-based entropy depends on the number of samples n because it is a summation instead of a mean, while the magnitude of sliced Wasserstein distance is mean and it is less affected by the number of samples. As mentioned in appendix G.4, the sum of PWSEP()) in eq. (76) can not serve as evaluation metrics like LSEPIN and WSEP, so we consider to evaluate with T(S: Z), LSEPIN and WSEP.
to mutual information
33

---

Published as a conference paper at ICLR 2024
Table 3: Metrics for the agents in fig. 7.
I(S; Z)
LSEPIN
WSEP
(a) MISL
787.9
161.5
15.3
(b) MISL+pLSEPIN
1060.6
244.6
23.5
(c) Parametric WSEP
1090.0
286.7
30,9
(d) MISL+nLSEPIN
965.6
243.4
29.3
(e) PWSEP
942.3
193.2
31.6
Parametric LSEPIN encourages the discriminability among skills, which shown by comparing figs. 7a and 7b. The agent learned with parametric LSEPIN has less overlapping among skills thus better discriminability. WSEP encourages more Wasserstein distances between skills and skills with more distances from each other are also more discriminable. Because particle-base based entropies are calculated by the distances between samples, which the same as the transportation cost of Wasserstein distance, the mutual information calculated with particle-based entropies measures not only discriminability
also some kind of "distance". This why the WSEP agent has high particle- based LSEPIN. PWSEP learns the skills by the order of the Matplotlib (Hunter, 2007) color cycle, so the red skill is learned in the end, so it goes away from other skills instead of going downwards to the undiscovered branch of the tree maze, which could be local optimum for PWSEP(i), In practice, because of local optimums and a limited number of skills, despite its favored theoretical properties, the PWSEP algorithm might perform the best.
By comparing figures in fig. 7 and table 3, we can see that LSEPIN and WSEP capture the diversity and separatability of skills.
I.2 ANT ENVIRONMENT
Ant is one of the most common environments for URL (Sharma et al., 2020; Park et al., 2022a; Kim al., 2021). It is challenging high-dimension continuous control but the skills be visualized by top-down view of x-y dimension. This environment has downstream tasks that require navigating the ant agent designated positions. We compare the top-down visualizations, downstream task performance, and our proposed metrics to see whether the correlation claimed by our theorems exists.
7.3
7.5
5.0
50
23
23
-23
-38
-00
-75
-75
-75
(a) Random
(b) AUSL+nI-SEPI
DADS origin
Figure Trajectory examples: Trajectories different skills are in different colors. (a) a randomly initialized agent. (b) is pre-trained by MISL with non non-parametric LSEPIN. (c) is pre-trained by the original implementation DADS with prior knowledge and specific tuning.
Figure 8 shows the top-down view of skills of three URL agents. Their skills are in different colors. Their corresponding metrics are shown in table DADS origin was implemented with prior knowledge and specific tuning. For example, there is specific logarithmic transformation of its intrinsic rewards which does not apply to our implementations.
"R" is the downstream task performance, MSEPIN is the median of I(S: Similarly. WSEP estimated by sliced Wasserstein distance. I(S: 1.) and I(S; Z) are estimated by particle-based entropy.
34

---

Published as a conference paper at ICLR 2024
Table 4: Metrics for the agents in fig. 8.
Agent
Dimension
[[S; 2)
LSEPIN
MSEPIN
WSEP
(a) Random
All
20.8
6.1
10.2
4.5
-0.95
x-y
129.2
12.3
24.5
10.4
(b) MISL w nLSEPIN
All
127.6
20.1
23.9
15.3
-0.46
x-y
288.7
53.6
60.6
50.4
(c) DADS
All
248.5
31.0
71.7
26.2
-0.37
x-y
639.4
55.9
178.8
95.1
The correlation coefficient between downstream task performance and the metrics are shown in table 5. We can see that there is a strong correlation between disentanglement metrics (LSEPINWSEPP) and downstream task performance. This empirically validates our proposed theorems which claim correlation between disentanglement metrics and downstream task performance.
Table Correlation coefficient between downstream task performance and the metrics
Dimension
I(S; 2)
LSEPIN
MSEPIN
WSEP
All
0.92
0.95
0.77
0.93
x-y
0.83
0.99
0.78
0.92
1.3 SCALABILITY OF PARAMETRIC OBJECTIVES
Although focusing the empirical performance, we find that the proposed practical algorithm design in appendix H has the scalability for a potentially larger number of skills.
(a) parametric WSEP (b) MISL+ (c) parametric WSEP
Figure 9: Trajectory examples: Trajectories of different skills are in different colors. (a) is learned by only parametric WSEP (b) is learned by MISI. with non-parametrit LSEPIN and state entropy. (c) is pre-trained by only parametric WSEP.
The agents in fig. are learned by 8 skills by our introduced practical algorithms. Empirical algorithm design is not our main focus but we can see that optimizing parametric WSEP (eq. (95)) alone has comparable state coverage MISL with state entropy for exploration
1.4 ABOUT UNSUPERVISED REINFORCEMENT LEARNING BENCHMARK
There exists an unsupervised reinforcement learning benchmark (URLB) (Laskin al., 2021) using DeepMind Control environments, but their downstream task setting is not suitable for evaluating vanilla MISL methods that approximate and maximize / (S; Z). MISL with ((S; Z) is essentially partitioning the state space and labeling each part with skill latent The downstream task in Laskin et al. (2021) focuses on movements like "walk", "run" and "flip", which require learning a specific sequence of state progression instead of reaching certain part of the state space. So it can be considered that the downstream task setting URLB is too complicated for current MISL methods. And the downstream finetuning performance of pure exploration methods like ProtoRI. Yarats et al.
35

---

Published as a conference paper at ICLR 2024
(2021), APT Liu & Abbeel (2021a), RND Burda et al. (2019) exceeds MISL methods from Lee et al. (2019); Eysenbach et al. (2019); Liu & Abbeel (2021b) for by large margin, this might be because the state distributions of the MISL learned skills cover fewer states that are useful for the downstream task than the pure exploration methods. The CIC method Laskin et al. (2022) optimizes modified MISI objective [T; 2) and outperforms the pure exploration methods, they claim CIC promotes better diversity and discriminability than previous MISL methods, which also accords with our results.
36


Reviews:
[{"review_id": "8JzhCCMTSC", "reviewer": "Reviewer_xY4U", "summary": "This paper analyzes unsupervised skill-learning through a rigorous mathematical lens, focusing on the properties of the learned skills and their usefulness for downstream tasks. Most existing works use mutual information as the skill-learning objective and Eysenbach et al. (2022) provides a mathematical analysis of the same. This work analyzes the mutual information-based skill learning paradigm with a focus on downstream task adaptability via worst-case adaptation cost. This work also introduces a complementary metric called LSEPIN to measure diversity of learned skills. The authors show that maximizing MI or LSEPIN is essentially optimizing the KL divergence between state distributions. As an alternative, they suggest using the Wasserstein metric owing to it being a proper metric, and propose a new skill learning objectives built upon Wasserstein distance.\n\n**References:**\n\nBenjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. The information geometry\nof unsupervised reinforcement learning. In The Tenth International Conference on Learning\nRepresentations, ICLR 2022, Virtual Event, April 25-29, 2022. [OpenReview.net](http://openreview.net/), 2022. URL\nhttps://openreview.net/forum?id=3wU2UX0voE.", "questions": "None", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "There is a long line of work on unsupervised skill learning based on mutual information maximization between states and skills, most of the work being motivated by intuition and empirical performance. This work complements Eysenbach et al. (2022) by providing a rigorous understanding of the properties of the learned skills and provides useful insights. The analysis presented in this work is novel, to the best of my knowledge and comprises a fairly significant advancement of our understanding of this sub-area of RL. The quality of analysis and writing is satisfactory, with sufficient background and context provided before explaining the main results of the paper.\n\n**References:**\n\nBenjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. The information geometry\nof unsupervised reinforcement learning. In The Tenth International Conference on Learning\nRepresentations, ICLR 2022, Virtual Event, April 25-29, 2022. [OpenReview.net](http://openreview.net/), 2022. URL\nhttps://openreview.net/forum?id=3wU2UX0voE.", "weaknesses": "This is a fairly strong submission which checks all the boxes. The only minor complaint is that 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.", "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": "gNzZtHOrvt", "reviewer": "Reviewer_ZXDZ", "summary": "The paper provides a theoretical analysis of learning skills using unsupervised reinforcement learning (URL), which serves as an initialization for learning a policy for a downstream task. The paper shows that Mutual Information Skill Learning (MISL) does not guarantee diversity and separability of learned skills, and proposes to replace the uniform distribution of skills objective with the Least SEParability and INformativeness (LSEPIN) metric to promote informativeness and separability. Moreover, the paper proposes to replace the KL divergence in MISL with Wasserstein distance that exploits better geometric properties. Finally, it proposes another Wasserstein distance-based algorithm (PWSEP) that can theoretically discover all optimal initial policies.\n\nThe authors show theoretically that LSEPIN bounds the Worst-case Adaptation Cost (WAC) and show that the Wasserstein distance has better geometrical properties (such as symmetry and triangle inequality) that leads to better skill separation. In addition, the authors provide experiments to validate the proposed theoretical results in the appendix.", "questions": "I would like to ask the following questions:\n\n1. 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?\n\n2. Is there a measure for adaptation other than Worst-case Adaptation Cost?\n\n3. Is there any advantage to the KL divergence over the Wasserstein distance? theoretically or computationally?\n\n4. What are the limitations of the PWSEP algorithms (specifically the WSEP objective)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper investigates an important topic and provides a rigorous analysis of the proposed ideas. In addition, it proposes a practical algorithm that was tested empirically and demonstrates superior results compared to existing MISL methods.", "weaknesses": "The paper is hard to read and follow, with lots of details.\n\nSince most of the contribution of the paper is placed in the appendix, including all experimental results, it is hard to understand and assess its contribution without reading carefully the appendix.", "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": "HT5XuqM0Bv", "reviewer": "Reviewer_j7aV", "summary": "This paper studies the geometry of state distributions learned with mutual information skill learning for the purpose of theoretical task adaptation analysis. The authors propose Least SEParability and INformativeness (LSEPIN) to measure the diversity and separability and show a relationship with worst-case adaptation cost (WAC). The authors theoretically prove the relationship between the optimization of LSEPIN and Similarly, the authors also propose Wasserstein based distance metric WSEP which is more suitable in symmetric polytope to measure the separability of skill. Similar to LSEPIN and WAC, the relationship between WSEP and Mean Adaptation Cost is studied.", "questions": "* Here are questions that we want to discuss with the authors.\n* We might incorrectly understand the contribution of this work. Why the theoretical derivation of Theorem 3.1 is important for a task adaptation?\n    * What can we additionally learn from the theorem 3.1 or how can we use the theorem for other task adaptation works? For my understanding, the relationship: increasing LSEPIN results in lower WAC.\n    * Why measuring $\\min_z I(S;z)$ can be used to measure diversity and separability? I guess two features should be computed among skills, but the $\\min_z I(S;z)$ is just a mutual information for a single code.\n    * 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?\n    * WAC is measured with the worst state distribution. how this could be meaningful and practical? That is, why 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.\n    * Could you please additionally describe the necessity of symmetry and triangle inequality?\n    * What is the limitation of WSEP?\n    * Could this study can be combined with in-distribution and out-distribution perspective of task adaptations?\n\n[Overall]\nAlthough the authors conducted several theoretical derivations and analysis, the choice of metrics and the relationship between them have little meaning. Mostly because the task adaptation is 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.*", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "* The geometry perspective of task adaptation is interesting.\n* Studies on the geometry is promising and can motivate readers.\n* The theoretical results are well justified.", "weaknesses": "* Some Definitions of words are not defined well (diversity and separability).\n* The flow of the paper is weakly organized and hurts readability (the first topic is LSEPIN and the second WSEP.) The paper have no discuss on other perspectives.\n* The findings from the theoretical derivation are not surprising. WAC and adaptation cost.\n* Most contents have high dependency with appendix. Although the authors studied various components, they are not well organized.", "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": "IMUlQohMBm", "reviewer": "Reviewer_LRDL", "summary": "This work targets on the Information Geometry in unsupervised Reinforcement Learning, especially Mutual Information Skill Learning (MISL). On the basis of previous work, this work first considers the diversity and separability of learned skills. As MISL can not guarantee these properties, this work proposes LSEPIN to measure the disentanglement, and then shows the connection between LSEPIN and downstream task adaptation cost. Moreover, this work investigates the information geometry of Wasserstein distance based skill learning methods. Finally, this work proposes PWSEP and theoretically shows that it can discover all optimal initial policies.", "questions": "- This paper claims that \"it is possible for WSEP to discover all vertices of the feasible polytope\". As the theoretical results can only show that all learned skills satisfy $p(z) > 0$ rather than WSEP can learn all $p(z) > 0$ skills, there is no direct evidence to suggest that WSEP can indeed find more skills than MISL (i.e., skills without maximum “distances”). Can authors show that WSEP can find more skills or even all skills? I think empirical results or theoretical results even in simple cases will be really helpful.\n\n- 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.\n\n- 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?\n\n- It's better to provide some proof sketch of theorems in the main text, like Thm 3.5.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Definitely an important problem to be tackling! Previous work[1] has established the connection between skill-based unsupervised RL and information geometry. [1] has shown that KL-based metric can only find skills with the ''largest radius'' and how to find all vertices of the skill polytope is still an open challenge. This work has elegantly solved this challenge and I believe it will be of interest to researchers of Unsupervised RL.\n\n- The description of the article is very clear, and the introduction to related work is very specific.\n\n- I have basically read all the proofs of the theorems, which are well-written.\n\nThanks to the authors for putting in the effort in doing this work!\n\n[1] Eysenbach, Benjamin, Ruslan Salakhutdinov, and Sergey Levine. \"The information geometry of unsupervised reinforcement learning.\" arXiv preprint arXiv:2110.02719 (2021).", "weaknesses": "Overall, I think this paper is well written and its contribution is solid. I do not find clear weaknesses but I still have some questions (see Questions). I will adjust my score accordingly based on the author's response and other reviewers' comments.", "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 (spotlight)