The Information Geometry of Unsupervised Reinforcement Learning
Authors: Benjamin Eysenbach
Ruslan Salakhutdinov
Sergey Levine
Venue: ICLR 2022 Oral (2022)
Keywords: unsupervised skill learning
reward-free RL
mutual information
DIAYN

Abstract: How can a reinforcement learning (RL) agent prepare to solve downstream tasks if those tasks are not known a priori? One approach is unsupervised skill discovery, a class of algorithms that learn a set of policies without access to a reward function. Such algorithms bear a close resemblance to representation learning algorithms (e.g., contrastive learning) in supervised learning, in that both are pretraining algorithms that maximize some approximation to a mutual information objective. While prior work has shown that the set of skills learned by such methods can accelerate downstream RL tasks, prior work offers little analysis into whether these skill learning algorithms are optimal, or even what notion of optimality would be appropriate to apply to them. In this work, we show that unsupervised skill discovery algorithms based on mutual information maximization do not learn skills that are optimal for every possible reward function. However, we show that the distribution over skills provides an optimal initialization minimizing regret against adversarially-chosen reward functions, assuming a certain type of adaptation procedure. Our analysis also provides a geometric perspective on these skill learning methods.

Published as a conference paper at ICLR 2022
THE INFORMATION GEOMETRY OF UNSUPERVISED REINEOR LEARNING
Benjamin Eysenbach'
Ruslan Salakhutdinov
Sergey Levine
CCarneiie Mellon University, beysenbalcs cmu. edu
Google Brain,
SUC Berkeley
ABSTRACT
How can a reinforcement learning (RL) agent prepare to solve downstream tasks if those tasks not known a priori? One approach is unsupervised skill discovery, class algorithms that learn a set of policies without access to a reward function. Such algorithms bear close resemblance (e.g., contrastive learning) in supervised learning, in that both are pretraining mutual information objective.
representation learning algorithms
algorithms that maximize some approximation While prior work has shown that the set of skills learned by such methods can accelerate downstream RL tasks, prior work offers little analysis into whether these skill learning algorithms are optimal, or even what notion of optimality would appropriate to apply to them. In this work, we show that unsupervised skill discovery algorithms based on mutual information maximization do not learn skills that are optimal for every possible reward function. However, we show that the distribution over skills provides optimal initialization minimizing regret against adversarially-chosen reward functions, assuming a certain type adaptation procedure. Our analysis also provides a geometrie perspective on these skill learning methods.
INTRODUCTION
The high sample complexity of reinforcement learning (RL) algorithms has prompted a large body of prior work to study pretraining of RL agents. During the pretraining stage, the agent collects unsu- pervised experience from the environment that is not labeled with any rewards. Prior methods have used this pretraining stage learn representations the environment that might assist the learning of downstream tasks. For example, some methods learn representations of the observations (Laskin al. 2020 Schwarzer et al. (02T) or representations of dynamics model (Ebert et al. 2018 8 Ha & Schmidhuber 2018 Sckar et al. 2020). In this work, we focus on methods that learn set of potentially usseul policies, often known as skills (Salge et al. 2014: Mohamed & Rezende 2015
Gregor al. 2016 Achiam al. 2018 Eysenbach et al.) 2018). That is, the learned representation corresponds to a reparametrization of policies. The aim of this unsupervised pretraining to learn skills that, when a reward function given, can quickly be combined or composed maximize this reward. Prior work has demonstrated that this general approach does accelerate learning downstream RL (Gregor et al. 2016 Eysenbach et al. 2018 Achiam al. 2018). However, prior work offers little analysis about when and where such methods are provably effective. Even simple questions, such what means for set of skills to be optimal, remain unanswered.
Algorithms for unsupervised skill learning are conceptually related to the representation learning methods used to improve supervised learning (Gutmann & Hyvärinen 2010, Belghazi al, 2018 Wu et al.) 2018 Oord et al. 2018 [Hjelm et al.) 2018 He et al., 2020). Both typically maximize a lower bound on mutual information, and the learned representations are often combined linearly to solve downstream tasks (Hjelm et al. 2018 Oord et al. 2018). However, whereas prior work in supervised learning has provided thorough analysis of when and where these representation learning methods produce useful features (Kraskov et al. [2004] Song & Ermon 2019; McAllester & Stratos 2020), there has been comparatively little analysis into when unsupervised skill learning methods produce skills that are useful for solving downstream RL. tasks.

---

Published as a conference paper at ICLR 2022
In this paper, we analyze when and where existing skill learning methods based on mutual information maximization are (or are not) optimal for preparing to solve unknown, downstream tasks. On the one hand, we show that the skills learned by these methods are not complete, in that they cannot be used to represent the solution to every RI. problem. This result implies that using the learned skills for hierarchical RI. may result in suboptimal performance, and suggests new opportunities for better skill learning algorithms. One the other hand, we show that existing methods acquire a policy initialization that is optimal for learning downstream tasks, if that adaptation is performed using an idealized adaptation procedure. To the best of our knowledge, this is the first result showing that unsupervised skill learning methods are optimal in any sense.
Our analysis also illuminates a number of properties of these methods. For example, we show that every skill is optimal for some reward function, and we provide a nontrivial upper bound on the number of unique skills learned. This result implies that these methods cannot learn an infinite number of unique skills, and instead will learn duplicate copies of some skills. The key to our analysis is to view RL algorithms and skill learning algorithms as geometric operations (see Fig. D Points correspond to distributions over states, and the set of all possible distributions is a convex polytope that lies probability simplex. We show that all this polytope and that maximizing mutual information corresponds to solving a facility assignment problem on the simplex.
policies lie at vertices of
The main contribution of this paper is a proof that skill learning algorithms based on mutual informa- tion are optimal for minimizing regret against unknown reward functions, assuming that adaptation is performed using a certain procedure. Our proof of optimality relies on certain problem assumptions, leaving the door open for future skill learning algorithms to perform better under different problem assumptions. This contribution provides a rigorous notion of what it means for an unsupervised RL algorithm be optimal, and also answers additional questions about unsupervised skill learning algorithms, such as whether the skills correspond to reward- maximmiini policies and how many unique skills will be learned.
PRIOR WORK
In the unsupervised RL setting, an agent interacts with the environment without access to a reward function, with the aim of learning some representation of the environment that will assist in learning downstream tasks. Prior work has proposed many approaches for this problem. including learning
a dynamics model (Ebert et al., 2018 Ha & Schmidhuber 2018 (Sekar et al.) 2020). learning
compact representations of observations (Laskin et al. 2020 Schwarzer et al.) (021), performing
goal-conditioned RI. without hand-specified rewards (Eysenbach et al.) 2020 Chebotar et al. 2021
Schwarzer et al. 2021). doing pure exploration (Salge et al.) 2014 Mohamed & Rezende 2015
Pathak et al.| 2017 Lee et al. 2019b Hazan et al. [2199] Seo et al. 2021) or learning collections of
skills (Gregor et al.
et al. 2018 Achiam et al. 2018 WardWWardeFare Farley et al., 2018
Sharma et al. 2019; Hansen et al. 2020; Campos et al. 2020). These prior methods are similar in that they all learn representations, with different methods learning representation of observations,
dynamics, or policies. While each approach works well in some setting (see Laskin et al. (202T) for recent benchmark), the precise connection between these pretraining methods and success on downstream tasks remains unclear. We will focus on unsupervised skill learning methods.
The problem of unsupervised pretraining for RL has also been studied in the RL theory commu- nity (Agarwal et al.) 2020; Misra et al. 2020 Modi et al.) 2021). where it is referred to as the reward-free setting. The algorithms proposed in these prior works use version of goal-conditioned RL for pretraining. However, the aim of these methods is different from ours: these methods aim to collect a sufficient breadth of data, whereas we focus on learning good representations.
Our paper is similar to prior work that uses geometric perspective to analyze RL algorithms DDdashi et al. 2019 Bellemare et al. 2019). For example, [Dadashi et al. (2019) visualize the value function as a point in a high-dimensional space, where each coordinate indicates the value of one state. In contrast, our analysis will visualize policies as points, where each coordinate indicates the probability of visiting that state. This difference allows us to analyze the unsupervised RL setting, where we cannot define a value function because the reward function is unknown. Our analysis uses ideas from the field of information geometry (Amari & Nagaoka 2000). Prior work has parametrized the RL problem in terms of the state distribution (Puterman 1990 Eq. 6.9.2), and proposed RL algorithms
that solve for the optimal state distribution (Ng et al.
1999 Wang al.) 2007 Nachum & Dai 2020).
2

