Bayesian Bellman Operators

1 Introduction . A Bayesian approach to reinforcement learning ( RL ) characterises uncertainty in the Markov decision process ( MDP ) via a posterior [ 35 , 78 ] . A great advantage of Bayesian RL is that it offers a natural and elegant solution to the exploration/exploitation problem , allowing the agent to explore to reduce uncertainty in the MDP , but only to the extent that exploratory actions lead to greater expected return ; unlike in heuristic strategies such as `` -greedy and Boltzmann sampling , the agent does not waste samples trying actions that it has already established are suboptimal , leading to greater sampling efficiency . Elementary decision theory shows that the only admissible decision rules are Bayesian [ 22 ] because a non-Bayesian decision can always be improved upon by a Bayesian agent [ 24 ] . In addition , pre-existing domain knowledge can be formally incorporated by specifying priors . In model-free Bayesian RL , a posterior is inferred over the Q-function by treating samples from the MDP as stationary labels for Bayesian regression . A major theoretical issue with existing model-free Bayesian RL approaches is their reliance on bootstrapping using a Q-function approximator , as samples from the exact Q-function are impractical to obtain . This introduces error as the samples are no long estimates of a Q-function and their dependence on the approximation is not accounted for . It is unclear what posterior , if any , these methods are inferring and how it relates to the RL problem . In this paper , we introduce Bayesian Bellman Operators ( BBO ) , a novel model-free Bayesian RL framework that addresses this issue and facilitates a theoretical exposition of the relationship between model-free Bayesian and frequentist RL approaches . Using our framework , we demonstrate that , by bootstrapping , model-free Bayesian RL infers a posterior over Bellman operators . For our main contribution , we prove that the BBO posterior concentrates on the true Bellman operator ( or the closest representation in our function space of Bellman operators ) . Hence a Bayesian method using the BBO posterior is consistent with the equivalent frequentist solution in the true MDP . We derive convergent gradient-based approaches for Bayesian policy evaluation and uncertainty estimation . Remarkably , our consistency and convergence results still hold when approximate inference is used . ⇤Correspondence to matthew.fellows @ cs.ox.ac.uk 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our framework is general and can recover empirically successful algorithms such as BootDQNprior+ [ 57 ] . We demonstrate that BootDQNprior+ ’ s lagged target parameters , which are essential to its performance , arise from applying approximate inference to the BBO posterior . Lagged target parameters can not be explained by existing model-free Bayesian RL theory . Using BBO , we extend BootDQNprior+ to continuous domains by developing an equivalent Bayesian actor-critic algorithm . Our algorithm can learn optimal policies in domains where state-of-the-art actor-critic algorithms like soft actor-critic [ 39 ] fail catastrophically due to their inability to properly explore . 2 Bayesian Reinforcement Learning . To aid the reader with notation , we provide a mathematical glossary in Appendix A . 2.1 Preliminaries . Formally , an RL problem is modelled as a Markov decision process ( MDP ) defined by the tuple hS , A , r , P , P0 , i [ 72 , 60 ] , where S is the set of states and A the set of available actions . At time t , an agent in state st 2 S chooses an action at 2 A according to the policy at ⇠ ⇡ ( ·|st ) . The agent transitions to a new state according to the state transition distribution st+1 ⇠ P ( ·|st , at ) which induces a scalar reward rt : = r ( st+1 , at , st ) 2 R with sups0 , a , s|r ( s0 , a , s ) | < 1 . The initial state distribution for the agent is s0 ⇠ P0 and the state-action transition distribution is denoted as P⇡ ( s0 , a0|s , a ) , which satisfies dP⇡ ( s0 , a0|s , a ) = d⇡ ( a0|s0 ) dP ( s0|s , a ) .. As the agent interacts with the environment it gathers a trajectory : ( s0 , a0 , r0 , s1 , a1 , r1 , s2 ... ) . We seek an optimal policy ⇡⇤ 2 argmax⇡ J⇡ that maximises the total expected discounted return : J⇡ : = E⇡ [ P1 t=0 trt ] where E⇡ is the expectation over trajectories induced by ⇡ . The Q-function is the total expected reward as a function of a state-action pair : Q⇡ ( s , a ) : = E⇡ [ P1 t=0 rt|s0 = s , a0 = a ] . Any Q-function satisfies the Bellman equation B [ Q⇡ ] = Q⇡ where the Bellman operator is defined as : B [ Q⇡ ] ( s , a ) : = EP⇡ ( s0 , a0|s , a ) [ r ( s0 , a , s ) + Q⇡ ( s0 , a0 ) ] . ( 1 ) 2.2 Model-based vs Model-free Bayesian RL . Bayes-adaptive MDPs ( BAMDPs ) [ 27 ] are a framework for model-based Bayesian reinforcement learning where a posterior marginalises over the uncertainty in the unknown transition distribution and reward functions to derive a Bayesian MDP . BAMDP optimal policies are the gold standard , optimally balancing exploration with exploitation but require learning a model of the unknown transition distribution which is typically challenging due to its high-dimensionality and multi-modality [ 67 ] . Furthermore , planning in BAMDPs requires the calculation of high-dimensional integrals which render the problem intractable . Even with approximation , most existing methods are restricted to small and discrete state-action spaces [ 6 , 38 ] . One notable exception is VariBAD [ 82 ] which exploits a meta learning setting to carry out approximate Bayesian inference . Unfortunately this approximation sacrifices the BAMDP ’ s theoretical properties and there are no convergence guarantees . Existing model-free Bayesian RL approaches attempt to solve a Bayesian regression problem to infer a posterior predictive over a value function [ 78 , 35 ] . Whilst foregoing the ability to separately model reward uncertainty and transition dynamics , modelling uncertainty in a value function avoids the difficulty of estimating high dimensional conditional distributions and mimics a Bayesian regression problem , for which there are tractable approximate methods [ 44 , 10 , 47 , 61 , 33 , 51 ] . These methods assume access to a dataset of N samples : DN : = { qi } i=1 : N from a distribution over the true Qfunction at each state-action pair : qi ⇠ PQ ( ·|si , ai ) . Each sample is an estimate of a point of the true Q-function qi = Q⇡ ( si , ai ) + ⌘i corrupted by noise ⌘i . By introducing a probabilistic model of this random process , the posterior over the Q-function P ( Q⇡|s , a , DN ) can be inferred , which characterises the aleatoric uncertainty in the sample noise and epistemic uncertainty in the model . Modeling aleatoric uncertainty is the goal of distributional RL [ 11 ] . In Bayesian RL we are more concerned with epistemic uncertainty , which can be reduced by exploration [ 57 ] . 2.3 Theoretical Issues with Existing Approaches . Unfortunately for most settings it is impractical to sample directly from the true Q-function . To obtain efficient algorithms the samples qi are approximated using bootstrapping : here a parametric function approximator Q̂ ! : S ⇥A ! R parametrised by ! 2 ⌦ is learnt as an approximation of the Q-function Q̂ ! ⇡ Q⇡ and then a TD sample is used in place of qi . For example a one-step TD estimate approximates the samples as : qi ⇡ ri + Q̂ ! ( si , ai ) , introducing an error that is dependent on ! . Existing approaches do not account for this error ’ s dependency on the function approximator . Samples are no longer noisy estimates of a point Q⇡ ( si , ai ) and the resulting posterior predictive is not P ( Q⇡|s , a , DN ) as it has dependence on Q̂ ! due to the dataset . This problem is made worse when a posterior is inferred over an optimal Q-function as it is impossible to sample from the optimal policy a priori to obtain unbiased samples . This is a major theoretical issue that raises the following questions : 1 . Do model-free Bayesian RL approaches that use bootstrapping still infer a posterior ? 2 . If it exists , how does this posterior relate to solving the RL problem ? 3 . What effect does approximate inference have on the solution ? 4 . Do methods that sample from an approximate posterior converge ? Contribution : Our primary contribution is to address these questions by introducing the BBO framework . In answer to Question 1 , BBO shows that , by introducing bootstrapping , we actually infer a posterior over Bellman operators . We can use this posterior to marginalise over all Bellman operators to obtain a Bayesian Bellman operator . Our theoretical results provide answers to Questions 2-4 , proving that the Bayesian Bellman operator can parametrise a TD fixed point as the number of samples N ! 1 and is analogous to the projection operator used in convergent reinforcement learning . Our results hold even under posterior approximation . Although our contributions are primarily theoretical , many of the benefits afforded by Bayesian methods play a significant role in a wide range of real-world applications of RL where identifying decisions that are being made under high uncertainty is crucial . We discuss the impact of our work further in Appendix B . 3 Bayesian Bellman Operators . Detailed proofs and a discussion of assumptions for all theoretical results are found in Appendix C. To introduce the BBO framework we consider the Bellman equation using a function approximator : B [ Q̂ ! ] = Q̂ ! . Using Eq . ( 1 ) , we can write the Bellman operator for Q̂ ! as an expectation of the empirical Bellman function b ! : B [ Q̂ ! ] ( s , a ) = EP⇡ ( s0 , a0|a , s ) [ b ! ( s0 , a0 , s , a ) ] , b ! ( s0 , a0 , s , a ) : = r ( s0 , a , s ) + Q̂ ! ( s0 , a0 ) . ( 2 ) When evaluating the Bellman operator to solve the Bellman equation , we can evaluate the function approximator Q̂ ! ( s , a ) but we can not evaluate B [ Q̂ ! ] ( s , a ) due to the uncertainty in reward function and transition distribution . In BBO we capture this uncertainty by treating the empirical Bellman function as a transformation of variables b ! ( · , s , a ) : S ⇥A ! R for each ( s , a ) . The transformed variable B : R ! R has a conditional distribution PB ( b|s , a , ! ) which is the pushforward of P⇡ ( s0 , a0|s , a ) under the transformation b ! ( · , s , a ) . For any PB-integrable function f : R ! R , the pushforward satisfies : EPB ( b|s , a , ! ) [ f ( b ) ] = EP⇡ ( s0 , a0|s , a ) [ f b ! ( s0 , a0 , s , a ) ] . ( 3 ) As the pushforward PB ( b|s , a , ! ) is a distribution over empirical Bellman functions , each sample b ⇠ PB ( ·|s , a , ! ) is a noisy sample of the Bellman operator at a point : bi = B [ Q̂ ! ] ( si , ai ) + ⌘i . To prove this , observe that taking expectations of b recovers B [ Q̂ ! ] ( s , a ) : EPB ( b|s , a , ! ) [ b ] =| { z } Eq . ( 3 ) EP⇡ ( s0 , a0|s , a ) [ b ! ( s0 , a0 , s , a ) ] =| { z } Eq . ( 2 ) B [ Q̂ ! ] ( s , a ) . si ai bi N Figure 1 : Graphical Model for BBO . As the agent interacts with the environment , it obtains samples from the transition distribution s0i ⇠ P ( ·|si , ai ) and policy a0i ⇠ ⇡ ( ·|s0i ) . From Eq . ( 3 ) a sample from the distribution bi ⇠ PB ( ·|si , ai , ! ) is obtained from these state-action pairs by applying the empirical Bellman function bi = ri+ Q̂ ! ( s0i , a0i ) . As we discussed in Section 2.3 , existing model-free Bayesian RL approaches incorrectly treat each bi as a sample from a distribution over the value function PQ ( Q⇡|s , a ) . BBO corrects this by modelling the true conditional distribution : PB ( b|s , a , ! ) that generates the data . The graphical model for BBO is shown in Fig . 1 . To model PB ( b|s , a , ! ) we assume a parametric conditional distribution : P ( b|s , a , ) with model parameters 2 and a conditional mean : EP ( b|s , a , ) [ b ] = B̂ ( s , a ) . It is also possible to specify a nonparametric model : P ( b|s , a ) . The conditional mean of the distribution B̂ defines a function space of approximators that represents a space of Bellman operators , each indexed by 2 . The choice of P ( b|s , a , ) should therefore ensure that the space of approximate Bellman operators characterised by B̂ is expressive enough to sufficiently represent the true Bellman operator . As we are not concerned with modelling the transition distribution in our model-free paradigm , we assume states are sampled either from an ergodic Markov chain , or i.i.d . from a buffer . Off-policy samples can be corrected using importance sampling . Assumption 1 ( State Generating Distribution ) . Each state si is drawn either i ) i.i.d . from a distribution ⇢ ( s ) with support over S or ii ) from an ergodic Markov chain with stationary distribution ⇢ ( s ) defined over a -algebra that is countably generated from S. We represent our preexisting beliefs in the true Bellman operator by specifying a prior P ( ) with a density p ( ) which assigns mass over parameterisations of function approximators 2 in accordance with how well we believe they represent B [ Q̂ ! ] . Given the prior and a dataset DN ! : = { bi , si , ai } i=1 : N of samples from the true distribution PB , we infer the posterior P ( |DN ! ) using Bayes ’ rule which has the density ( see Appendix D.1 for a derivation using both state generating distributions of Assumption 1 ) : p ( |DN ! ) = QN i=1 p ( bi|si , ai , ) p ( ) R QN i=1 p ( bi|si , ai , ) dP ( ) . ( 4 ) To be able to make predictions , we infer the posterior predictive : p ( b|DN ! , s , a ) : =R p ( b|s , a , ) dP ( |D N ! ) . Unlike existing approaches , our posterior density is a function of ! , which correctly accounts for the dependence on Q̂ ! in our data and the generating distribution PB ( b|s , a , ! ) . We highlight that it is possible to define a likelihood and prior that are functions of ! to encode any prior knowledge of how the underlying Bellman operator varies with ! , however this is not strictly necessary as the posterior automatically accounts for this dependence due to its conditioning on DN ! . As we anticipate that most applications of BBO will seek to learn an optimal policy , and hence and optimal Q-function , a prior that incorporates any knowledge available about the optimal Bellman operator will speed learning and give the agent an advantage . As our data depends on Q̂ ! , we must introduce a method of learning the correct Q-function approximator . As every Bellman operator characterises an MDP , the posterior predictive mean represents a Bayesian estimate of the true MDP by using the posterior to marginalise over all Bellman operators that our model can represent according to our uncertainty in their value : B ? ! , N ( s , a ) : =EP ( b|DN ! , s , a ) [ b ] = EP ( |DN ! ) h B̂ ( s , a ) i . ( 5 ) For this reason , we refer to the predictive mean B ? ! , N as the Bayesian Bellman operator and our Qfunction approximator should satisfy a Bellman equation using B ? ! , N . Our objective is therefore to find ! ? such that Q̂ ! ? = B ? ! ? , N . A simple approach to learn ! ? is to minimise the mean squared Bayesian Bellman error ( MSBBE ) between the posterior predictive and function approximator : MSBBEN ( ! ) : = Q̂ ! B ? ! , N 2 ⇢ , ⇡ ( 6 ) Here the distribution on the ` 2-norm is ⇢ ( s ) ⇡ ( a|s ) where recall ⇢ ( s ) is defined in Assumption 1 . Although the MSBBE has a similar form to a mean squared Bellman error with a Bayesian Bellman operator in place of the Bellman operator , our theoretical results in Section 3.1 show its frequentist interpretation is closer to the mean squared projected Bellman operator used by convergent TD algorithms [ 70 ] . We derive the MSBBE gradient in Appendix D.3 : r ! MSBBEN ( ! ) = E⇢ , ⇡ h⇣ Q̂ ! EP ( |DN ! ) h B̂ i⌘⇣ r ! Q̂ ! EP ( |DN ! ) h B̂ r ! log p ( |DN ! ) i⌘i . ( 7 ) If we can sample from the posterior then unbiased estimates of r ! MSBBEN ( ! ) can be obtained , hence minimising the MSBBE via a stochastic gradient descent algorithm is convergent if the standard Robbins-Munro conditions are satisfied [ 62 ] . When existing approaches are used , the posterior has no dependence on ! and the gradientr ! log p ( |DN ! ) is not accounted for , leading to gradient terms being dropped in the update . Stochastic gradient descent using these updates does not optimise any objective and so may not converge to any solution . The focus of our analysis in Section 4.1 is to extend convergent gradient methods for minimising the MSSBE to approximate inference techniques in situations where sampling from the posterior becomes intractable . Minimising the MSBBE also avoids the double sampling problem encountered in frequentist RL where to minimise the mean squared Bellman error , two independent samples from P ( s0|s , a ) are required to obtain unbiased gradient estimates [ 7 ] . In BBO , this issue is avoided by drawing two independent approximate Bellman operators B 1 and B 2 from the posterior 1 , 2 ⇠ P ( ·|DN ! ) instead .

Summary: The paper proposes a novel Bayesian treatment of the Bellman operator. The proposed formalism helps overcome significant limitations of previous approaches, where one typically infers a posterior over state-action values or transition functions. Instead, the authors focus on inferring posteriors over the Bellman operators (b = r_s’as + γ Q(s’, a’)). This formulation is attractive as it is free of strong assumptions and is malleable to theoretical analysis. The authors provide concrete (approximate and exact) instantiations of RL methods relying on the Bayesian Bellman operator (BBO), and demonstrate their empirical efficacy.