Skew-Fit: State-Covering Self-Supervised Reinforcement Learning

1 INTRODUCTION Reinforcement learning ( RL ) provides an appealing formalism for automated learning of behavioral skills , but separately learning every potentially useful skill becomes prohibitively time consuming , both in terms of the experience required for the agent and the effort required for the user to design reward functions for each behavior . What if we could instead design an unsupervised RL algorithm that automatically explores the environment and iteratively distills this experience into general-purpose policies that can accomplish new user-specified tasks at test time ? For an agent to learn autonomously , it needs an explo- ration objective . In the absence of any prior knowledge about which states are more useful , an effective exploration scheme is one that visits as many states as possible , allowing a policy to autonomously prepare for user-specified task that it might see at test time . This objective has been formalized as maximizing the entropy of the learned policy ’ s visited state distribution 1 H ( S ) ( Hazan et al. , 2018a ) , since a policy that maximizes this objective should approach a uniform distribution over valid states . Unfortunately , directly optimizingH ( S ) requires an accurate model of the policy and environment ( Hazan et al. , 2018a ) . Moreover , even if this optimization were tractable , another short-coming of this objective is that the resulting policy can not be used to solve new tasks : it only knows how to maximize state entropy . In other words , to develop principled unsupervised RL algorithms that result in useful policies , maximizingH ( S ) is not enough . We need a mechanism that allows us to control the resulting policy to achieve new tasks at test-time . 1We consider the distribution over terminal states in a finite horizon task and believe this work can be extended to infinite horizon stationary distributions . We argue that this can be accomplished by performing goal-directed exploration . In addition to maximizing the state entropy , we should be able to control where the policy goes by giving it a goal G that corresponds to a state that it must reach . Mathematically , a goal-conditioned policy should minimize the conditional entropy over the states given a goal , H ( S | G ) . This objective provides us with a principled way for training a policy to explore all states , by maximizingH ( S ) , such that the state that is reached can be controlled by commanding goals , which means minimizingH ( S | G ) . Directly using this objective is often intractable , since it requires optimizing the entropy of the marginal state distribution of the policy , H ( S ) . However , we can sidestep this issue by noting that the objective is the mutual information between the state and the goal , I ( S ; G ) , which can be written as : H ( S ) −H ( S|G ) = I ( S ; G ) = H ( G ) −H ( G|S ) . ( 1 ) Equation 1 thus gives an equivalent objective for an unsupervised RL algorithm : the agent should set diverse goals , maximizingH ( G ) , and learn how to reach them , minimizingH ( G | S ) . While the second term is the typical objective studied in goal-conditioned RL ( Kaelbling , 1993 ; Andrychowicz et al. , 2017 ) , maximizing the diversity of goals is crucial for effectively learning to reach all possible states . In a new environment , acquiring such a maximum-entropy goal distribution is challenging : how can an agent set diverse goals when it does not even know what states exist ? In this paper , we address this question via a new algorithm , Skew-Fit , which learns to model the uniform distribution over states , given only access to data collected by an autonomous goalconditioned policy . Our paper makes the following contributions . First , we propose a principled objective for unsupervised RL , based on Equation 1 . While a number of prior works ignore theH ( G ) term , we argue that jointly optimizing the entire quantity is needed to develop effective and useful exploration . Second , we propose a method called Skew-Fit and prove that , under some regularity conditions , it learns a generative model that converges to a uniform distribution over the goal space , even when the set of valid states is unknown ( e.g. , as in the case of images ) . Third , we empirically demonstrate that , when combined with goal-conditioned RL , Skew-Fit allows us to autonomously train goal-conditioned policies that reach diverse states . We test this method on a variety of simulated vision-based robot tasks without any task-specific reward function . In these experiments , Skew-Fit reaches substantially better final performance than prior methods , and learns much more quickly . We also demonstrate that our approach solves a real-world manipulation task , which requires a robot to learn to open a door from scratch in about five hours , directly from images , and without any manually-designed reward function . 2 PROBLEM FORMULATION . To ensure that an unsupervised reinforcement learning agent learns to reach all possible states in a controllable way , we maximize the mutual information between the state S and the goal G , I ( S ; G ) , as stated in Equation 1 . This section discusses how to optimize Equation 1 by splitting the optimization into two parts : minimizingH ( G | S ) and maximizingH ( G ) . 2.1 MINIMIZING H ( G | S ) : GOAL-CONDITIONED REINFORCEMENT LEARNING Standard RL considers a Markov decision process ( MDP ) , which has a state space S , action space A , and unknown dynamics ρ ( st+1 | st , at ) : S×S×A 7→ [ 0 , +∞ ) . Goal-conditioned RL also includes a goal space G. For simplicity , we will assume in our derivation that the goal space matches the state space , such that G = S , though the approach extends trivially to the case where G is a hand-specified subset of S , such as the global x-y position of a robot . A goal-conditioned policy π ( a | s , g ) maps a state s ∈ S and goal g ∈ S to a distribution over actions a ∈ A , and its objective is to reach the goal , i.e. , to make the current state equal to the goal . Goal-reaching can be formulated as minimizing H ( G | S ) , and many practical goal-reaching algorithms ( Kaelbling , 1993 ; Lillicrap et al. , 2016 ; Schaul et al. , 2015 ; Andrychowicz et al. , 2017 ; Nair et al. , 2018 ; Pong et al. , 2018 ; Florensa et al. , 2018a ) can be viewed as approximations to this objective by observing that the optimal goal-conditioned policy will deterministically reach the goal , resulting in a conditional entropy of zero : H ( G | S ) = 0 . See Appendix E for more details . Our method may thus be used in conjunction with any of these prior goal-conditioned RL methods in order to jointly minimizeH ( G | S ) and maximizeH ( G ) . 2.2 MAXIMIZING H ( G ) : SETTING DIVERSE GOALS We now turn to the problem of setting diverse goals or , mathematically , maximizing the entropy of the goal distributionH ( G ) . Let US be the uniform distribution over S , where we assume S has finite volume so that the uniform distribution is well-defined . Let pφ be the goal distribution from which goals G are sampled . Our goal is to maximize the entropy of pφ , which we write asH ( G ) . Since the maximum entropy distribution over S is the uniform distribution US , maximizingH ( G ) may seem as simple as choosing the uniform distribution to be our goal distribution : pφ = US . However , this requires knowing the uniform distribution over valid states , which may be difficult to obtain when S is a subset of Rn , for some n. For example , if the states correspond to images viewed through a robot ’ s camera , S corresponds to the ( unknown ) set of valid images of the robot ’ s environment , while Rn corresponds to all possible arrays of pixel values of a particular size . In such environments , sampling from the uniform distribution Rn is unlikely to correspond to a valid image of the real world . Sampling uniformly from S would require knowing the set of all possible valid images , which we assume the agent does not know when starting to explore the environment . While we can not sample arbitrary states from S , we can sample states by performing goal-directed exploration . To derive and analyze our method , we introduce a simple model of this process : a goal G ∼ pφ is sampled from the goal distribution pφ , and then the agent attempts to achieve this goal , which results in a distribution of states S ∈ S seen along the trajectory . We abstract this entire process by writing the resulting marginal distribution over S as p ( S | pφ ) . We assume that p ( S | pφ ) has full support , which can be accomplished with an epsilon-greedy goal reaching policy in a communicating MDP . We also assume that the entropy of the resulting state distribution H ( p ( S | pφ ) ) is no less than the entropy of the goal distributionH ( pφ ( S ) ) . Without this assumption , a policy could ignore the goal and stay in a single state , no matter how diverse and realistic the goals are . Note that this assumption does not require that the entropy of p ( S | pφ ) is strictly larger than the entropy of the goal distribution , pφ . This simplified model allows us to analyze the behavior of our goal-setting scheme separately from any specific goal-reaching algorithm . We will however show in Section 6 that we can instantiate this approach into a practical algorithm that jointly learns the goal-reaching policy . In summary , our goal is to acquire a maximum-entropy goal distribution pφ over valid states S , while only having access to state samples from p ( S | pφ ) . 3 SKEW-FIT : LEARNING A MAXIMUM ENTROPY GOAL DISTRIBUTION . Our method , Skew-Fit , learns a maximum entropy goal distribution pφ using samples collected from a goal-conditioned policy . We analyze the algorithm and show that Skew-Fit maximizes the entropy of the goal distribution , and present a practical instantiation for unsupervised deep RL . 3.1 SKEW-FIT ALGORITHM . To learn a uniform distribution over valid goal states , we present a method that iteratively increases the entropy of a generative model pφ . In particular , given a generative model pφt at iteration t , we would like to train a new generative model pφt+1 such that pφt+1 has higher entropy than pφt over the set of valid states . While we do not know the set of valid states S , we can sample states from p ( S | pφt ) , resulting in an empirical distribution pempt over the states pempt ( s ) , 1 N N∑ n=1 1 { s = Sn } , Sn ∼ p ( S | pφt ) , ( 2 ) and use this empirical distribution to train the next generative model pφt+1 . However , if we simply train pφt+1 to model this empirical distribution , it may not necessarily have higher entropy than pφt . The intuition behind our method is quite simple : rather than fitting a generative model to our empirical distribution , we skew the empirical distribution so that rarely visited states are given more weight . See Figure 2 for a visualization of this process . How should we skew the empirical distribution if we want to maximize the entropy of pφt+1 ? If we had access to the density of each state , pempt ( S ) , then we could simply weight each state by 1/pempt ( S ) . We could then perform maximum likelihood estimation ( MLE ) for the uniform distribution by using the following loss to train φt+1 : L ( φ ) = ES∼US [ log pφ ( S ) ] = ES∼pempt [ US ( S ) pempt ( S ) log pφ ( S ) ] ∝ ES∼pempt [ 1 pempt ( S ) log pφ ( S ) ] where we use the fact that the uniform distribution US ( S ) has constant density for all states in S. However , computing this density pempt ( S ) requires marginalizing out the MDP dynamics , which requires an accurate model of both the dynamics and the goal-conditioned policy . We avoid needing to model the entire MDP process by approximating pempt ( S ) with our previous learned generative model : pempt ( S ) ≈ p ( S | pφt ) ≈ pφt ( S ) . We therefore weight each state by the following weight function wt , α ( S ) , pφt ( S ) α , α < 0 . ( 3 ) where α is a hyperparameter that controls how heavily we weight each state . If our approximation pφt was exact , we could choose α = −1 and recover the exact importance sampling procedure described above . If α = 0 , then this skew step has no effect . By choosing intermediate values of α , we can trade off the reliability of our estimate pφt ( S ) with the speed at which we want to increase the entropy of the goal distribution . Variance Reduction As described , this procedure relies on importance sampling ( IS ) , which can have high variance , particularly if pφt ( S ) ≈ 0 . We therefore choose a class of generative models where the probabilities are prevented from collapsing to zero , as we will describe in Section 4 . To further reduce the variance , we train pφt+1 with sampling importance resampling ( SIR ) ( Rubin , 1988 ) . Rather than sampling from pempt and weighting the update from each sample by wt , α , SIR explicitly defines a skewed distribution as pskewedt ( s ) , 1 Zα pempt ( s ) wt , α ( s ) , Zα = N∑ n=1 pempt ( Sn ) wt , α ( Sn ) , ( 4 ) where Zα is the normalizing coefficient and pempt is given by Equation 2 . We note that computing Zα adds little computational overhead , since all of the weights already need to be computed . We then fit the generative model at the next iteration pφt+1 to pskewedt using standard MLE . We found that using SIR resulted in significantly lower variance than IS . See Appendix B.3 for this comparision . Goal Sampling Alternative Because pφt+1 ≈ pskewedt , at iteration t + 1 , one can sample goals from either pφt+1 or pskewedt . Sampling goals from pskewedt may be preferred if sampling from the learned generative model pφt+1 is computationally or otherwise challenging . In either case , one still needs to train the generative model pφt to create pskewedt . In our experiments , we found that both methods perform well . Summary Overall , Skew-Fit samples data from the environment and weights different samples by their density under the generative model pφt . We prove in the next section conditions under which this weighting makes the generative model at the next iteration pφt+1 have higher entropy . With higher entropy , the pφt+1 is more likely to generate goals at the frontier of unseen states , which results in more uniform state coverage . Skew-Fit is shown in Figure 2 and summarized in Algorithm 1 . Algorithm 1 Skew-Fit 1 : for Iteration t = 1 , 2 , ... do 2 : Collect N states { Si } Ni=1 by sampling goals from pφt ( or pskewedt ) and running goal- conditioned policy . 3 : Construct skewed distribution pskewedt ( Equation 3 and Equation 4 ) . 4 : Fit pφt+1 to skewed distribution pskewedt using MLE . 5 : end for

Summary: This paper introduced a very interesting idea to facilitate exploration in goal-conditioned reinforcement learning. The key idea is to learn a generative model of goal distribution to match the weighted empirical distribution, where the rare states receive larger weights. This encourages the model to generate more diverse and novel goals for goal-conditioned RL policies to reach.