Year: 2019

Abstract:
Most of existing advantage function estimation methods in reinforcement learning suffer from the problem of high variance, which scales unfavorably with the time horizon. To address this challenge, we propose to identify the independence property between current action and future states in environments, which can be further leveraged to effectively reduce the variance of the advantage estimation. In particular, the recognized independence property can be naturally utilized to construct a novel importance sampling advantage estimator with close-to-zero variance even when the Monte-Carlo return signal yields a large variance. To further remove the risk of the high variance introduced by the new estimator, we combine it with existing Monte-Carlo estimator via a reward decomposition model learned by minimizing the estimation variance. Experiments demonstrate that our method achieves higher sample efficiency compared with existing advantage estimation methods in complex environments.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND

heading: 2.1 NOTATIONS & PROBLEM SETTINGS
text: 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.

heading: 2.2 ADVANTAGE FUNCTION ESTIMATORS
text: 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).

heading: 3 UTILIZING INDEPENDENCE PROPERTY IN ADVANTAGE ESTIMATION
text: 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.

heading: 4 OPTIMAL COMBINATION WITH MONTE-CARLO ESTIMATOR
text: 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)

heading: 5 DEPENDENCY FACTOR ESTIMATION
text: The final challenge in our method is to estimate the dependency factor Cπ , which is crucial to make the advantage estimator low-biased. In this section, we will introduce our modeling and training method, which is able to give accurate dependency factor estimation in experiments.
It is hard to estimate the transition probability in equation 2 because of the high dimensionality of state space. Here we derive the ratio between two transition probability into a form which can be represented by an action classifier by equation 13, whose proof is shown in Appendix A.3.
Pπ(st+k, at+k|st, at) Pπ(st+k, at+k|st) = Pπ(at|st, st+k, at+k) π(at|st) = Pπ(at|st, st+k) π(at|st) , if k ≥ 1. (13)
In practice, we use an action classification modelPφ(at|st, st+k, k) to representPπ(at|st, st+k), using Cφ(st, at, st+k) :=
Pφ(at|st,st+k,k) π(at|st) to give an approximation of C π . We call Pφ(at|st, st+k, k) dependency model in later discussions.
Similar to the derivation in previous work (Liu et al., 2018), we show that the dependency model can be learned in an temporal difference manner:
Pπ(at|st, st+k2) = Est+k1∼Pπ(st+k1 |st+k2 ,st) [P π(at|st, st+k1)] , if k2 > k1 ≥ 1. (14)
Detailed proof of equation 14 is provided in Appendix A.4. Although the case is different from the common cases in temporal difference learning, it is clear that the samples of distribution Pπ(st+k1 |st+k2 , st) can be collected by directly rolling out policy π, then we are able to train model Pφ by minimizing temporal difference error. We use a mixture of temporal difference target and Monte-Carlo target to train model Cφ, which is detailed in Appendix B. Empirically, we demonstrate the effectiveness of our approach to accurately estimate dependency factor in section 8.2.

heading: 6 ALGORITHM
text: In this section, we present crucial algorithms details of the method presented in Section 3, 4 and 5.

heading: 6.1 MODEL DEPENDENCY AND PSEUDO-CODE
text: We illustrate the general framework of our algorithm shown by Figure 1. The model of dependency factor Cπ , lying in the basic part of our algorithm, can be directly trained by samples of current policy π. Two value functions Iw1 and Vw2 are trained by SGD on MSE loss, shown in equation 8 and equation 9, given the the reward decomposition rψ as input. The reward decomposition model rψ is also trained by SGD shown in equation 11, given the output of Iw1 and Vw2 . The pseudo-code of our algorithm is provided by Algorithm 1.
Algorithm 1 Policy Optimization Algorithm by Independence-aware Advantage Estimation 1: Initialize policy parameters θ, dependency model parameters φ, decomposition model parame-
ters ψ, and value function parameters w1, w2. 2: for each iteration do 3: Sample trajectory τ with policy πθ 4: Update φ by temporal difference learning 5: Update w1, w2 to minimize mean squared error by equation 8 and equation 9 6: Update ψ to minimize variance upper bound by equation 11 7: Compute advantage estimation ÂIAEt at each time-step t by equation 12 8: Update θ by Proximal Policy Optimization 9: end for

