Soft Q-Learning with Mutual-Information Regularization

1 INTRODUCTION . Reinforcement Learning ( RL ) ( Sutton & Barto , 1998 ) is a framework for solving sequential decision-making problems under uncertainty . Contemporary state-of-the-art RL methods often use an objective that includes an entropy regularization term ( Haarnoja et al. , 2018c ; 2017 ; Teh et al. , 2017 ) . Entropy regularized RL has been shown to capture multi-modal behaviour , as well as exhibiting superior exploration ( Haarnoja et al. , 2017 ) . Additionally , the learned policies are more robust , as the entropy bonus accounts for future action stochasticity ( Grau-Moya et al. , 2016 ) and reduces value overestimation ( Fox et al. , 2016 ) . While encouraging high-entropy policies can provide several benefits , it is possible to devise examples where entropy regularization actually impedes exploration . Since high-entropy policies tend to spread the probability mass across all actions equally , they can perform poorly when the RL problem contains actions that are rarely useful . In this paper , we propose to overcome the previous limitation by designing a reinforcement learning algorithm that dynamically adjusts the importance of actions while learning . We motivate our algorithm by phrasing RL as an inference problem with an adaptive prior action distribution . Where previous work assumes a uniform prior distribution over the actions ( Rawlik et al. , 2012 ; Levine , 2018 ) , we generalize the formulation by optimizing the prior . We show that this optimization process leads to an RL objective function with a regularizer based on the mutual information between states and actions . Additionally , we develop a novel algorithm that uses such mutual-information regularization to obtain an optimal action-prior for better performance and exploration in high-dimensional state spaces . This novel regularizer for RL encourages policies to be close to the marginal distribution over actions . This results in assigning higher probability to actions frequently used by the optimal policy , while actions that are used infrequently have lower probability under the prior . We demonstrate significant improvements on 19 Atari games over a deep Q-network ( Mnih et al. , 2015 ) ( DQN ) baseline without any regularization and over soft Q-learning ( Schulman et al. , 2017 ; Leibfried et al. , 2018 ) ( SQL ) that employs standard entropy regularization without prior adaptation . 2 BACKGROUND . 2.1 REINFORCEMENT LEARNING . We consider the standard Markov decision process ( MDP ) setting . Formally , an MDP is defined as the tuple 〈S , A , P , R , γ〉 where S is the state space , A the action space , and P : S × A × S → [ 0 , 1 ] denotes the state transition function . Upon taking action at ∈ A in state st ∈ S , the agent transitions to st+1 with probability P ( st+1|st , at ) . The reward function R : S × A → R quantifies the agent ’ s performance . The goal is to find a policy that maximizes the value function , i.e . π∗ ( a|s ) = argmaxπ V π ( s ) , where V π ( s ) = E [ ∑T t=0 γ tr ( st , at ) |s0 = s ] . Here , γ is a discount factor ( 0 < γ < 1 ) that allows to account for the future in different ways . The policy-dependent state transition probabilities are defined as Pπ ( s′|s ) : = ∑ a P ( s ′|a , s ) π ( a|s ) which can be written in matrix notation asPπ ∈ R|S|×R|S| where the rows are indexed by s and the columns by s′ . This allows us to conveniently define the agent ’ s stationary distribution over states µπ ( s ) and actions ρπ ( a ) as follows : Definition 1 ( Stationary distribution over states ) . The stationary distribution over states ( assumed to exist and to be unique ) is defined in vector form as µ > π : = limt→∞ ν > 0 P t π with ν0 being an arbitrary vector of probabilities over states at time t = 0 . The stationary distribution satisfies µπ ( s ′ ) = ∑ s Pπ ( s ′|s ) µπ ( s ) and therefore is a fixed point under the state transition probabilities µ > π = µ > πPπ . Definition 2 ( Stationary distribution over actions ) . Let µπ ( s ) be the stationary distribution over states induced by the policy π . Then the stationary distribution over actions under the policy π is defined as ρπ ( a ) : = ∑ s∈S µπ ( s ) π ( a|s ) . 2.2 MAXIMUM ENTROPY REINFORCEMENT LEARNING . Maximum entropy reinforcement learning augments the standard RL reward objective with an additional policy entropy term . The optimal value function under entropy regularization ( Haarnoja et al. , 2017 ) is defined as : V∗ ( s ) = max π E [ ∞∑ t=0 γt ( r ( st , at ) − 1 β log π ( at|st ) ) ∣∣∣∣s0 = s ] , ( 1 ) where 1β trades off between reward and entropy maximization , and the expectation operation is over state-action trajectories . The optimal policy that solves ( 1 ) can be written in closed form as : π∗ ( a|s ) = e βQ∗ ( s , a ) ∑ a∈A e βQ∗ ( s , a ) , where Q∗ ( s , a ) : = r ( s , a ) + ∑ s′∈S P ( s ′|s , a ) V∗ ( s′ ) . Note that the above represents a generalization of standard RL settings , where β →∞ corresponds to a standard RL valuation ( limβ→∞ V∗ ( s ) = maxπ V π ( s ) ) , while for β → 0 we recover the valuation under a random uniform policy . For intermediate values of β , we can trade off between reward maximization and entropy maximization . Interestingly , one can formulate the maximum entropy RL objective as an inference problem ( Levine , 2018 ) by specifying a prior distribution over trajectories that assumes a fixed uniform distribution over actions . Precisely this assumption is what encourages the policies to maximize entropy . As outlined in the introduction , encouraging policies to be close to a uniform distribution might be undesirable when some actions are simply non-useful or not frequently used . In Section 3 , we show that when relaxing the previous assumption , i.e . allowing for prior optimization , we obtain a novel variational inference formulation of the RL problem that constrains the policy ’ s mutual-information between states and actions . We show that such policies must be close to the marginal distribution over actions which automatically assigns high probability mass to overall useful actions and low probability to infrequently used actions . Before proceeding , however , it is insightful to show how prior optimization bridges the gap between entropy regularization and mutual-information regularization in a non-sequential decision-making scenario that considers one time step only . 2.3 MUTUAL-INFORMATION REGULARIZATION FOR ONE-STEP DECISION-MAKING . In a one-step decision-making scenario , entropy regularization assumes the following form max π ∑ s , a p ( s ) π ( a|s ) ( r ( s , a ) − 1 β log π ( a|s ) ) , where p ( s ) is some arbitrary distribution over states and the optimal policy balances expected reward maximization versus expected entropy maximization . Entropy regularization discourages deviations from a uniform prior policy . In a more general setting , when discouraging deviations from an arbitrary prior ρ ( a ) , a similar objective can be written as max π ∑ s , a p ( s ) π ( a|s ) ( r ( s , a ) − 1 β log π ( a|s ) ρ ( a ) ) = max π ∑ s , a p ( s ) π ( a|s ) r ( s , a ) − 1 β ∑ s p ( s ) KL ( π ( ·|s ) ||ρ ( · ) ) , ( 2 ) where KL refers to the Kullback-Leiber ( KL ) divergence . This framework has been proposed before in the literature under the name information-theory for decision-making ( Ortega & Braun , 2013 ) . Going one step further , one can also optimize for ρ in addition to π , which essentially means that policies are discouraged to deviate from an optimal prior distribution , leading to the following optimization problem max π ∑ s , a p ( s ) π ( a|s ) r ( s , a ) − 1 β min ρ ∑ s p ( s ) KL ( π ( ·|s ) ||ρ ( · ) ) , ( 3 ) where we utilize the fact that only the expected KL penalty depends on ρ . The minimum expected KL relates to the mutual information as follows : Proposition 1 ( Mutual Information ) . Let If be a functional , in particular : If ( pX , pY |X , qY ) : =∑ x pX ( x ) KL ( pY |X ( ·|x ) ||qY ( · ) ) , where pX ( x ) is the distribution of the input , pY |X ( y|x ) is the conditional distribution of the output conditioned on the input , and qY ( y ) a variational distribution of the output . Then , the mutual information 1 is recovered with I [ X , Y ] = min qY If ( pX , pY |X , qY ) , where the optimal variational distribution is q ? Y ( y ) = ∑ x pX ( x ) pY |X ( y|x ) , i.e . the true marginal distribution . See e.g . ( Cover & Thomas , 2006 , Lemma 10.8.1 ) for details . This allows us to rewrite the problem from Equation ( 3 ) as max π ∑ s , a p ( s ) π ( a|s ) r ( s , a ) − 1 β I [ S ; A ] , ( 4 ) yielding a penalty on the mutual information between states and actions . Notice that this problem is mathematically equivalent to rate-distortion theory from informationtheory ( Shannon , 1959 ) which formulates how to efficiently send information over an informationtheoretic channel with limited transmission rate . This framework has also been used to describe decision-making problems with limited information budgets ( Sims , 2011 ; Genewein et al. , 2015 ; Leibfried & Braun , 2015 ; 2016 ; Peng et al. , 2017 ; Hihn et al. , 2018 ) . In a decision-making context , the agent is considered as information-theoretic channel π ( a|s ) where s is the channel input and a the channel output . The agent aims at maximizing expected reward under the constraint that the information transmission rate is limited , where the transmission rate is given by the mutual-information between states and actions ( Cover & Thomas , 2006 ) . Intuitively , this means that the agent has to discard reward-irrelevant information in s to not exceed the limits in information transmission . In the following section , we generalize the rate-distortion formulation for decision-making to be applicable to a sequential decision-making scenario , i.e . the RL setting . We propose an inferencebased formulation where the mutual information arises as a consequence of allowing for optimizing the action-prior distribution . 1I ( X ; Y ) : = ∑ x , y p ( x , y ) log p ( x , y ) pY ( y ) pX ( x ) = ∑ x pX ( x ) KL ( pY |X ( ·|x ) ||pY ( · ) )

Summary: The authors take the control-as-inference viewpoint and learn a state-independent prior (which is typically held fixed). They claim that this leads to better exploration when actions have different importance. They relate this objective to a mutual information constrained RL objective in a limiting case. They then propose a practical algorithm, MIRL and compare their algorithm against DQN and Soft Q-learning (SQL) on 19 Atari games and demonstrate improvements over both.