Assumed Density Filtering Q-learning

1 Introduction

Bayesian reinforcement learning is a classic reinforcement learning (RL) technique that utilizes Bayesian inference to integrate new experiences with prior information about the problem in a probabilistic distribution. It explicitly quantifies the uncertainty of the learning parameters unlike standard RL approaches in which uncertainty is unaccounted for. Explicit quantification of the uncertainty can help guide policies that consider the exploration-exploitation trade-off by exploring actions with higher uncertainty more often [Osband et al., 2013; Osband et al., 2014]. Moreover, it can also regularize posterior updates by properly accounting for uncertainty.

∗Contact Author

Motivated by these advantages, a number of algorithms have been proposed in both model-based [Dearden et al., 1999; Duff, 2002; Guez et al., 2012; Poupart et al., 2006] and modelfree Bayesian RL [Dearden et al., 1998; Engel et al., 2003; Engel et al., 2005; Geist and Pietquin, 2010; Chowdhary et al., 2014]. However, Bayesian approaches to off-policy temporal difference (TD) learning have been less studied compared to alternative methods due to difficulty in handling the max nonlinearity in the Bellman optimality equation. Previous studies such as Dearden's Bayesian Q-learning [1998] and Kalman Temporal Difference Q-learning (KTD-Q) [Geist and Pietquin, 2010] suffer from their computational complexity and scalability. Yet off-policy TD methods such as Q-learning [Watkins and Dayan, 1992] have been widely used in standard RL, including extensions integrating neural network function approximations such as Deep Q-Networks (DQN) [Mnih et al., 2013].

In this paper, we introduce a novel approximate Bayesian Q-learning algorithm, denoted as ADFQ, which updates belief distributions of Q (action-value function) and approximates their posteriors using an online Bayesian inference algorithm known as assumed density filtering (ADF). In order to reduce the computational burden of estimating parameters of the approximated posterior, we propose a method to analytically estimate the parameters. Unlike Q-learning, ADFQ executes a non-greedy update by considering all possible actions for the next state and returns a soft-max behavior and regularization determined by the uncertainty measures of the Q-beliefs. This alleviates overoptimism and instability issues from the greedy update of Q-learning which have been discussed in a number of papers [Tsitsiklis, 2002; Harutyunyan et al., 2016; Hasselt, 2010; Hasselt et al., 2016]. We prove the convergence of ADFQ to the optimal Q-values by showing that ADFQ becomes identical to Q-learning as all state and action pairs are visited infinitely often.

