Making Sense of Reinforcement Learning and Probabilistic Inference

1 INTRODUCTION . Probabilistic inference is a procedure of making sense of uncertain data using Bayes ’ rule . The optimal control problem is to take actions in a known system in order to maximize the cumulative rewards through time . Probabilistic graphical models ( PGMs ) offer a coherent and flexible language to specify causal relationships , for which a rich literature of learning and inference techniques have developed ( Koller & Friedman , 2009 ) . Although control dynamics might also be encoded as a PGM , the relationship between action planning and probabilistic inference is not immediately clear . For inference , it is typically enough to specify the system and pose the question , and the objectives for learning emerge automatically . In control , the system and objectives are known , but the question of how to approach a solution may remain extremely complex ( Bertsekas , 2005 ) . Perhaps surprisingly , there is a deep sense in which inference and control can represent a dual view of the same problem . This relationship is most clearly stated in the case of linear quadratic systems , where the Ricatti equations relate the optimal control policy in terms of the system dynamics ( Welch et al. , 1995 ) . In fact , this connection extends to a wide range of systems , where control tasks can be related to a dual inference problem through rewards as exponentiated probabilities in a distinct , but coupled , PGM ( Todorov , 2007 ; 2008 ) . A great benefit of this connection is that it can allow the tools of inference to make progress in control problems , and vice-versa . In both cases the connections provide new insights , inspire new algorithms and enrich our understanding ( Toussaint & Storkey , 2006 ; Ziebart et al. , 2008 ; Kappen et al. , 2012 ) . Reinforcement learning ( RL ) is the problem of learning to control an unknown system ( Sutton & Barto , 2018 ) . Like the control setting , an RL agent should take actions to maximize its cumulative rewards through time . Like the inference problem , the agent is initially uncertain of the system dynamics , but can learn through the transitions it observes . This leads to a fundamental tradeoff : The agent may be able to improve its understanding through exploring poorly-understood states and actions , but it may be able to attain higher immediate reward through exploiting its existing knowledge ( Kearns & Singh , 2002 ) . In many ways , RL combines control and inference into a general framework for decision making under uncertainty . Although there has been ongoing research ∗These authors contributed equally to this work . †DeepMind , London , UK , { bodonoghue , iosband , cdi } @ google.com in this area for many decades , there has been a recent explosion of interest as RL techniques have made high-profile breakthroughs in grand challenges of artificial intelligence research ( Mnih et al. , 2013 ; Silver et al. , 2016 ) . A popular line of research has sought to cast ‘ RL as inference ’ , mirroring the dual relationship for control in known systems . This approach is most clearly stated in the tutorial and review of Levine ( 2018 ) , and provides a key reference for research in this field . It suggests that a generalization of the RL problem can be cast as probabilistic inference through inference over exponentiated rewards , in a continuation of previous work in optimal control ( Todorov , 2009 ) . This perspective promises several benefits : A probabilistic perspective on rewards , the ability to apply powerful inference algorithms to solve RL problems and a natural exploration strategy . In this paper we will outline an important way in which this perspective is incomplete . This shortcoming ultimately results in algorithms that can perform poorly in even very simple decision problems . Importantly , these are not simply technical issues that show up in some edge cases , but fundamental failures of this approach that arise in even the most simple decision problems . In this paper we revisit an alternative framing of ‘ RL as inference ’ . In fact , we show that the original RL problem was already an inference problem all along.1 Importantly , this inference problem includes inference over the agent ’ s future actions and observations . Of course , this perspective is not new , and has long been known as simply the Bayes-optimal solution , see , e.g. , Ghavamzadeh et al . ( 2015 ) . The problem is that , due to the exponential lookahead , this inference problem is fundamentally intractable for all but the simplest problems ( Gittins , 1979 ) . For this reason , RL research focuses on computationally efficient approaches that maintain a level of statistical efficiency ( Furmston & Barber , 2010 ; Osband et al. , 2017 ) . We provide a review of the RL problem in Section 2 , together with a simple and coherent framing of RL as probabilistic inference . In Section 3 we present three approximations to the intractable Bayes-optimal policy . We begin with the celebrated Thompson sampling algorithm , then we review the popular ‘ RL as inference ’ framing , as presented by Levine ( 2018 ) , and highlight a clear and simple shortcoming in this approach . Finally , we review K-learning ( O ’ Donoghue , 2018 ) , which we re-interpret as a modification to the RL as inference framework that provides a principled approach to the statistical inference problem , as well as a presenting a relationship with Thompson sampling . In Section 4 we present computational studies that support our claims . 2 REINFORCEMENT LEARNING . We consider the problem of an agent taking actions in an unknown environment in order to maximize cumulative rewards through time . For simplicity , this paper will model the environment as a finite horizon , discrete Markov Decision Process ( MDP ) M = ( S , A , R , P , H , ρ ) .2 Here S = { 1 , .. , S } is the state space , A = { 1 , .. , A } is the action space and each episode is of fixed length H ∈ N. Each episode ` ∈ N begins with state s0 ∼ ρ then for timesteps h = 0 , .. , H − 1 the agent selects action ah , observes transition sh+1 with probability P ( sh+1 , sh , ah ) ∈ [ 0 , 1 ] and receives reward rh+1 ∼ R ( sh , ah ) , where we denote by µ ( sh , ah ) = Erh+1 the mean reward . We define a policy π to be a mapping from S to probability distributions over A and write Π for the space of all policies . For any timestep t = ( ` , h ) , we define Ft = ( s00 , a00 , r01 , .. , s ` h−1 , a ` h−1 , r ` h ) to be the sequence of observations made before time t. An RL algorithm maps histories to policies πt = alg ( S , A , Ft ) . Our goal in the design of RL algorithms is to obtain good performance ( cumulative rewards ) for an unknown M ∈ M , whereM is some family of possible environments . Note that this is a different problem from typical ‘ optimal control ’ , that seeks to optimize performance for one particular known MDP M ; although you might still fruitfully apply an RL algorithm to solve problems of that type . For any environment M and any policy π we can define the action-value function , QM , πh ( s , a ) = Eπ , M H∑ j=h+1 rj | sh = s , ah = a . ( 1 ) 1Note that , unlike control , connecting RL with inference will not involve a separate ‘ dual ’ problem . 2This choice is for clarity ; continuous , infinite horizon , or partially-observed environments do not alter our narrative . Where the expectation in ( 1 ) is taken with respect to the action selection aj for j > h from the policy π and evolution of the fixed MDP M . We define the value function VM , πh ( s ) = Eα∼πQ M , π h ( s , α ) and writeQM , ? h ( s , a ) = maxπ∈ΠQ M , π h ( s , a ) for the optimal Q-values over policies , and the optimal value function is given by VM , ? h ( s ) = maxaQ M , ? h ( s , a ) . In order to compare algorithm performance across different environments , it is natural to normalize in terms of the regret , or shortfall in cumulative rewards relative to the optimal value , Regret ( M , alg , L ) = EM , alg [ L∑ ` =1 ( VM , ? 0 ( s ` 0 ) − H∑ h=1 r ` h ) ] . ( 2 ) This quantity depends on the unknown MDP M , which is fixed from the start and kept the same throughout , but the expectations are taken with respect to the dynamics of M and the learning algorithm alg . For any particular MDP M , the optimal regret of zero can be attained by the nonlearning algorithm algM that returns the optimal policy for M . In order to assess the quality of a reinforcement learning algorithm , which is designed to work across some family of M ∈ M , we need some method to condense performance over a set to a single number . There are two main approaches to this : BayesRegret ( φ , alg , L ) = EM∼φRegret ( M , alg , L ) , ( 3 ) WorstCaseRegret ( M , alg , L ) = max M∈M Regret ( M , alg , L ) , ( 4 ) where φ is a prior over the family M. These differing objectives are often framed as Bayesian ( average-case ) ( 3 ) and frequentist ( worst-case ) ( 4 ) RL 3 . Although these two settings are typically studied in isolation , it should be clear that they are intimately related through the choice ofM and φ . Our next section will investigate what it would mean to ‘ solve ’ the RL problem . Importantly , we show that both frequentist and Bayesian perspectives already amount to a problem in probabilistic inference , without the need for additional re-interpretation . 2.1 SOLVING THE RL PROBLEM THROUGH PROBABILISTIC INFERENCE . If you want to ‘ solve ’ the RL problem , then formally the objective is clear : find the RL algorithm that minimizes your chosen objective , ( 3 ) or ( 4 ) . To anchor our discussion , we introduce a simple decision problem designed to highlight some key aspects of reinforcement learning . We will revisit this problem setting as we discuss approximations to the optimal policy . Problem 1 ( One unknown action ) . FixN ∈ N ≥ 3 , > 0 and defineMN , = { M+N , , M − N , } . Both M+ and M− share S = { 1 } , H = 1 and A = { 1 , .. , N } ; they only differ through their rewards : R+ ( 1 ) = 1 , R+ ( 2 ) = +2 , R+ ( a ) = 1− for a = 3 , .. , N , R− ( 1 ) = 1 , R− ( 2 ) = −2 , R− ( a ) = 1− for a = 3 , .. , N. WhereR ( a ) = x ∈ R is a shorthand for deterministic reward of x when choosing action a . Problem 1 is extremely simple , it involves no generalization and no long-term consequences : It is an independent bandit problem with only one unknown action . For known M+ , M− the optimal policy is trivial : Choose at = 2 in M+ and at = 1 in M− for all t. An RL agent faced with unknown M ∈ M should attempt to optimize the RL objectives ( 3 ) or ( 4 ) . Unusually , and only because Problem 1 is so simple , we can actually compute the optimal solutions to both in terms of L ( the total number of episodes ) and φ = ( p+ , p− ) where p+ = P ( M = M+ ) , the probability of being in M+ . For L > 3 an optimal minimax ( minimizing the worst-case regret ) RL algorithm is to first choose a0 = 2 and observe r1 . If r1 = 2 then you know you are in M+ so pick at = 2 for all t = 1 , 2 .. , for Regret ( L ) = 0 . If r1 = −2 then you know you are in M− so pick at = 1 for all t = 1 , 2 .. , for Regret ( L ) = 3 . The worst-case regret of this algorithm is 3 , which can not be bested by any algorithm . 3Some frequentist results are high-probability bounds on the worst case rather than true worst-case bounds , but this distinction is not important for our purposes Actually , the same RL algorithm is also Bayes-optimal for any φ = ( p+ , p− ) provided p+L > 3 . This relationship is not a coincidence . All admissible solutions to the worst-case problem ( 4 ) are given by solutions to the average-case ( 3 ) for some ‘ worst-case ’ prior φ̃ ( Wald , 1950 ) . As such , for ease of exposition , our discussion will focus on the Bayesian ( or average-case ) setting . However , readers should understand that the same arguments apply to the worst-case objective . In Problem 1 , the key probabilistic inference the agent must consider is the effects of it own actions upon the future rewards , i.e. , whether it has chosen action 2 . Slightly more generally , where actions are independent and episode length H = 1 , the optimal RL algorithm can be computed via Gittins indices , but these problems are very much the exception ( Gittins , 1979 ) . In problems with generalization or long-term consequences , computing the Bayes-optimal solution is computationally intractable . One example of an algorithm that converges to Bayes-optimal solution in the limit of infinite computation is given by Bayes-adaptive Monte-Carlo Planning ( Guez et al. , 2012 ) . The problem is that , even for very simple problems , the lookahead tree of interactions between actions , observations and algorithmic updates grows exponentially in the search depth ( Strehl et al. , 2006 ) . Worse still , direct computational approximations to the Bayes-optimal solution can fail exponentially badly should they fall short of the required computation ( Munos , 2014 ) . As a result , research in reinforcement learning amounts to trying to find computationally tractable approximations to the Bayes-optimal policy that maintain some degree of statistical efficiency .

Summary: This paper criticizes the ‘RL as inference’ paradigm by highlighting its limitations and shows that a variant of this framework - the K-learning algorithm (O'Donoghue et. al., 2018) does not have these limitations. The paper first clarifies some points of confusion regarding RL as inference, namely the fact that RL was originally an inference problem all along. A simple example is used to demonstrate that the RL as inference framework (Levine, 2018) fails to choose the optimal actions that resolve epistemic uncertainty, whereas the K-learning algorithm does select the optimal action. Further, a connection is made which reveals that K-learning is an approximate version of Thompson sampling - the strategy of using as single posterior sample of parameters given data for greedy actions which originated in bandit settings. Some empirical results are provided highlighting the cases where Soft Q-learning (Levine, 2018) fails but Thompson sampling and K-learning do not.