---

Published as a conference paper at ICLR 2022
Our analysis similar procedure for adapting new reward functions after the unsupervised learning stage.
3 PRELIMINARIES
We focus on infinite-horizon MDPs with discrete states S and actions A, initial state distribu- tion (o())), dynamics p(5++1 St, at), and discount factor y E (0, We assume Marko- vian policy π (a 8), and define the discounted state occupancy measure of this policy as p (5) = (1 7) -1= 7 PP where (s) is the probability that policy visits state time t. We define to be the set of state marginal distributions that are feasible under the environment dynamics. The RL objective can be expressed using the reward function r(s) and the state marginal
distribution:
(1-7) (rr(se) = (s) [r(s)]
The factor of 1 Y accounts for the fact that the sum of discount factors + + +22 + 1-7' Without loss of generality, we focus on dependent reward functions; action retion-dependent reward functions can be handled by modifying the state to include the previous action. While our analysis will ignore function approximation error, it applies to any MDP, including MDPs where observations correspond to features (e.g., activations of a frozen neural network) rather than original states.
3.1
UNSUPERVISED SKILI. LEARNING ALGORITHMS
Prior skill learning algorithms learn policy (a(a (a s, Zipptt) with parameters 0 and conditioned on an additional input Zioput (ZZmpu Let pT* (s Zinput) denote the state marginal distribution of policy To s, Zinpui). We will focus methods (Florensa et al. 2017 Eysenbach et al. 2018) that maximize the mutual information between the representation Zinput and the states s visited by policy (a 2:npur):
max 8, (Zp(ze)
(ss; Zinputt A Ep(zagu) Ep u (s) (ispu) [log pT* (s Zinpur) log pTa (s)].
(I)
We will refer to such methods as mutual information skill learning (MISL., noting that our analysis might be extended to other mutual information objectives (Salge et al., 2014 Mohamed & Rezende 2015 Achiam et al. 20188 [Sharma et al. 2019). Because these methods have two learnable compo- nents, O and p( Zinput). ) analyzing them directly can be challenging. Instead,
will use simplified abstract model skill learning where both the policy parameters 0 and the latent variable Zinput are part of single representation, a single distribution p(z):
20). Then, we define skill learning algorithm as learning
max (s: z) A Bptx) (log p(s z) log p(s)]. p(z)
We can recover Eq. I by choosing distribution (z) that factors p(z (0, Zinpat) (())pnZ))) Whereas skills learning algorithms are often viewed as optimizing individual skills, this simplified abstract model lets think about these algorithms as assigning different probabilities to different policies a probability of zero to almost every skill. We will refer to the small number of policies that are given non-zero probability as the "skills." We will show that, in general, the learned skills fail to cover all possible optimal policies. However, we show that the distribution over skills that maximizes mutual information, when converted into distribution over states, provides the best state distribution reward function using an idealized optimizer. In general, this
Zinpur)).
for optimizing an adversarially-choss. idealized optimizer is infeasible to implement, suggesting that existing skill learning methods do not provide the best initialization for practical optimization methods.
4 OVERVIEW OF MAIN RESULTS
Our main results analyze the relationship between the skills learned via the mutual information objective in Equation 2 and optimal policies. We will show that, in the general case, maximizing mutual information alone does result in skills that cover the set of all optimal policies, malter
3

---

Published as a conference paper at ICLR 2022
how many skills are learned. Perhaps supprisigly, when the number of skills is large enough, additional skills no longer differentiate from the existing ones, despite some potentially optimal policies not being covered. However, will show that the the distribution over skills learned by MISL is optimal in a different sense: when this distribution over skills is converted into a distribution over states, it provides prior over states that is close to the state distributions of reward-man maarddmmaiiiizing policies. This prior is optimal in the sense that, when combined with particular adaptation procedure, minimizes regret against the worst-case reward function. We formally state this result below:
Theorem 4.1. The optimal initialization for adapting to an unknown reward function is the average state marginal given by maximizing mutual information (Eg. 2
min max ADAPTATIONNODIECCVVVEESS... = max [(s; z), p(z)
p ((S)
where the adaptation objective is defined
ADAPTATIONODBICCTVVE(SS)),1
min max P p'(e)EC (5)EC (ptt)
(o) ()))s))]]))))))))]]]])))))))[[ ( (p"(()) p(s))
regret
(3)
Moreover, the optimal prior can be written as the average state distribution of the skills: (*(s) = (xy z)].
Note that the optimal prior, (*() is not hard to compute: it corresponds to sampling one of the skills z (z() and then execute the corresponding policy, (aa s, s,
The regularized regret objective, ADAPTA- TIONOBJECTIVE, can be interpreted in several different ways. The regularization can be moti- vated as preventing overfitting to limited num- ber of samples of a noisy reward function. We could also interpret this term as reflecting the cost of adapting from to the new task, but under somewhat idcalized model of the learning pro- cess. Precisely, the adaptation objective is equiv- alent to performing one step of natural gradient descent (Amari) 1998; Kakade 2001; Martens]
Table 1: Notation for analysis.
r(s)
Reward function.
((s)
Initialization.
Optimal state marginal distribution reward r(s). Used to define regret.
P² (s)
Optimal state marginal distribution af- ter adaptation, which maximizes reward and minimizes information cost.
2020). where the underlying metric is the state marginal distribution (Raskutti & Mukherjee] 2015 Theorem 1). This natural gradient is different from the standard natural policy gradient (Kakade 2001), which does not depend on the environment dynamics and much easier estimate. This equivalence suggests that our notion of adaptation is idealized: it assumes that this adaptation step can be performed exactly, without considering how data is collected, how state marginal distributions are estimated, or how the underlying policy parameters should be updated. Nonetheless, this paper is (to the best of our knowledge) the first to provide any formal connection between the mutual information objective optimized by unsupervised skill learning algorithms and performance on downstream tasks.
RL AS GEOMETRY ON A SIMPLEX
In order to prove our main result and further analyze the kinds of policies learned by unsupervised skill learning methods, we will first introduce a geometric perspective on RL. This perspective will allow us to better understand how unsupervised skill learning relates to the space of potentially optimal policies (for state-dependent reward functions).
We visualize policies by their discounted state occupancy measure, |S -dimensional point lying on the probability simplex (see Fig. []. Note that this a distribution over states, not actions, so a policy that chooses actions uniformly will not necessarily lie at the center of the simplex. Of course, not all state marginals are feasible, as the MDP dynamics might preclude certain distributions. We will be interested in where skill learning algorithmms place skills this probability simplex, and whether these skills are optimal for learning downstream tasks. Note that not all points on the simplex correspond to realizable policies, and multiple policies could map to the same state occupancy measure.
4

---

Published as a conference paper at ICLR 2022
[1S
pls)) p(s))
Figure 1: Polytope of achievable policies: (Left) For 3-state MDP, the set of feasible state marginal distributions descrihed by a triangle [(1, ((1,0,0, 0), (0, 0), (0, 0, 1)] in R We plot state occupancy measure of 6 differrnt policies. (Center) We plot those same policies. but remove the coordinate axes to improve visual clarity. (Right) For any convex combination of valid state occupancy measures, there exists a Markovian policy that has this state occupancy measure. Thus, policy search happens inside convex polytope.
5.1 WHICH STATE MARGINALS ARE ACHIEVABLE?
Understanding which state marginals achievable is important for analyzing what unsupervised skill learning algorithms do, and how the learned skills relate to downstream state-dependent reward functions. While the of achievable state marginal distributions can be described by set of linear equations (Puterman 1990, Eq. 6.9.2) (see Appendix A.II, we offer different and complementary description based on the following property (Ziebart 2010, Theorem 2.8):
Proposition Given set of feasible state marginals, any convex combination is also feasible.
This property is illustrated in Fig. (right): if we plot the state marginals for six policies, then any point inside the convex hull of those state marginals is also feasible. That is, there exists Markovian policy that achieves that state marginal distribution. In general, we can obtain the set of all feasible state marginals by taking the convex hull of the state marginals of all deterministic policies. We will refer to this set of achievable state marginals as polytope. Every vertex of the state marginal polytope will contain deterministic policy and the state marginal of every policy can be represented as convex combination of the state marginals at the vertices.
5.2 WHICH POLICIES ARE OPTIMAL?
We now analyze where wwwar-maxximizin policies lie within this state marginal polytope. Again taking a geomet- ric perspective, we visualize reward functions as vectors starting from the origin, so that the expected return objec- tive can be expressed as the inner product between the state marginal distribution and the reward vector (see Fig. 2 This visualization helps illustrate two facts about the state marginal distributions of reward-maximizing maximizing policies:
For every state deteedeeenndnt reward func- tion, among the set of policies that maximize that reward function is one that lies at vertex of the state marginal polytope.
Proposition 2.
Proposition 3. For every vertex the the state marginal poly- tope, there exists a reward function for which that vertex corresponds to a reward-maximizing policy.
Figure 2: Maximizing Rewards: We vi- sualize reward functions as vectors starting the origin. Maximizing expected retum corresponds maximizing the dot produet between the state marginal distribution and this reward vector.
Proposition is an application of the maximum princi- ple (Rockafellar) 2015 Chpt. 32); Proposition is an ap- plication of the strong separation theorem (Hiriart- Urruty Lemaréchal 2004 Theorem 3.2.2). These observations suggest that, for the purpose of RL,
Note, however, that some deterministic policies may not lie vertices of this polytope.
5

---

Published as a conference paper at ICLR 2022
28%
28%
(a) Examples where MISI. recovers 3 |S] skills.
(b) Example where MISI. learns only unique skills.
Figure 3:
Information Geometry of Skill Learning: Maximizing mutual information is equivalent to mini- mizing the maximum divergence between prior distribution states p(s) (green square) and any achievable state marginal distribution (Lemma 6.5). MISL learns distribution over state distributions, p(z). that assigns non-zero probability to only small number state distributions (solid orange circles).
sufficient to search among policies whose state marginals lie at vertices These observations also suggests an objective for unsupervised skill learning:
Definition 5.1 (Vertex discovery problem). Given controlled Markov process (i.e., MDP without a reward function), find the smallest set policies such that every vertex of the state marginal polytope contains at least one policy.
While we unaware of any tractable algorithm for exactly solving this problem, the next section discusses how a prior skill learning algorithm provides an efficient approximation to this problem.
THE GEOMETRY OF SKILL LEARNING
Now that we know where reward-maximiza maximizing policies lie on the polytope, we now analyze where the state marginal distributions of skills lie on the polytope. Our main result is that mutual information skill learning recovers some, but not all, of the wward-maximizing policies. We show that the skills do not cover every vertex, and thus insufficient for representing the solution to every downstream RL task. We then will show that the average state distribution the skills provides an optimal initialization for solving downstream tasks using certain adaptation procedure.
The analysis in this section determines where skills lie on the polytope by analyzing the KL divergence between points the polytope. This analysis does not assume that the KL divergence is symmetric (which it is not) or that it obeys the triangle inequality (which it does not).
6.1 WHERE DO SKILLS LEARNED BY MUTUAL INFORMATION FALL ON THE POLYTOPE?
We analyze where the skills learned by MISL lie on the state marginal polytope by dissecting the mutual information objective. The mutual information can be written the average divergence between each policy's 's state distribution p(s 2) and the average state distribution, p(s):
T(s: z) Ep(x)) 1Kk(pps z) p(s)))
To optimize this mutual information, skill learning algorithm would choose to assign higher probability to some policies and lower probability (perhaps zero) to other policies. In particular, this mutual information suggests that skill learning algorithms should assign higher probabilities to policies whose state distributions are further from the average state distribution. In fact, one policy has state distribution that is further the average state distribution than another policy, then the skill learning algorithm should assign a probability of zero to this second policy. Thus, the net result of optimizing mutual information is place probability only policies that far away from their average as possible:
While it is well known that reward function has deterministic optimal policy PPuterman 1990 Sec. 4.4), our visualization adds nuance the story, suggesting that searching in the space of all deterministic policies be wasteful, least the case of rewards that depend only on state.

---

Published as a conference paper at ICLR 2022
Lemma 6.1. The MISL skills must lie vertices of the state marginal polytope.
This result, which we confirm experimentally in Fig. 3. answers the question of where the skills learned by mutual information fall on the polytope. Recalling that the vertices of the state marginal polytope correspond to optimal policies (Proposition 2 this result implies that the skills learned by MISL are subset of the set of optimal policies. Thus,
Corollary 6.1.1. MISL skilly are all optimal for some downstream state-dependent reward functions.
6.2 How MANY UNIQUE SKILLS ARE LEARNED?
A simple way to use skills acquired during pretraining to solve a downstream RL task is to simply search among the skills to find the one that achieves the highest reward. So, it is natural to ask whether skill learning algorithms learn skills that are optimal for every possible reward function. Indeed, this is natural notion of what it might mean for a skill learning algorithm to be optimal. Prior work motivates the use of larger dimensional z or continuous-valued z, arguing that these design decisions will allow the method to Iearn larger (perhaps infinite) number of skills (Eysenbach al., 2018 Hausman et al. 2018). However, it has remained unclear whether such design decisions are useful and whether the goal of acquiring an infinite or exponential number of skills is even possible. Among prior skill learning methods, we are unaware of prior work that characterizes the number of skills that these methods will learn. In this section we show that existing skill learning algorithms based on mutual information are not optimal in this sense. Our proof will follow simple counting argument based on the number of unique skills.
To count the number of unique skills, start by analyzing the distance between each skill and the state distribution averaged over skills, measured as a KL divergence between state marginals. We can use the KKT conditions to prove that all skills with non-zero support (i.e., for which p(z) 0) have an equal divergence from the average state marginal:
Lemma 6.2. Let p(s z) be given (not learned). Let p(z) be the (learned) solution to maximizing mutual information (Eq. 2 and let p(s) = [p(s)]p(( 2)] be the average state marginal. Then the following holds:
p(z)
(pl( ())p((s)) max (p(s) |22222")))))) ())))((s)) dmax-
z"
We visualize this result in Fig. 3 the dashed line denotes a state marginals p(s z) that are equidistant from the average state marginal (green square). Note that all skills where MISL places non-zero probability mass (denoted by solid orange dots instead of transparent orange dots) lie on this dashed circle.
This lemma allows us to count how many unique skills MISL acquires. To state the result formally, we need to assume that the vertices of the state marginal polytope do not end up in a special concyclic arrangement. This assumption, which we formally define in Appendix AA77 is similar to the collinearity assumption in linear regression (Mitchell & Beauchamp 1988) and occurs with measure
Lemma 6.3. Under Assumption 7 MISL will recover at most S distinct skills.
Proof. The proof follows a simple counting argument. Without loss of generality, we only consider vertices of the state marginal polytope candidate skills (Lemmal6.1) The location of the average state marginal (p(s)) on the probability simplex is determined by the location of the skills and the distance the skills. The average state marginal is vector in |S] -iddimensional space with [S] degrees of freedom; we subtract one to account for the constraint that the state marginal sum to 1.
Consider adding skills one by one. When z). we add an additional constraint to the location of this average state marginal; namely, that (p((p(() 2) || p(z))= dmax. The value of dmax adds one additional degree of freedom. Thus, after adding (IS) - 1) + skills, the average state marginal is fully specified. Adding additional skills would make the average state marginal ill-defined, as would have violate one of the distance constraints. Thus, we can have at most (S] unique skills.
add new skill (i.e., set p(z) > 0 for some
*Code repraduce: ittpa://g1tte http://wwww.....eaaa///////eee/a
1pynb
7

---

Published as a conference paper at ICLR 2022
We verify this result with simple experiment. We randomly generate tabular MDPs and learn 100 (possibly redundant) skills using MISL. As shown in Fig. 4 the number of unique skills for MDP never greater than the number of states in that MDP, supporting Lemma 6.3
20
theoretical
Number of unique skills
maximum
15-
empirical maximum
10
MISL need not learn exactly |S| skills. In some MDPs, such as the one illustrated in Fig. (3b] MISL acquires only unique skills for 3-state MDP. If an additional skill (at vertex) were added, then the average distance from the skills to the average state marginal would decrease, worsening the mutual information objective. Additionally, when the regularity assumption is violated, there could be more than [S] skills that all have KL divergence of dmax from the initialization p(s). in which case MISL could learn more than |S] skills.
10
15
20
State dimension
Figure 4: As predicted by our theory, the number of unique skills learned by MISL upper bounded by the number of states.
If MISL can only learn up to |S] unique skills, will it learn a skill for maximizing every possible reward function? While some environments have exactly S optimal policies, general environments can have up to [A|||S] optimal policies (see Appendix A.8). Thus, MISL does not solve the vertex discovery problem:
Theorem 6.4. In general, skill learning algorithms that maximize the mutual information [(s: z) will not recover the set of all optimal policies.
This result is important for determining how to use the output of an unsupervised skill learning algorithm to solve downstream tasks. One approach to adaptation is to find whichever skill receives the highest return. This result indicates that this approach will produce suboptimal performance on a large number of reward functions. In fact, for every reward maximizing policy which skill learning algorithm can represent, there can be an exponential number that the skill learning algorithm can fail represent. This result casts doubt on the use of the skills themselves for solving downstream tasks, at least those skills learned by the existing generation of skill learning algorithms.
Might other simple algorithms for skill discovery recover the entire set of optimal policies? We currently unaware of any algorithm that does solve the vertex discovery problem. In Appendix (A.5) we show that (1) including actions in the mutual information objective or (2) maximizing collection of indicator reward functions also fail to discover all optimal policies.
Unsupervised skill learning algorithms learn a full distribution over skills. Instead of just looking at the set of policies that are given non-zen zero probability (i.e., skills), do the numerical values of these probabilities give us useful information to help solve downstream tasks? the next section, we discuss how our main result (Theorem 4. helps to address this equestion.
6.3 AN OPTIMAL INITIALIZATION FOR AN UNKNOWN REWARD FUNCTION
We now discuss a different perspective on optimizing mutual information, thinking about the average state distribution learned by MISL and whether it is optimal for solving downstream tasks. This perspective provides a proof sketch of our main result, Theorem 4.1
Intuitively, the difficulty of solving downstream tasks depends on how much the agent must adapt its initialization to learn the optimal policy for downstream tasks. Different initializations will make some tasks easier and some tasks more difficult. If we think about adaptation happening in the space of state marginal distributions, then there a precise relationship between maximizing mutual information and minimizing the amoung of adaptation required.
Lemma 6.5 (Theorem 13.11 from Cover& & Thomas (2006), based on Gallager 1979); Ryabko (1979). Let p(s) z) be given. Maximizing mutual information is equivalent minimizing the divergence between the average state distribution p(s) and the state distribution of the furthest skill:
max (s; min max DRL. 2) p(s)) p(z)
p(s)
This characterization can be understood as a facility location problem (Farahani & Hekmatfar 2009)
or I-center problem (Minieka 1970, Garfinkel et al. 1977). Since maximizing mutual information
8

---

Published as a conference paper at ICLR 2022
learns average state distribution that is close other policies, we expect that this average state distribution will be a good initialization for learning these other policies. This result provides the intuition behind the proof of Theorem in Appendix A.6
Relationship with meta-learning The pretraining objective (Eq. 3) bears a resemblance to the objective for meta-learning algorithms (Franceschi et al. 2018: Lee et al. 2019a Rajeswaran et al.) 2019): both aim to find parameter initialization for which an inner optimization problem can be efficiently solved. Whereas these methods assume access to a training distribution, skill learning algorithms used for pretraining do not have information about the downstream task. Moreover, meta-learning typically studies the average performance at test time, our objective measures the worst-case performance at test time. From this perspective, unsupervised skill learning could be interpreted as optimal pretraining for an idealized downstream adaptation procedure that searches over policies in the space of state marginals. Of course, this is abstraction of practical RL methods, which must search in the space of policy parameters. This connection suggests that it might be possible to design meta-learning algorithms using the tools of mutual information skill learning, but the underlying mutual information objective might need to be modified so that adaptation can happen in the space of policy parameters, rather than state marginal distributions.
6.4 IMPLICATIONS FOR FUTURE SKILL LEARNING ALGORIHTM
Our analysis sheds light on fundamental limitations of skill learning algorithms based on mutual information. First, explained in Sec. 6.22 existing algorithms based on mutual information do not learn skills that are sufficient for representing every potentially optimal policy. Future skill learning algorithms that can represent every policy may prove better at adapting to new tasks.
Our analysis shows that MISL produces an optimal initialization for a particular adaptation procedure: natural gradient descent in the space of state marginal distribution. This natural gradient is not the standard natural policy gradient (Kakade [2001) which is performed in policy parameter space, but rather the natural gradient with respect the state marginal distribution. This is a very strong (hypothetical) reinforcement learning algorithm, because it explores directly in the space of state marginals. This gradient implicitly depends on the environment dynamics, and computing this gradient would likely require exploring the environment to determine how changes in policy's s parameters affect changes in the states visited. Said in other words, our analysis suggests that existing unsupervised skill learning methods are optimal for a somewhat unrealistic and idealized model of how downstream policy learning is performed. This suggests an opportunity to develop significantly more powerful algorithms that serve effective initializations for more realistic adaptation procedures, such as those studied in meta-RL. (Wang et al. 2016 Finn et al. 2017 Humplik et al. 2019).
6.5 LIMITATIONS OF ANALYSIS
In our analysis, we used a simplified model of skill learning algorithms, one that viewed policy parameters and latent codes as one monolithic representation, with different skills using entirely different sets of parameters. Practical implementations of skill learning algorithms are less expressive, which may change the geometry of the state marginal polytope. For example, a limited policy class may be insufficient to represent some state marginal distributions, introducing "holes" in the state marginal polytope, or splitting the polytope into two disjoint regions. It remains to be seen whether these practical implementations actually learn fewer unique skills, and how these expressivity constraints affect the average state distribution.
7 CONCLUSION
We have used geometric perspective to answer number of open questions in the study of skill dis- covery algorithms and unsupervised RL. We also showed that the skills learned by unsupervised skill learning algorithms are not sufficient for maximizing every downstream reward function. However, we do show that the average state distribution of the skills provides an initialization that optimal for solving unknown downstream tasks using particular adaptation procedure. Our analysis showed that the number of skills is bounded by the number of states, suggesting that scaling these methods to learn more skills will hit a theoretical upper bound. Taken together. we believe that this analysis sheds light on what prior methods are doing, when and where they should work, and what limitations could be addressed in the next generation of skill learning algorithms.
9

---

Published as a conference paper at ICLR 2022
We thank Abhishek Gupta, Alex Alemi, Archit Sharma, Jonathan Ho, Michael Janner, Ruosong Wang, Shane Gu and anonymous reviewers for discussions and feedback the paper. This material supported by the Fannic and John Hertz Foundation and the NSF GRFP IDGE1745016)
Acknowledgements.
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint arXiv: 10299, 2018.
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in Neural Information Processing Systems, 33:20095- 20107, 2020.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251-21) 1998.
Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191. American Mathematical Soc., 2000.
Mohamed Ishmacl Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning, pp. 531-540, 2018.
Marc Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. geometric perspective on optimal representations for reinforcement learning. Advances in Neural Information Processing Systems, 32:4358 4369, 2019.
Víctor Campos, Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giró-i Nieto, and Jordi Torres. Explore, discover and Iearn: Unsupervised discovery of state-covering skills. In International Conference Machine Learning, pp. 1317-1327. PMLR, 2020.
Yevgen Chebotar, Karol Hausman, Lu, Ted Xiao, Dmitry Kalashnikov, Jacob Varley, Alex Irpan, Benjamin Eysenbach, Ryan C Julian, Chelsea Finn, et al. Actionable models: Unsupervised offine reinforcement learning of robotic skills. In International Conference on Machine Learning, pp. 1518-151 1528. PMLR, 2021.
John Co-Reyes, YuXuan Liu. Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey Levine. Self- consistent trajectory autoencoder: Hierarchical reinforcement learning with trajectory embeddings. In International Conference on Machine Learning, pp. 1009-1018. PMLR, 2018.
Thomas M. Cover and Joy A. Thomas. Elements of Information Theory (Wiley Series Telecommunications and Signal Processing). Wiley-Interscience, USA, 2006. ISBN 047 0471241954,
Robert Dadashi. Adrien Ali Taiga, Nicolas Le Roux, Dale Schuurmans, and Marc G. Bellemare. The value function polytope in reinforcement learning. In International Conference Machine Learning, pp. 1486- 1495, 2019.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annic Xic, Alex Lec, and Sergey Levine, Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint 1rXiv:18/2.005555 2018.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. International Conference on Learning Representations, 2018.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-Icarning Learning to achieve goals via recursive classification. In International Conference on Learning Representations, 2020.
Reza Zanjirani Farahani and Masoud Hekmatfar. Facility location: concepts, models, algorithms and case studies. Springer Science & Business Media, 2009.
Eugene A Feinberg and Adam Shwartz. Handbook Markov decision processes: methods and applications, volume 40. Springer Science & Business Mcdia, 2012.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning, 1126-1135, PMLR, 2017.
Carlos Florensa, Yan Duan, and Picter Abbeel. Stochastic neural networks for hierarchical reinforcement learning. arXiv preprint arXiv: /004.0700..000 03012, 2017.
10

---

Published as a conference paper at ICLR 2022
Luca Franceschi, Paolo Frascomi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-lcariing International Conference on Machine Learning, pp. 1568-1577 PMLR, 2018.
Robert G Gallager. Source coding with side information universal coding. 1979.
RS Garfinkel, AW Neebe, and MR Rao. The m-center problem: Minimax facility location. Management Science, 1331021133-1112, 1977.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint 1r 2016.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: new estimation principle unnor- malized statistical models. International Conference on Artificial Intelligence and Statisticx, pp. 297-304. JMLR Workshop and Conference Proceedings, 2010.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cosa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, 2018.
Steven Hansen, Will Dabney, Andre Barreto, David Warde-Farley, Tom Van de Wiele, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. In International Conference on Learning Representations, 2020. URL SttLihttsss//peeeviiiete wwt/ Bew.oorummndBBBBAABBeeee BJBJHAk
Karol Hausman, Jost Tobias Springenberg. Ziyu Wang. Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. In International Conference on Learning Representations, 2018.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In International Conference on Machine Learning, pp. 2681-2691. PMLR, 2019.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Computer Vision and Puttern Recognition, pp. 9729-9738, 2020.
Caan-Baptiste Hiriart- Hiriart-Urruty and Claude Lemaréchal. Fundamentals of convex analysis. Springer Science Business Media, 2004.
Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchiln
Marchildon, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In International Conference on Learning Representations, 2018.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv: htt::/905.0422. 2019.
Sham M Kakade. A natural policy gradient. Advances in Neural Information Processing Systems, 14, 2001.
Alexander Kraskov, Harald Stöghauer, and Peter Grassberger. Estimating mutual information. Physical review E, 69(6))066118 2004.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference Muchine Learning, pp. 5639-5650. PMLR, 2020.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert 7han, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. In Thirty-ffthh Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-leamiig with differentiable convex optimization. In Computer Vision and Pattern Recognition, pp. 10657-10665, 2019a.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov. Efficient exploration via marginal matching. arXiv preprint arXiv: 01009990600744 2019b.
James Martens. New insights and perspectives the natural gradient method. Journal of Machine Learning Research, 21:1-76. 2020.
David McAllester and Karl Stratos. Formal limitations the measurement of mutual information. In Interna- tiomal Conference on Artificial Intelligence and Statistics, pp. 875-884. PMLR, 2020.
Edward Minieka. The m-center problem. Siom Review, 12(1):138-139 1970.
11

---

Published as a conference paper at ICLR 2022
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efticient rich-observation reinforeement learning. International Conference on Machine Learning. pp. 6961-6971. PMLR, 2020.
Toby Mitchell and John Beauchamp. Bayesian variable selection in linear regression. Journal of the american statistical association, 83(404): 1988.
Aditya Modi, Jinglin Chen, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. Model-free representation learning and exploration in low-rank mdps. arXiv preprint arXiv:2 2021.
Shakir Mohamed and Danilo Jimenez Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. Advances in Neural Information Processing Systems, 28:2125-2133, 2015.
Ofir Nachum Bo Dai. arXiv:2001.0 2020.
Reinforcement learning fenchel-rockafellar duality.
arXiv preprint
Andrew Y Ng. Ronald Parr, and Daphne Koller. Policy search via density estimation. In Advances in Neural Information Processing Systems, pp. 1022-1028, 1999.
Aaron van den Oord, Yazhe Li. and Oriol Vinyals. Representation learning with contrastive predictive coding. 2018.
arXiv preprint
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-
supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
Martin I. Puterman. Markov decision processes. Handbooks in operations research and management science, 2: 331-434, 1990.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. Advonces in Neural Information Processing Systems, 32:1 13-124, 2019.
Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent, IEEE Transactions on Information Theory, 61()):1451-14)) 2015.
Ralph Tyrell Rockafellar. Convex analysis. Princeton university press, 2015.
Boris Yakovlevich Ryabko. Coding of a source with unknown but ordered probabilities. Prublems of Information Trunsmission, 15(2):134-133 1979.
Christoph Salge, Cornelius Glackin, and Daniel Polani, Empowerment-an introduction. In Guided Self-
Organization: Inception, pp. 67-114.55 Springer, 2014.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron C Courville. Pretraining representations for data-efficient reinforcement learning. Advances Neural Information Processing Systems, 34, 2021.
Ramanan Sekar, Olch Rybkin, Kostas Danilidiis, Picter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. Intermational Conference on Machine Learning, pp. 8583-8592 3592, 2020.
Younggyo Seo, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration. In International Conference on Machine Learning, pp. 9443-9454. PMI.R, 2021.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-awee aware unsupervised discovery of skills. In International Conference on Learning Representations, 2019.
Satinder P Singh. Tommi Jaakkola, and Michael I Jordan. Reinforcement learning with soft state aggregation. Advances in Neural Information Processing Systems, pp. 361-368, 1995.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators. In International Conference on Learning Representations, 2019.
Benjamin Van Roy. Performance loss bounds for approximate value iteration with state aggregation. Mathematics of Operutions Research, 31(2):234-244 2006.
Jane X Wang. Zeb Aurth-Nelson Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXiv: /6/ /6/ 105763, 2016.
12

