Title: SOFT Q-LEARNING WITH MUTUAL-INFORMATION REGULARIZATION

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.

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 a t ∈ A in state s t ∈ S, the agent transitions to s t+1 with probability P (s t+1 |s t , a t ). 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) = arg max π V π (s), where V π (s) = E T t=0 γ t r(s t , a t )|s 0 = 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 as P π ∈ 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 µ π := lim t→∞ ν 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:

EQUATION

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

EQUATION

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

EQUATION

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 I f be a functional, in particular:

I f (p X , p Y |X , q Y ) := x p X (x)KL(p Y |X (•|x)||q Y (•))

, where p X (x) is the distribution of the input, p Y |X (y|x) is the conditional distribution of the output conditioned on the input, and q Y (y) a variational distribution of the output. Then, the mutual information1 is recovered with

I[X, Y ] = min q Y I f (p X , p Y |X , q Y ),

where the optimal variational distribution is q Y (y) = x p X (x)p Y |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) -

EQUATION

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.

3. MUTUAL-INFORMATION REGULARIZATION IN RL:

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.

3.1. VARIATIONAL INFERENCE RL FORMULATION WITH OPTIMAL ACTION-PRIORS:

The RL problem can be expressed as an inference problem by introducing a binary random variable R that denotes whether the trajectory τ := (s 0 , a 0 , . . . s T , a T ) 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(s t , a t )) (Levine, 2018) . We additionally introduce a scaling factor β > 0 into the exponential, i.e. p(R = 1|τ ) ∝ exp(β T t=0 r(s t , a t )). This will allow us to trade off reward and entropy maximization2 . 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 :

EQUATION

The generative model is written as p(τ ) = p(s 0 )

T -1 t=0 ρ(a t )P (s t+1 |s t , a t ) and the variational distribution as q(τ ) = p(s 0 )

T -1 t=0 π(a t |s t )P (s t+1 |s t , a t ). 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. ρ(a t ) = 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 π,ρ E q T t=0 r(s t , a t ) -1 β log π(at|st) ρ(at)

. Since we are interested in infinite horizon problems, we introduce a discount factor and take the limit lim T →∞ (Haarnoja et al., 2017) . This leads to the optimization objective that we use in our experiments:

EQUATION

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.

3.2. RECURSION, OPTIMAL POLICIES AND OPTIMAL PRIORS:

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(s t , a t ) -1 β log π(at|st) ρ(at)

|s 0 = s where one can show that V π,ρ satisfies a recursion similar to the Bellman equation:

EQUATION

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

EQUATION

with Z = a ρ(a) exp(βQ π * ,ρ (s, a)), and the the soft Q-function is defined as

EQUATION

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(s 0 ) t-2 t =0

π(a t |s t )P (s t +1 |s t , a t ) π(a t-1 |s t-1 )P (s|s t-1 , a t-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) .

4. MIRL: A PRACTICAL MUTUAL INFORMATION RL ALGORITHM:

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.

4.1. MIRL FOR TABULAR Q-LEARNING:

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,

EQUATION

) where s i ∼ ν 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

