Maximum Reward Formulation In Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) algorithms typically try to maximize the cumulative finite horizon undiscounted return , R ( τ ) = ∑T t=0 rt , or the infinite horizon discounted return , R ( τ ) = ∑∞ t=0 γ trt . rt is the reward obtained at time step t , γ is the discount factor in the range [ 0 , 1 ) , and τ is the agent ’ s trajectory . τ consists of actions ( a ) sampled from the policy ( π ( · | s ) ) and states ( s′ ) sampled from the probability transition function P ( s′|s , a ) of the underlying Markov Decision Process ( MDP ) . The action-value function Qπ ( s , a ) for a policy π is given by Qπ ( s , a ) = E τ∼π [ R ( τ ) | ( s0 , a0 ) = ( s , a ) ] The corresponding Bellman equation for Qπ ( s , a ) with the expected return defined as R ( τ ) =∑∞ t=0 γ trt is Qπ ( st , at ) = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st+1 ) [ r ( st , at ) + γQ π ( st+1 , at+1 ) ] This Bellman equation has formed the foundation of RL . However , we argue that optimizing for only the maximum reward achieved in an episode is also an important goal . Reformulating the RL problem to achieve the largest reward in an episode is the focus of this paper , along with empirical demonstrations in one toy and one real-world domain . In the de novo drug design pipeline , molecule generation tries to maximize a given reward function . Existing methods either optimize for the expected cumulative return , or for the reward at the end of the episode , and thus fail to optimize for the very high reward molecules that may be encountered in the middle of an episode . This limits the potential of several of these reinforcement learning based drug design algorithms . We overcome this limitation by proposing a novel functional formulation of the Bellman equation : Qπmax ( st , at ) = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st+1 ) [ max ( r ( st , at ) , γQ π max ( st+1 , at+1 ) ) ] ( 1 ) Other use cases of this formulation ( i.e. , situations where the single best reward found , rather than the total rewards , are important ) are - symbolic regression ( Petersen ( 2020 ) , Udrescu & Tegmark ( 2020 ) ) which is interested in finding the single best model , active localization ( Chaplot et al . ( 2018 ) ) must find the robot ’ s one most likely pose , green chemistry ( Koch et al . ( 2019 ) ) wants to identify the one best product formulation , and other domains that use RL for generative purposes . This paper ’ s contributions are to : • Derive a novel functional form of the Bellman equation , called max-Bellman , to optimize for the maximum reward in an episode . • Introduce the corresponding evaluation and optimality operators , and prove the convergence of Q-learning with the max-Bellman formulation . • Test on a toy environment and draw further insights with a comparison between Q-learning and Q-learning with our max-Bellman formulation . • Use this max-Bellman formulation to generate synthesizable molecules in an environment that mimics the real drug discovery pipeline , and demonstrate significant improvements over the existing state-of-the-art methods . 2 RELATED WORK . This section briefly introduces fundamental RL concepts and the paper ’ s main application domain . 2.1 REINFORCEMENT LEARNING . Bellman ’ s dynamic programming paper ( Bellman , 1954 ) introduced the notions of optimality and convergence of functional equations . This has been applied in many domains , from control theory to economics . The concept of an MDP was proposed in the book Dynamic Programming and Markov Processes ( Howard , 1960 ) ( although some variants of this formulation already existed in the 1950s ) . These two concepts of Bellman equation and MDP are the foundations of modern RL . Q-learning was formally introduced in ( Watkins & Dayan , 1992 ) and different convergence guarantees were further developed in ( Jaakkola et al. , 1993 ) and ( Szepesvári , 1997 ) . Q-learning convergence to the optimal Q-value ( Q ? ) has been proved under several important assumptions . One fundamental assumption is that the environment has finite ( and discrete ) state and action spaces and each of the states and actions can be visited infinitely often . The learning rate assumption is the second important assumption , where the sum of learning rates over infinite episodes is assumed to go to infinity in the limit , whereas the sum of squares of the learning rates are assumed to be a finite value ( Tsitsiklis , 1994 ; Kamihigashi & Le Van , 2015 ) . Under similar sets of assumptions , the on-policy version of Q-learning , known as Sarsa , has also been proven to converge to the optimal Q-value in the limit ( Singh et al. , 2000 ) . Recently , RL algorithms have seen large empirical successes as neural networks started being used as function approximators ( Mnih et al. , 2016 ) . Tabular methods can not be applied to large state and action spaces as these methods are linear in the state space and polynomial in the action spaces in both time and memory . Deep reinforcement learning ( DRL ) methods on the other hand , can approximate the Q-function or the policy using neural networks , parameterized by the weights of the corresponding neural networks . In this case , RL algorithms easily generalize across states , which improves the learning speed ( time complexity ) and sample efficiency of the algorithm . Some popular Deep RL algorithms include DQN ( Mnih et al. , 2015 ) , PPO ( Schulman et al. , 2017 ) , A2C ( Mnih et al. , 2016 ) , SAC ( Haarnoja et al. , 2018 ) , TD3 ( Fujimoto et al. , 2018b ) , etc . 2.2 DE NOVO DRUG DESIGN . De novo drug design is a well-studied problem and has been tackled by several methods , including evolutionary algorithms ( Brown et al . ( 2004 ) ; Jensen ( 2019 ) ; Ahn et al . ( 2020 ) ) , generative models ( Simonovsky & Komodakis ( 2018 ) ; Gómez-Bombarelli et al . ( 2018 ) ; Winter et al . ( 2019 ) ; Jin et al . ( 2018 ) ; Popova et al . ( 2018 ) ; Griffiths & Hernández-Lobato ( 2020 ) ; Olivecrona et al . ( 2017 ) ) , and reinforcement learning based approaches ( You et al . ( 2018a ) ; Zhou et al . ( 2018 ) ) . While the effectiveness of the generated molecules using these approaches has been demonstrated on standard benchmarks such as Guacamol ( Brown et al . ( 2019 ) ) , the issue of synthesizability remains a problem . While all the above approaches generate molecules that optimize a given reward function , they do not account for whether the molecules can actually be effectively synthesized , an important practical consideration . Gao & Coley ( 2020 ) further highlighted this issue of synthesizability by using a synthesis planning program to quantify how often the molecules generated using these existing approaches can be readily synthesized . To attempt to solve this issue , Bradshaw et al . ( 2019 ) used a variational auto-encoders based approach to optimize the reward function with single-step reactions . Korovina et al . ( 2019 ) employed a random selection of reactants and reaction conditions at every time step of a multi-step process . PGFS ( policy gradient for forward synthesis ) from Gottipati et al . ( 2020 ) generates molecules via multi-step chemical synthesis and simultaneously optimized for the given reward function . PGFS leveraged TD3 algorithm ( Fujimoto et al . ( 2018a ) ) , and like existing approaches , optimizes for the usual objective of total expected discounted return . 3 METHOD . This section briefly describes the previous attempts at optimizing for the maximum reward in an episode , defines the Q-function and new functional form of the Bellman equation , defines the corresponding max-Bellman operators , and proves its convergence properties . Quah & Quek ( 2006 ) also try to formulate a maximum reward objective function . They define the Q-function and derive the max-Bellman equation in the following way ( rephrased according to this paper ’ s notation ) : Qπmax ( st , at ) , E [ max t′≥t ( γt ′−tr ( st′ , at′ ) ) | ( st , at ) , π ] , ∀ ( st , at ) ∈ S ×A ( 2 ) Then , the corresponding functional form of Bellman equation ( i.e. , the max-Bellman formulation ) was derived as follows : Qπmax ( st , at ) = E [ max t′≥t ( γt ′−tr ( st′ , at′ ) ) | ( st , at ) , π ] = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st′ ) max ( r ( st , at ) , E [ max t′≥t+1 ( γt ′−tr ( st′ , at′ ) ) | ( st+1 , at+1 ) , π ] ) = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st+1 ) max ( r ( st , at ) , γQ π max ( st+1 , at+1 ) ) ( 3 ) However , the second equality is incorrect , which can be shown with a counter example . Consider an MDP with four states and a zero-reward absorbing state . Assume a fixed deterministic policy and γ close to 1 . Let s1 be the start state , with deterministic reward r ( s1 ) = 1 . With probability 1 , s1 goes to s2 . Let r ( s2 ) = 0 , and with equal probability 0.5 it transitions to either s3 or s4 . The immediate reward r ( s3 ) = 2 and r ( s4 ) = 0 , after that it always goes to the absorbing state with 0 rewards thereafter . Since there can be only 2 trajectories , one with a maximum reward of 2 and the other with a maximum reward of 1 , we have Qπmax ( s1 ) = 1.5 , but Q π max ( s2 ) = 1 , and max ( r ( s1 ) , Q π max ( s2 ) ) = 1 6= 1.5 . This is because the expectation and max operators are not interchangeable . While the counter example above is not applicable in the current setting of PGFS ( as it has a deterministic environment and uses a deterministic policy gradient algorithm ( TD3 ) ) , the definition of the Q-function given in equation-2 does not generalize to stochastic environments . Moreover , in order to be able to accommodate multiple possible products of a chemical reaction and to truly leverage a stochastic environment ( and stochastic policy gradient algorithms ) , the Q-function should be able to recursively optimize for the expectation of maximum of reward at the current time step and future expectations . Therefore , we define the Q-function as : Qπmax ( st , at ) = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st+1 ) [ max ( r ( st , at ) , γEst+2∼P ( ·|st+1 , at+1 ) at+2∼π ( ·|st+2 ) [ max ( r ( st+1 , at+1 ) , ... ) ] ) ] ( 4 ) Then , the corresponding functional form of Bellman equation ( i.e. , max-Bellman formulation ) can be obtained as follows : Qπmax ( st , at ) = Est+1∼P ( ·|st , at ) at+1∼π ( ·|st+1 ) max ( r ( st , at ) , γQ π max ( st+1 , at+1 ) ) ( 5 ) While the work of Quah & Quek ( 2006 ) is focused on deriving the various forms of the max-Bellman equations , the major contributions of our paper are to 1 ) propose the max-Bellman equation with a particular motivation towards drug discovery , 2 ) define the corresponding operators , 3 ) provide their convergence proofs , and 4 ) validate the performance on a toy domain and a real world domain ( de-novo drug design ) . Other related ideas and formulations are described in Appendix section-B . Based on Equation 5 , we can now define the max-Bellman evaluation operator and the max-Bellman optimality operator . For any function Q : S ×A→ R and for any state-action pair ( s , a ) , ( MπQ ) ( s , a ) , max ( r ( s , a ) , γEs′∼P ( ·|s , a ) a′∼π ( ·|s′ ) [ Q ( s′ , a′ ) ] ) ( M ? Q ) ( s , a ) , max ( r ( s , a ) , γEs′∼P ( ·|s , a ) [ max a′∈A Q ( s′ , a′ ) ] ) M ? and Mπ : ( S ×A → R ) → ( S ×A → R ) are operators that takes in a Q function and returns another modified Q function by assigning the Q value of the state s , action a , to be the maximum of the reward obtained at the same state and action ( s , a ) and the discounted future expected Q value . Proposition 1 . The operators have the following properties . • Monotonicity : let Q1 , Q2 : S ×A → R such that Q1 ≥ Q2 element-wise . Then , MπQ1 ≥MπQ2 and M ? Q1 ≥M ? Q2 • Contraction : both operators are γ-contraction in supremum norm i.e. , for any Q1 , Q2 : S ×A → R , ‖MπQ1 −MπQ2‖∞ ≤ γ‖Q1 −Q2‖∞ ( 6 ) ‖M ? Q1 −M ? Q2‖∞ ≤ γ‖Q1 −Q2‖∞ ( 7 ) Proof . We will provide a proof only for M ? . The proof of Mπ is similar and is provided in Section A of the Appendix . Monotonicity : Let Q1 and Q2 be two functions such that Q1 ( s , a ) ≥ Q2 ( s , a ) for any state-action pair ( s , a ) ∈ S ×A . We then have : max a′∈A Q1 ( s , a ′ ) ≥ Q2 ( s , a ) , ∀ ( s , a ) By the definition of max , we obtain : max a′∈A Q1 ( s , a ′ ) ≥ max a′∈A Q2 ( s , a ′ ) , ∀s and by linearity of expectation , we have : γEs′∼P ( ·|s , a ) [ max a′∈A Q1 ( s ′ , a′ ) ] ≥ γEs′∼P ( ·|s , a ) [ max a′∈A Q2 ( s ′ , a′ ) ] , ∀ ( s , a ) Since , M ? Q1 ( s , a ) ≥ γEs′∼P ( ·|s , a ) [ max a′∈A Q1 ( s ′ , a′ ) ] we get : M ? Q1 ( s , a ) ≥ γEs′∼P ( ·|s , a ) [ max a′∈A Q2 ( s ′ , a′ ) ] Moreover , because M ? Q1 ( s , a ) ≥ r ( s , a ) , we obtain : M ? Q1 ( s , a ) ≥ max ( r ( s , a ) , γEs′∼P ( ·|s , a ) [ max a′∈A Q2 ( s ′ , a′ ) ] ) = M ? Q2 ( s , a ) which is the desired result . Contraction : Denote fi ( s , a ) = γEs′∼P ( ·|s , a ) [ maxa′∈AQi ( s′ , a′ ) ] , the expected action-value function of the next state . By using the fact that max ( x , y ) = 0.5 ( x+ y + |x− y| ) , ∀ ( x , y ) ∈ R2 , we obtain : max ( r , f1 ) −max ( r , f2 ) = 0.5 ( r + f1 + |r − f1| ) − 0.5 ( r + f2 + |r − f2| ) = 0.5 ( f1 − f2 + |r − f1| − |r − f2| ) ≤ 0.5 ( f1 − f2 + |r − f1 − ( r − f2 ) | ) = 0.5 ( f1 − f2 + |f1 − f2| ) ≤ |f1 − f2| Therefore ‖M ? Q1 −M ? Q2‖∞ = ‖max ( r , f1 ) −max ( r , f2 ) ‖∞ ≤ ‖f1 − f2‖∞ ≤ γ‖max a′ Q1 ( · , a′ ) −max a′ Q2 ( · , a′ ) ‖∞ ≤ γ‖Q1 −Q2‖∞ ( max is a non-expansion ) The left hand side of this equation represents the largest difference between the two Q-functions . Recall that γ lies in the range [ 0 , 1 ) and hence all differences are guaranteed go to zero in the limit . The Banach fixed point theorem then lets us conclude that the operator M ? admits a fixed point . We denote the fixed point of M ? by Q ? max . Based on equation 5 , we see that Qπmax is the fixed point of Mπ . In the next proposition , we will prove that Q ? max corresponds to the optimal action-value function in the sense that Q ? max = maxπ Q π max . Proposition 2 ( Optimal policy ) . The deterministic policy π ? is defined as π ? ( s ) = arg maxa∈AQ ? max ( s , a ) . π ? , the greedy policy with respect to Q ? max , is the optimal policy and for any stationary policy π , Qπ ? max = Q ? max ≥ Qπmax . Proof . By the definition of greediness and the fact that Q ? is the fixed point of M ? , we have : π ? ( s ) = arg maxa∈AQ ? max ( s , a ) ⇒ Mπ ? Q ? = M ? Q ? max = Q ? max . This proves that Q ? max is the fixed point of the evaluation operator Mπ ? , which implies that Q ? max is the action-value function of π ? . For any function Q : S × A → R and any policy π , we have M ? Q ≥ MπQ . Using monotonicity , we have ( M ? ) 2Q = M ? ( M ? Q ) ≥M ? ( MπQ ) ≥Mπ ( MπQ ) = ( Mπ ) 2Q . ( 8 ) We then use induction to show that for any n ≥ 1 , ( M ? ) nQ ≥ ( Mπ ) nQ . As both operators are contractions , by the fixed-point theorem , ( M ? ) nQ and ( Mπ ) nQ converge to Q ? max and Q π max , respectively , when n goes to infinity . We conclude then that Q ? max ≥ Qπmax . Thus , we have shown that and optimality operator defined based on our novel formulation of Bellman equation converges to a fixed point ( Proposition 1 ) and that fixed point is the optimal policy ( Proposition 2 ) .

Summary: This paper proposes a modified bellman equation for reinforcement learning that optimizes the maximum expected single step reward along a trajectory, instead of the maximum cumulative reward. This formulation is applied to the generation of molecules with optimized properties of interest. A recently published molecule generation algorithm, that constructs molecules step wise via the (predicted) chemical reactions of building blocks, is modified with this new bellman formulation, and shows modest improvements in optimizing for some HIV activity targets.