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: The paper describes how to adapt the Q-learning algorithm so that it is compatible with Bayesian deep learning, as opposed to the standard (semi-)gradient-based deep learning implementation. It explains how to propagate uncertainties about network parameters through the neural network computations and ultimately to the q-values, which enables action selection using Thompson sampling. Finally, it presents empirical results on CartPole, Lunar Lander, and five Atari games.