Information asymmetry in KL-regularized RL

1 INTRODUCTION . For many interesting reinforcement learning tasks , good policies exhibit similar behaviors in different contexts , behaviors that need to be modified only slightly or occasionally to account for the specific task at hand or to respond to information becoming available . For example , a simulated humanoid in navigational tasks is usually required to walk – independently of the specific goal it is aiming for . Similarly , an agent in a simulated maze tends to primarily move forward with occasional left/right turns at intersections . This intuition has been explored across multiple fields , from cognitive science ( e.g . Kool & Botvinick , 2018 ) to neuroscience and machine learning . For instance , the idea of bounded rationality ( e.g . Simon , 1956 ) emphasizes the cost of information processing and the presence of internal computational constraints . This implies that the behavior of an agent minimizes the need to process information , and more generally trades off task reward with computational effort , resulting in structured repetitive patterns . Computationally , these ideas can be modeled using tools from information and probability theory ( e.g . Tishby & Polani , 2011 ; Ortega & Braun , 2011 ; Still & Precup , 2012 ; Rubin et al. , 2012 ; Ortega & Braun , 2013 ; Tiomkin & Tishby , 2017 ) , for instance , via constraints on the channel capacity between past states and future actions in a Markov decision process . In this paper we explore this idea , starting from the KL regularized expected reward objective ( e.g . Todorov , 2007 ; Toussaint , 2009 ; Kappen et al. , 2012 ; Rawlik et al. , 2012 ; Levine & Koltun , 2013 ; Teh et al. , 2017 ) , which encourages an agent to trade off expected reward against deviations from a prior or default distribution over trajectories . We explore how this can be used to inject subjective knowledge into the learning problem by using an informative default policy that is learned alongside the agent policy This default policy encodes default behaviours that should be executed in multiple contexts in absence of addi- tional task information and the objective forces the learned policy to be structured in alignment with the default policy . To render this approach effective , we introduce an information asymmetry between the default and agent policies , preventing the default policy from accessing certain information in the state . This prevents the default policy from collapsing to the agent ’ s policy . Instead , the default policy is forced to generalize across a subset of states , implementing a form of default behavior that is valid in the absence of the missing information , and thereby exerting pressure that encourages sharing of behavior across different parts of the state space . Figure 1 illustrates the proposed setup , with asymmetry imposed by hiding parts of the state from the default policy . We investigate the proposed approach empirically on a variety of challenging problems including both continuous action problems such as controlling simulated high-dimensional physical embodied agents , as well as discrete action visual mazes . We find that even when the agent and default policies are learned at the same time , significant speed-ups can be achieved on a range of tasks . We consider several variations of the formulation , and discuss its connection to several ideas in the wider literature , including information bottleneck , and variational formulations of the EM algorithm for learning generative models . 2 KL AND ENTROPY REGULARIZED REINFORCEMENT LEARNING . Throughout this paper we use st and at to denote the state and action at time step t , and rps , aq the instantaneous reward for the agent if it executes action a in state s. We denote the history up to time t by xt “ ps1 , a1 , . . . , stq , and the whole trajectory by τ “ ps1 , a1 , s2 , . . .q . Our starting point is the KL regularized expected reward objective Lpπ , π0q “ Eπτ “ ř t γ trpst , atq ´ αγtKL “ πpat|xtq } π0pat|xtq ‰‰ , ( 1 ) where π is the agent policy ( parameterized by θ and to be learned ) , π0 the default policy , and Eπτ r¨s is taken with respect to the distribution πτ over trajectories defined by the agent policy and system dynamics : πτ pτq “ pps1q ś t πpat|xtqppst ` 1|st , atq . Note that our policies are history-dependent . KLrπpat|xtq } π0pat|xtqs is the Kullback-Leibler ( KL ) divergence between the agent policy π and a default or prior policy π0 given history xt . The discount factor is γ P r0 , 1s and α is a hyperparameter scaling the relative contributions of both terms . Intuitively , this objective expresses the desire to maximize the reward while also staying close to a reference behaviour defined by π0 . As discussed later , besides being a convenient way to express a regularized RL problem , it also has deep connections to probabilistic inference . One particular instantiation of eq . ( 1 ) is when π0 is the uniform distribution ( assuming a compact action space ) . In this case one recovers , up to a constant , the entropy regularized objective ( e.g . Ziebart , 2010 ; Fox et al. , 2015 ; Haarnoja et al. , 2017 ; Schulman et al. , 2017a ; Hausman et al. , 2018 ) : LHpπq “ Eπτ r ř t γ trpst , atq ` αγtHrπpat|xtqss . ( 2 ) This objective has been motivated in various ways : it prevents the policy from collapsing to a deterministic solution thus improving exploration , it encourages learning of multiple solutions to a task which can facilitate transfer , and it provides robustness to perturbations and model mismatch . One approximation of the entropy regularized objective is for the history dependent entropy to be used as an additional ( auxiliary ) loss to the RL loss ; this approach is widely used in the literature ( e.g . Williams & Peng , 1991 ; Mnih et al. , 2016 ) . While the motivations for considering the entropy regularized objective are intuitive and reasonable , the choice of regularizing towards an uniform policy is less obvious , particularly in cases with large or high dimensional action spaces . In this work we explore whether regularization towards more sophisticated default policies can be advantageous . Both objectives ( 1 ) and ( 2 ) can be generalized beyond the typical Markov assumption in MDPs . In particular , additional correlations among actions can be introduced , e.g . using latent variables Hausman et al . ( 2018 ) . This can be useful when , as discussed below , either π0 or π are not given full access to the state , rendering the setup partially observed . In the following we will not explore such extensions , though note that we do work with policies πpat|xtq and π0pat|xtq that depend on history xt . 3 LEARNING DEFAULT POLICIES . Many works that consider the KL regularized objective either employ a simple or fixed default policy or directly work with the entropy formulation ( e.g . Rubin et al. , 2012 ; Fox et al. , 2015 ; Haarnoja et al. , 2017 ; Hausman et al. , 2018 ) . In contrast , here we will be studying the possibility of learning the default policy itself , and the form of the subjective knowledge that this introduces to the learning system . Our guiding intuition , as described earlier , is the notion of a default behaviour that is executed in the absence of additional goal-directed information . Instances which we explore in this paper include a locomotive body navigating to a goal location where the locomotion pattern depends largely on the body configuration and less so on the goal , and a 3D visual maze environment with discrete actions , where the typical action includes forward motion , regardless of the specific task at hand . To express the notion of a default behavior , which we also refer to as “ goal-agnostic ” ( although the term should be understood very broadly ) , we consider the case where the default policy π0 is a function ( parameterized by φ ) of a subset of the interaction history up to time t , i.e . π0pat|xtq “ π0pat|xDt q , where xDt is a subset of the full history xt and is the goal-agnostic information that we allow the default policy to depend on . We denote by xGt the other ( goal-directed ) information in xt and assume that the full history is the disjoint union of both . The objective ( 1 ) specializes to : Lpπ , π0q “ Eπτ “ ř t γ trpst , atq ´ αγtKL “ πpat|xtq } π0pat|xDt q ‰‰ , ( 3 ) To give a few examples : If xDt is empty then the default policy does not depend on the history at all ( e.g . uniform policy ) . If xDt “ a1 : t´1 then it depends only on past actions . In multitask learning xGt can be the task identifier , while x D t the state history . And finally , in continuous control x D t can contain proprioceptive information about the body , while xGt contains exteroceptive ( goal-directed ) information ( e.g . vision ) . By hiding information from the default policy , the system forces the default policy to learn the average behaviour over histories xt with the same value of xDt . If x D t hides goal-directed information , the default policy will learn behaviour that is generally useful regardless of the current goal . We can make this precise by noting that optimizing the objective ( 1 ) with respect to π0 amounts to supervised learning of π0 on trajectories generated by πτ , i.e . this is a distillation process from πτ to π0 ( Hinton et al. , 2015 ; Rusu et al. , 2016 ; Parisotto et al. , 2016 ; Teh et al. , 2017 ) . In the nonparametric case , the optimal default policy π0˚ can be derived as : π0˚pat|xDt q “ ř t̃ γ t̃ ş ` 1pxDt “ x̃Dt̃ qπpat|x̃t̃q ˘ πτ px̃t̃qdx̃t̃ ř t̃ γ t̃ ş ` 1pxDt “ x̃Dt̃ q ˘ πτ px̃t̃qdx̃t̃ , ( 4 ) where πτ px̃t̃q is the probability of seeing history x̃t̃ at time step t̃ under the policy π , and the indicator 1pxDt “ x̃Dt̃ q is 1 if the goal-agnostic information of the two histories matches and 0 otherwise . It is also worth considering the effect of the objective eq . ( 3 ) on the learned policy π . Since π0 is learned alongside π and not specified in advance , this objective does not favor any particular behavior a priori . Instead it will encourage a solution in which similar behavior will be executed in different parts of the state space that are similar as determined by xDt , since the policy π is regularized towards the default policy π0 . More generally , during optimization of π the default policy effectively acts like a shaping reward while the entropy contained in the KL discourages deterministic solutions . 4 CONNECTION TO INFORMATION BOTTLENECK AND VARIATIONAL EM . 4.1 INFORMATION BOTTLENECK . Reinforcement learning objectives with information theoretic constraints have been considered by multiple authors ( Tishby & Polani , 2011 ; Still & Precup , 2012 ; Tiomkin & Tishby , 2017 ) . Such constraints can be motivated by the internal computational limitations of the agent , which limit the rate with which information can be extracted from states ( or observations ) and translated into actions . Such capacity constraints can be expressed via an information theoretic regularization term that is added to the expected reward . Specializing to our scenario , where the “ information flow ” to be controlled is between the goal-directed history information xGt and action at ( so that the agent prefers default , goal-agnostic , behaviour ) , consider the objective : LI “ Eπτ “ ř t γ trpst , atq ´ αγtMIrxGt , at|xDt s ‰ , ( 5 ) where MIrxGt , at|xDt s is the conditional mutual information between xGt and at given xDt . The conditional mutual information can be upper bounded : Eπτ rMIrxGt , at|xDt ss “ Eπτ „ log πτ pxGt |xDt qπpat|xGt , xDt q πτ pxGt |xDt qπτ pat|xDt q ď Eπτ „ log πpat|xGt , xDt q π0pat|xDt q ( 6 ) where the inequality is from the fact that the KL divergence KLrπτ pat|xDt q } π0pat|xDt qs is positive ( see Alemi et al. , 2016 ) . Re-introducing this into ( 5 ) we find that the KL regularized objective in eq . ( 3 ) can be seen as a lower bound to eq . ( 5 ) , where the agent has a capacity constraint on the channel between goal-directed history information and ( future ) actions . See section A in the appendix for a generalization including latent variables . In this light , we can see our work as a particular implementation of the information bottleneck principle , where we penalize the dependence on the information that is hidden from the default policy .

Summary: This is a very interesting piece of work. We know from cognitive science literature, that there are 2 distinct modes of decision making - habit based and top-down control (goal directed) decision making. The paper proposes to use this intuition by using information theoretic objective such that the agent follows "default" policy on average and agent gets penalized for changing its "default" behaviour, and the idea is to minimize this cost on average across states.