---

Published as a conference paper at ICLR 2022
Tao Wang, Michael Bowling, and Dale Schourmans. Dual representations for dynamic programming and reinforcement learning. 2007 /EEE International Symposium on Approximate Dynamic Programming and Reinforcement Learning, pp. 44-51, 2007.
David Warde-Farley, Tom Van de Wiele. Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through non-parametric discriminative rewards. In International Conference on Learning Representations, 2018.
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised feature learning via non non-parametric instance discrimination. In Computer Vision and Pattern Recognition, pp. 3733-3742, 2018.
Brian D Zicbart. Modeling purposeful adaptive behavior with the principle of maximum causal entroyy. Carnegic Mellon University, 2010.
A PROOFS AND ADDITIONAL ANALYSIS
A.1 MATHEMATICAL DEFINITION OF THE STATE MARGINAL POLYTOPE
We provide formal description of the constraint set. Precisely, the set of achievable state marginal distributions are those for which the state-action marginal distribution p(s, a) satisfies the following linear constraints (Puterman 1990 Eq. 6.9.2):
Σ pls's = W'A
+7
Σ sES,aEA
pls' |ss ssa))s(s a) Vs' E ES
((s. a) 2 0 Vs E
The constraint that p be a valid probability distribution (i.e., Ex.a p a) redundant with the constraints above. These constraints can also be written in matrix form (I PPp)pp - -1p where i is a binary matrix that sums the action probabilities for each state (i.e., ip Ps), the transition matrix, and pi is the initial state distribution.
A.2 ADDITIONAL INTCITION FOR THE STATE MARGINAL POLYTOPE
This section provides some more intuition for what the state marginal polytope represents by ex- amining a number of properties. First, note that a policy that chooses actions uniformly does not necessarily achieve uniform state occupancy measure. For example, if actions directly control the next state, then duplicating one action would cause a random policy to shift from the center of the simplex. As second example,
state may simply not be reachable.
Second, note that two policies can have the same state occupancy measure. This case can happen if (say) multiple actions have similar effects Note that, from the perspective of reward maximizing, the state occupancy measure is a sufficient representation of policies (i.e., is an effect method for state aggregation (Van Roy 2006 Singh et al. 1995)).
agents visit similar states in different orders.
Finally, because some states may be measures are realizable: some states may be challenging to reach, meaning that the agent cannot spend more than small fraction of the episode that state. consequence this fact is that we will restrict our attention to subset of the state occupancy measures on the probability simplex. Our next section analyzes which state occupancy measures are achievable.
challenging to reach than others, not all state occupancy
A.3 CONVEX COMBINATIONS OF FEASIBLE STATE MARGINALS ARE FEASIBLE
This section proves that convex combinations of feasible state marginal distributions are also feasible. First, let two policies (a s) and (2 8) be given and consider their corresponding state marginals, (s) and (2()). We illustrate these marginals in Fig. (center). Any state marginal on the line between the two state marginals P1 and Pa is also achievable. We express all such state marginals as Py = Ap1 + (1 A)p2. One way to construct the policy that achieves the state marginal PA is to form non-Markovian mixture policy: sample i BERNOULLLI()) the start of
13