EQUATION

) where α Q is a learning rate and T ρ soft Q is the empirical soft-operator defined as (T ρ soft Q)(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 a i is obtained by

a i = arg max a π i (a|s i ), if u > a ∼ ρ i (•) if u ≤ , where π i (a|s) = 1 Z ρ i (a) exp(β i Q i (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.

4.2. MIRL WITH PARAMETRIC FUNCTION APPROXIMATORS:

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

EQUATION

)

where M 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 ρ soft Q is the empirical soft-operator from the tabular setting, here repeated for convenience (T ρ soft Q)(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 + α β ( 1 L(θ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 s i and apply action a i ∼ π b (•|s i ) 4:

Get r i , s i+1 and store (s i , a i , r i , s i+1 ) in replay memory M 5: 

Update prior ρ i+1 (•) = ρ i (•)(1 -α ρ ) + α ρ π i (•|s i ) 6: if i mod update frequency == 0 then 7: Update Q-function θ i+1 = θ i -α Q ∇ θ L(θ, ρ i+1

5. EXPERIMENTS:

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.

5.1. GRID WORLD:

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 • 10 5 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.

5.2. ATARI:

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 • 10 7 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 10 5 environment steps. Evaluating a single snapshot offline is done by running the policy for 30 episodes that last at most 4.5 • 10 3 environment steps but terminate earlier in case of a terminal event. When evaluating the agents, the epsilon value is = 0.05, whereas in training is linearly annealed over the first 10 6 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 z normalized = z-zrandom zhuman-zrandom • 100%, where z stands for the score obtained from our agent at test time, z random stands for the score that a random agent obtains and z human 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.

6. RELATED WORK:

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. The mutual information, central to our approach, is a basic quantity in information theory to measure the statistical dependence between two random variables. Machine learning applications that use the mutual information are numerous including the information-bottleneck method (Tishby et al., 1999) , rate-distortion theory (Cover & Thomas, 2006; Tishby & Polani, 2011) , clustering (Still & Bialek, 2004 ) and curiosity driven exploration (Still & Precup, 2012) .

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 Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In AAAI, volume 8, pp. 1433-1438. Chicago, IL, USA, 2008.

A APPENDIX:

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

EQUATION

where µ π and ρ π are the stationary distributions induced by the policy π over states and actions, respectively, and thus I f (µ π , π, ρ π ) 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

EQUATION

where K(γ) := C 1-γ (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 s t at time t following the Markovian dynamics is written as in Equation (10).

A standard result in the MDP literature (Bertsekas, 1995) 

is that arg max π lim γ→1 (1 -γ)E q ∞ t=0 γ t r(s t , a t ) ⇐⇒ arg max π E s∼µπ 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 γ t I f (ν t , π, ρ) ≤ K(γ) ⇐⇒ min ρ (1 -γ) ∞ t=0 γ t I f (ν t , π, ρ) ≤ C.

Taking the limit γ → 1 in the last inequality and interchanging the limit and the min operators5 , we obtain the constraint min ρ lim γ→1 (1 -γ) ∞ t=0 γ t I f (ν 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 I f (µ π , π, ρ π ) := min ρ I f (µ π π, ρ) can also be written as

EQUATION

where ν t (s) is defined as in (10).

Proof. Following similar steps as in (Bertsekas, 1995, p.186) we have

I f (µ π , π, ρ π ) = min ρ lim N →∞ 1 N N -1 t=0 I f (ν t , π, ρ) = min ρ lim N →∞ lim γ→1 N -1 t=0 γ t I f (ν t , π, ρ) N t=0 γ t = min ρ lim γ→1 lim N →∞ N t=0 γ t I f (ν t , π, ρ) N t=0 γ t = min ρ lim γ→1 (1 -γ) ∞ t=0 γ t I f (ν 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 I f (µ π , π, ρ π ) can also be written as

EQUATION

where ν t (s) is defined as in (10). 

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) = I f (µ π , π, ρ π ),

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.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. 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

:

I(X; Y ) := x,y p(x, y) log p(x,y) p Y (y)p X (x) = x pX (x)KL(p Y |X (•|x)||pY (•))

Other authors absorb this scaling factor into the reward, but we keep it as an explicit hyperparameter.

This 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 ).

Pracitcally, we replace the squared loss with the Huber lossMnih et al. (2015).

× 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.

We can interchange the operators because the l.h.s. of the inequality is finite for any 0 < γ < 1.

Figures and Tables:

Figure 1: Grid world experiments. Left column:The top shows a corridor grid world with 3 × 20 cells where the goal is on the right. The bottom shows an 8 × 8 grid world where an important action (left) has to be made exactly once to arrive at the goal. Middle column: Evaluation of Qlearning (QL), SQL with standard uniform exploration and with marginal exploration (SQL m), and MIRL. We clearly see that MIRL outperforms the baselines on the corridor, and is comparable to the baselines on the 8 × 8 world. Right column: We see that MIRL is able to identify the correct action (go right) faster than the baselines in the corrider (top). The bottom reports how having infrequent but important actions does not affect the performance of MIRL.

Figure3: Normalized score for eight games comparing MIRL against standard SQL and a modified version of SQL that explores with the marginal distribution over actions (SQL m). The exploration method slightly improves SQL but not sufficiently enough to achieve MIRL's performance. See individual plots and games in the Appendix.

Let I ρ (s) := KL(π(•|s)||ρ(•)) and I ρπ (s) := KL(π(•|s)||ρ π (•)). Note that both previous quantities are bounded for all t.

Figure5: Scores for all games on the evaluation snapshots.

Figure 6: Comparison between standard SQL (SQL u), SQL with marginal exploration (SQL m), and MIRL.

Left panel: Median normalized score across 19 Atari games. Comparison between our method mutual information RL (MIRL), SQL and DQN, demonstrating MIRL's superior performance. Right panels: Top figures show the raw score for 2 example games reporting MIRL's superior performance on RoadRunner and Seaquest. The bottom plots show the evolution of the estimated prior over actions. For RoadRunner the prior converges to stable values during training. In Seaquest, the algorithm seems not to have converged yet after 50 million environment steps which is why the prior probabilities have not converged yet either (however, the formation of separate trajectory clusters towards the end of training indicates the ongoing process of convergence). See Appendix for details and plots for all environments. The curves are smoothed with an exponential moving average with effective window size of 10 6 environment steps.

Table content: <table><tr><td>Normalized Score (%)</td><td>0 50 100 150 200</td><td>0</td><td>1 Normalized Score (Median across games) 2 3 4 Environment Steps DQN SQL MIRL</td><td>1e7 5</td><td>0.0 50000 40000 0 10000 20000 30000 0.1 0.2 0.3 0.4 Prior Probabilities Reward</td><td>0 0</td><td>1 RoadRunnerNoFrameskipv4 2 3 4 Environment Steps 1 2 3 4 Environment Steps RoadRunnerNoFrameskip-v4</td><td>5 1e7 5 1e7 DQN SQL MIRL</td><td>0.000 8000 0 2000 4000 6000 0.025 0.050 0.075 0.100 0.125 Prior Probabilities Reward</td><td>0</td><td>0</td><td>1 SeaquestNoFrameskipv4 2 3 4 Environment Steps 1 2 3 4 Environment Steps SeaquestNoFrameskip-v4</td><td>1e7 5 5 1e7 DQN SQL MIRL</td></tr><tr><td colspan="3">Figure 2:</td><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/><td/></tr></table>

Mean Normalized score in 19 Atari games for DQN, SQL and our approach MIRL.

Table content: <table><tr><td>Game</td><td colspan="3">DQN (%) SQL (%) MIRL (%)</td></tr><tr><td>Alien</td><td>101.58</td><td>51.02</td><td>40.23</td></tr><tr><td>Assault</td><td>250.61</td><td>283.62</td><td>357.40</td></tr><tr><td>Asterix</td><td>166.32</td><td>242.73</td><td>330.19</td></tr><tr><td>Asteroids</td><td>9.74</td><td>8.57</td><td>7.80</td></tr><tr><td>BankHeist</td><td>97.12</td><td>94.62</td><td>166.26</td></tr><tr><td>BeamRider</td><td>99.16</td><td>113.64</td><td>117.21</td></tr><tr><td>Boxing</td><td>2178.57</td><td>2283.33</td><td>2338.89</td></tr><tr><td>ChopperCommand</td><td>72.71</td><td>26.37</td><td>65.03</td></tr><tr><td>DemonAttack</td><td>350.95</td><td>451.78</td><td>469.30</td></tr><tr><td>Gopher</td><td>474.18</td><td>538.87</td><td>429.44</td></tr><tr><td>Kangaroo</td><td>351.48</td><td>393.16</td><td>405.9</td></tr><tr><td>Krull</td><td>843.16</td><td>886.68</td><td>1036.04</td></tr><tr><td>KungFuMaster</td><td>122.14</td><td>142.04</td><td>121.41</td></tr><tr><td>Riverraid</td><td>77.21</td><td>109.37</td><td>76.02</td></tr><tr><td>RoadRunner</td><td>548.90</td><td>613.62</td><td>695.88</td></tr><tr><td>Seaquest</td><td>21.95</td><td>36.00</td><td>64.86</td></tr><tr><td>SpaceInvaders</td><td>166.62</td><td>200.38</td><td>164.79</td></tr><tr><td>StarGunner</td><td>653.44</td><td>681.12</td><td>574.89</td></tr><tr><td>UpNDown</td><td>183.19</td><td>230.82</td><td>394.21</td></tr><tr><td>Mean</td><td>356.26</td><td>388.83</td><td>413.46</td></tr></table>

). Hyperparameters for tabular experiments.

Table content: <table><tr><td>Parameter</td><td>Value</td></tr><tr><td>γ</td><td>0.99</td></tr><tr><td>α ρ</td><td>2 • 10 -3</td></tr><tr><td>ω</td><td>0.8</td></tr></table>

Peer review (iclr2019):
 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. Generally I found the idea interesting and at a high level the deficiency of entropy regularization makes sense. However, I had great trouble understanding the reasoning behind their method and did not find the connection to mutual information helpful. Furthermore, I had a number of questions about the experiments. If the authors can clarify their motivation and reasoning and strengthen the experiments, I*d be happy to raise my score. In Sec 3.1, why is it sensible to optimize the prior? Can the authors give intuition for maximizing .log p(R = 1) wrt to the prior? This is critical for justifying their approach. Currently, the authors provide a connection to MI, but don*t explain why this matters. Does it justify the method? What insight are we supposed to take away from that? The experiments could be strengthened by addressing the following: * What was epsilon during training? Why was epsilon = 0.05 in evaluation? This is quite high compared to previous work, and it makes sense that this would degrade MIRLs performance less than DQN and SQL. * What is the performance of SQL if we use .rho as the action selector in .epsilon-greedy. This would help understand if the performance gains are due to the impact on the policy or due to the changes in the behavior policy. * Plotting beta over time * Comparing the action distributions for SQL and MIRL to understand the impact of the penalty. In general, a deeper analysis of the impact on the policy is important. * Are their environments we would expect MIRL to outperform SQL based on your theoretical understanding? Does it? * How many seeds were run per game? * How and why were the 19 games selected from the full set? Comments: The abstract claims state-of-the-art performance, however, what is actually shown is that MIRL outperforms DQN and SQL. With a fixed prior, the action prior can be absorbed into the reward (e.g., Levine 2018), so it is of no loss of generality to assume a uniform prior. Could state that the stationary distribution is assumed to exist and be unique. In Sec 3.1, why is the prior state independent? In Sec 3.1, p(R = 1|.tau) is defined to be proportional to exp(.beta .sum_t r_t). Is this well-specified? How would we compute the normalizing constant since p(R = 0 | .tau) is not defined? Throughout, I suggest that the authors not use the phrases "closed form" and "analytic" for expressions that are in terms of intractable quantities. It should be noted that Sec 3.2 Optimal policy for a fixed prior .rho follows from Levine 2018 and others by transforming the fixed prior into a reward bonus. In Sec 3.2, the last statement does not appear to be necessary for the next subsection. Remove or clarify? I believe that the connection to MI can be simplified. Plugging in the optimal .rho into Eq 3, we can see that Eq 3 simplifies to .max_.pi E_q[ .sum_t .gamma^t r_t] - (1 - gamma)/.beta MI_p(s, a) where p(s, a) = d^.pi(s) * .pi(a | s) and d^.pi is the discounted state visitation distribution. Thus Eq 3 can be thought of as a lower bound on the MI regularized objective. In Sec 4, the authors state the main difference between their soft operator and the typical soft operator. What other differences are there? Is that the only one? Sec 5 references the wrong Haarnoja reference in the first paragraph. In Sec 5, alpha_beta = 3 * 10^5. Is that correct? ===== 11/26 At this time, the authors have not responded to the reviews. I have read the other reviews and comments, and I*m not inclined to change my score. ==== 12/7 The authors have addressed most of my concerns, so I have raised my score. I*m still concerned that the exploration epsilon is quite different than existing work (e.g., https://github.com/google/dopamine/tree/master/baselines).