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; Fujimoto et al., 2018]. The large variance associated with vanilla policy gradient estimator has prompted a series of previous works to use advantage function estimation, due to its variance-minimized form [Bhatnagar et al., 2008], to get a stable policy gradient estimation [Mnih et al., 2016; Schulman et al., 2015a; Schulman et al., 2015b; Schulman et al., 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.

∗Corresponding author.

Meanwhile, in many real-world reinforcement learning applications, we observe that not all future rewards have a 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, where each round lasts constant timesteps. 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 the 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 MonteCarlo 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 settings on environments with high-dimensional observations, showing that our method outperforms other advantage estimation methods in sample efficiency. Results of our experiments are reported in Section 7.

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 Monte-Carlo estimator which do not contribute to advantage function. Moreover, we propose a practical advantage estimation method to identify and utilize the independence property in environments, which achieves better performance than other advantage estimation methods.

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 needs to be estimated.

2.2 Advantage Function Estimators

Monte-Carlo estimator ˆ AMC t of advantage function Aπ (st, at) is formalized below:

Here Vθ(st) denotes the function approximator of the 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 works focus on reducing the variance of ˆ AMC t 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):

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 estimator. 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 constant with respect to different choices of 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, inspired by the cases we described above. In later discussions, we will name the proposed estimator as importance sampling advantage estimator.

By importance sampling, we present our way to derive Aπ (st, at) into a form which utilizes independence property:

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)1 .

For the simplicity of discussion, we will use the following definition:

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) that future states st+k that has nearly zero dependency factor 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.

Finally, we define the form of the importance sampling advantage estimator as follows:

where τt follows distribution P π (τt|st). By previous analysis, we have E[ ˆ AIS t ] = Aπ (st, at). 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.

1 It is worth noting that when k ≥ 1, we are able to omit At+k in the importance ratio in equation (1), since At+k is independently sampled by St+k.

4 Optimal Combination withMonte-Carlo Estimator