heading: 6.2 COMPUTATIONAL COMPLEXITY
text: There might be concerns about computational complexity for training dependency model and reward decomposition model, since the training requires O(T 2) number of (st, st+k) pairs as training samples. We apply two techniques, trajectory truncation and block-wise training, enabling our algorithm to have a comparable computational complexity with PPO algorithm.
Trajectory Truncation. We truncate the trajectory into slices of 128 timesteps similar to the approach in PPO algorithm. We apply our method into the truncated trajectories by considering a reward Qθ(st, at) at the last step. The reward decomposition model also considers the decomposition over the final reward Qθ(st, at).
Block-wise Training. When training dependency model and reward decomposition model, we provide a total of 1282 numbers of (st, st+k) pairs, and accumulate gradient for every valid (st, st+k) pair. Only 128 × 2 times of forward and backward passes in CNN feature extractor are computed though we deal with a squared number of training data, making our method computationally efficient. Our method requires less than 3 times training time compared to PPO algorithm.

heading: 7 RELATED WORK
text: Policy gradient (Sutton et al., 2000) provides the basic form to optimize a parameterized policy in expected returns. Generalized Advantage Estimation (Kimura et al., 2000; Schulman et al., 2015b) replaces Monte-Carlo estimator by the mixture of N-step temporal difference estimator, reducing the variance of policy gradient estimator while introducing bias. Another series of previous works focus on how to optimize parameterized policy without focusing on the selection of estimator. Among these works, TRPO (Schulman et al., 2015a) and PPO (Schulman et al., 2017) are part of the recent works that reaches state-of-the-art performance on a variety of tasks.
Some of previous works have shown that RL algorithm tends to be unstable when planning horizon is long (Jiang et al., 2015). Another series of works (François-Lavet et al., 2015; Xu et al., 2018) focus on how to select discount factor γ to improve the performance of RL algorithm. Our method suggests another solution to the problem in RL, since IAE leads to substantial variance reduction in return signal even if the planning horizon is long, which improves the performance while keeping the long planning horizon.
Our method performs gradient descent on estimation variance to improve the estimator as training proceeds. This method has been used in recent works on various applications. One previous work (Hanna et al., 2017) focuses on optimizing a behaviour policy to minimize the variance of off-policy value estimation; another previous work (Grathwohl et al., 2017) focuses on getting the optimal variance balance between REINFORCE estimator and reparameterized gradient estimator by minimizing estimation variance.
In the problem of density ratio estimation, our method is similar to one previous work (Liu et al., 2018) which transforms the density ratio estimation problem into an temporal difference learning problem. Our method is different that we focus on estimating the dependency between current actions and future states in a fixed policy, while this previous work focuses on estimating the ratio of the future state reaching probability between two different policies.

heading: 8 EXPERIMENTS
text: In our experiments, we provide empirical results to answer the following questions:
• Can our dependency model training method in section 5 precisely estimates the dependency factor Cπ , and captures the independence property in environments?
• Can our method utilize the independence property to reduce the variance in advantage estimation, and further give more accurate advantage estimation than other advantage estimation methods?
• Can IAE improve the overall performance of policy optimization algorithm, for instance, PPO algorithm?
To answer the first question, we train the dependency model by method in section 5 and compare the prediction with ground-truth value, proving the capacity of our training method to model dependency factor Cπ . This part of results are detailed in section 8.2.
For the second question, we show that IAE method gives advantage estimation with less variance in tabular settings, and reduces value function training error in function approximation settings. In the Pixel Grid World environment, we further show that our method gives advantage estimation closer to ground-truth advantage function than MC and GAE method under cosine similarity metric. This part of results are detailed in section 8.3. To demonstrate how IAE method utilize independence property, we also provide case studies in Appendix C.3 to visualize the effect of variance reduction by low dependency factor in Pixel Grid World environment.
For the last question, we provide training curves in section 8.4 in Pixel Grid Worlds. Compared with PPO algorithm with Monte-Carlo advantage estimation and generalized advantage estimation, IAE method makes policy optimization process more sample-efficient.

