MUTUAL-INFORMATION REGULARIZATION
Year: 2019

Abstract:
We propose a reinforcement learning (RL) algorithm that uses mutual-information regularization to optimize a prior action distribution for better performance and exploration. Entropy-based regularization has previously been shown to improve both exploration and robustness in challenging sequential decision-making tasks. It does so by encouraging policies to put probability mass on all actions. However, entropy regularization might be undesirable when actions have significantly different importance. In this paper, we propose a theoretically motivated framework that dynamically weights the importance of actions by using the mutualinformation. In particular, we express the RL problem as an inference problem where the prior probability distribution over actions is subject to optimization. We show that the prior optimization introduces a mutual-information regularizer in the RL objective. This regularizer encourages the policy to be close to a nonuniform distribution that assigns higher probability mass to more important actions. We empirically demonstrate that our method significantly improves over entropy regularization methods and unregularized methods.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND

heading: 2.1 REINFORCEMENT LEARNING
text: 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).

heading: 2.2 MAXIMUM ENTROPY REINFORCEMENT LEARNING
text: 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.

heading: 2.3 MUTUAL-INFORMATION REGULARIZATION FOR ONE-STEP DECISION-MAKING
text: 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 (·))

heading: 3 MUTUAL-INFORMATION REGULARIZATION IN RL
text: In this section, we first derive mutual-information regularization for the RL setting from a variational inference perspective. Subsequently, we derive expressions for the optimal policy and the optimal prior that are useful for constructing a practical algorithm.

heading: 3.1 VARIATIONAL INFERENCE RL FORMULATION WITH OPTIMAL ACTION-PRIORS
text: The RL problem can be expressed as an inference problem by introducing a binary random variable R that denotes whether the trajectory τ := (s0, a0, . . . sT , aT ) is optimal (R = 1) or not (R = 0). The likelihood of an optimal trajectory can then be expressed as p(R = 1|τ) ∝ exp( ∑T t=0 r(st, at)) (Levine, 2018). We additionally introduce a scaling factor β > 0 into the exponential, i.e. p(R = 1|τ) ∝ exp(β ∑T t=0 r(st, at)). This will allow us to trade off reward and entropy maximization
2. Next, we can define the posterior trajectory probability assuming optimality, i.e. p(τ |R = 1). Here we treat τ as a latent variable with prior probability p(τ), and we specify the log-evidence as log p(R = 1) = log ∫ p(R = 1|τ)p(τ)dτ . We now introduce a variational distribution q(τ) to approximate the posterior p(τ |R = 1). This leads to an Evidence Lower BOund (ELBO) of the previous expression (scaled by 1β ) 3:
1 β log p(R = 1) = 1 β log
∫ p(R = 1|τ)p(τ)dτ
≥ 1 β Eτ∼q(τ)
[ log
p(R = 1|τ)p(τ) q(τ)
] (5)
The generative model is written as p(τ) = p(s0) ∏T−1 t=0 ρ(at)P (st+1|st, at) and the variational
distribution as q(τ) = p(s0) ∏T−1 t=0 π(at|st)P (st+1|st, at). The RL problem can now be stated as a maximization of the ELBO w.r.t π. The maximum entropy RL objective is recovered when assuming a fixed uniform prior distribution over actions, i.e. ρ(at) = 1|A| for all t.
We obtain a novel variational RL formulation by introducing an adaptive prior over actions ρ. Contrary to maximum entropy RL, where the prior of the generative model is fixed and uniform, here the prior over actions is subject to optimization. Starting from Equation (5) and substituting p(τ) and q(τ) we obtain the following ELBO: maxπ,ρ Eq [∑T t=0 ( r(st, at)− 1β log π(at|st) ρ(at) )] . Since we are interested in infinite horizon problems, we introduce a discount factor and take the limit limT→∞ (Haarnoja et al., 2017). This leads to the optimization objective that we use in our experiments:
max π,ρ
Eq [ ∞∑ t=0 γt ( r(st, at)− 1 β log π(at|st) ρ(at) )] , (6)
where 0 < γ < 1 is the discount factor. In the following, we show that the the solution for the prior and the policy can be expressed in a concise form giving rise to a novel RL regularization scheme.