---

Published as a conference paper at ICLR 2022
each episode and then sample actions iT(aa s) for each step of that episode. By construction, this mixture policy will have a state occupancy measure of px-
It turns out that there also exists a Markovian policy xa that achieves the state marginal px:
Ap1 (s)
(a s) A(s) (a 8) + (1 - (22)
where
((s)
(4)
Api (s) + A)p2(s)
While the existence proof somewhat involved (Ziehart 20T0, Theorem 2.8) FFeineegg & Shwartz/
2012 Theorem 6.1), the construction of the policy is surprisingly simple. First, note that it is sufficient
to ensure px (s, a) (s, a) (1 1))22s, a). Then, we apply Bayes' rule to determine the
corresponding policy:
()(s.
Apı a) + (1 ())2(ss.)
xa s) Ps(s)
(s)(s) + (1 A)p2(s)
Amı(a| s))
A) s)p2(s)
(p)(s) + +(1 - A)p
(( (a s) + (1 (x()))222( ()) where X(s) - Ap1(s)
Ap1 (s) + (1 - {)p)(s)
By induction, one can also show that convex combinations of multiple state marginals also yields achievable state marginal distributions. We illustrate this observation in Fig. [ (right). The construc- tion of the corresponding policy is analogous to Eq.
A.4 PROOF OF LEMMA 6.1
We now prove Lemma 6.1
Proof. If p(z) were non-zero for a policy where the state marginal p(s z) did not lie vertex, then the mutual information objective could be improved by shifting p(z) to place probability mass on the vertex rather than interior point p(s z).
A.S ALTERNATIVE STRATEGIES ALSO FAIL TO DISCOVERY ALL VERTICES
Indicator reward functions. Another simple method for skill discovery, based in prior work (Misra et al. 2020, Sec. 3.1), to define sparse reward function for visiting a particular state and action: rs.o (s', a²) A (s' 8, a' a). This method also fails learn all vertices of the state marginal polytope, we show in the following worked example.
We define the dynamics to be uniform over the three states. This counterexample has deterministic dynamics, represented by arrows below:
3-state, 2-action MDP, set y = 0.5, and define the initial state distribution
None of the policies that are optimal for reward functions {Ts.a) are optimal for the reward r(s, a) (ss E (1, The intuition behind this counterexample that there are multiple ways to maximize this reward function: remaining at state I and remaining at state 2. However, none of the optimal policies for (ra.a)} employ this strategy; rather, all these policies deterministically attempt repeat single transition.
Including actions in the mutual information (L((s, a)) Prior work on skill learning has considered many alternative definitions of the mutual information objective, experimenting with conditioning on actions (Eysenbach et al.) 2018), initial states (Gregor et al. 2016). and entire trajectories RCo-Reyes et 2018). We refer the reader to Acciia al. (2018) for survey of these methods. Consider, for instance, the objective [(ss, a): z). Our reasoning from Lemma 6.. suggests that this alternative objective would result in learning at most [S] [A] skills, more than the [ skills
14

---

Published as a conference paper at ICLR 2022
(a) [(s; z)
(b) ((( a):
Figure 5: Counterexample 1: Maximizing [((s, z); does not result in discovering vertices of the state marginal polytope.
learned when maximizing [ss; z). This objective results discovering vertices of the state-action marginal polytope. However, these additional skills do not necessarily cover additional vertices of the state marginal polytope. We now present a worked example where maximizing [((s, a); z) only covers 3 of 6 vertices, the same number covered by maximizing (s;;))
We describe Counterexample from Sec. i more detail. This example has states, actions, and 7
skills. The matrix describing the ittteacction marginals p(s. z) is given as follows:
0.0716829
0. 0.18424628
0.1795034
0.13672623
0.00295358
0.095103347
0.03845856
0.05906211
0.08669371
0.01741351
0.02568876
0.10246764
0.11750427
0.14607
0.11480248
0.00287273
0.00954288
0.10701075
0.02102705
0.02261918
0.14991443
0.23005368
0.06343457
0.01514798
0.13386153
0.06171002
0.09080011
0.16650628
0.05625721
0.06271223
0.10353768
0.09575361
0.02849315
0.00634354
0.13433545
0.05968919
0.05926209
0.10606773
0.04541879
0.05518233
0.14630736
0.12098995
P.o
0.05692926
0.0745452
0.04428352
0.12954463
0.0261651
0.13530253
0.13704367
0.01333569
0.10392952
0.04119388
0.11826303
0.12777541
0.12092377
0.0703527
0.08563638
0.08155578
0.04348018
0.05650998
0.06520597
0.00354202
0.11542116
0.09079959
0.14099273
0.05765374
0.04477772
0.03235234
0.07636344
0.00824066
0.09351822
0.27113241
0.20516233
0.19498386
0.03909264
0.11229717
0.06929778
0.09957651
(0.02888635
0.08654659
0.00456412
0.19995985
0.0328871
0.0167457
The matrix describing the state marginals computed as follows:
I.
0.
0.
1.
0.
0.
1.
0.
1.
0.
0.
0.
1.
0. 1. O.
Pa P.
1.
0. 1. 0.
0. 0. 1.
0. 0. 1.
0.
0.
1.
0. 0.
1.
We illustrate the skills found by maximizing [(s; z) and a() z) in Fig. 5 Note that including the
actions in the mutual information does not result discovering more vertices of the state marginal polytope.
15

---

Published as a conference paper at ICLR 2022
A.6 PROOF OF THEOREM 4.1
Proof. Our proof proceeds in two steps. First, we will solve the adaptation problem, assuming that both the reward function r(s) and the initialization p(s) given. We then substitute the optimal policy p (s) into the pretraining problem and solve for the optimal initialization p(s).
First, we solve for the optimal policy p" (() of the adaptation step. Using calculus of variations, the optimal policy can be written as
p(s)er(s)
p (*(s)
I p(s))err ds'
We can then express the cost of adaptation as follows:
ADAPTATIONOBECCIV(((()).11
max B (x](s)) p+(eE
E, (w)[(()) + DKI. (s) ll p(s))
J (s
I (()))))())
max B [[](()])]]]] (a)EC
+o8pts)++(1s) log pls pls')e 108
I p(s')er
((ss))rr
S (())
max ([)(s)][s)]] p+(s)EC
L r()dd log ((s)c
-f dertedds
max B, [r-ss] - log pls)e ds. (+(eE
We can note write the pretraining problem (Eq. 3 as follows:
min max ADAPIATIONNODICCTVV(((())s11) r(s)) p(a)EC (())
min max max (p((())(s))) log p(s)e ds.
p(A)EC (())R) (+
We now aim to solve this optimization problem for p(s). We start by noting that the max operator is commutative, so we can swap the order of the maximization of r(s) and p (s). That is. instead of the adversary first choosing an reward function and then choosing an optimal policy for that reward function, we can consider the adversary first choosing an optimal policy and then choosing a reward function for which that policy is optimal.
min max ADAPPATOONORICCTIVV((SS)) r())eC(
min max M(-)E (elec (())
max ())(()))))) log ((s)e
Using calculus of variations, we determine that the optimal reward function satisfies
r(s) log px (s) p(s) +
where à E IR is constant scalar. There are many optimal reward functions, one for each choice of b. However, as the scalar b will cancel out in the next step of our proof, this multiplicity not concern. Substituting this worst-case reward function, we can write the overall pretraining objective as
min max ADAPRAATOONOBBECTTVE((SI1,,,,, r(s) (())
a6s)nn(s)) pGST c'ds
min max (p+(s)llog flo (s)) p(s) + bb - log
p(s)eC p*)
min max RBptss) fog pt (s) logg(s) + loget
(t)
= min max Ept(s log p² (s) log p(s)]
plx)eC (+)E
= min max DRL. (p (s) p(s)).
A(a)eC (*)EE
Applying Lemma 6.5 completes the proof.
A.7 ASSUMPTION ON CONCYCLIC VERTICES
Lemma 6.3 requires the following regularity condition:
16

---

Published as a conference paper at ICLR 2022
Assumption 1. Assume that any subset of [S] vertices is not concyclic with respect the KL divergence. That is, for all subsets of vertices V S VERTICES(C) with size [S, there does not exist a state distribution p(s) that has equal divergence all vertices in the subset:
DKL p(s)) = (p) p(s))
VijjV.
This assumption can understood visually using Fig. 3a (Left). A simplex in [S] has S - 1 degrees of freedom. So, given |S| 1 points, these |S] - points. If we add an additional point to the simplex, it is very unlikely (measure zero) be the same distance from this center.
uniquely determine center that is equidistant from
A.8 EXAMPLE WHERE ALL DETERMINISTIC POLICIES ARE UNIQUE VERTICES
@1
01
81
$2
85
0/2
0/2
02
0/2
Figure 6
B EXPERIMENTAL DETAILS AND ADDITIONAL. PLOTS
For all figures in the paper, the dynamics were randomly generated. One way generate the state marginal polytope is to randomly sample the initial state distribution and the dynamics distribution. We employed simpler, but mathematically equivalent, sampling procedure. We first sampled of state marginal distributions, and then computed the convex hull. Note that, given set of state marginal distributions, always possible to construct an MDP whose state marginal polytope is the convex hull of this set. Assign one action to each skill, and define the dynamics as
p(s++1 St: at) = p(s z = a) for all states St-
Then, define the initial state distribution to uniform over all states.
Please see the supplemental Jupyter notebook for code reproduce the figures. Fig. plots the same experiment as Fig. 3 in the main text, repeated for different mudommlyssample dynamics.
17

