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 novel advantage estimate for reinforcement learning based on estimating the extent to which past actions impact the current state. More precisely, the authors train a classifier to predict the action taken k-steps ago from the state at time t, the state at t-k and the time gap k. The idea is that when it is not possible to accurately predict the action, the action choice had no impact on the current state, and thus should not be assigned credit for the current reward, they refer to this as the "independence property" between current action and future states. Based on this idea, the authors introduce a "dependency factor", using the ratio P(s_{t+k},a_{t+k}|s_t,a_t)/P(s_{t+k},a_{t+k}|s_t). They later show that this can be reworked using Bayes theorem into a ratio of the form P(a_t|s_t,s_{t+k})/\pi(a_t|s_t) which is more convenient to estimate. The authors show mathematically that, when the dependency factor is computed with the true probabilities and use to weight each reward in a trajectory, the result is an unbiased advantage estimator. Importantly the expectation, in this case, is taken over trajectories sampled according to the policy pi conditioned only on S_t. This is distinct from the Monte-Carlo estimator which is based only on samples in which A_t, the action whose advantage is being estimated, was selected.