Title: INDEPENDENCE-AWARE ADVANTAGE ESTIMATION

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.

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 Monte-Carlo estimator which do not contribute to advantage function. • 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 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 S 0 , γ ∈ (0, 1] is the discount factor, T is the total time steps. We denote S t , A t , R t as the random variable of state, action, reward at time t, and τ t := (S t , A t , R t , S t+1 , ..., S T , A T , R T ) as trajectory starting from time t.

We denote π : S × A → R as a stochastic policy, and use the notation of Q π (s t , a t ), V π (s t ), A π (s t , a t ) as state-action value function, state value function and advantage function respectively. In the following discussions, we will recognize (s t , a t ) as a constant state-action pair whose advantage function needed to be estimated.

2.2. ADVANTAGE FUNCTION ESTIMATORS:

Monte-Carlo estimator ÂMC t of advantage function A π (s t , a t ) is formalized below:

ÂMC t := -V θ (s t ) + T -t k=0 γ k R t+k , where τ t ∼ P π (τ t |s t , a t ).

Here V θ (s t ) denotes the function approximator of value function V π (s t ). We use τ t ∼ P π (τ t |s t , a t ) to denote that trajectory τ t is generated by policy π from s t , a t .

Some previous work focuses on reducing the variance of ÂMC 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 π (s t , a t ):

ÂTD(n) t :=        -V θ (s t ) + n-1 k=0 γ k R t+k + γ n V θ (S t+n ), if n < T -t ÂMC t , if n ≥ T -t ÂGAE t := (1 -λ) ∞ n=0 λ n ÂTD(n+1) t

, where τ t ∼ P π (τ t |s t , a t ).

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 s t whose advantage functions with respect to all actions are needed to be estimated. For a set of s t+k which can be reached from s t , we have independence property such that the probability P π (s t+k |s t , a t ) is a constant with respect to a t . Although the Monte-Carlo return estimator from s t+k may have large variance, it is clear that the return after reaching s t+k gives no contribution to A π (s t , a t ) 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 s t+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 π (s t , a t ) into a form which utilizes independence property:

EQUATION

To briefly summarize our derivation, we perform importance sampling in every future time t + k, estimating the discounted reward γ k R t+k in distribution P π (S t+k , A t+k |s t , a t ) by sampling on distribution P π (S t+k , A t+k |s t ). It is worth noting that when k ≥ 1, A t+k is independently sampled by S t+k , and we are able to omit A t+k in the probability ratio.

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