heading: 3.2 RECURSION, OPTIMAL POLICIES AND OPTIMAL PRIORS
text: Crucial for the construction of a practical algorithm are concise expressions for the optimal policy and the prior. More concretely, the optimal policy takes the form of a Boltzmann distribution weighted by the prior ρ. When fixing the policy, the optimal prior is the marginal distribution over actions under the discounted stationary distribution over states. This finding is important to devise a method for efficiently learning an optimal prior in practice.
Optimal policy for a fixed prior ρ: We start by defining the value function with the information cost as Vπ,ρ(s) := E [∑∞ t=0 γ t ( r(st, at)− 1β log π(at|st) ρ(at) ) |s0 = s ] where one can show that Vπ,ρ
2Other authors absorb this scaling factor into the reward, but we keep it as an explicit hyperparameter. 3This is obtained by multiplying and dividing by q(τ) inside the integral and applying Jensen’s inequality
f (Ex∼q[x]) ≥ Ex∼q[f(x)] (for any concave function f ).
satisfies a recursion similar to the Bellman equation: Vπ,ρ(s) = Eπ [ r(s, a)− 1
β log π(a|s) ρ(a)
+ γEs′ [Vπ,ρ(s′)] ] . (7)
When considering a fixed ρ, the problem of maximizing Equation (7) over the policy can be solved analytically by standard variational calculus (Rubin et al., 2012; Genewein et al., 2015). The optimal policy is then given by
π∗(a|s) := 1 Z ρ(a) exp(βQπ∗,ρ(s, a)) (8) with Z = ∑ a ρ(a) exp(βQπ∗,ρ(s, a)), and the the soft Q-function is defined as Qπ,ρ(s, a) := r(s, a) + γEs′ [Vπ,ρ(s′)]. (9) Being able to write the optimal policy in this way as a function of Q-values is needed in order to estimate the optimal prior as we show next.
Optimal prior for a fixed policy: In order to solve for the optimal prior, we rewrite the problem in Equation (6) as
max π,ρ ∞∑ t=0 ∑ s γtνt(s) ∑ a π(a|s) ( r(s, a)− 1 β log π(a|s) ρ(a) ) ,
where we have defined the marginal distribution over states at time t as
νt(s) := ∑
s0,a0,...,st−1,at−1
p(s0) ( t−2∏ t′=0 π(at′ |st′)P (st′+1|st′ , at′) ) π(at−1|st−1)P (s|st−1, at−1).
(10) For fixed π, we eliminate the max operator for π and all components that do not depend on ρ:
max ρ − 1 β ∞∑ t=0 ∑ s γtνt(s)KL(π(·|s)||ρ(·)).
Swapping the sums and letting p(s) := ∑∞ t=0 γ
tνt(s) be the unnormalized discounted marginal distribution over states, we obtain maxρ− 1β ∑ s p(s)KL(π(·|s)||ρ(·)). The solution to the latter,
ρ?(a) = ∑ s p(s)π(a|s)∑ s,a p(s)π(a|s) , can easily be obtained by adding the constraint that the action-prior is
a valid distribution (i.e., ∑ a ρ(a) = 1 and ρ(a) > 0 ∀a), and using the method of Lagrange multipliers and standard variational calculus. The connection to the mutual information becomes clear when plugging ρ? back into the objective yielding −k · 1β I(S,A) scaled by a positive constant k (because p(s) is not normalized) that can be absorbed into β. Additionally, we also formalize the connection to the stationary mutual-information for the limit case of γ → 1 in the Appendix. With the form of the optimal prior for a fixed policy at hand, one can easily devise a stochastic approximation method (e.g. ρi+1(a) = (1 − αρ)ρi(a) + αρπ(a|s) with αρ ∈ [0, 1] and s ∼ p(·) to estimate the optimal ρ using the current estimate of the optimal policy from Equation (8). We note that here we sample from the undiscounted distribution over states instead, rather than the true discounted state-distribution in the equation above. This is a common practice used in actor-critic RL that results in a biased estimator of the objective (Thomas, 2014).

heading: 4 MIRL: A PRACTICAL MUTUAL INFORMATION RL ALGORITHM
text: In this section, we present the MIRL agent. We focus on the tabular setting first and then port our algorithm to high-dimensional state spaces that require parametric function approximators.

heading: 4.1 MIRL FOR TABULAR Q-LEARNING
text: Our tabular MIRL agent is a modification of an ordinary Q-learning agent with a different update scheme for Q-values. In parallel to updating Q-values, the MIRL agent needs to update the prior ρ as well as the parameter β. The behavioral policy π is also different and utilizes soft Q-values. This is outline in more detail below.
Prior Updates: We approximate the optimal prior by employing the following update equation, ρi+1(a) = (1− αρ)ρi(a) + αρπi(a|si) (11)
where si ∼ νi(·) and αρ is a learning rate. Assuming a fixed policy, it is easy to show that this iterative update converges to ρπ(a) = ∑ s µπ(s)π(a|s), thus estimating correctly the optimal prior.
Q-function Updates: Concurrently to learning the prior, MIRL updates the tabular Q-function as Q(s, a)← Q(s, a) + αQ ((T ρsoftQθ̄)(s, a, s
′)−Q(s, a)) (12) where αQ is a learning rate and T ρ softQ is the empirical soft-operator defined as (T ρ softQ)(s, a, s
′) := r(s, a)+γ 1β log ∑ a′ ρ(a
′) exp (βQ(s′, a′)). Importantly, this operator differs from other soft operators arising when employing entropy regularization. Entropy regularization assumes a fixed uniform prior, whereas in our case, the optimal prior is estimated in the course of learning.
Behavioural policy: Since the Q-function can be learned off-policy, the experience samples can conveniently be drawn from a behavioural policy πb different from the current estimate of the optimal policy. As such, the behavioural policy used in our experiments is similar in spirit to an -greedy policy but it better exploits the existence of the estimated optimal prior when both exploring and exploiting. When exploring, MIRL’s behavioural policy samples from the current estimate of the optimal prior ρi which has adjusted probabilities, in contrast to vanilla -greedy that samples all actions with equal frequency. Additionally, when exploiting, MIRL selects the maximum probability action that depends not only on the Q-values but also on the current estimate of the optimal action-prior, instead of selecting the action with highest Q-value as in traditional -greedy. More formally, given a random sample u ∼ Uniform[0,1] and epsilon , the action ai is obtained by
ai = { argmaxa πi(a|si), if u > a ∼ ρi(·) if u ≤ ,
where πi(a|s) = 1Z ρi(a) exp(βiQi(s, a)), see Equation (8).
Parameter β Updates: The parameter β can be seen as a Lagrange multiplier that quantifies the magnitude of penalization for deviating from the prior. As such, a small fixed value of β would restrict the class of available policies and evidently constrain the asymptotic performance of MIRL. In order to remedy this problem and obtain better asymptotic performance, we use the same adaptive β-scheduling over rounds i from (Fox et al., 2016) in which βi is updated linearly according to βi+1 = c · i with some positive constant c. This update favours small values of β at the beginning of training and large values towards the end of training when the error over Q-values is small. Therefore, towards the end of training when β is large, MIRL recovers ordinary Q-learning without a constraint. This ensures that the asymptotic performance of MIRL is not hindered.

heading: 4.2 MIRL WITH PARAMETRIC FUNCTION APPROXIMATORS
text: For parametric function approximators, the scheme for updating the prior and the behavioural policy is the same as in the tabular setting but Q-function updates and β-scheduling need to be adjusted for high-dimensional state spaces. The pseudocode of our proposed algorithm is outlined in Algorithm 1 and follows standard literature for parametric value learning, see e.g. Mnih et al. (2015).
Q-function Updates: Q-function parameters are obtained by minimizing the following loss L(θ, ρ) := Es,a,r,s′∼M [( (T ρsoftQθ̄)(s, a, s ′)−Qθ(s, a) )2] (13)
whereM is a replay memory (Mnih et al., 2015),Qθ̄ is a target network that is updated after a certain number of training iterations and T ρsoftQ is the empirical soft-operator from the tabular setting, here repeated for convenience (T ρsoftQ)(s, a, s ′) := r(s, a) + γ 1β log ∑ a′ ρ(a ′) exp (βQ(s′, a′)).
Parameter β Updates: We use the same adaptive β-scheduling from Leibfried et al. (2018) in which βi is updated according to the inverse of the empirical loss of the Q-function, i.e. βi+1 = (1−αβ)βi+αβ( 1L(θi,ρi+1) ). This provides more flexibility than the linear scheduling scheme from the tabular setting, more suitable for high-dimensional state spaces where it is impossible to visit all state-action pairs.
Algorithm 1 MIRL 1: Input: the learning rates αρ, αQ and αβ , a Q-network Qθ(s, a), a target network Qθ̄(s, a), a
behavioural policy πb, an initial prior ρ0 and parameters θ0 at t = 0. 2: for i = 1 to N iterations do 3: Get environment state si and apply action ai ∼ πb(·|si) 4: Get ri, si+1 and store (si, ai, ri, si+1) in replay memoryM 5: Update prior ρi+1(·) = ρi(·)(1− αρ) + αρπi(·|si) 6: if i mod update frequency == 0 then 7: Update Q-function θi+1 = θi − αQ∇θL(θ, ρi+1)|θi according to Equation (13) 8: Update parameter βi+1 = (1− αβ)βi + αβ ( 1
L(θi,ρi+1) ) 9: end if
10: end for

heading: 5 EXPERIMENTS
text: We evaluate our MIRL agent both in the tabular setting using a grid world domain, and in the parametric function approximator setting using the Atari domain.

heading: 5.1 GRID WORLD
text: As an intuitive example, we evaluate our method in a grid world domain where the agent has to reach a goal. Reaching the goal gives a reward of 9 but each step yields a reward of −1. After the end of an episode (when reaching the goal), the agent’s location is randomly re-sampled uniformly over the state space. We compare against two baselines, Q-learning without any regularization, and SQL (Fox et al., 2016) which employs entropy regularization with the dynamic β-scheduling scheme outlined earlier. We train the agents for 2.5 · 105 environment steps following the procedure outlined in Fox et al. (2016). Both SQL and MIRL update the Lagrange multiplier β over time by using a linear scheduling scheme with a constant c = 10−3. MIRL additionally updates the estimate of the optimal prior by using a learning rate αρ = 2 · 10−3. In all experiments, we use an adaptive
learning rate for Q-values αQ = n(s, a)−ω that depends on the state-action-visitation frequencies n(s, a) (Fox et al., 2016). See Appendix for further details.
During the training of each algorithm, snapshots of the Q-tables (and the estimate of the prior in the case of MIRL) are stored every 100 environment steps for evaluation. The evaluation for a single snapshot is conducted by running the policy for 30 episodes lasting at most 100 environment steps. The epsilon value when in evaluation mode is set to = 0.05 (same as in training). Every individual experiment is repeated with 10 different initial random seeds and results are averaged across seeds.
Figure 1 summarizes our grid world experiments on two instances: a corridor where the goal is to the right, and a square world where one important action has to be executed exactly once. In the corridor, MIRL clearly outperforms competing approaches (Q-learning and SQL), whereas in the square world, MIRL attains comparable performance as the baselines. Note that these results remain valid when equipping SQL with the same marginal exploration scheme as MIRL.

heading: 5.2 ATARI
text: We conduct experiments on 19 Atari games (Brockman et al., 2016) with Algorithm 1 (MIRL), and compare against DQN (Mnih et al., 2015) and SQL (Haarnoja et al., 2018c) with a dynamic β-scheduling scheme based on loss evolution that leads to improved performance over vanilla SQL with fixed β (Leibfried et al., 2018). The three algorithms use a neural network for the estimation of Q-values as in (Mnih et al., 2015). The network receives as an input the state s which is composed of the last four frames of the game with some extra pre-processing (see Appendix), and it outputs a vector of Q-values, i.e. one value for each valid action. We train the network for 5 ·107 environment steps, where a training iteration is performed every four steps. The target network Qθ̄ is updated every 104 training iterations. Both SQL and MIRL update the Lagrange multiplier β over time by using an exponential moving average of the inverse loss 4 (Leibfried et al., 2018) with αβ = 3·10−5. In addition, MIRL updates the estimate of the optimal prior by using a learning rate αρ = 5 · 10−5. Additional details can be found in the Appendix.
For evaluation, we create snapshots of the network agents every 105 environment steps. Evaluating a single snapshot offline is done by running the policy for 30 episodes that last at most 4.5 · 103 environment steps but terminate earlier in case of a terminal event. When evaluating the agents, the
4Pracitcally, we replace the squared loss with the Huber loss Mnih et al. (2015).
epsilon value is = 0.05, whereas in training is linearly annealed over the first 106 steps (Mnih et al., 2015) from 1.0 to 0.1.
To summarize the results across all games, we normalize the episodic rewards obtained in the evaluation. The normalized episodic rewards are computed as follows znormalized = z−zrandomzhuman−zrandom · 100%, where z stands for the score obtained from our agent at test time, zrandom stands for the score that a random agent obtains and zhuman for the score a human obtains. Random and human scores are taken from Mnih et al. (2015) and Van Hasselt et al. (2016). As seen in Figure 2, our algorithm significantly outperforms the baselines in terms of the median normalized score. In particular, after 50 million interactions we obtain about 30% higher median normalized score compared to SQL and 50% higher score compared to DQN. MIRL attains the final performance of SQL in about half the amount of interactions with the environment and, similarly, it attains DQN’s final performance in about five times less interactions.
In Table 1 we show the comparison between best-performing agents for all the environments, where the best-performing agent is the agent that achieves the best score in evaluation mode considering all snapshots. Although this measure is not very robust, we include it since it is a commonly reported measure of performance in the field. MIRL outperforms the other baselines in 11 out of 19 games compared to SQL and DQN that are best on 5 and 3 games respectively.
In Figure 3, we conduct additional experiments on a subset of eight Atari games comparing MIRL and DQN with two different ablations of SQL: one ablation using uniform exploration and another ablation using the same marginal exploration scheme as MIRL (denoted SQL m). These experiments confirm the importance of the difference in values between MIRL and SQL rather than the difference in the exploration protocol.

heading: 6 RELATED WORK
text: The connection between reinforcement learning and inference is well established (Dayan & Hinton, 1997; Levine, 2018). Several authors have proposed RL algorithms based on optimizing the ELBO in Equation (5). In policy search, a popular approach is to optimize the lower bound using an iterative procedure similar to expectation maximization (Deisenroth et al., 2013). Different approximations can then be used for the trajectory distributions, resulting in different algorithms (Kober & Peters, 2009; Peters et al., 2010; Hachiya et al., 2011). The recent maximum a posteriori policy optimisation (MPO) (Abdolmaleki et al., 2018) framework uses a similar approach and combines an expectation maximization style update with off-policy estimation of a regularized Q-function. A key difference between MPO and our method is that MPO treats the approximate distribution q(τ) as an auxiliary distribution used to optimize the policy that generates p(τ). In our method, as well as in the maximum entropy based methods discussed below, we optimize the policy used to generate q(τ), while using the distribution p(τ) generated by the prior as an auxiliary distribution. While both approaches can be related to optimizing the ELBO, they can be shown to optimize different versions of the KL constraint on the target policy (Levine, 2018).
Maximum entropy reinforcement learning represents another family of inference-based RL methods. This formulation can be derived from the same evidence lower bound in Equation (5), by fixing the generative policy for p(τ) to be a uniform prior. Maximum entropy RL has been derived under
different conditions in many different settings. Ziebart et al. (2008) proposed maximum entropy learning for inverse reinforcement learning problems. Several authors have applied the same principles for trajectory optimization (Kappen, 2005; Todorov, 2008; Levine & Koltun, 2013). These maximum entropy methods typically assume the availability of some sort of transition model. More recently, maximum entropy learning has also been studied in model-free settings by introducing alternative soft-operators (Asadi & Littman, 2017) or soft Q-learning approaches (Rawlik et al., 2012; Fox et al., 2016; Haarnoja et al., 2017). Soft Q-learning learns a softened value function by replacing the hard maximum operator in the Q-learning update with a softmax operator. Several authors have discussed the benefits of this approach and provided generalizations under linear programming formulations (Neu et al., 2017). In particular, Fox et al. (2016) and Haarnoja et al. (2017) show that maximum entropy learning improves exploration and robustness. Furthermore, Haarnoja et al. (2018b) show that the resulting policies are composable and can be used to directly build solutions to unseen problems. Additionally, entropy-regularization has shown to be crucial to prove convergence guarantees on value learning with non-linear function approximators (Dai et al., 2018). The soft Qlearning framework has also been used in actor-critic settings (Haarnoja et al., 2018c) and to show a connection between value-based and policy gradient methods (Schulman et al., 2017; Nachum et al., 2017). The method has also been extended to hierarchical settings (Florensa et al., 2017; Haarnoja et al., 2018a). In the multi-task setting, the Distral framework (Teh et al., 2017) combines entropy regularization with an additional KL regularization used to transfer knowledge between tasks.
7 DISCUSSION AND CONCLUSION
Using a variational inference perspective, we derived a novel RL objective that allows optimization of the prior over actions. This generalizes previous methods in the literature that assume fixed uniform priors. We show that our formulation is equivalent to applying a mutual-information regularization and derive a novel algorithm (MIRL) that learns the prior over actions. We demonstrate that MIRL significantly improves performance over SQL and DQN.
We recognize that our approach might fail under certain conditions. For example, in the case when there is an action that is useful only once (similar to our
8× 8 grid world example) but is most of the times penalized with a negative reward. When the negative reward is too strong, MIRL might assign very low probability to that action and never explore it. However, this problem might be alleviated by a weighted mixing of our exploration policy with a uniform distribution.
An interesting direction for future work is to investigate the convergence properties of the alternating optimization problem presented here. We believe that, at least in the tabular case, the framework of stochastic approximation for two timescales (Borkar, 2009) is sufficient to prove convergence. On the experimental side, one could also investigate how our approach can be combined with the Rainbow framework (Hessel et al., 2017) which is the current state of the art in performance.

heading: A APPENDIX
text: A.1 CONNECTION TO MUTUAL INFORMATION FOR γ → 1.
The goal of this section is to show that when γ → 1, Equation (6) can be expressed as the following average-reward formulation (Puterman, 1994) with a constraint on the stationary mutual information
max π Es∼µπ [∑ a π(a|s)r(s, a) ] s.t. If (µπ, π, ρπ) ≤ C, (14)
where µπ and ρπ are the stationary distributions induced by the policy π over states and actions, respectively, and thus If (µπ, π, ρπ) is defined as the stationary mutual-information. Note that for a fixed stationary distribution over states, this problem coincides exactly with the well-known ratedistortion problem (Cover & Thomas, 2006).
We start by expressing (6) as a constrained problem
max π
Eq [ ∞∑ t=0 γtr(st, at) ] s.t. min ρ ∞∑ t=0 γtIf (νt, π, ρ) ≤ K(γ), (15)
where K(γ) := C1−γ (although we have set K(·) as a function of γ, it is without loss of generality since we can always obtain a desired K(·) by choosing an appropriate C for a given γ) and the marginal probability of state st at time t following the Markovian dynamics is written as in Equation (10).
A standard result in the MDP literature (Bertsekas, 1995) is that
argmax π lim γ→1
(1− γ)Eq [ ∞∑ t=0 γtr(st, at) ] ⇐⇒ argmax π Es∼µπ [∑ a π(a|s)r(s, a) ]
which basically says that the optimal policy for the limit γ → 1 of an infinite horizon problem is equivalent to the average reward formulation. Now it only remains to make explicit a similar equivalence on the constraint.
We rewrite the constraint by multiplying on both sides by (1− γ) assuming γ ∈ (0, 1)
min ρ ∞∑ t=0 γtIf (νt, π, ρ) ≤ K(γ) ⇐⇒ min ρ (1− γ) ∞∑ t=0 γtIf (νt, π, ρ) ≤ C.
Taking the limit γ → 1 in the last inequality and interchanging the limit and the min operators 5, we obtain the constraint minρ limγ→1(1 − γ) ∑∞ t=0 γ
tIf (νt, π, ρ) ≤ C. Then, we see the connection between the last inequality and the constraint on Equation (14) using the following Proposition 2.
Proposition 2. Let µπ(s) and ρπ(a) be the stationary distribution over states and actions under policy π according to Definitions (1) and (2). Then the stationary mutual information defined as If (µπ, π, ρπ) := minρ If (µππ, ρ) can also be written as
If (µπ, π, ρπ) = min ρ lim γ→1 (1− γ) ∞∑ t=0 γtIf (νt, π, ρ) (16)
where νt(s) is defined as in (10).
Proof. Following similar steps as in (Bertsekas, 1995, p.186) we have
If (µπ, π, ρπ) = min ρ lim N→∞
1
N N−1∑ t=0 If (νt, π, ρ)
= min ρ lim N→∞ lim γ→1
∑N−1 t=0 γ
tIf (νt, π, ρ)∑N t=0 γ t
= min ρ lim γ→1 lim N→∞
∑N t=0 γ
tIf (νt, π, ρ)∑N t=0 γ t
= min ρ lim γ→1 (1− γ) ∞∑ t=0 γtIf (νt, π, ρ),
where we used Proposition 3 (shown next) in the first equality and where the limits in the third equality can be interchanged due to the monotone convergence theorem.
Proposition 3. Let µπ(s) and ρπ(a) be the stationary distribution over states and actions under policy π according to Definitions (1) and (2). Then the stationary mutual information defined as If (µπ, π, ρπ) can also be written as
If (µπ, π, ρπ) = min ρ lim N→∞
1
N N−1∑ t=0 ∑ st νt(st)KL(π(·|st)||ρ(·)), (17)
where νt(s) is defined as in (10).
5We can interchange the operators because the l.h.s. of the inequality is finite for any 0 < γ < 1.
Proof. Let Iρ(s) := KL(π(·|s)||ρ(·)) and Iρπ (s) := KL(π(·|s)||ρπ(·)). Note that both previous quantities are bounded for all t. Then
min ρ lim N→∞
1
N N−1∑ t=0 ∑ s νt(s)Iρ(s) =
= min ρ lim N→∞
1
N ( K∑ t=0 ∑ s νt(s)Iρ(s) + N−1∑ t=K+1 ∑ s νt(s)Iρ(s) )
= min ρ lim N→∞
1
N N−1∑ t=K+1 ∑ s νt(s)Iρ(s)
= min ρ lim N→∞
1
N N−1∑ t=K+1 ∑ s µπ(s)Iρ(s)
= min ρ ∑ s µπ(s)Iρ(s) lim N→∞ 1 N N−1∑ t=K+1 1
= ∑ s µπ(s)Iµ(s) lim N→∞ N −K N
= ∑ s µπ(s)Iµ(s)
= If (µπ, π, ρπ),
where we assumed that νt(s) = µπ(s) for all t > K and finite but large enough K.
Since in practice we use a discount factor γ / 1, our original problem formulation in (6) can be seen as an approximation to the problem with stationary mutual-information constraints in (14).
The conclusion of this section is that we have established a clear link between the ELBO with optimizable priors and the average reward formulation with stationary mutual-information constraints.
A.2 HYPERPARAMETERS
Here we describe the hyperparameters used for both the Grid World experiments (see Table 2) and the Atari experiments (see Table 3).
A.3 SPECIFIC PLOTS FOR INDIVIDUAL GAMES
0 1 2 3 4 5 Environment Steps 1e7
0.00
0.05
0.10
0.15
0.20
Pr io
r P ro
ba bi
lit ie
s
AlienNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.2
0.4
0.6
Pr io
r P ro
ba bi
lit ie
s
AssaultNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
AsterixNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.000
0.025
0.050
0.075
0.100
0.125
0.150
Pr io
r P ro
ba bi
lit ie
s
AsteroidsNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
BankHeistNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
Pr io
r P ro
ba bi
lit ie
s
BeamRiderNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
BoxingNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
ChopperCommandNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
0.5
0.6
Pr io
r P ro
ba bi
lit ie
s
DemonAttackNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
GopherNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
KangarooNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
Pr io
r P ro
ba bi
lit ie
s
KrullNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.2
0.4
0.6
Pr io
r P ro
ba bi
lit ie
s
KungFuMasterNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s
RiverraidNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
0.4
Pr io
r P ro
ba bi
lit ie
s RoadRunnerNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.000
0.025
0.050
0.075
0.100
0.125
Pr io
r P ro
ba bi
lit ie
s
SeaquestNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
Pr io
r P ro
ba bi
lit ie
s
SpaceInvadersNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.0
0.1
0.2
0.3
Pr io
r P ro
ba bi
lit ie
s
StarGunnerNoFrameskip-v4
0 1 2 3 4 5 Environment Steps 1e7
0.1
0.2
0.3
0.4
0.5
0.6
Pr io
r P ro
ba bi
lit ie
s
UpNDownNoFrameskip-v4
Figure 4: Prior Evolution for all games. We can see that MIRL’s prior has fully converged for some games whereas for other games it is still about to converge.
A.4 ABLATION STUDY
In the ablation study summarized in Figure 3, we show how marginal exploration affects SQL. In Figure 6, we show the same plot for individual games. We clearly see that the marginal exploration improves performance, but is not the defining factor for all the improvements obtained by MIRL.
A.5 EVOLUTION OF THE LAGRANGE MULTIPLIER
In Figure 7, we show the evolution of β over time (environment steps) for the MIRL agent. As we can see, the β-values usually start at a high value (not shown for visual reasons) and typically go down and stabilize at some value. At first sight, this might be seen as a negative side effect since
lower β values imply a stronger constraint. However, we note that the constraint is highly dependent on the scale of the reward (or its sparsity), and therefore, the β value is not meaningful without a proper specification of this reward scaling.
Consequently, given that β-values are not meaningful here, we propose to instead show the multiplication of β times the current maximum Q-value estimates denoted as β × maxQ. Note that βQ(s, a) appears on the exponential term of the policy, i.e. π(a|s) = 1Z ρ(a) exp(βQ(s, a)), and therefore, is the term that shapes the deviation from the action-prior distribution. Additionally, the Q-values serve us as proper scaling for each game and account also for the learning of the agent. In particular, while the agent is learning and increasing its reward acquisition, the Q-values are going to be higher, thus, effectively needing a smaller β to shape the policy probabilities.
On Figure 8, we show the evolution of the term β ×maxQ for all the games. As we can see for the majority of games, β ×maxQ increases over time or has high value. This is important since a high value denotes that the policy is highly affected by this term.

Review summaries:
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.