Diverse Exploration via InfoMax Options

1 INTRODUCTION . Abstracting a course of action as a higher-level action , or an option ( Sutton et al. , 1999 ) , is a key ability for reinforcement learning ( RL ) agents in several aspects , including exploration . In RL problems , an agent learns to approximate an optimal policy only from experience , given no prior knowledge . This leads to the necessity of exploration : an agent needs to explore the poorly known states for collecting environmental information , sometimes sacrificing immediate rewards . For statistical efficiency , it is important to explore the state space in a deep and directed manner , rather than taking uniformly random actions ( Osband et al. , 2019 ) . Options can represent such directed behaviors by capturing long state jumps from their starting regions to terminating regions . It has been shown that well-defined options can facilitate exploration by exploiting an environmental structure ( Barto et al. , 2013 ) or , more generally , by reducing decision steps ( Fruit and Lazaric , 2017 ) . A key requirement for such explorative options is diversity . If all options have the same terminating region , they will never encourage exploration . Instead , options should lead to a variety of regions for encouraging exploration . However , automatically discovering diverse options in a scalable , online manner is challenging due to two difficulties : generalization and data limitation . Generalization with function approximation ( Sutton , 1995 ) is important for scaling up RL methods to large or continuous domains . However , many existing option discovery methods for exploration are graph-based ( e.g. , Machado et al . ( 2017 ) ) and incompatible with function approximation , except for that by Jinnai et al . ( 2020 ) . Discovering options online in parallel with polices requires us to work with limited data sampled from the environment and train the model for evaluating the diversity in a data-efficient manner . To address these difficulties , we introduce the infomax termination objective defined as the mutual information ( MI ) between options and their corresponding state transitions . This formulation reflects a simple inductive bias : for encouraging exploration , options should terminate in a variety of regions per starting regions . Thanks to the information-theoretical formulation , this objective is compatible with function approximation and scales up to continuous domains . A key technical contribution of this paper is the optimization scheme for maximizing this objective . Specifically , we employ a simple classification model over options as a critic for termination conditions , which makes our method data-efficient and tractable in many domains . The paper is organized as follows . After introducing background and notations , we present the infomax termination objective and derive a practical optimization scheme using the termination gradient theorem ( Harutyunyan et al. , 2019 ) . We then implement the infomax objective on the option-critic architecture ( OC ) ( Bacon et al. , 2017 ) with algorithmic modifications , yielding the InfoMax Option Critic ( IMOC ) algorithm . Empirically , we show that ( i ) IMOC improves exploration in structured environments , ( ii ) IMOC improves exporation in lifelong learning , ( iii ) IMOC is scalable to MuJoCo continuous control tasks , and ( iv ) the options learned by IMOC are diverse and meaningful . We then relate our method to other option-learning methods and the empowerment concept ( Klyubin et al. , 2005 ) , and finally give concluding remarks . 2 BACKGROUND AND NOTATION . We assume the standard RL setting in the Markov decision process ( MDP ) , following Sutton and Barto ( 2018 ) . An MDPM consists of a tuple ( X , A , p , r , γ ) , where X is the set of states , A is the set of actions , p : X ×A×X → [ 0 , 1 ] is the state transition function , r : X ×A → [ rmin , rmax ] is the reward function , and 0 ≤ γ ≤ 1 is the discount factor . A policy is a probability distribution over actions conditioned on a state x , π : X ×A → [ 0 , 1 ] . For simplicity , we consider the episodic setting where each episode ends when a terminal state xT is reached . In this setting , the goal of an RL agent is to approximate a policy that maximizes the expected discounted cumulative reward per episode : JRL ( π ) = Eπ , x0 [ T−1∑ t=0 γtRt ] , ( 1 ) where Rt = r ( xt , at ) is the reward received at time t , and x0 is the initial state of the episode . Relatedly , we define the action-value function Qπ ( xt , at ) def = Ext , at , π [ ∑T−1 t′=t γ t′−tRt′ ] and the state-value function V π ( xt ) def = Ext , π ∑ a π ( a|xt ) Qπ ( xt , a ) . Assuming that π is differentiable by the policy parameters θπ , a simple way to maximize the objective ( 1 ) is the policy gradient method ( Williams , 1992 ) that estimates the gradient by : ∇θπJRL ( π ) = Eπ , xt [ ∇θπ log π ( at|xt ) Â ( xt , at ) ] , ( 2 ) where Â ( xt , at ) is the estimation of the advantage function Aπ ( xt , at ) def = Qπ ( xt , at ) − V π ( xt ) . A common choice of Â ( xt , at ) is N -step TD error ∑N i=0 γ iRt+i + γ N V̂ ( xt+N ) − V̂ ( xt ) , where N is a fixed rollout length ( Mnih et al. , 2016 ) . 2.1 OPTIONS FRAMEWORK . Options ( Sutton et al. , 1999 ) provide a framework for representating temporally abstracted actions in RL . An option o ∈ O consists of a tuple ( Io , βo , πo ) , where Io ⊆ X is the initiation set , βo : X → [ 0 , 1 ] is a termination function with βo ( x ) denoting the probability that option o terminates in state x , and πo is intra-option policy . Following related studies ( Bacon et al. , 2017 ; Harutyunyan et al. , 2019 ) , we assume that Io = X and learn only βo and πo . Letting xs denote an option-starting state and xf denote an option-terminating state , we can write the option transition function as : P o ( xf |xs ) = βo ( xf ) Ixf=xs + ( 1− βo ( xs ) ) ∑ x pπ o ( x|xs ) P o ( xf |x ) , ( 3 ) where I is the indicator function and pπo is the policy-induced transition function pπo ( x′|x ) def=∑ a∈A π o ( a|x ) p ( x′|x , a ) . We assume that all options eventually terminate so that P o is a valid probability distribution over xf , following Harutyunyan et al . ( 2019 ) . To present option-learning methods , we define two option-value functions : QO and UO , where QO is the option-value function denoting the value of selecting an option o at a state xt defined by QO ( xt , o ) def = Eπ , β , µ [ ∑T−1 t′=t γ t′−tRt ] . Analogously to Qπ and V π , we let VO denote the marginalized option-value function VO ( x ) def = ∑ o µ ( o|x ) QO ( x , o ) , where µ ( o|xs ) : X ×O → [ 0 , 1 ] is the policy over options . Function UO ( x , o ) def = ( 1− βo ( x ) ) QO ( x , o ) + βo ( x ) VO ( x ) is called the option-value function upon arrival ( Sutton et al. , 1999 ) and denotes the value of reaching a state xt with o and not having selected the new option . 2.2 OPTION CRITIC ARCHITECTURE . OC ( Bacon et al. , 2017 ) provides an end-to-end algorithm for learning πo and βo in parallel . To optimize πo , OC uses the intra-option policy gradient method that is the option-conditional version of the gradient estimator ( 2 ) , ∇θπoJRL ( πo ) = E [ ∇θπo log πo ( at|xt ) Âo ( xt , at ) ] , where Âo is an estimation of the option-conditional advantage Aπ o . For optimizing βo , OC directly maximizes QO using the estimated gradient : ∇θβoQO ( x , o ) = γE [ −∇θβoβ o ( x ) ( QO ( x , o ) − VO ( x ) ) ] . ( 4 ) Intuitively , this decreases the termination probability βo ( x ) when holding an o is advantageous , i.e. , QO ( x ) − VO ( x ) is positive , and vice versa . Our method basically follows OC but has a different objective for learning βo . 2.3 TERMINATION CRITIC . Recently proposed termination critic ( TC ) ( Harutyunyan et al. , 2019 ) optimizes βo by maximizing the information-theoretic objective called predictability : JTC ( P o ) = −H ( Xf |o ) , ( 5 ) where H denotes entropy and Xf is the random variable denoting the option-terminating states . Maximizing −H ( Xf |o ) makes the terminating region of an option smaller and more predictable . In other words , we can compress terminating regions by optimizing the objective ( 5 ) . To derivate this objective by the beta parameters θβo , Harutyunyan et al . ( 2019 ) introduced the termination gradient theorem : Theorem 1 . Let βo be parameterized with a sigmoid function and ` βo denote the logit of βo . We have ∇θβP o ( xf |xs ) = ∑ x P o ( x|xs ) ∇θβ ` βo ( x ) ( Ixf=x − P o ( xf |x ) ) , ( 6 ) Leveraging the theorem 1 , TC performs gradient ascent using the estimated gradient : ∇θβoJ TC ( P o ) = −Exs , x , xf [ ∇θβ ` βo ( x ) βo ( x ) ( ( logP oµ ( x ) − logP oµ ( xf ) ) + ( 1− P o ( xf |xs ) P oµ ( x ) P oµ ( xf ) P o ( x|xs ) ) ) ] . where P oµ ( x ) is the marginalized distribution of option-terminating states . Contrary to the termination objective of OC ( 4 ) , this objective does not depend on state values , making learned options robust against the reward structure of the environment . Our method is inspired by TC and optimizes a similar information-theoretic objective , not for predictability but for diversity . Also , our infomax objective requires an estimation of p̂ ( o|xs , xf ) instead of the option transition model P o ( xf |xs ) , which makes our method tractable in more environments . 3 INFOMAX OPTION CRITIC . We now present the key idea behind the InfoMax Option Critic ( IMOC ) algorithm . We first formulate the infomax termination objective based on the MI maximization , then derive a practical gradient estimation for maximizing this objective on βo , utilizing the termination gradient theorem 1 . To evaluate the diversity of options , we use the MI between options and option-terminating states conditioned by option-starting states : J IMOC = I ( Xf ; O|Xs ) = H ( Xf |Xs ) −H ( Xf |Xs , O ) , ( 7 ) where I denotes conditional MI I ( A ; B|Z ) = H ( A|Z ) −H ( A|B , Z ) , Xs is the random variable denoting an option-starting state , and O is the random variable denoting an option . We call this objective the infomax termination objective . Let us interpret Xf |Xs as the random variable denoting a state transition induced by an option . Then maximizing the MI ( 7 ) ( i ) diversifies a state transition Xf |Xs and ( ii ) makes an option-conditional state transition Xf |Xs , o more deterministic . Note that the marginalized MI I ( Xf ; O ) also makes sense in that it prevents the terminating region of each option from being too broad , as predictability ( 5 ) does . However , in this study , we focus on the conditional objective since it is easier to optimize . To illustrate the limitation of infomax options , we conducted an analysis in a toy four-state deterministic chain environment , which has four states and two deterministic actions ( go left and go right ) per each state . Since deriving the exact solution is computationally difficult , we searched options that maximize H ( Xf |Xs ) from deterministic options that has deterministic option-policies and termination functions ( thus has the minimum H ( Xf |Xs , O ) ) . Among multiple solutions , Figure 1 shows two interesting instances of deterministic infomax options when |O| = 2 . The left options enable diverse behaviors per state , although they fail to capture long-term behaviors generally favorable in the literature ( e.g. , Mann et al . ( 2015 ) ) . On the other hand , the right options enable relatively long , two step state transitions , but they are the same and the rightmost state and the next one . Furthermore , an agent can be caught in a small loop that consists of the leftmost state and the next one . This example shows that ( i ) we can obtain short and diverse options with only a few options , ( ii ) to obtain long and diverse options , we need sufficiently many options , and ( iii ) an agent can be caught in a small loop with only a few options , failing to visit diverse states . As we show in Appendix A , this ’ small loop ’ problem can not happen with four options . Thus , the number of options is important when we are to maximize the MI ( 7 ) and a limitation of this method . However , in experiments , we show that we can practically learn diverse options with relatively small number of options . For maximizing the MI by gradient ascent , we now derive the gradient of the infomax termination objective ( 7 ) . First , we estimate the gradient of the objective using the option transition model P o and marginalized option-transition model P ( xf |xs ) = ∑ o µ ( o|xs ) P o ( xf |xs ) . Proposition 1 . Let βo be parameterized with a sigmoid function . Given a trajectory τ = xs , . . . , x , . . . , xf sampled by πo and βo , we can obtain unbiased estimations of∇θβH ( Xf |Xs ) and ∇θβH ( Xf |Xs , O ) by ∇θβH ( Xf |Xs ) = Exs , x , xf , o [ −∇θβ ` βo ( x ) βo ( x ) ( logP ( x|xs ) − logP ( xf |xs ) ) ] ( 8 ) ∇θβH ( Xf |Xs , O ) = Exs , x , xf , o [ −∇θβ ` βo ( x ) βo ( x ) ( logP o ( x|xs ) − logP o ( xf |xs ) ) ] ( 9 ) where ` βo ( x ) denotes the logit of βo ( x ) . Note that the additional term βo is necessary because x is not actually a terminating state . The proof follows section 4 in Harutyunyan et al . ( 2019 ) and is given in Appendix B.1 . The estimated gradient of the infomax termination objective ( 7 ) can now be written as : ∇θβI ( Xf ; O|Xs ) = ∇θβH ( Xf |Xs ) −∇θβH ( Xf |Xs , O ) = Exs , x , xf , o [ −∇θβ ` βo ( x ) βo ( x ) ( logP ( x|xs ) − logP ( xf |xs ) − ( logP o ( x|xs ) − logP o ( xf |xs ) ) ) ] , ( 10 ) which means that we can optimize this objective by estimating P o and P . However , estimating the probability over the state space can be difficult , especially when the state space is large , as common in the deep RL setting . Hence , we reformulate the gradient using Bayes ’ rule in a similar way as Gregor et al . ( 2017 ) . The resulting term consists of the reverse option transition p ( o|xs , xf ) that denotes the probability of having an o given a state transition xs , xf . Proposition 2 . We now have ∇θβI ( Xf ; O|Xs ) =∇θβH ( Xf |Xs ) −∇θβH ( Xf |Xs , O ) =Exs , x , xf , o [ ∇θβ ` βo ( x ) βo ( x ) ( log p ( o|xs , x ) − log p ( o|xs , xf ) ) ] ( 11 ) The proof is given in Appendix B.2 . In the following sections , we estimate the gradient ( 11 ) by learning a classification model over options p̂ ( o|xs , xf ) from sampled option transitions .

Summary: The authors propose a modification of the option-critic algorithm for hierarchical reinforcement learning. The proposed algorithm modifies how the termination conditions of the options are improved by experience. Specifically, the algorithm aims to maximize the mutual information between the options and their termination states. The authors develop an optimization scheme for achieving this objective and provide empirical results in a number of domains