heading: 8.1 ENVIRONMENT SETTINGS
text: We perform experiment on two types of environments: finite-state MDPs and Pixel Grid World.
Finite-state MDP settings. To evaluate the quality of advantage estimation of our method in tabular cases, we construct different 3-state MDPs with different transition probability and reward functions. We categorize state transition probability settings into connected settings and isolated settings, and categorize reward settings into high-variance settings and low-variance settings. Detailed descriptions of transition probability and reward settings are provided in Appendix C.1.
Pixel Grid World Environment. To evaluate our method in function approximation settings, we build Pixel Grid World environment where observations are provided by high-dimensional images.
As illustrated in Figure 2a, the blue square represents the position of the agent and the yellow square represents the position of current goal. The agent gets positive reward for reaching the goal. To make the problem harder, the environment will do periodical reset multiple times in an episode, by which the agent and the goal are randomly repositioned. We use two different reward settings: per-step punishment setting and no punishment setting, which is detailed in Appendix C.1.
The hyper-parameters and network architectures we use in our experiments are presented in Appendix C.2.

heading: 8.2 DEPENDENCY FACTOR MODELING
text: In this section, we investigate our estimation of dependency factor Cπ , and show the general similarity between our estimation and ground-truth Cπ . We train our model Cφ by data generated by a fixed uniform random policy π. Figure 2a and 2b show the case where the dependency is precisely captured: given the future state st+k shown in Figure 2a, the model Pφ(at|st, st+k, k) correctly predicts down and right action that more likely lead current state st to future state st+k. We also build a dataset consisting of 300 random (st, st+k) pairs, where k is uniformly sampled from 1 to 30. We evaluate the mean KL divergence between ground-truth Pπ(at|st, st+k) and prediction Pφ(at|st, st+k, k) averaged in 10 runs, as shown in Figure 2c. The mean KL divergence decrease to a relatively small value during training, showing that the dependency model Pφ leads to a generally precise estimation of dependency factor Cπ .

heading: 8.3 THE VARIANCE AND ACCURACY OF INDEPENDENCE-AWARE ADVANTAGE ESTIMATION
text: We evaluate the variance of IAE estimator on a variety of finite-state MDP settings. We train r(st, st+k) for 10000 episodes and then test the advantage estimator by performing advantage estimation multiple times to get the estimation variance. We compare the variance of IAE estimator with Monte-Carlo advantage estimator (MC) and importance sampling advantage estimator (IS). In this experiment, we use the precise value of dependency factor for IS and IAE estimator. Table 1
demonstrates the standard derivation of advantage estimation. In both environments suitable for MC estimation and ones suitable for IS estimation, our method gives estimation with less variance than both MC and IS method. In some cases, IAE estimation dramatically reduces the variance of both MC and IS estimation.
On function approximation settings, we show that our method dramatically reduces the mean squared error in training value function approximators, as shown in Figure 3a. We initialize reward decomposition model rψ to be zero for all inputs, which constructs a precise Monte-Carlo advantage estimator initially, and compare the value function training error with or without training rψ . When the reward decomposition model rψ is fixed, the loss of value function training keeps to be high because of the stochasticity of Monte-Carlo return signal, while in our method, the rψ makes reward to be distributed into importance sampling advantage estimator, making the mean squared error of value function to reduce. In Figure 3b, we show that IAE estimation has much higher cosine similarity to ground-truth advantage function, compared with Monte-Carlo and GAE estimation.
Besides, we investigate how independence property reduces the variance of estimation by visualizing the reward decomposition in Pixel Grid World Environment, which is detailed in Appendix C.3.

heading: 8.4 PERFORMANCE OF POLICY OPTIMIZATION
text: We run Proximal Policy Optimization algorithm with IAE advantage estimation method, and compare the result to PPO algorithm with Monte-Carlo and GAE advantage estimation. Figure 4 shows
the result on two different reward settings. Compared with Monte-Carlo and GAE estimation, IAE estimation makes the policy improvement process more sample-efficient.