---

Published as a conference paper at ICLR 2022
Figure 7: Information Geometry of Skill Learning: More examples of the skills learned by MISL.
18


Reviews:
[{"review_id": "w3uNCspHI1S", "reviewer": "Reviewer_tCVb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "Overall, the paper presents an insightful analysis of some of the challenges of unsupervised RL. In general, I believe that the primary impact of this paper would be through its geometric approach of analyzing the space of learnable policies. The concrete theorems proven by the paper, while interesting, appear to have a relatively large set of conditions. It would be interesting to see how the analysis can be applied to more general settings.", "summary_of_the_paper": "The paper treats the problem of unsupervised RL, which it defines as the problem of pretraining a system, without having access to a reward function, to learn a collection of policies, that are labeled skills. The idea is that when the reward function is presented, the target policy can be assembled as a combination or composition of these skills. Such policies had been previously proposed to be learned using mutual information maximization approaches. \n\nThe primary contribution of this paper that it analyzes the space of policies/skills that are learnable using mutual information based approaches. The authors present this analysis through a geometric lens on the probability simplex of the possible states of the system. A given policy or skill can be associated to a point in this space by associating with a skill or policy its discounted state probability distribution. \n\nUsing these analytical tools, the authors prove several interesting results. For instance, they show that by maximizing mutual information alone, one cannot learn sufficient skills to cover all the set of optimal policies - the number of unique skills learned through this approach is bounded by the dimensionality of the state space. ", "summary_of_the_review": "The paper analyzes some of the challenges of unsupervised RL through a geometric approach of analyzing the space of learnable policies. The approach allows insights about the characteristics and certain optimality criteria that applies to skills learned through mutual information type approaches. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "uk5j00FOlOb", "reviewer": "Reviewer_Xkef", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "STRENGTHS\n\n(S1) The paper provides a very nice analysis for a growing class of methods that seek to solve a problem of growing interest in the community. The results are interesting--they say what unsupervised skill discovery methods _can_ and _can not_ do--and the discussion is timely as more and more researchers start to work in the space of unsupervised reinforcement learning.\n\n(S2) The geometric perspective presented is intuitive and does a nice job of communicating the main results of the paper.\n\nWEAKNESSES\n\n(W1) The paper lacks a good empirical confirmation of the claims made here. For example, if unsupervised skill discovery indeed hits the theoretical upper bound that the authors claim, shouldn't there be even a very small toy domain on which the authors could do some experiments to validate that claim? That said, I understand the contribution of the paper here is mostly theoretical in nature.", "summary_of_the_paper": "The authors provide an in-depth study of unsupervised skill discovery from the RL literature. In particular, the authors analyze such approaches from the perspective of \"information geometry,\" and show that these algorithms _can not_ learn skills that are optimal for all possible reward functions, but _can_ provide a good initialization for online learning approaches that seek to adapt the initial skills to find optimal policies for new reward functions.", "summary_of_the_review": "Overall, I really like the paper and think that the results are important and should be presented to the community. However, I do think the paper could be improved with even a small empirical study.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "qyHpW0fGSui", "reviewer": "Reviewer_r6uL", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "## Strengths:\n* This paper attempts to answer the fundamental question of why RL practitioners should care about unsupervised skill discovery. Providing a satisfying answer to this question should attract more attention to learning skills in this setting, as well as test and clarify the fundamentals on which the problem is based.\n* The high level results this paper presents also seem interesting.\n* Adapting the value function polytope to a polytope over state visitation frequencies and analyzing the problem from a geometric perspective is also a novel approach that deserves to be studied further.\n* The reframing of the unsupervised skill discovery problem as a vertex discovery problem is very interesting.\n\n## Weaknesses:\n* I am not completely convinced of the correctness of all the results. The paper introduces a lot of tools for its analysis. Specifically, the state marginal distribution polytope is an entirely new concept and its use in the analysis in this paper was nontrivial to follow. My questions in this regard are below.\n* The adaptation procedure that this paper refers to when it asserts that the mutual information maximizing skills are optimal initializations ignores the dynamics of the environment. As such, I am not sure if this adaptation procedure should even be considered feasible. The paper should reframe its statement to make the _strong_ assumption on this adaptation clear, rather than merely stating that the initialization is optimal \"under some assumption about how the adaptation is performed\".\n\n## Detailed Comments and Questions:\n* The problem setup specifies $\\gamma \\in [0, 1]$. Doesn't $\\gamma=1$ not work for the state marginal distribution and the subsequent RL objective?\n* The value function polytope is a closely related to the value function polytope approach (Dadashi et al., 2019). This related work seems almost hidden in section 2, and the contrast between the two approaches beyond the fact that one looks at value functions and the other at state marginals is not clarified. I felt like a paragraph clarifying the difference so that a reader familiar with that work might be able to grasp the state visitation polytope better would be useful. \n* In the 3-state MDP example considered for this paper, the start state distribution and the dynamics of the environment are not illustrated sufficiently. Does the agent start at all three states with equal probability?\n* When analysing this polytope in terms of state-only reward functions, does the origin indicate reward functions that do not prefer any particular state over the other? Are the rewards adjusted to be positive, so that they lie in this quadrant that is being analyzed?\n* In general, section 5 seems very condensed. I understand that the limits of the paper might be forcing the authors to be succinct, but I come away from this section feeling like I haven't completely understood the geometric view the paper is utilizing.\n* In Section 6, the notion of policies being closer or further in terms of the KL divergence is mentioned when analyzing the skills learned via mutual information maximization. I'm having a hard time wrapping my head around this concept, and perhaps it stems from not grasping the polytope in the previous section completely. Do points further away from each other on the polytope have higher a higher value of KL divergence? But if it was possible for two policies to spend all their time at individual vertices of the outer triangle (if the dynamics of the MDP allowed it), wouldn't the KL divergence be infinite? The notion of treating distance in this polytope in terms of KL divergence is something that I am not convinced of, and leads to me being unconvinced of the rest of the results of this paper.\n* In Lemma 6.2, does the probability of sampling a skill ($p(z)$) have to be uniform? I am unsure how skills with different probabilities could lead to state marginals that have the same KL divergence from the average. Do the skills with lower probability of being sampled travel farther?\n* For theorem 6.4, if the agent is able to learn $|\\mathcal{S}|$ skills that each maximize the time spent at individual states, then one of these skills should maximize the return for any state based reward function. What am I missing here? Why do we need to consider all $|\\mathcal{A}|^|\\mathcal{S}|$ policies?\n* On page 7, after the proof of Lemma 6.3, the paper mentions a regularity assumption. What is this assumption?\n* ICLR allows 10 pages of content this year. Given the density of Section 5, I feel like this paper wasted some space that could have been used for some clarifying exposition on the geometric view that the paper proposes.\n\n## Other Comments:\n* After Equation (3): \"... can be written is the  average state ...\", _is_ should be _as_.\n* Section 6.3: \"Maximizing mutual information pushes skills away from one another, not as finding a skill or policy that is close to every other policy\". I'm not sure what this sentence is trying to say. It is hard to parse.", "summary_of_the_paper": "This paper is trying to analyze whether unsupervised skill discovery is useful for more easily solving any possible downstream tasks in an MDP. It does so by adapting the idea of the value function polytope to a state visitation distribution polytope. It also specifies possible reward functions in this geometric setting and analyzes the connection between points on this polytope and returns with respect to the reward function.\n\nNext, it casts the mutual information based skill discovery problem in this geometric space and tries to analyze the skills learned. From their analysis, the paper suggests ways to infer how many skills can be learned, and whether those skills are optimal with respect to some downstream tasks.\n\nIt seems that these skills can be guaranteed to be the vertex of the above polytope, but not to be optimal with respect to all downstream reward functions. They then suggest that the skills learned might be useful for an adaptation procedure that ignores the dynamics of the environment.", "summary_of_the_review": "The problems that this paper attempts to solve are interesting and would add value to the setting of unsupervised skill discovery. Specifically, it considers the question of what the objective of such skill discovery should be.\nHowever, the geometric tools introduced in the paper and the subsequent analysis does not seem polished enough and is hard to follow. I am unconvinced of the analysis and thus scoring the paper conservatively.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Strong submission that analyses the unsupervised skill discovery setting from the perspective of information geometry, which leads to some interesting conclusions. In particular, it is shown that this does not lead to skills that are optimal for all reward functions, but does provide a good initialization for methods that aim to find optimal policies.
 
Across the board, the reviewers believe the analysis provided by this work is both important and novel. And while there were some initial concerns raised, such as lack of empirical confirmation of some of the claims and some questions about the analysis, the authors have addressed all of these concerns convincingly.
 
Hence, I strongly recommend acceptance of this submission.
Decision: Accept (Oral)