Analytically Tractable Bayesian Deep Q-Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has gained increasing interest since the demonstration it was able to reach human performance on video game benchmarks using deep Q-learning ( DQN ) ( Mnih et al. , 2015 ; Van Hasselt et al. , 2016 ) . Deep RL methods typically require an explicit definition of an exploration-exploitation function in order to compromise between using the current policy and exploring the potential of new actions . Such an issue can be mitigated by opting for a Bayesian approach where we estimate the uncertainty of the neural network ’ s parameters using Bayes ’ theorem , and select the optimal action given a state according to the parameter uncertainty using Thompson sampling technique ( Strens , 2000 ) . Bayesian deep learning ( BDL ) methods based on variational inference ( Kingma et al. , 2015 ; Hernández-Lobato & Adams , 2015 ; Blundell et al. , 2015 ; Louizos & Welling , 2016 ; Osawa et al. , 2019 ; Wu et al. , 2019 ) , Monte-Carlo dropout ( Gal & Ghahramani , 2016 ) , or Hamiltonian Monte-Carlo sampling ( Neal , 1995 ) have shown to perform well on regression and classification benchmarks , despite being generally computationally more demanding than their deterministic counterparts . Note that none of these approaches allow performing the analytical inference for the weights and biases defining the neural network . Goulet et al . ( 2021 ) recently proposed the tractable approximate Gaussian inference ( TAGI ) method that allows estimating the posterior of the neural network ’ s parameters using a closed-form analytical method . More specifically , TAGI leverages Gaussian conditional equations to analytically infer this posterior without the need of numerical approximations ( i.e. , sampling techniques or Monte Carlo integration ) or optimization on which existing BDL methods rely . In addition , TAGI is able to maintain the same computational complexity as deterministic neural networks based on the gradient backpropagation . For convolutional architectures applied on classification benchmarks , this approach was shown to exceed the performance of other Bayesian and deterministic approaches based on gradient backpropagation , and to do so while requiring a smaller number of training epochs ( Nguyen & Goulet , 2021 ) . In this paper , we propose a new perspective on probabilistic inference for Bayesian deep reinforcement learning which , up to now has been relying on gradient-based methods . More specifically , we present how can we adapt the temporal difference Q-learning framework ( Sutton , 1988 ; Watkins & Dayan , 1992 ) to make it compatible with TAGI . Section 2 first reviews the theory behind TAGI and the expected value formulation through the Bellman ’ s Equation . Then , we present how the actionvalue function can be learned using TAGI . Section 3 presents the related work associated with Bayesian reinforcement learning , and Section 4 compares the performance of the TAGI-DQN with the deterministic DQN on on- and off-policy RL approaches . 2 TAGI-DQN FORMULATION . This section presents how to adapt the DQN frameworks in order to make them compatible with analytical inference . First , Section 2.1 reviews the fundamental theory behind TAGI , and Section 2.1 reviews the concept of long-term expected value through the Bellman ’ s equation ( Sutton & Barto , 2018 ) . Then , Section 2.3 presents how to make the Q-learning formulation ( Watkins & Dayan , 1992 ) compatible with TAGI . 2.1 TRACTABLE APPROXIMATE GAUSSIAN INFERENCE . Bayesian deep learning aims to estimate the posterior of neural network ’ s parameters θ conditional on a given training dataset D = { ( xi , yi ) , ∀i ∈ 1 : D } using Bayes ’ theorem so that f ( θ|D ) = f ( θ ) f ( D|θ ) f ( D ) , ( 1 ) where f ( θ ) is the prior Probability Density Function ( PDF ) of parameters , f ( D|θ ) is the likelihood and f ( D ) is the marginal likelihood or evidence . Until recently , the posterior f ( θ|D ) in Equation 1 has been considered to be intractable ( Goodfellow et al. , 2016 ; Goan & Fookes , 2020 ) . TAGI has addressed this intractability issue by leveraging the Gaussian conditional equations ; For a random vector of parameter θ and observations Y for which the joint PDF satisfies f ( θ y ) = N ( [ θ y ] ; [ µθ µY ] , [ Σθ Σ ᵀ Yθ ΣYθ ΣY ] ) , ( 2 ) the Gaussian conditional equation describing the PDF of θ conditional on y is defined following f ( θ|y ) = N ( θ ; µθ|y , Σθ|y ) µθ|y = µθ + Σ ᵀ YθΣ −1 Y ( y − µY ) Σθ|y = Σθ −ΣᵀYθΣ−1Y ΣYθ . ( 3 ) In its simple form , the intractable computational requirements for directly applying Equation 3 limits it to trivial neural networks . In order to overcome this challenge , TAGI leverages the inherent conditional independence structure of hidden layers Z ( j−1 ) ⊥ Z ( j+1 ) |z ( j ) under the assumptions that parameters θ are independent between layers . This conditional independence structure allows breaking down the computations for Equation 3 into a layer-wise two-step procedure ; forward uncertainty propagation and backward update . The first forward uncertainty propagation step is intended to build the joint prior between the hidden states of a layer j , Z ( j ) , and the parameters θ ( j ) directly connecting into it . This operation is made by propagating the uncertainty from the parameters and the input layer through the neural network . In order to maintain the analytical tractability of the forward step , we must address two challenges : The first challenge is related to the product of the weights W and activation units A involved in the calculation of the hidden states Z so that Z ( j+1 ) i = A∑ k=1 W ( j ) i , k A ( j ) k +B ( j ) i , ( 4 ) where B is the bias parameter , ( j ) refers to the jth layer , i refers to a node in the current layer , k refers to a node from the previous layer , and A is the number of hidden units from the previous layer . The second challenge is to propagate uncertainty through the non-linearity when activating hidden states with an activation function φ ( . ) , A ( j+1 ) i = φ ( Z ( j+1 ) i ) . ( 5 ) In order to tackle these issues , TAGI made the following assumptions A1 . The prior for the parameters and variables in the input layer are Gaussians . A2 . The product of a weight and an activation unit in Equation 4 is approximated by a Gaussian random variable , WA ≈ N ( µWA , σ2WA ) , whose moments are computed analytically using Gaussian multiplicative approximation ( GMA ) . A3 . The activation function φ ( . ) in Equation 5 is locally linearized using a first-order Taylor expansion evaluated at the expected value of the hidden unit being activated . Note that the linearization procedure in the assumption A3 may seems to be a crude approximation , yet it has been shown to match or exceeds the state-of-the-art performance on fully-connected neural networks ( FNN ) ( Goulet et al. , 2021 ) , as well as convolutional neural networks ( CNN ) and generative adversarial networks ( Nguyen & Goulet , 2021 ) . In order to maintain a linear computational complexity with respect to the number of weights for the forward steps , TAGI relies on the following assumptions A4 . The covariance for all parameters in the network and for all the hidden units within a same layer are diagonal . A5 . The hidden states on a given layer are independent from the parameters connecting into the subsequent layer , i.e. , Z ( j ) ⊥ θ ( j ) . A6 . The hidden states at layer j + 1 depend only on the hidden states and parameters directly connecting into it . The empirical validity of those assumptions as well as the GMA formulations are provided by Goulet et al . ( 2021 ) . The second backward update-step consists in performing layer-wise recursive Bayesian inference going from hidden-layer to hidden-layer , and from hidden-layer to the parameters connecting into it ( see Figure 1 ) . The assumptions A1-A6 allow performing analytical inference while still maintaining a linear computational complexity with respect to the number of weight parameters in the network where only the posterior of hidden states for the output layer are computed using Equation 3 . For the remaining layers , those quantities are obtained using the Rauch-Tung-Striebel recursive procedure that was developed in the context of state-space models ( Rauch et al. , 1965 ) . For simplicity , we define the short-hand notation { θ+ , Z+ } ≡ { θ ( j+1 ) , Z ( j+1 ) } and { θ , Z } ≡ { θ ( j ) , Z ( j ) } so that the posterior for the parameters and hidden states are computed following f ( Z|y ) = N ( z ; µZ|y , ΣZ|y ) µZ|y = µZ + JZ ( µZ+|y − µZ+ ) ΣZ|y = ΣZ + JZ ( ΣZ+|y −ΣZ+ ) JᵀZ JZ = ΣZZ+Σ −1 Z+ , ( 6 ) f ( θ|y ) = N ( θ ; µθ|y , Σθ|y ) µθ|y = µθ + Jθ ( µZ+|y − µZ+ ) Σθ|y = Σθ + Jθ ( ΣZ+|y −ΣZ+ ) Jᵀθ Jθ = ΣθZ+Σ −1 Z+ . ( 7 ) Figure 1 presents a directed acyclic graph ( DAG ) describing the interconnectivity in such a neural network . TAGI allows inferring the diagonal posterior knowledge for weights and bias parameters , either using one observation at a time , or using mini-batches of data . As we will show in the next sections , this online learning capacity is best suited for RL problems where we experience episodes sequentially and where we need to define a tradeoff between exploration and exploitation , as a function of our knowledge of the expected value associated with being in a state and taking an action . 2.2 EXPECTED VALUE AND BELLMAN ’ S EQUATION . We define r ( s , a , s′ ) as the reward for being in a state s ∈ RS , taking an action a ∈ A = { a1 , a2 , · · · aA } , and ending in a state s′ ∈ RS . For simplicity , we use the short-form notation for the reward r ( s , a , s′ ) ≡ r ( s ) in order to define the value as the infinite sum of discounted rewards v ( s ) = ∞∑ k=0 γkr ( st+k ) . ( 8 ) As we do not know what will be the future states st+k for k > 0 , we need to consider them as random variables ( St+k ) , so that the value V ( st ) becomes a random variable as well , V ( st ) = r ( st ) + ∞∑ k=1 γkr ( St+k ) . ( 9 ) Rational decisions regarding which action to take among the set A is based the maximization of the expected value as defined by the action-value function q ( st , at ) = µV ≡ E [ V ( st , at , π ) ] = r ( st ) + E [ ∞∑ k=1 γkr ( St+k ) ] , ( 10 ) where it is assumed that at each time t , the agent takes the action defined in the policy π . In the case of episode-based learning where the agent interacts with the environment , we assume we know the tuple of states st and st+1 , so that we can redefine the value as V ( st , at ) = r ( st ) + γ ( r ( st+1 ) + ∞∑ k=1 γkr ( St+1+k ) ) = r ( st ) + γV ( st+1 , at+1 ) . ( 11 ) Assuming that the value V ∼ N ( v ; µV , σ2V ) in Equations 9 and 11 is described by Gaussian random variables , we can reparameterize these equations as the sum of the expected value q ( s , a ) and a zero-mean Gaussian random variable E ∼ N ( ; 0 , 1 ) , so that V ( s , a ) = q ( s , a ) + σV E , ( 12 ) where the variance σ2V and E are assumed here to be independent of s and a . Although in a more general framework this assumption could be relaxed , such an heteroscedastic variance term is outside from the scope of this paper . Using this reparameterization , we can write Equation 11 as the discounted difference between the expected values of two subsequent states q ( st , at ) = r ( st ) + γq ( st+1 , at+1 ) − σVtEt + γσVt+1Et+1 = r ( st ) + γq ( st+1 , at+1 ) + σV E . ( 13 ) Note that in Equation 13 , σVt and γσVt+1 can be combined in a single standard deviation parameters σV with the assumption that Ei ⊥ Ej , ∀i 6= j . In the case where at a time t , we want to update the Q-values encoded in the neural net only after observing n-step returns ( Mnih et al. , 2016 ) , we can reformulate the observation equation so that q ( st , at ) = n−t−1∑ i=0 γir ( st+i ) + γ n−tq ( sn , an ) + σV Et , ∀t = { 1 , 2 , · · · , n− 1 } . ( 14 ) Note that in the application of Equation 14 , we employ the simplifying assumption that Et ⊥ Et+i , ∀i 6= 0 , as Equation 13 already makes simplifying assumptions for the independence of σ2V and E . Note that in a more general framework , this assumption could be relaxed . An example of n-step returns is presented in the the algorithm displayed in §1 from Appendix A . The following subsections will present , for the case of categorical actions , how to model the deterministic action-value function q ( s , a ) using a neural network .

Summary: This paper proposes applying TAGI (Goulet et al, 2019), a framework for approximate Bayesian inference in deep neural networks, to deep q-learning. Using assumptions of Gaussianity on various distributions, TAGI infers both the neural network parameters and unit activations, without needing gradient descent. Because all distributions are assumed to be Gaussian, the means and variances can be updated analytically with linear computational complexity. TAGI is essentially used to infer the correct Q-value function, and the approach is applied to lunar lander, cartpole, and several Atari games.