heading: 9 CONCLUSIONS
text: In this work, we addressed the large variance problem in advantage estimation for policy gradient methods. We proposed a novel advantage estimation method by importance sampling, which identifies and utilizes the independence property, reducing the variance by ignoring those independent rewards. We further combined the proposed estimator with Monte-Carlo estimator in the optimal way, making the final IAE estimator to have low variance in general cases. The effectiveness of our method can be verified on Pixel-input environments compared with previous advantage estimation methods such as Monte-Carlo and Generalized Advantage Estimation.
There are a few directions to explore in the future. First, we only considered problems with discrete actions in this work, and we will extend our method to problems with continuous actions. Second, we will test our algorithm on more reinforcement learning environments, such as Atari games and continuous control tasks.

heading: A PROOFS
text: A.1 PROOF OF UNBIASEDNESS
In this section, we show the proof of theorem 1, which demonstrates that any choice of rt+k will not bias the combined estimator as long as rt+k is a function of st, at and τt+k.
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) ] , (1)
where Cπk (st, at, st+k, at+k) := Pπ(st+k, at+k|st, at) Pπ(st+k, at+k|st) − 1.
Proof. Denote r(st, at, τt+k) = E [rt+k|st, at, τt+k]. Since Aπ(st, at) can be written by
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 ] ,
We only need to prove that
Eτt∼Pπ(τt|st) [ T−t∑ k=0 γkrt+kC π k (st, at, St+k, At+k) ]
= Eτt∼Pπ(τt|st,at) [ T−t∑ k=0 γkrt+k ] − Eτt∼Pπ(τt|st) [ T−t∑ k=0 γkrt+k ] . (2)
We have derivation that
Eτt∼Pπ(τt|st) [ T−t∑ k=0 γkrt+kC π k (st, at, St+k, At+k) ] + Eτt∼Pπ(τt|st) [ T−t∑ k=0 γkrt+k ]
= Eτt∼Pπ(τt|st) [ T−t∑ k=0 γkrt+k(1 + C π k (st, at, St+k, At+k)) ]
= T−t∑ k=0 ∫ S ∑ at+k γkEτt+k∼Pπ(τt+k|st+k,at+k) [r(st, at, τt+k)] (1 + C π k (st, at, st+k, at+k))P π(st+k, at+k|st)dst+k
= T−t∑ k=0 ∫ S ∑ at+k γkEτt+k∼Pπ(τt+k|st+k,at+k) [r(st, at, τt+k)] Pπ(st+k, at+k|st, at) Pπ(st+k, at+k|st) Pπ(st+k, at+k|st)dst+k
= T−t∑ k=0 ∫ S ∑ at+k γkEτt+k∼Pπ(τt+k|st+k,at+k) [r(st, at, τt+k)]P π(st+k, at+k|st, at)dst+k
= Eτt∼Pπ(τt|st,at) [ T−t∑ k=0 γkrt+k ] ,
which proves equation 2 and further proves theorem 1.
A.2 UPPER BOUND OF VARIANCE
In this section, we show that the if we use the sum of three random variable in equation 3 as the advantage estimator, the variance of advantage estimation is bounded by three times the sum of the variance of estimation individual random variable inside expectation in equation 3. We begin with the following lemma:
Lemma 1. Suppose X1, X2, · · · , Xn are random variables. Then
Var [ n∑ i=1 Xi ] ≤ n n∑ i=1 Var[Xi].
Proof.
Var [ n∑ i=1 Xi ] = n∑ i=1 Var[Xi] + 2 ∑ 1≤i<j≤n Cov[Xi, Xj ]
≤ n∑ i=1 Var[Xi] + 2 ∑ 1≤i<j≤n √ Var[Xi]Var[Xj ]
Then we have
n n∑ i=1 Var[Xi]−Var [ n∑ i=1 Xi ]
≥ (n− 1) n∑ i=1 Var[Xi]− 2 ∑ 1≤i<j≤n √ Var[Xi]Var[Xj ]
= ∑
1≤i<j≤n
(√ Var[Xi]− √ Var[Xj ] )2 ≥ 0,
which proves lemma 1.
From lemma 1, it is clear that the variance of advantage estimation is bounded by three times the sum of three individual variances shown by equation 5.
A.3 ACTION PROBABILITY FORM FOR DEPENDENCY FACTOR
In this section, we give detailed proof of equation 13 in the main body, which is written by:
Pπ(st+k, at+k|st, at) Pπ(st+k, at+k|st) = Pπ(at|st, st+k, at+k) π(at|st) = Pπ(at|st, st+k) π(at|st) , if k ≥ 1. (3)
Proof. We separately prove the first and the second equation. The first equation is proved by:
Pπ(st+k, at+k|st, at) Pπ(st+k, at+k|st) = Pπ(st+k, at+k, st, at) Pπ(st, at) · P π(st) Pπ(st+k, at+k, st)
= Pπ(st+k, at+k, st, at) Pπ(st+k, at+k, st) · P π(st) Pπ(st, at) = Pπ(at|st, st+k, at+k)
π(at|st) .
The second equation is proved by:
Pπ(at|st, st+k, at+k) = Pπ(st, at, st+k, at+k)
Pπ(st, st+k, at+k)
= Pπ(st, at, st+k)P π(at+k|st, at, st+k) Pπ(st, st+k)Pπ(at+k|st, st+k) = Pπ(st, at, st+k)π(at+k|st+k) Pπ(st, st+k)π(at+k|st+k) = Pπ(st, at, st+k)
Pπ(st, st+k) = Pπ(at|st, st+k).
Thus, it is proved that at+k can be abandoned from the probability condition when k ≥ 1. When k = 0, the conditional probability can be directly given by a closed form solution:
Pπ(At = a|st, st+k, at+k) = {
1, a = at+k 0, a 6= at+k
A.4 DEPENDENCY FACTOR ESTIMATION
We prove that the dependency factor can be learned in an inverse temporal difference manner in this section, which is formulated by equation 14 in the main body:
Pπ(at|st, st+k2) = Est+k1∼Pπ(st+k1 |st+k2 ,st)P π(at|st, st+k1), if k2 > k1 ≥ 1.
Proof.
Est+k1∼Pπ(st+k1 |st+k2 ,st)P π(at|st, st+k1)
= ∫ S Pπ(st+k1 |st+k2 , st)Pπ(at|st, st+k1)dst+k1
= ∫ S Pπ(st+k1 , st+k2 , st) Pπ(st+k2 , st) Pπ(at, st, st+k1) Pπ(st, st+k1) dst+k1
= ∫ S Pπ(at, st, st+k1) Pπ(st+k2 , st) Pπ(st+k2 |st+k1 , st)dst+k1
= ∫ S Pπ(at, st, st+k1) Pπ(st+k2 , st) Pπ(st+k2 |st+k1)dst+k1
= ∫ S Pπ(at, st, st+k1 , st+k2) Pπ(st+k2 , st) dst+k1 = Pπ(at, st, st+k2)
Pπ(st+k2 , st) = Pπ(at|st, st+k2)

heading: B TEMPORAL DIFFERENCE LEARNING OF DEPENDENCY MODEL
text: We use a mixture of temporal difference and Monte-Carlo training target to train dependency model. For a trajectory sample τ = {si, ai, ri}Ti=1, we provide training target for Pφ(at|st, st+k, k) for every valid (t, k) pair, and use cross-entropy loss to update φ. For the Monte-carlo training target, we directly use one-hot probability distribution of action at as the ground-truth; for the temporal difference training target, we use distribution of Pφ(at|st, st+k−1, k − 1) as the ground-truth. We mix the two training target with 0.85k weight on Monte-Carlo target and 1 − 0.85k weight on temporal difference target to achieve the optimal bias-variance tradeoff.

heading: C EXPERIMENT
text: C.1 ENVIRONMENT SETTINGS
Finite-state MDP settings. We construct different 3-state MDPs with different transition probability and reward functions. We categorize transition probability settings into two types:
• Connected settings, where there is high mutual reaching probability between every state pair, causing the n-step transition probability to converge rapidly into stationary distribution. We use the following transition matrix for connected settings:
P = [ 0.2 0.5 0.3 0.5 0.3 0.2 0.2 0.2 0.6 ] • Isolated settings, where there is low mutual reaching probability between some state pairs,
causing the n-step transition probability to converge slowly into stationary distribution. We use the following transition matrix for isolated settings:
P = [ 0.9 0.05 0.05 0.45 0.1 0.45 0.05 0.05 0.9 ]
In connected settings, the n-step transition probability is very close into stationary distribution, resulting in a small dependency factor in future rewards. Thus, importance sampling estimator will have lower variance in connected settings than in isolated settings.
We also categorize reward settings into two types:
• High-variance settings, where the variance of Monte-Carlo return signal is large compared to mean total return. We use the following reward function for high variance settings:
R = [3 2 1]
• Low-variance settings, where the variance of Monte-Carlo return signal is small compared to mean total return. We use the following reward function for low variance settings:
R = [2.1 2 1.9]
It is clear that Monte-Carlo advantage estimator will have larger variance in high-variance settings than in low-variance settings.
Pixel Grid World environment. We build Pixel Grid World environment where observations are provided by high-dimensional images. As shown in Figure 5, the blue square represents the position of the agent and the yellow square represents the position of current goal. The size of grid world is 8 × 8, and observations are provided by 128 × 128 × 3 RGB pixels. There are 4 actions in Pixel Grid World, making the agent to move toward 4 different directions. We use two reward settings in experiments:
• Per-step punishment setting, where agent gets r = −0.03 reward in every step before reaching goal, and r = 0 reward in every step after reaching goal. Agent gets r = 1 reward when reaching goal for the first time. • No punishment setting, where agent gets r = 1 reward when reaching goal for the first
time, and gets r = 0 reward otherwise.
To make the problem harder, the environment will do periodical reset multiple times in an episode, by which the agent and the goal are randomly repositioned. The reset period is set to be 30 steps, and an episode contains 5 rounds.
C.2 HYPER-PARAMETER SETTINGS AND NETWORK ARCHITECTURE
We use γ = 0.999 for all experiments in policy optimization. For GAE estimator, we use the default settings λ = 0.95 in all experiments. We use 8 parallel environments to gather data, use the trajectory truncation length of 128, and train each truncated trajectory with 5 epochs. We use a global Adam optimizer for training both dependency model and reward decomposition model, and set learning rate to be 2.5 × 10−4. For initialization, we use orthogonal initialization with √ 2 scale for all layers, except that we initialize the last weight matrix of D to be zeros. For training Pφ(at|st, st+k, k), we weight the MC loss by 0.85k and TD loss by 1− 0.85k, and use the mixture loss to train dependency model. For PPO algorithm, we use the same hyper-parameter settings with OpenAI Baseline.
We use the network architecture shown in Figure 6 for dependency model and reward decomposition model. We use the same CNN feature extractor settings as the settings in the OpenAI Baseline. For time embedding, we store embedding vector of 128 dimensions for each k ∈ {1, 2, · · · , 128} separately.
C.3 CASE STUDY ON REWARD DECOMPOSITION MODEL
In Pixel Grid World environment, we fix a stochastic policy and train model rψ to minimize estimation variance. When the model rψ is near convergence, we visualize reward decomposition for every future reward after a state-action pair (st, at). These results are shown in Figure 7. It is clear that all the independent rewards after reset are assigned to importance sampling estimator, which is multiplied by a small dependency factor and barely contribute to advantage estimation. By contrast, rewards before reset are assigned to the Monte-Carlo estimator, which contribute to advantage estimation.

Review summaries:
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.

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.

This paper tries to control the variance of advantage function by utilizing the independence property between current action and future states. The practical approach they are using is to learn a dependency model of reward function as a control variate to lower the variance. Using the control variate technique they derive a (maybe complicated) algorithm to update the policy by PPO. Empirical results seems competitive.