Learning Diverse Options via InfoMax Termination Critic

1 INTRODUCTION . Behavior learning from environmental interaction is a fundamental problem in artificial intelligence and robotics . Recently , combined with deep neural networks ( DNN ) , reinforcement learning ( RL ) has successfully learned complex behaviors guided by human-designed reward functions ( Heess et al. , 2017 ; OpenAI et al. , 2019 ) . To reuse trained agents across multiple reward functions , or shortly tasks , abstracting a course of action as a higher-level building block ( Barto and Mahadevan , 2003 ; Barto et al. , 2013 ) would be useful . A representative formulation of action abstraction in RL is the options framework ( Sutton et al. , 1999 ) , where each option consists of a sub policy and its termination condition . Serving as reusable behavioral blocks , options have the potential of accelerating learning in new tasks ( Brunskill and Li , 2014 ) . For example , if a wheeled robot already had learned to lean and turn left and right , it can learn to drive on unfamiliar roads faster than learning from scratch . However , discovering reusable options remains challenging due to the difficulty of defining and measuring reusability a priori . Although we can measure the reusability of options when we know the task distribution ( e.g. , in a Bayesian sense ( Solway et al. , 2014 ) ) , it is difficult to measure it without prior knowledge . Hence , in such cases , we need a reasonable heuristic assumption about future tasks . In this paper , we follow a common assumption that diversifying options improves the reusability of options ( e.g. , Barto et al . ( 2004 ) ) . That is , we assume that tasks are uniformly distributed over the state space . Diverse options are expected to work to some extent for any task , thus we argue that this is a reasonable heuristics . Under this assumption , the next problem is to measure the diversity of options . The most intuitive strategy would be visiting all states as equally as possible . However , this can be computationally too expensive when the environment has continuous or large state spaces . As an alternative , in the context of learning skills without termination conditions , mutual information ( MI ) maximization has been widely applied for learning diverse policies ( Eysenbach et al. , 2019 ; Baumli et al. , 2021 ; Choi et al. , 2021 ) . An advantage of MI maximization is the existence of tractable approximations with DNN , scaling it up to more complex domains . For learning diverse options , we maximize MI between options and their terminating states conditioned by a starting state . This MI maximization diversifies the destinations of an agent when choosing different options at a state while keeping the transition as deterministic as possible , leading to diverse and meaningful options . Specifically , we propose to maximize this MI by optimizing the termination conditions of options . We derive an unbiased gradient estimator of this MI w.r.t . termination conditions using the termination gradient theorem ( Harutyunyan et al. , 2019 ) . As a key contribution , we reformulate the estimated gradient using Bayes ’ rule and derive a tractable approximation with a leaned classifier , yielding the InfoMax Termination Critic ( IMTC ) algorithm . The easiness of approximation is advantageous compared to the Termination Critic ( TC , Harutyunyan et al . ( 2019 ) ) . The same objective is used by Variational Intrinsic Control ( VIC , Gregor et al . ( 2016 ) ) for diversifying sub-policies , but the original VIC used a constant termination probability for all states and options . However , since VIC depends on relating starting and terminating states to an option , we argue that diversifying termination conditions is also important . We show that our method helps VIC learn clearer options . In this paper , we first introduce backgrounds and notations , followed by the derivation of IMTC . We implement IMTC on the Option Critic Architecture ( OC , Bacon et al . ( 2017 ) ) and Proximal Policy Optimization ( PPO , Schulman et al . ( 2017 ) ) . As a minor contribution , our implementation includes a newly proposed optimistic advantage estimation that speeds up option learning for a single task . In experiments , we qualitatively demonstrate that IMTC successfully learns diverse and meaningful options in a reward-free RL ( Jin et al. , 2020 ) setting . Figure 1 shows the learned options by IMTC and OC combined with VIC rewards in PointBilliard domain . We can see that options learned by IMTC are clearly separated and different . Then we test the reusability of learned options in task adaptation experiments : we first pre-train options without rewards and then transfer them on various tasks . We show that IMTC helps quick adaptation to specific tasks especially in complex tasks where object manipulation is required . Since IMTC is only for learning terminating conditions of options , we can combine it with other methods for learning diverse sub-policies . Our experiments observed that IMTC improves the performance of VIC and RVIC ( Baumli et al. , 2021 ) . 2 BACKGROUND AND NOTATION . We assume the standard RL setting in the Markov decision process ( MDP ) , following Sutton and Barto ( 2018 ) . MDP M consists of a tuple ( X , A , p , r ) , where X is the finite set of states , A is the finite set of actions , p : X × A × X → [ 0 , 1 ] is the state transition function , r : X × A → [ rmin , rmax ] is the reward function , where rmin , rmax are minimum and maximum . We use Xt and At for denoting random variables of the state and action experienced at time t. We define G def = ∑∞ t=0 γ tRt as discounted cumulative return , where Rt = r ( Xt , At ) is the reward received at time t and 0 < γ < 1 is the discount factor . We consider maximizing expected discounted return E [ G|π ] . As an agent , we consider a memoryless policy π : X × A → [ 0 , 1 ] . π induces two value functions : action-value function Qπ ( x , a ) def= Eπ [ G|X0 = x , A0 = a , π ] and state-value function V π ( x ) def = ∑ a π ( a|x ) Qπ ( x , a ) . We use dπ for denoting an agent ’ s average occupancy measure dπ ( x ) = lim n→∞ E [ ∑n t=0 IXt=x ] n , where I is the indicator function . Assuming that π is differentiable by the policy parameters θπ , the policy gradient ( PG ) method ( Williams , 1992 ) maximizes Gπ by updating θπ via gradient ascent . A common formulation of PG estimates the gradient by ∇θπGπ = Ex , a , π [ ∇θπ log π ( a|x ) Â ( x , a ) ] , where Â ( x , a ) is an estimation of the advantage function Aπ ( x , a ) def= Qπ ( x , a ) − V π ( x ) . Among many variants of PG methods , we implemented our method on PPO ( Schulman et al. , 2017 ) . At each gradient step , PPO updates θπ to maximize clip ( π ( a|x ) πold ( a|x ) Â , − , ) , where clip ( x , − , ) = max ( − , min ( , x ) ) and πold is the policy before being updated . This clipping heuristics prevents π from updating too rapidly . Options Framework Options ( Sutton et al. , 1999 ) provide a framework for formulating temporally abstracted actions in RL . An option o ∈ O consists of a tuple ( Io , βo , πo ) , where Io ⊆ X is the initiation βo : X → [ 0 , 1 ] is a termination function , and πo is an intra-option policy . Following related studies ( Bacon et al. , 2017 ; Harutyunyan et al. , 2019 ) , we assume that Io = X focuses on learning βo and πo . We let µ : X × O → [ 0 , 1 ] denote a policy over options . A typical RL agent sample At from π ( ·|Xt ) . Analogously , at time t , an RL agent with options first samples a termination Tt from βOt ( Xt ) . If Tt = 1 , Ot+1 is sampled from µ ( ·|Xt ) and if not , the current option remains the same . The next action At is sampled from πOt+1 ( ·|Xt ) . For option learning methods , we use π to denote the resulting policy induced by µ , ∪o∈Oπo , and ∪o∈Oβo . Option value functions With options , we have three option-value functions QO , VO , and UO . QO is the option-value function denoting the value of selecting an option o at state x defined by QO ( x , o ) def = E [ G|X0 = x , O0 = o ] . Similar to the relationship between Qπ and V π , we let VO denote the marginalized option-value function VO ( x ) def = ∑ o µ ( o|x ) QO ( x , o ) . 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 x with o and not having selected the new option . We use these notations in Section 4.2 . Termination gradient theorem Analogously to p , we let P o : X × O ×X → [ 0 , 1 ] denote the state transition probability induced by options . When an agent is at xs and having an option o , the probability that o ends at xf is given by : P o ( xf |xs ) = βo ( xf ) Ixf=xs + ( 1− βo ( xs ) ) ∑ x pπ o ( x|xs ) P o ( xf |x ) , ( 1 ) where pπ o is the policy-induced transition function pπ o ( x′|x ) def= ∑ a∈A π o ( a|x ) P ( x′|x , a ) . Here we assume that all options eventually terminate , such that P o is a valid probability distribution over xf . Interestingly , P o is differentiable with respect to the parameter of βo . Harutyunyan et al . ( 2019 ) introduced the termination gradient theorem : Theorem 1 . Let βo be parameterized by θβ , and let ` βo denote the logit of βo , i.e. , ` βo = log ( β o ( x ) 1−βo ( x ) ) . We have ∇θβP o ( xf |xs ) = ∑ x P o ( x|xs ) ∇θβ ` βo ( x ) ( Ixf=x − P o ( xf |x ) ) . ( 2 ) We use this theorem to derive an unbiased estimator of our target gradient . 3 INFOMAX TERMINATION CRITIC . We now present the InfoMax Termination Critic ( IMTC ) algorithm . To learn diverse options , we propose to maximize the following MI at each state xs : I ( Xf ; O|xs ) = H ( Xf |xs ) −H ( Xf |xs , O ) = H ( O|xs ) −H ( O|Xf , xs ) , ( 3 ) where I denotes the MI I ( A ; B|c ) = H ( A|c ) −H ( A|B , c ) , H denotes entropy , and O is a random variable denoting options experienced by an agent . Precisely , we let η denote the probability of having an option o when leaving a state x : η ( o|x ) = βo ( x ) µ ( o|x ) + ( 1− βo ( x ) ) ∑ x′∈X p πo ( x|x′ ) η ( o|x′ ) and let ηπ denote η marginalized over dπ : ηπ ( o ) = ∑ x∈X dπ ( x ) η ( o|x ) . Then , we define O as a random variable with ηπ . By decomposing this MI as I ( Xf ; O|xs ) = H ( Xf |xs ) − H ( Xf |xs , O ) , we can interprete this MI maximization as maximization of H ( Xf |xs ) and minimization of H ( Xf |xs , O ) . Maximizing H ( Xf |xs ) diversifies possible destinations of an agent . Thus , the resulting options are expected to be diverse in the sense that they likely lead to different destinations . On the other hand , minimizing H ( Xf |xs , O ) makes option state transitions more deterministic , leading to more meaningful options . To build intuition , we searched options that maximize the MI ( 3 ) by a simple brute-force search from a limited set of πo and βo in a 3 × 3 Gridworld environment . Specifically , we considered only deterministic intra-option polciies and 0.1 or 0.8 as the value of βo ( x ) . The number of options is fixed to |O| = 3 . In this environment , an agent can take four actions : Up , Down , Left , and Right . The taken action fails with probability 0.1 , and it takes a uniformly random action . Figure 2 shows the result . Arrows represent intra-option policies and colors of the room represent termination probabilities . Here , we can see the tendency that IMTC prefers options that have almost separated termination regions and different policies per option . We also visualized TC ( Harutyunyan et al. , 2019 ) options that maximize −H ( Xf |o ) . For TC options , we can see a similar tendency as IMTC options , but options 0 and 1 of TC share the intra-option policies on the left side of the room . Thanks to the diversity term H ( Xf |xs ) , IMTC successfully avoids such a policy overlap . In addition , we can see that IMTC has smaller termination regions than TC , which can enhance long-range abstraction . We propose to maximize the MI ( 3 ) by updating βo via gradient ascent w.r.t . θβ . For this purpose , we now derive an unbiased estimator of the gradient . First , we write the gradient of the MI using the option transition model P o and marginalized option-transition model P ( xf |xs ) =∑ o η π ( o ) P o ( xf |xs ) . Proposition 1 . Let βo be parameterized using 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 ) = Edπ , η [ −∇θβ ` βo ( x ) βo ( x ) ( logP ( x|xs ) − logP ( xf |xs ) ) ] ( 4 ) ∇θβH ( Xf |xs , O ) = Edπ , η [ −∇θβ ` βo ( x ) βo ( x ) ( logP o ( x|xs ) − logP o ( xf |xs ) ) ] , ( 5 ) 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 is based on Section 4 in Harutyunyan et al . ( 2019 ) and is given in Appendix A.1 . Now the targetting gradient can be written as : ∇θβI ( Xf ; O|xs ) = ∇θβH ( Xf |xs ) −∇θβH ( Xf |xs , O ) = Edπ , η [ −∇θβ ` βo ( x ) βo ( x ) ( logP ( x|xs ) − logP ( xf |xs ) − ( logP o ( x|xs ) − logP o ( xf |xs ) ) ) ] , ( 6 ) which means that we can approximate MI maximization by estimating P o and P . However , in RL , learning a density model over the state space in RL is often difficult , especially with large or continuous state spaces . For example , it has been tackled with data compression methods ( Bellemare et al. , 2016 ) and generative models ( Ostrovski et al. , 2017 ) . Hence , we reformulate the gradient using Bayes ’ rule to avoid estimating P o and P . The resulting term consists of the inverse option transition function pO ( o|xs , xf ) , which denotes the probability of having an o given a state transition xs , xf . Lemma 1 . We now have ∇θβI ( Xf ; O|xs ) = ∇θβH ( Xf |xs ) −∇θβH ( Xf |xs , O ) = Edπ , η [ ∇θβ ` βo ( x ) βo ( x ) ( log pO ( o|xs , x ) − log pO ( o|xs , xf ) ) ] . ( 7 ) The proof is provided in Appendix A.2 . Equation ( 7 ) requires estimating of pO per updating πo and βo , which is computational quite expensive . Thus , we approximate approximate the gradient ( 7 ) by regressing a classification model over options p̂O ( o|xs , xf ) on sampled option transition .

Summary: This paper proposed an algorithm, termed infomax Termination Critic (IMTC) to learn diversified options in RL. This algorithm learns termination conditions of options by maximizing mutual information between options and corresponding state transitions. The experiments demonstrate the IMTC algorithm learns diversified options and can be reused in various tasks.