Independence-aware Advantage Estimation

1 INTRODUCTION . Policy gradient method ( Sutton et al. , 2000 ) and its variants have demonstrated their success in solving a variety of sequential decision making tasks , such as games ( Mnih et al. , 2016 ) and continuous control ( Lillicrap et al. , 2015 ) . The large variance associated with vanilla policy gradient estimator has prompted a series of previous works to use advantage function estimation , due to its varianceminimized form ( Bhatnagar et al. , 2008 ) , to get a stable policy gradient estimation ( Mnih et al. , 2016 ; Schulman et al. , 2015a ; b ; 2017 ) . For a policy π and a state-action pair ( s , a ) , all these works estimate the advantage function Aπ ( s , a ) by subtracting an estimate of the value function V π ( s ) from the estimate of Q-value Qπ ( s , a ) . The estimation of Qπ ( s , a ) or V π ( s ) typically involves a discounted sum of future rewards , which still suffers from the high variance especially when facing the long time horizon . Meanwhile , in many real-world reinforcement learning applications , we observe that not all future rewards have dependency with the current action . For example , consider a simple multi-round game where at the end of each round of this game , the agent will be assigned a reward , representing whether it wins this round . An episode of the whole game consists of multiple independent rounds . In this example , an action in the current round will not affect the rewards in future rounds , and not all rewards received in future states do contribute to the advantage function of the current action . However , most of existing RL methods ( Sutton et al. , 2000 ; Mnih et al. , 2013 ; Schulman et al. , 2015b ) sum all future rewards to evaluate each action without considering their dependency . By identifying the independence between current action and future states in the environment , we are able to take advantage of such independence to reduce the variance of advantage estimation . In this paper , we propose Independence-aware Advantage Estimation ( IAE ) , an algorithm that can identify and utilize the independence property between current action and future states . We first introduce a novel advantage estimator that can utilize the independence property by importance sampling . The estimator formalizes a dependency factor Cπ , representing the contribution level of each future reward to advantage function estimation . For those states with no dependency on the current action , there will be a close-to-zero dependency factor Cπ , and the importance sampling estimator can reduce the variance of advantage estimation by ignoring the rewards on these states . For those states with a large dependency factor , the importance sampling estimator will potentially increase variance . In order to take advantage of variance reduction caused by small Cπ while removing the risk of increased variance by large Cπ , we further combine existing Monte-Carlo estimator with the proposed estimator by decomposing the reward into two estimators and learning the optimal decomposition by minimizing the corresponding estimation variance . Ideally , when facing states with zero dependency on the current action , our model can learn to distribute all the reward into the importance sampling estimator , where the reward can be ignored ; when those states yield extremely large Cπ , our model can learn to distribute part of rewards into the Monte-Carlo estimator to reduce the potential high variance caused by importance sampling . Details of our method are described in Section 3 , 4 and 5 . Empirically , we show that our estimated advantage function is closer to ground-truth advantage function Aπ than existing advantage estimation methods such as Monte-Carlo and Generalized Advantage Estimation ( Schulman et al. , 2015b ) . We also test IAE advantage estimation in policy optimization algorithms , showing that our method outperforms other advantage estimation methods in sample efficiency . Results of our experiments are reported in Section 8 . Our contributions can be summarized as follows : • As far as we know , we are the first to explore and utilize the independence property between current action and future states in environments to improve advantage estimation . The independence property can help us ignore the unnecessary high variance parts in MonteCarlo estimator which do not contribute to advantage function . • We propose a practical advantage estimation method to identify and utilize the indepen- dence property in environments , which achieves better performance than other advantage estimation methods in both tabular settings and function approximation settings . 2 BACKGROUND . 2.1 NOTATIONS & PROBLEM SETTINGS . We consider a finite-horizon Markov Decision Process defined by ( S , A , P , R , ρ0 , γ , T ) , where S is the set of states , A is the finite set of actions , P : S ×A×S → R denotes the transition probability , R : S × A → R denotes the reward function , ρ0 : S → R denotes the distribution of initial state S0 , γ ∈ ( 0 , 1 ] is the discount factor , T is the total time steps . We denote St , At , Rt as the random variable of state , action , reward at time t , and τt : = ( St , At , Rt , St+1 , ... , ST , AT , RT ) as trajectory starting from time t. We denote π : S × A → R as a stochastic policy , and use the notation of Qπ ( st , at ) , V π ( st ) , Aπ ( st , at ) as state-action value function , state value function and advantage function respectively . In the following discussions , we will recognize ( st , at ) as a constant state-action pair whose advantage function needed to be estimated . 2.2 ADVANTAGE FUNCTION ESTIMATORS . Monte-Carlo estimator ÂMCt of advantage function A π ( st , at ) is formalized below : ÂMCt : = −Vθ ( st ) + T−t∑ k=0 γkRt+k , where τt ∼ Pπ ( τt|st , at ) . Here Vθ ( st ) denotes the function approximator of value function V π ( st ) . We use τt ∼ Pπ ( τt|st , at ) to denote that trajectory τt is generated by policy π from st , at . Some previous work focuses on reducing the variance of ÂMCt at the cost of introducing bias ( Schulman et al. , 2015b ) , by using the n-step TD estimator and GAE estimator of advantage function Aπ ( st , at ) : Â TD ( n ) t : = −Vθ ( st ) + n−1∑ k=0 γkRt+k + γ nVθ ( St+n ) , if n < T − t ÂMCt , if n ≥ T − t ÂGAEt : = ( 1− λ ) ∞∑ n=0 λnÂ TD ( n+1 ) t , where τt ∼ Pπ ( τt|st , at ) . 3 UTILIZING INDEPENDENCE PROPERTY IN ADVANTAGE ESTIMATION . In many cases , we can utilize the independence between current action and future states to avoid unnecessary parts of variance in the Monte-Carlo estimators . Consider the example where we have a current state st whose advantage functions with respect to all actions are needed to be estimated . For a set of st+k which can be reached from st , we have independence property such that the probability Pπ ( st+k|st , at ) is a constant with respect to at . Although the Monte-Carlo return estimator from st+k may have large variance , it is clear that the return after reaching st+k gives no contribution to Aπ ( st , at ) in this case . In this section , we propose a new advantage estimator based on importance sampling , which removes the variance in Monte-Carlo return estimator after st+k by utilizing independence property , exactly as we described above . In later discussions , we will name the proposed estimator as importance sampling advantage estimator . By importance sampling approach , we present our way to derive Aπ ( st , at ) into a form which utilizes independence property : Aπ ( st , at ) = Eτt∼Pπ ( τt|st , at ) [ T−t∑ k=0 γkRt+k ] − Eτt∼Pπ ( τt|st ) [ T−t∑ k=0 γkRt+k ] = Eτt∼Pπ ( τt|st ) [ T−t∑ k=0 γkRt+k ( Pπ ( St+k , At+k|st , at ) Pπ ( St+k , At+k|st ) − 1 ) ] . ( 1 ) To briefly summarize our derivation , we perform importance sampling in every future time t + k , estimating the discounted reward γkRt+k in distribution Pπ ( St+k , At+k|st , at ) by sampling on distribution Pπ ( St+k , At+k|st ) . It is worth noting that when k ≥ 1 , At+k is independently sampled by St+k , and we are able to omit At+k in the probability ratio . For the simplicity of discussion , we will use the following definition : Cπk ( st , at , st+k , at+k ) : = Pπ ( st+k , at+k|st , at ) Pπ ( st+k , at+k|st ) − 1 . ( 2 ) where we call Cπk ( st , at , st+k , at+k ) the dependency factor , since the value captures how taking a specific action at changes the probability of reaching a future state-action pair ( st+k , at+k ) . It is clear from equation 1 and equation 2 that future state pair st+k that has dependency factor Cπk ( st , at , st+k , at+k ) close to 0 has small contribution to A π ( st , at ) , which further demonstrates that the rewards from independent future states do not contribute to advantage estimation , even if Monte-Carlo return signal has high variance . In practice , we face the challenge to estimate the dependency factor Cπ by data samples . We propose a novel modeling method and a temporal difference training strategy to solve this problem , which is detailed in Section 5 . 4 OPTIMAL COMBINATION WITH MONTE-CARLO ESTIMATOR . The advantage estimation method proposed in section 3 nicely deal with those future rewards which are independent on current action , since we have dependency factor close to zero and further ignore those rewards . However , the importance sampling advantage estimator may badly deal with those rewards with large dependency factor , which can increase the variance in estimation . To illustrate , consider we have st whose advantage functions with respect to all actions are needed to be estimated . The Monte-Carlo return starting from st following π is close to a constant q , while there is a large gap between Pπ ( St+k|st , at ) and Pπ ( St+k|st ) . This dependent case can cause high variance in importance sampling advantage estimator , even when Monte-Carlo estimation has low variance . To deal with the potential high variance problem , we seek to find the optimal combination between the proposed importance sampling estimator and Monte-Carlo estimator . There have been some previous works ( Grathwohl et al. , 2017 ) focusing on combining two estimators by optimizing a control variate , producing an estimator with less variance . Inspired by that , we decompose the reward into two estimators with a reward decomposition model , and learn the reward decomposition model by minimizing estimation variance . The following theorem demonstrates our derivation to combine two estimators : Theorem 1 . Suppose rt+k ∼ R̂ ( rt+k|st , at , τt+k ) , where R̂ is any probability distribution . Then Aπ ( st , at ) = Eτt∼Pπ ( τt|st , at ) [ T−t∑ k=0 γk ( Rt+k − rt+k ) ] − Eτt∼Pπ ( τt|st ) [ T−t∑ k=0 γk ( Rt+k − rt+k ) ] + Eτt∼Pπ ( τt|st ) [ T−t∑ k=0 γkrt+kC π k ( st , at , St+k , At+k ) ] ( 3 ) The proof of theorem 1 is provided in Appendix A.1 . The sum of all the terms containing rt+k in equation 3 constructs a control variate with zero expectations , while the expectation of rest of terms containing Rt+k is the value of advantage function Aπ ( st , at ) . From equation 3 , we find that rt+k determines the way in which rewards are divided into two estimators . If rt+k is close to 0 , rewards are divided into the Monte-Carlo estimator ; if rt+k is close to Rt+k , rewards are divided into the importance sampling advantage estimator . We parameterize rt+k as rt+k , ψ , seeking to optimize ψ to gain an advantage estimator with minimized variance . In practice , we represent rt+k , ψ by a neural network rψ ( st , at , St+k , At+k , Rt+k , k ) , whose architecture is shown by Figure 6 in the appendix . For simplicity , we will use Ĵψ ( τt ) and Îψ ( τt , at ) to denote two random variables inside of expectation in equation 3 , which is written by : Ĵψ ( τt ) : = T−t∑ k=0 γk ( Rt+k − rt+k ) ; Îψ ( τt , at ) : = T−t∑ k=0 γkrt+kC π k ( st , at , St+k , At+k ) ( 4 ) We use Varψ ( st , at ) to denote the variance of advantage estimator derived from equation 3 . Though we are not able to write the closed-form of variance in advantage estimation since we don ’ t know how three sampling processes in equation 3 are correlated , we are able to derive the variance upper bound of possible estimators by equation 5 , whose proof is shown in Appendix A.2 . Varψ ( st , at ) ≤ 3 [ Varτt∼Pπ ( τt|st , at ) [ Ĵψ ( τt ) ] +Varτt∼Pπ ( τt|st ) [ Ĵψ ( τt ) ] +Varτt∼Pπ ( τt|st ) [ Îψ ( τt , at ) ] ] ( 5 ) Based on equation 5 , we further derive the upper bound of variance which is friendly to optimize : Eat∼π ( at|st ) Varψ ( st , at ) ≤ 6Varτt∼Pπ ( τt|st ) [ Ĵψ ( τt ) ] + 3Eat∼π ( at|st ) Varτt∼Pπ ( τt|st ) [ Îψ ( τt , at ) ] ( 6 ) We seek to optimize ψ to minimize the variance upper bound shown in equation 6 . To achieve this , we firstly use two value function approximators Vw1 ( st ) , Iw2 ( st , at ) to approximate the expectation of two estimators respectively : Vw1 ( st ) ≈ Eτt∼Pπ ( τt|st ) [ Ĵψ ( τt ) ] ; Iw2 ( st , at ) ≈ Eτt∼Pπ ( τt|st ) [ Îψ ( τt , at ) ] ( 7 ) To approximate the expectation , we use SGD to minimize the mean squared error . The parameter update process can be finally expressed as follows , where αw represents the step-size : w′1 = w1 + αw ( Ĵψ ( τt ) − Vw1 ( st ) ) ∇w1Vw1 ( st ) , τt ∼ Pπ ( τt|st ) ( 8 ) w′2 = w2 + αw ( Îψ ( τt , at ) − Iw2 ( st , at ) ) ∇w2Iw2 ( st , at ) , τt ∼ Pπ ( τt|st ) ( 9 ) If we assume that Vw1 ( st ) and Iw2 ( st , at ) accurately represent the expectation , then we get the gradient to optimizeψ as follows , in order to minimize the upper bound of Eat∼π ( at|st ) Varψ ( st , at ) : ∇ψ ( 6Varτt∼Pπ ( τt|st ) [ Ĵψ ( τt ) ] + 3Eat∼π ( at|st ) Varτt∼Pπ ( τt|st ) [ Îψ ( τt , at ) ] ) = 3Eτt∼Pπ ( τt|st ) [ 2∇ψ ( Ĵψ ( τt ) − Vw1 ( st ) ) 2 + ∑ at π ( at|st ) ∇ψ ( Îψ ( τt , at ) − Iw2 ( st , at ) ) 2 ] ( 10 ) We can use SGD to optimize ψ expressed by equation 11 , where αψ represents the step-size : ψ′ = ψ + αψ [ 2 ( Ĵψ ( τt ) − Vw1 ( st ) ) ∇ψĴψ ( τt ) + ∑ at π ( at|st ) ( Îψ ( τt , at ) − Iw2 ( st , at ) ) ∇ψ Îψ ( τt , at ) ] , τt ∼ Pπ ( τt|st ) ( 11 ) Here we will illustrate why performing gradient descent on ψ leads to an advantage estimation with lower variance . For a single rt+k , ψ , the gradient component in the first term of equation 11 pushes the rt+k , ψ to change towards the direction to the mean total return ; meanwhile , the second term of equation 11 counteracts the gradient in the first term , preventing rt+k , ψ from constructing a constant return by the restriction in variance of importance sampling estimator∑T−t k=0 γ krt+k , ψC π k ( st , at , St+k ) . When we have the independence property in environments ( i.e . the value of Cπk is close to zero ) , the counteraction effect in the second term will disappear . With the gradient in the first term , rt+k , ψ will be rapidly optimized towards the mean total return , making the variance of advantage estimator to dramatically decrease along the training process . By replacing the expectation terms in equation 3 by function approximators , the form of independence-aware advantage estimator is given by : ÂIAEt : = Ĵψ ( τt ) − Vw1 ( st ) + Iw2 ( st , at ) , where τt ∼ Pπ ( τt|st , at ) . ( 12 )

Summary: This paper proposes a new advantage estimator in reinforcement learning based on importance sampling. This form allows for a significantly lower-variance estimator for situations where the current action "stops mattering" to the future state. A control variate, as in Grathwohl et al., is used to combine the importance sampling estimator with the "standard" estimator in a way that is always unbiased and attempts to minimize the overall variance.