C π k (s t ,

EQUATION

where we call C π k (s t , a t , s t+k , a t+k ) the dependency factor, since the value captures how taking a specific action a t changes the probability of reaching a future state-action pair (s t+k , a t+k ). It is clear from equation 1 and equation 2 that future state pair s t+k that has dependency factor C π k (s t , a t , s t+k , a t+k ) close to 0 has small contribution to A π (s t , a t ), 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 s t whose advantage functions with respect to all actions are needed to be estimated. The Monte-Carlo return starting from s t following π is close to a constant q, while there is a large gap between P π (S t+k |s t , a t ) and P π (S t+k |s t ). 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 r t+k ∼ R(r t+k |s t , a t , τ t+k ), where R is any probability distribution. Then

EQUATION

The proof of theorem 1 is provided in Appendix A.1. The sum of all the terms containing r t+k in equation 3 constructs a control variate with zero expectations, while the expectation of rest of terms containing R t+k is the value of advantage function A π (s t , a t ). From equation 3, we find that r t+k determines 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 R t+k , rewards are divided into the importance sampling advantage estimator. We parameterize r t+k as r t+k,ψ , seeking to optimize ψ to gain an advantage estimator with minimized variance. In practice, we represent r t+k,ψ by a neural network r ψ (s t , a t , S t+k , A t+k , R t+k , k), whose architecture is shown by Figure 6 in the appendix.

For simplicity, we will use Ĵψ (τ t ) and Îψ (τ t , a t ) to denote two random variables inside of expectation in equation 3, which is written by:

EQUATION

We use Var ψ (s t , a t ) 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.

EQUATION

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

EQUATION

We seek to optimize ψ to minimize the variance upper bound shown in equation 6. To achieve this, we firstly use two value function approximators V w1 (s t ), I w2 (s t , a t ) to approximate the expectation of two estimators respectively:

EQUATION

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:

EQUATION

EQUATION

If we assume that V w1 (s t ) and I w2 (s t , a t ) accurately represent the expectation, then we get the gradient to optimize ψ as follows, in order to minimize the upper bound of E at∼π(at|st) Var ψ (s t , a t ):

EQUATION

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

EQUATION

Here we will illustrate why performing gradient descent on ψ leads to an advantage estimation with lower variance. For a single r t+k,ψ , the gradient component in the first term of equation 11 pushes the r t+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 r t+k,ψ from constructing a constant return by the restriction in variance of importance sampling estimator

T -t k=0 γ k r t+k,ψ C π k (s t , a t , S t+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, r t+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:

ÂIAE t := Ĵψ (τ t ) -V w1 (s t ) + I w2 (s t , a t )

, where τ t ∼ P π (τ t |s t , a t ).

(12)

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.

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.

EQUATION

In practice, we use an action classification model P φ (a t |s t , s t+k , k) to represent P π (a t |s t , s t+k ), using C φ (s t , a t , s t+k ) :=

P φ (at|st,s t+k ,k) π(at|st)

to give an approximation of C π . We call P φ (a t |s t , s t+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:

EQUATION

)

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 π (s t+k1 |s t+k2 , s t ) 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.

6. ALGORITHM:

In this section, we present crucial algorithms details of the method presented in Section 3, 4 and 5. Update θ by Proximal Policy Optimization 9: end for

6.2. COMPUTATIONAL COMPLEXITY:

There might be concerns about computational complexity for training dependency model and reward decomposition model, since the training requires O(T 2 ) number of (s t , s t+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 θ (s t , a t ) at the last step. The reward decomposition model also considers the decomposition over the final reward Q θ (s t , a t ).

Block-wise Training. When training dependency model and reward decomposition model, we provide a total of 128 2 numbers of (s t , s t+k ) pairs, and accumulate gradient for every valid (s t , s t+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.

7. RELATED WORK:

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.

8. EXPERIMENTS:

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.

8.1. ENVIRONMENT SETTINGS:

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.

8.2. DEPENDENCY FACTOR MODELING:

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 s t+k shown in Figure 2a , the model P φ (a t |s t , s t+k , k) correctly predicts down and right action that more likely lead current state s t to future state s t+k . We also build a dataset consisting of 300 random (s t , s t+k ) pairs, where k is uniformly sampled from 1 to 30. We evaluate the mean KL divergence between ground-truth P π (a t |s t , s t+k ) and prediction P φ (a t |s t , s t+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 π .

8.3. THE VARIANCE AND ACCURACY OF INDEPENDENCE-AWARE ADVANTAGE ESTIMATION:

We evaluate the variance of IAE estimator on a variety of finite-state MDP settings. We train r(s t , s t+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.

8.4. PERFORMANCE OF POLICY OPTIMIZATION:

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.

9. 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 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.

A PROOFS:

A.1 PROOF OF UNBIASEDNESS

In this section, we show the proof of theorem 1, which demonstrates that any choice of r t+k will not bias the combined estimator as long as r t+k is a function of s t , a t and τ t+k .

Theorem 1. Suppose r t+k ∼ R(r t+k |s t , a t , τ t+k ), where R is any probability distribution. Then

EQUATION

where Since A π (s t , a t ) can be written by

A π (s t , a t ) = E τt∼P π (τt|st,at) T -t k=0 γ k R t+k -E τt∼P π (τt|st) T -t k=0 γ k R t+k ,

We only need to prove that

EQUATION

We have derivation that Thus, it is proved that a t+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:

E τt∼P π (τt|st) T -t k=0 γ k r t+k C π k (s t , a t , S t+k , A t+k ) + E τt∼P π (τt|st) T -t k=0 γ k r t+k = E τt∼P π (τt|st) T -t k=0 γ k r t+k (1 + C π k (s t , a t , S t+k , A t+k )) = T -t k=0 S a t+k γ k E τ t+k ∼P π (τ t+k |s t+k ,a t+k ) [r(s t , a t , τ t+k )] (1 + C π k (s t ,

P π (A t = a|s t , s t+k , a t+k ) = 1, a = a t+k 0, a = a t+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 π (a t |s t , s t+k2 ) = E s t+k 1 ∼P π (s t+k 1 |s t+k 2 ,st) P π (a t |s t , s t+k1 ), if k 2 > k 1 ≥ 1.

Proof. 

E

B TEMPORAL DIFFERENCE LEARNING OF DEPENDENCY MODEL:

We use a mixture of temporal difference and Monte-Carlo training target to train dependency model. For a trajectory sample τ = {s i , a i , r i } T i=1 , we provide training target for P φ (a t |s t , s t+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 a t as the ground-truth; for the temporal difference training target, we use distribution of P φ (a t |s t , s t+k-1 , k -1) as the ground-truth. We mix the two training target with 0.85 k weight on Monte-Carlo target and 1 -0.85 k weight on temporal difference target to achieve the optimal bias-variance tradeoff. 

C EXPERIMENT 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 φ (a t |s t , s t+k , k), we weight the MC loss by 0.85 k and TD loss by 1 -0.85 k , 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.

:



Figures and Tables:

Figure 1: Model Dependency Graph.

Figure 2: Results on dependency factor modeling. (a): The top and bottom image respectively illustrate s t and s t+k in the case study, and here we set k to be 7. (b): Dashed lines show the true distribution of P π (a t |s t , s t+k ); solid lines show the dependency model prediction P φ (a t |s t , s t+k , k); x-axis represents the number of training iterations of dependency model. Four different colors represent four different actions. Blue line shows the KL divergence between predicted distribution and true distribution. (c):The blue line shows the mean KL divergence between P π (a t |s t , s t+k ) and P φ (a t |s t , s t+k , k) over the dataset of random (s t , s t+k ) pairs, averaged in 10 runs.

Figure 3: (a): Mean squared error of two value functions V w1 and I w2 averaged in 10 runs on Pixel Grid World. Green and red lines show the MSE of V w1 and I w2 respectively, when r ψ is fixed; blue and orange lines show the MSE of V w1 and I w2 respectively, when r ψ is trained by our method. (b): The cosine similarity between advantage estimation and ground-truth advantage function. We compare IAE estimation, Monte-Carlo estimation and GAE estimation.

Figure 4: Overall performance curve. Figure (a) and (b) respectively show the training curve on Pixel Grid World environment in per-step punishment setting and no punishment setting, averaged in 10 random seeds. In per-step punishment setting agent gets negative rewards before reaching goals, while in no punishment setting agent gets no reward before reaching goals.

Figure 5: Observation in Pixel Grid World.

Figure 6: Network architecture for dependency model and reward decomposition model.

a t , s t+k , a t+k ) := P π (s t+k , a t+k |s t , a t ) P π (s t+k , a t+k |s t ) -1.

Table content: <table/>



Table content: <table><tr><td/><td colspan="4">Connected-low Isolated-low Connected-high Isolated-high</td></tr><tr><td colspan="2">MC 0.63</td><td>1.66</td><td>6.21</td><td>16.58</td></tr><tr><td>IS</td><td>1.43</td><td>8.68</td><td>1.45</td><td>8.63</td></tr><tr><td colspan="2">IAE 0.59</td><td>0.60</td><td>0.68</td><td>0.71</td></tr></table>

C π k (s t , a t , s t+k , a t+k ) := P π (s t+k , a t+k |s t , a t ) P π (s t+k , a t+k |s t ) -1. Proof. Denote r(s t , a t , τ t+k ) = E [r t+k |s t , a t , τ t+k ].

Table content: <table/>

a t , s t+k , a t+k ))P π (s t+k , a t+k |s t )ds t+k k E τ t+k ∼P π (τ t+k |s t+k ,a t+k ) [r(s t , a t , τ t+k )] P π (s t+k , a t+k |s t , a t ) P π (s t+k , a t+k |s t ) P π (s t+k , a t+k |s t )ds t+k k E τ t+k ∼P π (τ t+k |s t+k ,a t+k ) [r(s t , a t , τ t+k )] P π (s t+k , a t+k |s t , a t )ds t+k P π (a t |s t , s t+k , a t+k ) = P π (s t , a t , s t+k , a t+k ) P π (s t , s t+k , a t+k ) = P π (s t , a t , s t+k )P π (a t+k |s t , a t , s t+k ) P π (s t , s t+k )P π (a t+k |s t , s t+k ) = P π (s t , a t , s t+k )π(a t+k |s t+k ) P π (s t , s t+k )π(a t+k |s t+k ) = P π (s t , a t , s t+k ) P π (s t , s t+k ) = P π (a t |s t , s t+k ).

Table content: <table><tr><td>Under review as a conference paper at ICLR 2020</td></tr><tr><td>The second equation is proved by:</td></tr><tr><td>T -t</td></tr><tr><td>=</td></tr><tr><td>k=0 S a t+k</td></tr><tr><td>T -t</td></tr><tr><td>=</td></tr><tr><td>k=0 S a t+k</td></tr><tr><td>T -t</td></tr><tr><td>γ k r t+k ,</td></tr><tr><td>k=0</td></tr><tr><td>which proves equation 2 and further proves theorem 1.</td></tr></table>

(a t , s t , s t+k1 , s t+k2 ) P π (s t+k2 , s t ) ds t+k1 = P π (a t , s t , s t+k2 ) P π (s t+k2 , s t ) = P π (a t |s t , s t+k2 )

Table content: <table/>

Peer review (iclr2020):
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. They go on to say that this estimator will tend to have lower variance than the conventional Monte-Carlo estimator when future rewards are independent of current actions. However, the variance can actually be higher, due to the importance sampling ratio used, when future rewards are highly dependent on the current action. They propose a method to combine the two estimators on a per reward while maintaining unbiasedness using a control variate style decomposition. This introduces a tunable reward decomposition parameter which determines how to allocate each reward between the two estimators. The authors propose a method to tune this parameter by approximately optimizing an upper bound on the variance of the combined estimator. As a final contribution, the authors introduce a temporal-difference method of estimating the action probability P(a_t|s_t,s_{t+k}) required by their method. In the experiments, the authors provide empirical evidence that various aspects of their proposed method can work as suggested on simple problems. They also provide a simple demonstration where their advantage estimator is shown to improve sample efficiency in a control problem. This paper suffers from moderate clarity issues, but I lean toward acceptance primarily because I feel that the central idea is a solid contribution. The idea of improving credit assignment by explicitly estimating how much actions impact future states seems reasonable and interesting. The temporal difference method introduced for estimating P(a_t|s_t,s_{t+k}) is also interesting and clever. I*m less confident in the introduced method for trading off between the Monte Carlo and importance sampling estimators. The experiments seem reasonably well executed and do a fair job of highlighting different aspects of the proposed method. The derivation of the combined estimator was very confusing to me. It*s strange that the derivation of the variance lower bound includes terms which are drawn from both a state conditional and state-action conditional trajectory. You*re effectively summing variances computed with respect to two different measures, but the quantity being bounded is referred to as just the *variance of the advantage estimator*. What measure is this variance supposed to be computed with respect to? Especially given that as written the two estimators rely on samples drawn from two different measures. It doesn*t help that the advantage estimator whose variance is being constructed is never explicitly defined but just referred to as *advantage estimator derived from equation 3*. Nevertheless, if we ignore the details of what exactly it is a lower bound of, the sum of the three variances in equation 5 seems a reasonable surrogate to minimize. Related to the above point I don*t fully understand what the variances shown in table 1 mean in the experiments section. For the IAE estimator for example, is the variance computed based on each sample using three independent trajectories (one for each term) or is it computed from single trajectories? If it*s from single trajectories I can*t understand how the expression would be computed. Questions for the authors: -Could you please explicitly define the *advantage estimator derived from equation 3*? -Could you please explain precisely how the variance is computed in table 1? Update: Having read the other reviews and authors response I will maintain my score of a weak accept, though given more granularity I would raise my score to a 7 after the clarifications. I appreciate the authors* clarification of the advantage estimator and feel the related changes to the paper are very helpful. I still feel the central idea of the work is quite strong. However, I also feel the control variate part of the work is very loosely specified. In particular, given the use of function approximation in practice instead of actually sampling 3 trajectories the validity of the control variate method applied is questionable. As the authors say *if the random variable in either term has high variance, function approximators will tend to have large error*, This may be true initially but the function approximator can already reduce variance over time by learning, so it*s not clear how the function approximators and control variate complement each other. This is something I feel would be worthwhile to explore more in future work. Also, I feel it*s worth pointing out that a concurrent paper presenting a very similar idea is scheduled to be presented at NeurIPS 2020, which can be found here: https://papers.nips.cc/paper/9413-hindsight-credit-assignment. I don*t feel this in any way undermines the contribution of the work presented here, but merely wanted to make the meta reviewer aware in case it was relevant to their decision. In fact, I feel this work complements that one in a number of ways, including the presentation of the temporal difference method for learning the action probabilities.