The Information Geometry of Unsupervised Reinforcement Learning

1 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 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. , 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 . 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 can not 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 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 can not 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 . 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 . 2 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. , 2021 ) , performing goal-conditioned RL without hand-specified rewards ( 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 sufficient 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 can not 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 et al. , 2007 ; Nachum & Dai , 2020 ) . Our analysis uses a similar procedure for adapting to 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 distribution p0 ( s0 ) , dynamics p ( st+1 | st , at ) , and discount factor γ ∈ ( 0 , 1 ) . We assume a Markovian policy π ( a | s ) , and define the discounted state occupancy measure of this policy as ρπ ( s ) = ( 1− γ ) ∑∞ t=0 γ tPπt ( s ) , where P π t ( s ) is the probability that policy π visits state s at time t. We define 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 : ( 1− γ ) Eπ [ ∞∑ t=0 γtr ( st ) ] = Eρπ ( s ) [ r ( s ) ] . 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.1 UNSUPERVISED SKILL LEARNING ALGORITHMS . Prior skill learning algorithms learn a policy πθ ( a | s , zinput ) with parameters θ and conditioned on an additional input zinput ∼ p ( zinput ) . Let ρπθ ( s | zinput ) denote the state marginal distribution of policy πθ ( a | s , zinput ) . We will focus on methods ( Eysenbach et al. , 2018 ; Florensa et al. , 2017 ) that maximize the mutual information between the representation zinput and the states s visited by policy πθ ( a | s , zinput ) : max θ , p ( zinput ) I ( s ; zinput ) , Ep ( zinput ) Eρπθ ( s|zinput ) [ log ρ πθ ( s | zinput ) − log ρπθ ( s ) ] . ( 1 ) 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 ( zinput ) , analyzing them directly can be challenging . Instead , we will use a simplified abstract model of skill learning where both the policy parameters θ and the latent variable zinput are part of a single representation , z = ( θ , zθ ) . Then , we can define a skill learning algorithm as learning a single distribution p ( z ) : max p ( z ) I ( s ; z ) , Ep ( z ) Eρ ( s|z ) [ log ρ ( s | z ) − log ρ ( s ) ] . ( 2 ) We can recover Eq . 1 by choosing a distribution ( z ) that factors as p ( z = ( θ , zinput ) = δ ( θ ) p ( zinput ) . Whereas skills learning algorithms are often viewed as optimizing individual skills , this simplified 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 .

Summary: 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.   The 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.   Using 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.