The advantage estimation method proposed in section 3 nicely deals with those future rewards which are independent with current action, since the dependency factors are close to zero and those rewards are ignored by importance sampling. However, the importance sampling advantage estimator may badly deal with those rewards with large dependency factors, which can increase the variance in estimation. To illustrate, consider the case where Monte-Carlo return starting from st following π is always 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 the Monte-Carlo estimator. There have been some previous works [Grathwohl et al., 2017; Liu 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 the two estimators:

Theorem 1. Suppose R′t+k ∼ ˆR(R′t+k |st, at, τt+k), where ˆR is any probability distribution. Then

The proof of theorem 1 is not hard after realizing the sum of three terms including R′t+k is zero, which can be explained by the correctness of importance sampling. In equation (4), the first and second expectation can be estimated by Monte-Carlo estimator, and the last expectation can be estimated by the importance sampling estimator. As a bridge to connect these two parts, R′t+k aims to determine the way in which rewards are divided into two estimators. If R′t+k is close to 0, rewards are divided into the Monte-Carlo estimator; if R′t+k is close to Rt+k, rewards are divided into the importance sampling advantage estimator. Further, we parameterize R′t+k as R′t+k,ψ with a deep neural network R′ψ (st, at, St+k, At+k, Rt+k, k), which is composed of feature extractors of st and St+k, a concatenation layer that merges features of current state and future state, and a multi-layer perceptron which produces the outputs.

In order to learn such R′t+k,ψ , we first define the advantage estimator derived from equation (4) and derive the form of variance (section 4.1). We then demonstrate our method to optimize ψ by minimizing the variance of this estimator (section 4.2). Finally, we give the practical form of independence-aware advantage estimator with function approximations (section 4.3).

4.1 Definition of Advantage Estimator

In this section, we precisely define the combined estimator and derive its form of variance. For simplicity, we will use Jψ(τt) and Iψ(τt, at) to denote two terms inside expectation in equation (4), which is written by:

Derived from equation (4), we define the form of the combined estimator as follows:

By equation (4), we have E[ ˆ ACombined t ] = Aπ (st, at). The variance of ˆ ACombined t can be directly derived as follows:

4.2 Optimize ψ for Variance Minimization

Based on equation (7), we further derive the upper bound of variance which is friendly to optimize:

We use L(ψ) as the objective function for optimizing ψ, since it provides the variance upper bound of the advantage estimator ˆ ACombined t . To minimize the objective function, we firstly use two value function approximators Vw1 (st), Iw2 (st, at) to approximate the expectation of two estimators respectively:

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 stepsize:

where τt follows distribution P π (τt|st).

If we assume that the two expectation terms in equation (9) can be precisely captured by Vw1 (st) and Iw2 (st, at), we can get the gradient of the L(ψ) with respect to ψ as follows 1 :

We can use SGD to optimize ψ expressed by equation (13), where αψ represents the step-size:

here τt follows distribution P π (τt|st). Note that the two gradient terms of equation (13) are optimizing the same MSE loss function as equation (10) and (11), but here the different parameter ψ is being optimized.

Here we will illustrate why performing gradient descent on ψ leads to an advantage estimation with lower variance. The gradient component in the first term of equation (13) will adapt reward decomposition to make Jψ(τt) move towards the mean value, and further reduce its variance; meanwhile, the second term of equation (13) counteracts the gradient in the first term, preventing the Jψ(τt) to become constant by the restriction in variance of importance sampling estimator Iψ(τt, at) = ∑T-t k=0 γ kR′t+k,ψ C π k (st, at, St+k, At+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 termwill disappear. With the gradient in the first term, R′t+k,ψ will be rapidly optimized to reduce the variance of Jψ(τt), making the variance of advantage estimator to dramatically decrease along the training process.

4.3 Advantage Estimator with Function Approximators

In regular reinforcement learning settings, we cannot sample multiple trajectories from the same state st. To handle this issue, we replace two terms in equation (6) by function approximators Vw1 and Iw2 defined in equation (9). This leads to the form of independence-aware advantage estimator:

where τt follows distribution P π (τt|st, at).

5 Dependency Factor Estimation

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.

1 We have the form of Var[Xψ] in the objective function L(ψ). Since we have Var[Xψ] = E[(Xψ - E[Xψ])2 ], it can be derived that ∇ψVar[Xψ] = E[∇ψ(Xψ - E[Xψ])2 ]. By replacing E[Xψ] with the approximated value Vw1 (st) and Iw2 (st, at), we can get the gradient with respect to ψ in equation (12).

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 probabilities into a form which can be represented by an action classifier by equation (15), whose proof can be directly obtained by the definition of conditional probability.

In order to approximate and learn the value of the dependency factor, we use an action classification model Pφ(at|st, st+k, k) to approximate P π (at|st, st+k). If the action classifier is learned accurately, the approximated dependency factor Cφ(st, at, st+k) := Pφ(at|st,st+k,k) π(at|st) - 1 equals the dependency factor C π (st, at, st+k). We call Pφ(at|st, st+k, k) dependency model in later discussions.

Inspired by the derivation in previous work [Liu et al., 2018], we can prove that the probability P π (at|st, st+k) has temporal difference property as follows:

when k2 > k1 ≥ 1.

Moreover, if equation (16) still holds after substituting the probability P π (at|st, st+k) by the prediction of the dependency model Pφ(at|st, st+k, k), and this dependency model has accurate prediction when k = 1, this dependency model must have accurate prediction for any k ≥ 1 on any current state st and reachable future state st+k.

With the above analysis, we are able to train the model Pφ by minimizing the difference between two sides of equation (16) after substituting P π by the prediction of Pφ. Practically, we use a mixture of temporal difference target Pφ(at|st, st+k, k) and the ground truth at as the training target for the dependency model Pφ(at|st, st+k+1, k +1) in the next step. We demonstrate the effectiveness of our approach to accurately estimate dependency factors in section 7.2.

6 RelatedWork

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.

Some of previous works [Liu et al., 2017; Wu et al., 2018; Papini et al., 2018] discuss other approaches to reduce variance in policy gradient estimation. Comparing our work to this series of work, there are significant differences in the estimator being used and the cases where variance is reduced. Our method relies on independence property to reduce variance; in contrast, previous works use Stein's identity [Liu et al., 2017] or the property that each dimension of action is individually sampled [Wu et al., 2018] to reduce variance, where fundamental difference exists.

Some previous works discuss how properties of future states can be leveraged to enable better credit assignment. HCA [Harutyunyan et al., 2019] proposes a form of advantage estimator similar to the form of equation (3) in our work, but their approach can yield an estimator with larger variance than Monte-Carlo estimator as we discuss in section 4. Beyond their work, we propose a novel reward decomposition model and a learning approach to effectively reduce the variance of the proposed estimator, and make IAE work well in environments with high-dimensional observation. Another previous work [Mesnard et al., 2020] extracts the independent features of future observation to serve as value baseline function in policy gradient, yielding another form of policy gradient estimator. Their work focuses on how to extract trajectory features and how to make them independent with actions, instead of leveraging future states with independence property to decrease the variance of advantage estimation, which forms a different direction from our work.

Our method performs gradient descent on estimation variance to improve the estimator as training proceeds. Similar approaches have 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.

7 Experiments

In our experiments, we provide empirical results to answer the following questions:

• Can our dependency model training method in section 5 precisely estimate the dependency factor C π , and capture 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 algorithms, for instance, PPO algorithm?

To answer the first question, we train the dependency model by the method in section 5 and compare the prediction with ground-truth value, proving the capability of our training method to model the dependency factor C π . This part of the results is detailed in section 7.2.

For the second question, we show that IAE 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 groundtruth advantage function than MC and GAE method under cosine similarity metric. This part of the results is detailed in section 7.3.

For the last question, we provide training curves in section 7.4 in Pixel Grid World environment. Compared with the PPO algorithm with Monte-Carlo advantage estimation and generalized advantage estimation, IAE makes the policy optimization process more sample-efficient.

7.1 Environment Settings

We perform experiments 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. In isolated transition settings, there are some state pairs with low mutual reaching probability; in connected transition settings, all state pairs have high mutual reaching probability. In high-variance reward settings, the variance of Monte-Carlo return signal is large compared to the average total return; in the low-variance reward setting, the variance of Monte-Carlo return signal is small compared to the average total return.

Pixel Grid World environment. To evaluate our method in function approximation settings, we build Pixel Grid World environment where observations are provided by 128×128× 3 RGB pixels. As illustrated in Figure 1a, the blue square represents the position of the agent and the yellow square represents the position of the current goal. The agent gets a positive reward for reaching the goal. To make the problem harder, the environment will do periodic resets 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. In the per-step punishment setting, the agent gets r = -0.03 reward in every step before reaching its goal, r = 1 reward when reaching its goal for the first time, and r = 0 reward for every step after reaching its goal. In no punishment setting, the agent gets r = 1 reward when reaching its goal for the first time, and gets r = 0 reward otherwise.

7.2 Dependency Factor Modeling

In this section, we investigate our estimation of the dependency factor C π , and show the general similarity between our estimation and the ground-truth C π . We train our model Cφ with data generated by a fixed uniform random policy π. Figure 1a and 1b show the case where the dependency is pre-

(a)

cisely captured: given the future state st+k shown in Figure 1a, the model Pφ(at|st, st+k, k) correctly predicts Right and Up actions that more likely lead the current state st to the 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 the true value of P π (at|st, st+k) and the prediction from Pφ(at|st, st+k, k) averaged in 10 runs, as shown in Figure 1c. The mean KL divergence decreases to a relatively small value during training, showing that the dependency model Pφ gives a generally precise estimation of the dependency factor.

7.3 Variance and Accuracy of Independence-aware Advantage Astimation

We evaluate the variance of IAE estimator on a variety of finite-state MDP settings. We train tabular reward decomposition 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) on the same state-action pair. For IAE, we individually sample three trajectories for each estimation of advantage function, use these three trajectories as samples of the three random variables in equation (4). In this experiment, we use the precise value of dependency factors for IS and IAE estimators. 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

(a)

less variance than both MC and IS methods. 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 2a. We initialize the 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 the reward decomposition. When the reward decomposition model is fixed, the loss of value function training keeps being high because of the high variance of Monte-Carlo return signal, while in our method, the reward decomposition helps reward to be distributed into the importance sampling advantage estimator, reducing the mean squared error of value function. In Figure 2b, we show that IAE estimation has much higher cosine similarity to groundtruth advantage function, compared with Monte-Carlo (MC) and GAE estimation. For GAE, we use λ = 0.95.

7.4 Performance of Policy Optimization

We run Proximal Policy Optimization algorithm [Schulman et al., 2017] with IAE advantage estimation method, and compare the result to PPO algorithm with Monte-Carlo (MC) and GAE advantage estimation. For GAE, we use λ = 0.95.

Figure 3 shows the result on two different reward settings. Compared with two existing advantage estimation methods, Monte-Carlo and GAE, IAE makes the policy improvement process more sample-efficient.

8 Conclusions

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 an 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 existing advantage estimation methods such as Monte-Carlo and GAE.