T HE I NFORMATION G EOMETRY OF U NSUPERVISED R EINFORCEMENT L EARNING
A BSTRACT
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.
I NTRODUCTION
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 unsupervised experience from the environment that is not labeled with any rewards.
Prior methods have used this pretraining stage to learn representations of the environment that might assist the learning of downstream tasks.
For example, some methods learn representations of the observations (Laskin et al., 2020; Schwarzer et al., 2021) or representations of the dynamics model (Ebert et al., 2018;
Ha & Schmidhuber, 2018; Sekar et al., 2020).
In this work, we focus on methods that learn a set of potentially-useful policies, often known as skills (Salge et al., 2014; Mohamed & Rezende, 2015;
Gregor et al., 2016; Achiam et al., 2018; Eysenbach et al., 2018).
That is, the learned representation corresponds to a reparametrization of policies.
The aim of this unsupervised pretraining is to learn skills that, when a reward function is given, can quickly be combined or composed to 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 et al., 2018).
However, prior work offers little analysis about when and where such methods are provably effective.
Even simple questions, such as what it means for a 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 et al., 2018;
Wu et al., 2018;
Oord et al., 2018; Hjelm et al.,
; 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.
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 RL problem.
This result implies that using the learned skills for hierarchical RL 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 ﬁrst 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 inﬁnite 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. 1).
Points correspond to distributions over states, and the set of all possible distributions is a convex polytope that lies on a probability simplex.
We show that all reward-maximizing policies lie at vertices of this polytope and that maximizing mutual information corresponds to solving a facility assignment problem on the simplex.
The main contribution of this paper is a proof that skill learning algorithms based on mutual information 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 to be optimal, and also answers additional questions about unsupervised skill learning algorithms, such as whether the skills correspond to reward-maximizing policies and how many unique skills will be learned.
P RIOR W ORK
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., 2021), performing goal-conditioned RL without hand-speciﬁed 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., 2019; Seo et al., 2021) or learning collections of skills (Gregor et al., 2016;
Eysenbach et al., 2018; Achiam et al., 2018; Warde-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.
(2021) for a 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 community (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 a version of goal-conditioned
RL for pretraining.
However, the aim of these methods is different from ours: these methods aim to collect a sufﬁcient breadth of data, whereas we focus on learning good representations.
Our paper is similar to prior work that uses a geometric perspective to analyze RL algorithms (Dadashi 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 deﬁne a value function because the reward function is unknown.
Our analysis uses ideas from the ﬁeld 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 et al., 2007; Nachum & Dai, 2020).
Our analysis uses a similar procedure for adapting to new reward functions after the unsupervised learning stage.
P RELIMINARIES
We focus on inﬁnite-horizon MDPs with discrete states S and actions A , initial state distribution p 0 ( s 0 ) , dynamics p ( s t + 1 | s t , a t ) , and discount factor γ ∈ (0 , 1) .
We assume a Markovian policy π ( a | s ) , and deﬁne the discounted state occupancy measure of this policy as ρ π ( s ) =
(1 − γ ) (cid:80) ∞
t =
γ t P πt ( s ) , where P πt ( s ) is the probability that policy π visits state s at time t .
We deﬁne C 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:
[Equation]
The factor of 1 − γ accounts for the fact that the sum of discount factors 1 + γ + γ 2 + · · · = 11 − γ .
Without loss of generality, we focus on state-dependent reward functions ; action-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. U NSUPERVISED S KILL L EARNING A LGORITHMS
Prior skill learning algorithms learn a policy π θ ( a | s , z input ) with parameters θ and conditioned on an additional input z input ∼ p
( z input ) .
Let ρ π θ ( s | z
input
) denote the state marginal distribution of policy π θ ( a | s , z input ) .
We will focus on methods (Florensa et al., 2017; Eysenbach et al., 2018) that maximize the mutual information between the representation z input and the states s visited by policy π θ ( a | s , z input ) :
[Equation]
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., 2018; Sharma et al., 2019).
Because these methods have two learnable components, θ and p ( z input ) , analyzing them directly can be challenging.
Instead, we will use a simpliﬁed abstract model of skill learning where both the policy parameters θ and the latent variable z input are part of a single representation,
z = ( θ, z θ ) .
Then, we can deﬁne a skill learning algorithm as learning a single distribution p ( z ) :
[Equation]
We can recover Eq. 1 by choosing a distribution ( z ) that factors as p ( z = ( θ, z input ))
= δ ( θ ) p ( z input ) .
Whereas skills learning algorithms are often viewed as optimizing individual skills, this simpliﬁed abstract model lets us 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 a distribution over states, provides the best state distribution for optimizing an adversarially-chosen reward function using an idealized optimizer.
In general, this idealized optimizer is infeasible to implement, suggesting that existing skill learning methods do not provide the best initialization for practical optimization methods.
O VERVIEW OF M AIN R ESULTS
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 not result in skills that cover the set of all optimal policies, no matter how many skills are learned .
Perhaps surprisingly, 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, we 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 a prior over states that is close to the state distributions of reward-maximizing policies.
This prior is optimal in the sense that, when combined with a particular adaptation procedure, it 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 (Eq. 2):
[Equation]
where the adaptation objective is deﬁned as
[Equation]
Moreover, the optimal prior can be written as the average state distribution of the skills: ρ ∗ ( s ) =
E p ∗ ( z ) [ ρ ( s | z )] .
Note that the optimal prior, ρ ∗ ( s ) is not hard to compute: it corresponds to sampling one of the skills z ∼ p ∗ ( z ) and then execute the corresponding policy, π ( a | s , z ) .
[Table]
af- ter adaptation, which maximizes reward and minimizes information cost.
The regularized regret objective, A DAPTATION O BJECTIVE , can be interpreted in several different ways.
The regularization can be motivated as preventing overﬁtting to a limited number of samples of a noisy reward function.
We could also interpret this term as reﬂecting the cost of adapting from to the new task, but under a somewhat idealized model of the learning process.
Precisely, the adaptation objective is equivalent to performing one step of natural gradient descent (Amari, 1998; Kakade, 2001; Martens, 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 is much easier to 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 ﬁrst to provide any formal connection between the mutual information objective optimized by unsupervised skill learning algorithms and performance on downstream tasks.
RL AS G EOMETRY ON A S IMPLEX
In order to prove our main result and further analyze the kinds of policies learned by unsupervised skill learning methods, we will ﬁrst 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, a |S| -dimensional point lying on the probability simplex (see Fig. 1).
Note that this is 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 algorithms place skills on 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.
5. W HICH S TATE M ARGINALS ARE A CHIEVABLE ?
Understanding which state marginals are 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 set of achievable state marginal distributions can be described by a set of linear equations (Puterman, 1990, Eq.
(see Appendix A.1), we offer a different and complementary description based on the following property (Ziebart, 2010, Theorem 2.8):
Proposition 1.
Given a set of feasible state marginals, any convex combination is also feasible.
This property is illustrated in Fig.
1 (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 a 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 a polytope .
Every vertex of the state marginal polytope will contain a deterministic policy, 1 and the state marginal of every policy can be represented as a convex combination of the state marginals at the vertices.
5. W HICH P OLICIES ARE O PTIMAL ?
We now analyze where reward-maximizing policies lie within this state marginal polytope.
Again taking a geometric perspective, we visualize reward functions as vectors starting from the origin, so that the expected return objective 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 policies:
Proposition 2.
For every state-dependent reward function, among the set of policies that maximize that reward function is one that lies at a vertex of the state marginal polytope.
Proposition 3.
For every vertex of the state marginal polytope, there exists a reward function for which that vertex corresponds to a reward-maximizing policy.
Proposition 2 is an application of the maximum principle (Rockafellar, 2015, Chpt. 32); Proposition 3 is an application of the strong separation theorem (Hiriart-Urruty & Lemaréchal, 2004, Theorem 3.2.2).
These observations suggest that, for the purpose of RL, it is
[Figure]
sufﬁcient to search among policies whose state marginals lie at vertices.
These observations also suggests an objective for unsupervised skill learning:
Deﬁnition 5.1 (Vertex discovery problem) .
Given a controlled Markov process (i.e., an MDP without a reward function), ﬁnd the smallest set of policies such that every vertex of the state marginal polytope contains at least one policy.
While we are unaware of any tractable algorithm for exactly solving this problem, the next section discusses how a prior skill learning algorithm provides an efﬁcient approximation to this problem.
T HE G EOMETRY OF S KILL L EARNING
Now that we know where reward-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 reward-maximizing policies.
We show that the skills do not cover every vertex, and are thus insufﬁcient for representing the solution to every downstream
RL task.
We then will show that the average state distribution of the skills provides an optimal initialization for solving downstream tasks using a certain adaptation procedure.
The analysis in this section determines where skills lie on the polytope by analyzing the KL divergence between points on 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. W HERE DO S KILLS L EARNED BY M UTUAL I NFORMATION F ALL ON THE P OLYTOPE ?
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 as the average divergence between each policy’s state distribution ρ
( s | z ) and the average state distribution, ρ ( s ) :
[Equation]
To optimize this mutual information, a 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, if one policy has a state distribution that is further to 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 to place probability only on policies that are as far away from their average as possible:
Lemma 6.1.
The MISL skills must lie at vertices of the state marginal polytope.
This result, which we conﬁrm experimentally 3 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 a subset of the set of optimal policies.
Thus,
Corollary 6.1.1.
MISL skills are all optimal for some downstream state-dependent reward functions.
6. H OW M ANY U NIQUE S KILLS ARE L EARNED ?
A simple way to use skills acquired during pretraining to solve a downstream RL task is to simply search among the skills to ﬁnd 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 a 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 learn a larger (perhaps inﬁnite) number of skills (Eysenbach et al., 2018;
Hausman et al., 2018).
However, it has remained unclear whether such design decisions are useful and whether the goal of acquiring an inﬁnite 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 a simple counting argument based on the number of unique skills.
To count the number of unique skills, we 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 ρ ( s | z ) be given (not learned).
Let p ( z ) be the (learned) solution to maximizing mutual information (Eq. 2) and let ρ ( s ) =
E p ( z ) [ ρ ( s | z )] be the average state marginal.
Then the following holds:
[Equation]
We visualize this result in Fig.
3: the dashed line denotes a state marginals ρ ( 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 deﬁne in Appendix A.7, is similar to the collinearity assumption in linear regression (Mitchell & Beauchamp, 1988) and occurs with measure zero.
Lemma 6.3.
Under Assumption 1, 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 as candidate skills (Lemma 6.1).
The location of the average state marginal ( ρ ( s ) ) on the probability simplex is determined by the location of the skills and the distance to the skills.
The average state marginal is a vector in |S| -dimensional 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 we add a new skill (i.e., set p ( z ) > 0 for some z ), we add an additional constraint to the location of this average state marginal; namely, that
D KL ( ρ ( s | z ) (cid:107) p ( z )) = d max .
The value of d max adds one additional degree of freedom.
Thus, after adding ( |S| − 1) + 1 skills, the average state marginal is fully speciﬁed.
Adding additional skills would make the average state marginal ill-deﬁned, as it would have to violate one of the distance constraints.
Thus, we can have at most |S| unique skills.
We verify this result with a 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 an MDP is never greater than the number of states in that MDP, supporting Lemma 6.3.
MISL need not learn exactly |S| skills.
In some
MDPs, such as the one illustrated in Fig.
3b,
MISL acquires only 2 unique skills for a 3-state
MDP.
If an additional skill (at a 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 a KL divergence of d max from the initialization ρ ( s ) , in which case MISL could learn more than |S| skills.
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 I ( 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 ﬁnd 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 a skill learning algorithm can represent, there can be an exponential number that the skill learning algorithm can fail to 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 are 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 a 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-zero probability (i.e., skills), do the numerical values of these probabilities give us useful information to help solve downstream tasks?
In the next section, we discuss how our main result (Theorem 4.1) helps to address this equestion.
6. A N O PTIMAL I NITIALIZATION FOR AN U NKNOWN R EWARD F UNCTION
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 difﬁculty 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 difﬁcult.
If we think about adaptation as happening in the space of state marginal distributions, then there is 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 ρ ( s | z ) be given.
Maximizing mutual information is equivalent to minimizing the divergence between the average state distribution ρ ( s ) and the state distribution of the furthest skill:
[Equation]
This characterization can be understood as a facility location problem (Farahani & Hekmatfar, 2009) or a 1-center problem (Minieka, 1970; Garﬁnkel et al., 1977).
Since maximizing mutual information learns an average state distribution that is close to 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 4.1 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 ﬁnd a parameter initialization for which an inner optimization problem can be efﬁciently 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 an 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 modiﬁed so that adaptation can happen in the space of policy parameters, rather than state marginal distributions.
6. I MPLICATIONS FOR F UTURE S KILL L EARNING A LGORITHMS
Our analysis sheds light on fundamental limitations of skill learning algorithms based on mutual information.
First, as explained in Sec. 6.2, existing algorithms based on mutual information do not learn skills that are sufﬁcient 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 to 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 a policy’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 signiﬁcantly more powerful algorithms that serve as 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. L IMITATIONS OF A NALYSIS
In our analysis, we used a simpliﬁed 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 insufﬁcient 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.
C ONCLUSION
We have used a geometric perspective to answer a number of open questions in the study of skill discovery algorithms and unsupervised RL.
We also showed that the skills learned by unsupervised skill learning algorithms are not sufﬁcient for maximizing every downstream reward function.
However, we do show that the average state distribution of the skills provides an initialization that is optimal for solving unknown downstream tasks using a 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.
Acknowledgements.
We thank Abhishek Gupta, Alex Alemi, Archit Sharma, Jonathan Ho, Michael Janner,
Ruosong Wang, Shane Gu and anonymous reviewers for discussions and feedback on the paper.
This material is supported by the Fannie and John Hertz Foundation and the NSF GRFP (DGE1745016)
A P ROOFS AND A DDITIONAL A NALYSIS A. M ATHEMATICAL D EFINITION OF THE S TATE M ARGINAL P OLYTOPE
We provide a formal description of the constraint set.
Precisely, the set of achievable state marginal distributions are those for which the state-action marginal distribution ρ ( s , a ) satisﬁes the following linear constraints (Puterman, 1990, Eq.
[Equation]
The constraint that ρ be a valid probability distribution (i.e., (cid:80) s , a ρ ( s , a ) = 1 ) is redundant with the constraints above.
These constraints can also be written in matrix form as (˜ I − γP ) (cid:126)ρ s,a = (1 − γ ) (cid:126)p 1 , where ˜ I is a binary matrix that sums the action probabilities for each state (i.e., ˜ I(cid:126)ρ s,a = (cid:126)ρ s ), P is the transition matrix, and (cid:126)p 1 is the initial state distribution.
A. A DDITIONAL I NTUITION FOR THE S TATE M ARGINAL P OLYTOPE
This section provides some more intuition for what the state marginal polytope represents by examining a number of properties.
First, note that a policy that chooses actions uniformly does not necessarily achieve a 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 a second example, one 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 or two agents visit similar states in different orders.
Note that, from the perspective of reward maximizing, the state occupancy measure is a sufﬁcient representation of policies (i.e., it is an effect method for state aggregation (Van Roy, 2006; Singh et al., 1995)).
Finally, because some states may be more challenging to reach than others, not all state occupancy measures are realizable: some states may be challenging to reach, meaning that the agent cannot spend more than a small fraction of the episode at that state.
A consequence of this fact is that we will restrict our attention to a subset of the state occupancy measures on the probability simplex.
Our next section analyzes which state occupancy measures are achievable.
A. C ONVEX C OMBINATIONS OF F EASIBLE S TATE M ARGINALS ARE F EASIBLE
This section proves that convex combinations of feasible state marginal distributions are also feasible.
First, let two policies π 1 ( a | s ) and π 2 ( a | s ) be given and consider their corresponding state marginals, ρ 1 ( s ) and ρ 2 ( s ) .
We illustrate these marginals in Fig.
1 (center).
Any state marginal on the line between the two state marginals ρ 1 and ρ 2 is also achievable.
We can express all such state marginals as ρ λ = λρ 1 + (1 − λ ) ρ 2 .
One way to construct the policy that achieves the state marginal
ρ λ is to form a non-Markovian mixture policy: sample i ∼ B ERNOULLI ( λ ) at the start of each episode and then sample actions a ∼ π i ( a | s ) for each step of that episode.
By construction, this mixture policy will have a state occupancy measure of ρ λ .
It turns out that there also exists a Markovian policy π λ that achieves the state marginal
ρ λ : π λ
[Equation]
) where λ ( s ) (cid:44)
[Equation]
s
[Equation]
While the existence proof is somewhat involved (Ziebart, 2010, Theorem 2.8)(Feinberg & Shwartz, 2012, Theorem 6.1), the construction of the policy is surprisingly simple.
First, note that it is sufﬁcient to ensure ρ λ ( s , a ) = λρ 1 ( s , a ) + (1 − λ )
ρ 2 ( s , a ) .
Then, we apply Bayes’ rule to determine the corresponding policy:
[Equation]
s )
[Equation]
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.
1 (right).
The construction of the corresponding policy is analogous to Eq. 4.
A. P ROOF OF L EMMA 6.
We now prove Lemma 6.1:
Proof.
If p ( z ) were non-zero for a policy where the state marginal ρ ( s | z ) did not lie at a vertex, then the mutual information objective could be improved by shifting p ( z ) to place probability mass on the vertex rather than interior point
ρ
( s | z ) .
A. A LTERNATIVE S TRATEGIES A LSO F AIL TO D ISCOVERY A LL V ERTICES
Indicator reward functions.
Another simple method for skill discovery, based in prior work (Misra et al., 2020, Sec. 3.1), is to deﬁne a sparse reward function for visiting a particular state and action: r s,a ( s (cid:48) , a (cid:48) )
(cid:44)
( s (cid:48)
= s, a (cid:48) = a ) .
This method also fails to learn all vertices of the state marginal polytope, as we show in the following worked example.
We deﬁne the dynamics for a 3-state, 2-action MDP, set γ = 0 . 5 , and deﬁne the initial state distribution to be uniform over the three states.
This counterexample has deterministic dynamics, represented by arrows below:
[Equation]
None of the policies that are optimal for reward functions { r s , a } are optimal for the reward r ( s , a ) =( s ∈ { 1 , 2 } ) .
The intuition behind this counterexample is that there are multiple ways to maximize this reward function: remaining at state 1 and remaining at state 2.
However, none of the optimal policies for { r s,a ) } employ this strategy; rather, all these policies deterministically attempt to repeat a single transition.
Including actions in the mutual information ( I (( s , a ); z ) ).
Prior work on skill learning has considered many alternative deﬁnitions of the mutual information objective, experimenting with conditioning on actions (Eysenbach et al., 2018), initial states (Gregor et al., 2016), and entire trajectories (Co-Reyes et al., 2018).
We refer the reader to Achiam et al.
(2018) for a survey of these methods.
Consider, for instance, the objective I (( s , a ); z ) .
Our reasoning from Lemma 6.3 suggests that this alternative objective would result in learning at most |S| · |A| skills, more than the |S| skills
[Figure]
learned when maximizing I ( s ; z ) .
This objective results in 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 I (( s , a ); z ) only covers 3 of 6 vertices, the same number as covered by maximizing I ( s ; z ) .
We describe Counterexample 1 from Sec.
6 in more detail.
This example has 3 states, 4 actions, andskills.
The matrix describing the state-action marginals p ( s , a | z ) is given as follows:
[Equation]
 0 .0 .0 .0 .
00456412 0 .0 .0 .

The matrix describing the state marginals is computed as follows:
P a =
P sa      1 . 0 .
    
We illustrate the skills found by maximizing I ( s ; z ) and I (( s , a ); z ) in Fig.
Note that including the actions in the mutual information does not result in discovering more vertices of the state marginal polytope.
A. P ROOF OF T HEOREM 4.
Proof.
Our proof proceeds in two steps.
First
[Equation]
we will solve the adaptation problem, assuming that both the reward function r ( s ) and the initialization ρ ( s ) are given.
We then substitute the optimal policy ρ ∗ ( s ) into the pretraining problem and solve for the optimal initialization ρ ( s )
[Equation]
First
[Equation]
we solve for the optimal policy ρ ∗ ( s ) of the adaptation step.
Using calculus of variations
[Equation]
the optimal policy can be written as
[Equation]
We can then express the cost of adaptation as follows:
[Equation]
[Table]
[Equation]
We can note write the pretraining problem (Eq. 3) as follows:
[Equation]
[Table]
[Equation]
[Table]
We now aim to solve this optimization problem for ρ ( s ) .
We start by noting that the max operator is commutative, so we can swap the order of the maximization of r ( s ) and ρ + ( s ) .
That is, instead of the adversary ﬁrst choosing an reward function and then choosing an optimal policy for that reward function, we can consider the adversary ﬁrst choosing an optimal policy and then choosing a reward function for which that policy is optimal.
[Equation]
Using calculus of variations
[Equation]
we determine that the optimal reward function satisﬁes
[Equation]
where b ∈ R is a 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 is not a concern.
Substituting this worst-case reward function
[Equation]
we can write the overall pretraining objective as
[Equation]
Applying Lemma 6.5 completes the proof
[Equation]
A. A SSUMPTION ON C ONCYCLIC V ERTICES
Lemma 6.3 requires the following regularity condition:
Assumption 1.
Assume that any subset of |S| vertices is not concyclic with respect to the KL divergence.
That is, for all subsets of vertices V ⊆
V ERTICES ( C ) with size |S| , there does not exist a state distribution ρ ( s ) that has an equal divergence to all vertices in the subset:
[Equation]
This assumption can be understood visually using Fig.
3a (Left) .
A simplex in |S| has |S| − 1 degrees of freedom.
So, given |S| − 1 points, we can uniquely determine a center that is equidistant from these |S| − 1 points.
If we add an additional point to the simplex, it is very unlikely (measure zero) to be the same distance from this center.
A.
E XAMPLE WHERE ALL DETERMINISTIC POLICIES ARE UNIQUE VERTICES
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
B E XPERIMENTAL D ETAILS AND A DDITIONAL P LOTS
For all ﬁgures in the paper, the dynamics were randomly generated.
One way to generate the state marginal polytope is to randomly sample the initial state distribution and the dynamics distribution.
We employed a simpler, but mathematically equivalent, sampling procedure.
We ﬁrst sampled a set of state marginal distributions, and then computed the convex hull.
Note that, given a set of state marginal distributions, it is always possible to construct an MDP whose state marginal polytope is the convex hull of this set.
Assign one action to each skill, and deﬁne the dynamics as
[Equation]
for all states
[Equation]
Then, deﬁne the initial state distribution to be uniform over all states.
Please see the supplemental Jupyter notebook for code to reproduce the ﬁgures.
Fig.
7 plots the same experiment as Fig.
3 in the main text