Efficient Exploration via State Marginal Matching

1 INTRODUCTION . Reinforcement learning ( RL ) algorithms must be equipped with exploration mechanisms to effectively solve tasks with limited reward signals . These tasks arise in many real-world applications where providing human supervision is expensive . The inability of current RL algorithms to adequately explore limits their applicability to long-horizon control tasks . A wealth of prior work has studied exploration for RL . While , in theory , the Bayes-optimal exploration strategy is optimal , it is intractable to compute exactly , motivating work on tractable heuristics for exploration . Exploration methods based on random actions have limited ability to cover a wide range of states . More sophisticated techniques , such as intrinsic motivation , accelerate learning in the single-task setting . However , these methods have two limitations . First , they do not explicitly define an objective to quantify “ good exploration , ” but rather argue that exploration arises implicitly through some iterative procedure . Lacking a well-defined optimization objective , it remains challenging to understand what these methods are doing and why they work . Similarly , the lack of a metric to quantify exploration , even if only for evaluation , makes it challenging to compare exploration methods and assess progress in this area . The second limitation is that these methods target the single-task setting . Because these methods aim to converge to the optimal policy for a particular task , it is challenging to repurpose these methods to solve multiple tasks . We address these shortcomings by considering a multi-task setting , where many different reward functions can be provided for the same set of states and dynamics . Rather than exploring from scratch for each task , we aim to learn a single , task-agnostic exploration policy that can be adapted to many possible downstream reward functions , amortizing the cost of learning to explore . This exploration 1Videos and code : https : //sites.google.com/view/smm-anonymous policy can be viewed as a prior on the policy for solving downstream tasks . Learning will consist of two phases : during training , we acquire this task-agnostic exploration policy ; during testing , we use this exploration policy to quickly explore and maximize the task reward . Learning a single exploration policy is considerably more difficult than doing exploration throughout the course of learning a single task . The latter is done by intrinsic motivation ( Pathak et al. , 2017 ; Tang et al. , 2017 ; Oudeyer et al. , 2007 ) and count-based exploration methods ( Bellemare et al. , 2016 ) , which can effectively explore to find states with high reward , at which point the agent can decrease exploration and increase exploitation of those high-reward states . While these methods perform efficient exploration for learning a single task , the policy at any particular iteration is not a good exploration policy . For example , the final policy at convergence would only visit the high-reward states discovered for the current task . What objective should be optimized to obtain a good exploration policy ? We recast exploration as a problem of State Marginal Matching : given a desired state distribution , we learn a mixture of policies for which the state marginal distribution matches this desired distribution . Without any prior information , this objective reduces to maximizing the marginal state entropyH [ s ] , which encourages the policy to visit as many states as possible . The distribution matching objective also provides a convenient mechanism to incorporate prior knowledge about the task , whether in the form of safety constraints that the agent should obey ; preferences for some states over other states ; reward shaping ; or the relative importance of each state dimension for a particular task . We also propose an algorithm to optimize the State Marginal Matching ( SMM ) objective . First , we reduce the problem of SMM to a two-player , zero-sum game between a policy player and a density player . We find a Nash Equilibrium for this game using fictitious play ( Brown , 1951 ) , a classic procedure from game theory . Our resulting algorithm iteratively fits a state density model and then updates the policy to visit states with low density under this model . Our analysis of this approach sheds light on prior work on exploration . In particular , while the policy learned by existing exploration algorithms does not perform distribution matching , the replay buffer does , an observation that potentially explains the success of prior methods . On both simulated and real-world tasks , we demonstrate that our algorithm explores more effectively and adapts more quickly to new tasks than state-of-the-art baselines . 2 RELATED WORK . Most prior work on exploration has looked at exploration bonuses and intrinsic motivation . One class of exploration methods uses prediction error of some auxiliary task as an exploration bonus , which provides high ( intrinsic ) reward in states where the predictive model performs poorly ( Pathak et al. , 2017 ; Oudeyer et al. , 2007 ; Schmidhuber , 1991 ; Houthooft et al. , 2016 ; Burda et al. , 2018 ) . Another set of approaches ( Tang et al. , 2017 ; Bellemare et al. , 2016 ; Schmidhuber , 2010 ) directly encourage the agent to visit novel states . While all methods effectively explore during the course of solving a single task ( Taïga et al. , 2019 ) , the policy obtained at convergence is often not a good exploration policy ( see Section 4 ) . In contrast , our method converges to a highly-exploratory policy by maximizing state entropy in the training objective ( Eq . 2 ) . Many exploration algorithms can be classified by whether they explore in the space of actions , policy parameters , goals , or states . Common exploration strategies including -greedy and Ornstein–Uhlenbeck noise ( Lillicrap et al. , 2015 ) , as well as standard MaxEnt RL algorithms ( Ziebart , 2010 ; Haarnoja et al. , 2018 ) , explore in the action space . Recent work ( Fortunato et al. , 2017 ; Plappert et al. , 2017 ) shows that adding noise to the parameters of the policy can result in good exploration . Most closely related to our work are methods that perform exploration in the space of states or goals ( Colas et al. , 2018 ; Held et al. , 2017 ; Nair et al. , 2018 ; Pong et al. , 2019 ; Hazan et al. , 2018 ) . In fact , Hazan et al . ( 2018 ) consider the same State Marginal Matching objective that we examine and propose a similar algorithm . In relation to Hazan et al . ( 2018 ) , our main contributions are ( 1 ) empirically showing that exploration based on state-entropy is competitive with existing state-of-the-art exploration methods , and ( 2 ) explaining how existing exploration methods based on prediction error are implicitly maximizing this state-entropy objective . In Appendix C.1 , we also discuss how goal-conditioned RL ( Kaelbling , 1993 ; Schaul et al. , 2015 ) can be viewed as a special case of State Marginal Matching when the goal-sampling distribution is learned jointly with the policy . The problems of exploration and meta-reinforcement learning are tightly coupled . Meta-reinforcement learning algorithms ( Duan et al. , 2016 ; Finn et al. , 2017 ; Rakelly et al. , 2019 ; Mishra et al. , 2017 ) must perform effective exploration if they hope to solve a downstream task . Some prior work has explicitly looked at the problem of learning to explore ( Gupta et al. , 2018 ; Xu et al. , 2018 ) . Our problem statement is similar to meta-learning , in that we also aim to learn a policy as a prior for solving downstream tasks . However , whereas meta-RL requires a distribution of task reward functions , our method will require only a single target state marginal distribution . Due to the simpler problem assumptions and training procedure , our method may be easier to apply in real-world domains . Related to our approach are standard maximum action entropy algorithms ( Haarnoja et al. , 2018 ; Kappen et al. , 2012 ; Rawlik et al. , 2013 ; Ziebart et al. , 2008 ; Theodorou & Todorov , 2012 ) . While these algorithms are referred to as MaxEnt RL , they are maximizing entropy over actions , not states . These algorithms can be viewed as performing inference on a graphical model where the likelihood of a trajectory is given by its exponentiated reward ( Toussaint & Storkey , 2006 ; Levine , 2018 ; Abdolmaleki et al. , 2018 ) . While distributions over trajectories induce distributions over states , computing the exact relationship requires integrating over all possible trajectories , an intractable problem for most MDPs . A related but distinct class of relative entropy methods use a similar entropy-based objective to limit the size of policy updates ( Peters et al. , 2010 ; Schulman et al. , 2015 ) . Finally , the idea of distribution matching has been employed successfully in imitation learning ( Ziebart et al. , 2008 ; Ho & Ermon , 2016 ; Finn et al. , 2016 ; Fu et al. , 2017 ) . Similar to some inverse RL algorithms ( Ho & Ermon , 2016 ; Fu et al. , 2018 ) , our method iterates between learning a policy and learning a reward function , though our reward function is obtained via a density model instead of a discriminator . While inverse RL algorithms assume access to expert trajectories , we instead assume access to the density of the target state marginal distribution . In many realistic settings , such as robotic control with many degrees of freedom , providing fully-specified trajectories may be much more challenging than defining a target state marginal distribution . The latter only requires some aggregate statistics about expert behavior , and does not even need to be realizable by any policy . In summary , our work unifies prior exploration methods as performing approximate distribution matching , and explains how state distribution matching can be performed properly . This perspective provides a clearer picture of exploration , and this observation is useful particularly because many of the underlying ingredients , such as adversarial games and density estimation , have seen recent progress and therefore might be adopted to improve exploration methods . 3 STATE MARGINAL MATCHING . In this section , we propose the State Marginal Matching problem as a principled objective for learning to explore , and offer an algorithm for optimizing it . We consider a parametric policy πθ ∈ Π , { πθ | θ ∈ Θ } that chooses actions a ∈ A in a Markov Decision Process ( MDP ) M with fixed episode lengths T , dynamics distribution p ( st+1 | st , at ) , and initial state distribution p0 ( s ) . The MDPM together with the policy πθ form an implicit generative model over states . We define the state marginal distribution ρπ ( s ) as the probability that the policy visits state s : ρπ ( s ) , E s1∼p0 ( S ) , at∼πθ ( A|st ) st+1∼p ( S|st , at ) [ 1 T T∑ t=1 1 ( st = s ) ] We emphasize that ρπ ( s ) is not a distribution over trajectories , and is not the stationary distribution of the policy after infinitely many steps , but rather the distribution over states visited in a finite-length episode.2 We also note that any trajectory distribution matching problem can be reduced to a state marginal matching problem by augmenting the current state to include all previous states . We assume that we are given a target distribution p∗ ( s ) over states s ∈ S that encodes our belief about the tasks we may be given at test-time . For example , a roboticist might assign small values of p∗ ( s ) to states that are dangerous , regardless of the desired task . Alternatively , we might also learn p∗ ( s ) from data about human preferences ( Christiano et al. , 2017 ) . For goal-reaching tasks , we can analytically derive the optimal target distribution ( Appendix C ) . Given p∗ ( s ) , our goal is to find a 2ρπ ( s ) approaches the policy ’ s stationary distribution in the limit as the episodic horizon T →∞ . parametric policy that is “ closest ” to this target distribution , where we measure discrepancy using the Kullback-Leibler ( KL ) divergence : max π∈Π F ( ρπ ( s ) , p∗ ( s ) ) , max π∈Π −DKL ( ρπ ( s ) ‖ p∗ ( s ) ) = max π∈Π Es∼ρπ ( s ) [ log p ∗ ( s ) − log ρπ ( s ) ] ( 1 ) = max π∈Π Es∼ρπ ( s ) [ log p ∗ ( s ) ] +Hπ [ s ] ( 2 ) This is the same objective as in Hazan et al . ( 2018 ) . Note that we use the reverse-KL ( Bishop , 2006 ) , which is mode-seeking ( i.e. , exploratory ) . We show in Appendix C that the policies obtained via State Marginal Matching provide an optimal exploration strategy for a particular distribution over reward functions . To gain intuition for the State Marginal Matching objective , we decomposed it in two ways . In Equation 2 , we see that State Marginal Matching is equivalent to maximizing the reward function r ( s ) , log p∗ ( s ) while simultaneously maximizing the entropy of states . Note that , unlike traditional MaxEnt RL algorithms ( Ziebart et al. , 2008 ; Haarnoja et al. , 2018 ) , we regularize the entropy of the state distribution , not the conditional distribution of actions given states , which results in exploration in the space of states rather than in actions . Moreover , Equation 1 suggests that State Marginal Matching maximizes a pseudo-reward r ( s ) , log p∗ ( s ) − log ρπ ( s ) , which assigns positive utility to states that the agent visits too infrequently and negative utility to states visited too frequently ( see Figure 1 ) . We emphasize that maximizing this pseudo-reward is not a RL problem because the pseudo-reward depends on the policy .

Summary: The paper proposes to frame exploration in reinforcement learning as a distribution matching problem. More specifically, the proposed method (SMM) aims to minimize the reverse KL between the state distribution induced by the policy and a desired state distribution. The desired state distribution can be used to guide exploration (e.g. by penalizing bad states) or can be chosen uniformly (resulting in a policy that maximizes state entropy).