ADFQ is computationally efficient and is extended to complex environments with a neural network. There are previous works that implement Bayesian approaches to Deep RL by using uncertainty in the neural network weights and show promising performance in several Atari games [Azizzadenesheli et al., 2018; O'Donoghue et al., 2017; Osband et al., 2016]. However, these approaches only focus on exploration and uncertainty information does not directly applied to updating RL parameters. Our method differs from these approaches as it explicitly computes the variances of the Q-beliefs and uses them both for exploration and in the value update. Bellemare et al. [2017] proposed a gradient-based categorical DQN algorithm using a distributional perspective. The value distribution in their work represents the inherent randomness of the agent's interactions with its environment. In contrast, the Q-belief defined in ADFQ is a belief distribution of a learning agent on a certain state-action pair. Therefore, only ϵ-greedy is used in their experiments. We evaluate ADFQ with Thompson sampling (TS) [Thompson, 1933] as well as ϵ-greedy methods in various Atari games and they outperform DQN and Double DQN [Hasselt, 2010]. Particularly, the non-greedy update in ADFQ dramatically improves the performance in domains with a large number of actions and higher stochasticity. Example source code is available online 1 .

2 Background

2.1 Assumed Density Filtering

Assumed Density Filtering (ADF) is a general technique for approximating the true posterior with a tractable parametric distribution in Bayesian networks. It has been independently rediscovered for a number of applications and is also known as moment matching, online Bayesian learning, and weak marginalization [Opper, 1999; Boyen and Koller, 1998; Maybeck, 1982]. Suppose that a hidden variable x follows a tractable parametric distribution p(x|θt) where θt is a set of parameters at time t. In the Bayesian framework, the distribution can be updated after observing some new data (Dt) using Bayes' rule, ˆ p(x|θt,Dt) ∝ p(Dt|x, θt)p(x|θt). In online settings, a Bayesian update is typically performed after a new data point is observed, and the updated posterior is then used as a prior for the following iteration.

When the posterior computed by Bayes' rule does not belong to the original parametric family, it can be approximated by a distribution belonging to the parametric family. In ADF, the posterior is projected onto the closest distribution in the family chosen by minimizing the reverse Kullback-Leibler divergence denoted as KL(ˆ p||p) where ˆ p is the original posterior distribution and p is a distribution in a parametric family of interest. Thus, for online Bayesian filtering, the parameters for the ADF estimate is given by θt+1 = argminθ KL(ˆ p(·|θt, Dt)||p(·|θ)).

2.2 Q-learning

RL problems can be formulated in terms of an MDP described by the tuple, M = 〈S,A,P,R, γ〉 where S and A are the state and action spaces, respectively, P : S × A × S → [0, 1] is the state transition probability kernel, R : S × A → IR is a reward function, and γ ∈ [0, 1) is a discount factor. The value function is defined as V π (s) = Eπ[∑∞t=0 γ tR(st, at)|s0 = s] for all s ∈ S, the expected value of cumulative future rewards starting at a state s and following a policy π thereafter. The state-action value (Q) function is defined as the value for a state-action pair, Qπ (s, a) = Eπ[∑∞t=0 γ tR(st, at)|s0 = s, a0 = a] for all s ∈ S, a ∈ A. The objective of a learning agent in RL is to find an optimal policy π∗ = argmaxπ V π .

1 https://github.com/coco66/ADFQ

Finding the optimal values, V ∗ (·) and Q∗ (·, ·), requires solving the Bellman optimality equation:

and V ∗ (s) = maxa∈A(s) Q∗ (s, a) ∀s ∈ S where s′ is the subsequent state after executing the action a at the state s. Qlearning is the most popular off-policy TD learning technique due to its relatively easy implementation and guarantee of convergence to an optimal policy [Watkins and Dayan, 1992; Kaelbling et al., 1996]. At time step t, it updates Q(st, at) after observing a reward rt and the next state st+1 (one-step TD learning). The update is based on the TD error - a difference between the TD target, rt + γ maxb Q(st+1, b), and the current estimate on Q(st, at) with a learning rate α ∈ (0, 1]:

3 Bayesian Q-learning with ADF

3.1 Belief Updates on Q

We define Qs,a as a Gaussian random variable with mean µs,a and variance σ2 s,a corresponding to the action value function Q(s, a) for all s ∈ S and a ∈ A. We assume that the random variables for different states and actions are independent and have different means and variances, Qs,a ∼ N(µs,a, σ2 s,a ) where µs,a ̸= µs′ ,a′ if s ̸= s′ or a ̸= a′ ∀s ∈ S, ∀a ∈ A. According to the Bellman optimality equation in Eq.1, we can define a random variable for V (s) as Vs = maxa Qs,a. In general, the probability density function for the maximum of Gaussian random variables, M = max1≤k≤N Xk where Xk ∼ N(µk, σ2 k ), is no longer Gaussian (see the appendix).

For one-step Bayesian TD learning, the beliefs on Q = {Qs,a}∀s∈S,∀a∈A can be updated at time t after observing rt and st+1 using Bayes' rule. In order to reduce notation, we drop the dependency on t denoting st = s, at = a, st+1 = s′ , rt = r, yielding the causally related 4-tuple τ =< s, a, r, s′ >. We use the one-step TD target with a small Gaussian white noise, r + γVs′ +W where W ∼ N(0, σ2w ), as the likelihood for Qs,a. The noise parameter, σw, reflects stochasticity of an MDP. We will first derive the belief updates on Q-values with σw = 0 for simplicity and then extend the result to the general case. The likelihood distribution can be represented as a distribution over Vs′ as p(r + γVs′ |q, θ) = pVs′ ((q - r)/γ|s′ , θ) where q is a value corresponding to Qs,a and θ is a set of mean and variance of Q. From the independence assumptions on Q, the posterior update is reduced to an update for the belief on Qs,a:

The resulting posterior distribution is given as follows where φ(·) is the standard Gaussian probability density function (PDF) and Φ(·) is the standard Gaussian cumulative distribution function (CDF) (derivation details in the appendix):

Z is a normalization constant and

Note that all next actions are considered in Eq.2 unlike the conventional Q-learning update which only considers the subsequent action resulting in the maximum Q-value at the next step (maxb Q(s′ , b)). This can lead to a more stable update rule as updating with only the maximum Q-value has inherent instability. The Bayesian update considers the scenario where the true maximum Q-value may not be the one with the highest estimated mean, and weights each subsequent Q-value accordingly. Each term for action b inside the summation in Eq.2 has three important features. First of all, ¯µτ,b is an inversevariance weighted (IVW) average of the prior mean and the TD target mean. Therefore, the Gaussian PDF part becomes closer to the TD target distribution if it has a lower uncertainty than the prior, and vice versa as compared in the first row of Fig.1. Next, the TD error, δτ,b = (r + γµs′ ,b) - µs,a, is naturally incorporated in the posterior distribution with the form of a Gaussian PDF in the weight cτ,b. Thus, a subsequent action which results in a smaller TD error contributes more to the update. The sensitivity of a weight value is determined by the prior and target uncertainties. An example case is described in the second row of Fig.1 where δτ,0 < δτ,1 and σs′ ,0 > σs′ ,1. Finally, the product of Gaussian CDFs provides a soft-max operation. The red curve with dots in the third row of Fig.1 represents ∏b′̸=b Φ(q|r + γµτ,b′ , γστ,b′ ) for each b. For a certain q value (x-axis), the term returns a larger value for a larger µs′ ,b as seen in the black circles. This results has a similarity with the soft Bellman equation [Ziebart, 2010], but the degree of softness in this case is determined by the uncertainty measures rather than a hyperparameter.

3.2 ADF on Q-Belief Updates

The posterior distribution in Eq.2, however, is no longer Gaussian. In order to continue the online Bayesian update, we approximate the posterior with a Gaussian distribution using ADF. When the parametric family of interest is spherical Gaussian, it is shown that the ADF parameters are obtained by matching moments. Thus, the mean and variance of the approximate posterior are given by those of the true posterior, Eˆ pQs,a [q] and Varˆ pQs,a [q], respectively. It is fairly easy to derive the mean and variance when |A| = 2. The derivation is presented in the appendix. However, to our knowledge, there is no analytically tractable solution for |A| > 2.

When σw > 0, the expected likelihood is obtained by solving ∫R p(r + γVs′ + w|q, θ)pW (w)dw which is an integral of a similar form with the posterior in Eq.2. Therefore, a closedform expression is also not available in general except when |A| = 2 (see the appendix).

In the next sections, we prove the convergence of the means to the optimal Q-values for the case |A| = 2 with the exact solutions for the ADF parameters. Then, we show how to derive an analytic approximation for the ADF parameters which becomes exact in the small variance limit.

3.3 Convergence to Optimal Q-values

The convergence theorem of the Q-learning algorithm has previously been proven [Watkins and Dayan, 1992]. We, therefore, show that the online Bayesian update using ADF with the posterior in Eq.2 converges to Q-learning when |A| = 2. We apply an approximation from Lemma 1 in order to prove Theorem 1. Proofs are presented in the appendix.

Lemma 1. Let X be a random variable following a normal distribution, N(µ, σ2 ). Then we have:

where [x]+ = max(0, x) is the ReLU nonlinearity.

Theorem 1. Suppose that the mean and variance of Qs,a ∀s ∈ S, ∀a ∈ A are iteratively updated by the mean and variance of ˆ pQs,a after observing r and s′ at every step. When |A| = 2, the update rule of the means is equivalent to the Q-learning update if all state-action pairs are visited infinitely often and the variances approach 0. In other words, at the kth update on µs,a:

Interestingly, ατ approaches 1 when σs,a/σs′ ,b+ → ∞ and 0 when σs,a/σs′ ,b+ → 0 for σw = 0. Such behavior remains when σw > 0 but ατ eventually approaches 0 as the number of visits to (s, a) goes to infinity. This not only satisfies the convergence condition of Q-learning but also provides a natural learning rate - the smaller the variance of the TD target (the higher the confidence), the more Qs,a is updated from the target information rather than the current belief. We show empirical evidence that the contraction condition on variance in Theorem 1 holds in the appendix.

4 Analytic ADF Parameter Estimates

When |A| > 2, the update can be solved by numerical approximation of the true posterior mean and variance using a number of samples. However, its computation becomes unwieldy due to the large number of samples needed for accurate estimates. This becomes especially problematic with small variances as the number of visits to corresponding state-action pairs grows. In this section, we show how to accurately estimate the ADF parameters using an analytic approximation. This estimate becomes exact in the small variance limit.

4.1 Analytic Approximation of Posterior

Applying Lemma 1 to the Gaussian CDF terms in Eq.2, the posterior is approximated to the summation over b ∈ A of the following term:

Similar to Laplace's method, we approximate each term as a Gaussian distribution by matching the maximum values as well as the curvature at the peak of the distribution. In other words, the maximum of the distribution is modeled locally near its peak by the quadratic concave function:

We find µ∗τ,b and σ∗ τ,b by matching the first and the second derivatives, respectively:

where H(·) is a Heaviside step function. The RHS of the selfconsistent piece-wise linear equation Eq.8 is an IVW average mean of the prior, the TD target distribution of b, and other TD target distributions whose means are larger than µ∗τ,b . The height of the peak is computed as,

Algorithm 1 ADFQ algorithm

• 1: Initialize randomly µs,a, σs,a ∀s ∈ S and ∀a ∈ A 2: for each episode do 3: Initialize s0 4: for each time step t do 5: Choose an action, at ∼ πaction (st; θt) 6: Perform the action and observe rt and st+1 7: for each b ∈ A do 8: Compute µ∗τ,b, σ∗ τ,b, k∗ τ,b using Eq.7-9 9: end for 10: Update µst,at and σst,at using Eq.10 and Eq.11 11: end for 12: end for

The final approximated distribution is a Gaussian mixture model with µ∗τ,b , σ∗ τ,b , w∗ τ,b as mean, variance, and weight, respectively, for all b ∈ A where w∗ τ,b = k∗ τ,b / ∑b′ k∗ τ,b′ . Therefore, we update the belief distribution over Qs,a with the mean and variance of the Gaussian mixture model:

The final mean is the weighted sum of each individual mean with a weight from k∗ τ,b and the final variance is the weighted sum of each individual variance added to a non-negative term accounting for the dispersion of the means. As shown in Eq.9, the weights are determined by TD errors, variances, and relative distances to larger TD targets. Each weight includes the TD error penalizing term, cτ,b, and also decreases as the number of TD targets larger than µ∗τ,b increases. Therefore, the weight provides a softened maximum property over b. The algorithm is summarized in Algorithm 1. Its space complexity is O(|S||A|). The computational complexity of each update is O(|A|2 ) which is higher than Q-learning but only by a factor of |A| and constant in the number of states.

4.2 Approximate Likelihood for Stochastic MDPs

In an asymptotic limit of σw/σs′ ,b → 0, ∀b ∈ A and |A| = 2, the expected likelihood distribution for σw > 0 is similar to p(r + γVs′ |q, θ) but the variance of its Gaussian PDF term is γ 2σ2 s′ ,b + σ2w instead of γ 2σ2 s′ ,b (see the appendix for details). Extending this result to the general case (|A| = n for n ∈ N), the posterior distribution for σw > 0 is same with Eq.2 but γ 2σ2 s′ ,b is replaced by γ 2σ2 s′ ,b+σ2w in cτ,b, ¯µτ,b, and ¯στ,b (Eq.35). Therefore, µ∗τ,b , σ∗ τ,b , and k∗ τ,b in the ADFQ algorithm (Table.1) are also changed accordingly. In practice, the nonzero noise parameter is needed when an MDP is stochastic.

4.3 Convergence of ADFQ

Theorem 1 extends to the ADFQ algorithm (Proof in the appendix). The contraction behavior of the variances in the case of Theorem 1 is also empirically observed in ADFQ.

Theorem 2. The ADFQ update on the mean µs,a ∀s ∈ S, ∀a ∈ A for |A| = 2 is equivalent to the Q-learning update if the variances approach 0 and if all state-action pairs are visited infinitely often. In other words, we have :

As we have observed the behavior of ατ in Theorem 1, the learning rate ατ again provides a natural learning rate with the ADFQ update. We can therefore think of Q-learning as a special case of ADFQ.

5 Demonstration in Discrete MDPs

To demonstrate the behavior of the ADFQ update, we look at the simple MDP (γ = 0.9) in Fig.2 at a specific iteration. An episode starts at s0 and terminates at either s2 or s3. At s1, each action returns a stochastic reward with p = 0.2. The optimal deterministic policy at s1 is a1. Suppose an RL learner has already visited (s1, a1) 3 times and obtained a reward of r = 5 every time. Now it is on the t-th iteration with (s1, a1) The plots in Fig.2 show the ADFQ update for Qs0,a0 at t + 1 when rt = +5 (left) and rt = -5 (right). When it receives a less expected reward, -5, at t, σs1,a1;t is updated to a larger value than the one in the rt = +5 case. Then, the episode is terminated and the next episode starts at st+1 = s0, at+1 = a0. ADFQ considers both Qs1,a0 and Qs1,a1 for updating Qs0,a0 . Due to the relatively large TD error and variance of Qs1,a1 , a lower value is assigned to w∗ τ,b=1 . In this same scenario, Q-learning would update Q(s0, a0) only from Q(s1, a0) and regulate the update amount with the learning rate which is usually fixed or determined by the number of visits.

In order to show the benefits of the update rule, we examined Q-learning, ADFQ, and a numerical approximation of the mean and variance of Eq.2 (denoted as ADFQ-Numeric) for the convergence to the optimal Q-values in the presented MDP and a similar MDP but with 10 terminating states and 10 actions. Random exploration is used in order to evaluate only the update part of each algorithm. During learning, we computed the root mean square error (RMSE) between the estimated Q-values (or means) and Q∗ , and plotted the averaged results over 5 trials in Fig.3. As shown, ADFQ converged to the optimal Q-values quicker than Q-learning in both cases and showed more stable performance. ADFQ-Numeric suffers from correctly estimating the parameters when its variances become small as it is previously pointed out, and resulted a poor convergence result in the large MDP.

6 ADFQwith Neural Networks

In this section, we extend our algorithm to complex environments with neural networks similar to Deep Q-Networks (DQN) proposed in [Mnih et al., 2013]. In the Deep ADFQ model with network parameters ξ, the output of the network is mean µ(s, a; ξ) and variance σ2 (s, a; ξ) of each action for a given state s as shown in Fig.4. In practice, we use - log(σs,a) instead of σ2 s,a for the output to ensure positive values for the variance. As in DQN, we have a train network (ξ) and a target network (ξ ′ ). Mean and variance for s and s′ from the target network are used as inputs into the ADFQ algorithm to compute the desired mean, µADFQ, and standard deviation, σADFQ for the train network. We used prioritized experience replay [Schaul et al., 2015] and a combined Huber loss functions of mean and variance.

In order to demonstrate the effectiveness of our algorithm, we tested on six Atari games, Enduro (|A| = 9), Boxing (|A| = 18), Pong (|A| = 6), Asterix (|A| = 9), Kung-Fu Master (|A| = 14), and Breakout (|A| = 4), from the OpenAI gym simulator [Brockman et al., 2016]. For baselines, we used DQN and Double DQN with prioritized experience replay implemented in OpenAI baselines 2 with their default hyperparameters for all games. We used ϵ-greedy action policy with ϵ annealed from 1.0 to 0.01 for the baselines as well as ADFQ. In ADFQ, the greedy selection is performed on the mean values instead of Q-values. Additionally, we examined Thompsing Sampling (TS) for ADFQ which selects at = argmaxa qst,a where qst,a ∼ pQst,a (·|θt).The algorithms were evaluated for TH = 10M training steps (5M for Pong). Each learning was greedily evaluated at every epoch (= TH/100) for 3 times, and their averaged results are presented in Fig.5. The entire experiment was repeated for 3 random seeds. Rewards were normalized to {-1, 0, 1} and different from raw scores of the games. Both ADFQ with TS and with ϵ-greedy notably surpassed DQN and Double DQN in Enduro, Boxing, Asterix, and Kung-Fu Master and showed similar results in Pong. The performance of ADFQ in Breakout is explained as Breakout is the only tested domain where there is no dynamic object interrupting the learning agent. As the demonstration in Sec.5 and the additional experiments in the appendix show, improvements of ADFQ from Q-learning is more significant when an experimental domain has high stochasticity and its action space is large due to the non-greedy update with uncertainty measures. Additionally, ADFQ showed more stable performance in all tested domains overcoming DQN's instability. ADFQ with TS achieved slightly higher performance than the ϵ-greedy method utilizing the uncertainty in exploration.

7 Discussion

We proposed an approach to Bayesian off-policy TD method called ADFQ. ADFQ demonstrated that it could improve some of the issues from the greedy update of Q-learning by showing the quicker convergence to Q∗ than Q-learning and surpassing DQN and Double DQN in various Atari games. The presented ADFQ algorithm demonstrates several intriguing results.

2 https://github.com/openai/baselines

First, unlike the conventional Q-learning algorithm, ADFQ incorporates the information of all available actions for the subsequent state in the Q-value update. Each subsequent stateaction pair contributes to the update based on its TD target mean and variance as well as its TD error. Particularly, we make use of our uncertainty measures not only in exploration but also in the value update as natural regularization. The advantages of this non-greedy update are noticeable in highly stochastic domains and domains with a large action space in the experiment. Next, we prove that ADFQ converges to Q-learning as the variances decrease and can be seen as a more general form of Q-learning. Last, one of the major drawbacks of Baysian RL approaches is their high computational complexity and poor scalability. ADFQ is computationally efficient and is extended to Deep ADFQ with a neural network.

We would like to highlight the fact that ADFQ is a Bayesian counterpart of Q-learning and is orthogonal to most other advancements made in Deep RL. Deep ADFQ merely changes the loss function and we compare with basic architectures here to provide insight as to how it may improve the performance. ADFQ can be used in conjunction with other extensions and techniques applied to Q-learning and DQN.

Acknowledgments

This work was supported by funding from the National Institutes of Health U19 program.