Direct Advantage Estimation
Abstract
The predominant approach in reinforcement learning is to assign credit to actions based on the expected return.
However, we show that the return may depend on the policy in a way which could lead to excessive variance in value estimation and slow down learning.
Instead, we show that the advantage function can be interpreted as causal effects and shares similar properties with causal representations.
Based on this insight, we propose Direct Advantage Estimation (DAE), a novel method that can model the advantage function and estimate it directly from on-policy data while simultaneously minimizing the variance of the return without requiring the (action- )value function.
We also relate our method to Temporal Difference methods by showing how value functions can be seamlessly integrated into DAE.
The proposed method is easy to implement and can be readily adapted by modern actor-critic methods.
We evaluate DAE empirically on three discrete control domains and show that it can outperform generalized advantage estimation (GAE), a strong baseline for advantage estimation, on a majority of the environments when applied to policy optimization.
Introduction
Reinforcement learning (RL) methods aim to maximize cumulative rewards in sequential decision making problems [Sutton and Barto, 2018].
Through interactions with the environment, agents learn to identify which actions lead to the highest return.
One major difﬁculty of this problem is that, typically, a vast amount of decision makings are involved in determining the rewards, making it complex to identify which decisions are crucial to the outcomes.
This is also known as the credit assignment problem [Minsky, 1961].
A straightforward approach is to assign credit according to the expected return, which is aligned with the RL objective[Sutton and Barto, 2018].
However, the expected return accumulates all the future rewards and does not directly reﬂect the effect of an action.
Consider, for example, the environment illustrated in Figure 1.
Since the transitions at each state are independent of the actions chosen, the effect of any action should be immediate.
The expected return conditioned on any action, however, depends on the policy of all subsequent states.
This property could be undesirable because the expected return would vary signiﬁcantly whenever the policy changes during training, making it a difﬁcult target to learn.
A natural question to ask is, then, how should we deﬁne the effect of an action?
In this work, we take inspiration from the causality literature, where the notion of causal effects is deﬁned.
We show that the advantage function can be interpreted as the causal effect of an action on the expected return.
Additionally, we show that it shares similar properties with causal representations, which are expected to improve generalization and transfer learning [Schölkopf et al., 2021].
We prove that the advantage function minimizes the variance of the return under a constraint, and propose Direct Advantage Estimation (DAE), a method that can model and estimate the advantage sss... r = u r =d r =u r =d r =u r =d function directly from sampled trajectories unlike previous methods that rely on learning either the value function or the Q -function beforehand.
Furthermore, we show that DAE can be combined with a value function approximator and be updated in a way similar to Temporal Difference methods [Sutton, 1988].
We then illustrate how DAE can be integrated into modern deep actor-critic methods using the example of Proximal Policy Optimization (PPO)
[Schulman et al., 2017].
Finally, we test our method empirically on three discrete domains including (1) a synthetic environment, (2) the MinAtar suite [Young and Tian, 2019] and (3) the Arcade Learning Environment [Bellemare et al., 2013], and demonstrate that DAE outperforms Generalized Advantage Estimation (GAE)
[Schulman et al., 2015b] on most of them.
The main contributions of this paper are
• We analyze the properties of the advantage function and show how they are linked to causality.
• We develop DAE, a novel on-policy method for estimating the advantage function directly, and compare it empirically with GAE using PPO.
Background
We consider a discounted Markov Decision Process ( S , A , P, r, γ ) with ﬁnite state space S , ﬁnite action space A , transition probability P ( s (cid:48) | s, a ) , expected reward function r : S ×
A → R , and discount factor γ ∈ [0 , 1) .
A policy is a function π : S ×
A → R with π ( a | s ) denoting the probability of action a given state s .
A trajectory τ = ( s 0 , a 0 , s 1 , a 1 , ... ) is said to be sampled from a policy π if a t ∼ π ( ·| s t ) and s t +1 ∼ P ( ·| s t , a t ) .
We denote the return of a trajectory by G ( τ ) = (cid:80) t ≥ {+
+} γ t r t , where r t = r ( s t , a t ) .
The goal of an RL agent is to ﬁnd a policy π ∗ that maximizes the expected return, i.e., π ∗ = arg max π E τ ∼ π
[ G ( τ )] .
Given a policy π , we deﬁne the value function by V π ( s ) =
E τ ∼ π
[ G ( τ ) | s
=
s ] , the action-value function by Q π ( s, a ) =
E τ ∼ π
[ G ( τ ) | s
= s, a 0 = a ] , and the advantage function by A π ( s, a ) =
Q π ( s, a ) − V π ( s )
[Sutton and Barto, 2018, Baird, 1993].
For simplicity, we shall omit the superscript.
These functions quantify how good a certain state (state-action pair) is, which can be used to improve the policy.
In practice, however, they are usually unknown to the agent, and have to be learned from its interactions with the environment.
Below, we brieﬂy summarize how they are typically learned.
Monte Carlo (MC) method.
By deﬁnition, an unbiased way to estimate V ( s ) is by sampling trajectories starting from state s and averaging the sample returns V ( s ) ≈ 1 n (cid:80) ni
=1 G ( τ i ) .
While being conceptually simple and unbiased, MC methods are known to suffer from high variance, which can cause slow learning.
Temporal Difference (TD) learning.
A different approach based on bootstrapping is proposed by
Sutton [1988].
In this approach, we begin with an initial guess of the value function V and improve our estimates iteratively by sampling transitions ( s, a, r, s (cid:48) ) and updating our current estimates to match the bootstrapping targets via V ( s ) ← V ( s ) + α ( r + V ( s (cid:48) )
− V ( s )) .
This is also known as the
TD(0) method.
Unlike MC methods, TD methods have lower variance at the cost of being biased due to their use of bootstrapping targets.
Methods based on bootstrapping have been shown to converge faster than MC methods empirically and have been widely adapted in recent deep RL methods [Mnih et al., 2015, 2016, Schulman et al., 2017].
One disadvantage of the TD(0) method is that it only propagates the reward one step at a time, which can be slow when the rewards are sparse.
Instead, one can utilize the rewards that are up to n -steps ahead and update value estimates by r t + γr t {++
+} + ... + γ n V ( s t + n ) .
This allows the reward information to propagate faster to relevant state-action pairs.
Additionally, this view uniﬁes both methods in the sense that TD(0) methods correspond to the case when n = 1 while n → ∞ recovers
MC methods.
Choosing the parameter n can be seen as a bias and variance trade-off, where lower n relies more on bootstrapping which could be biased while higher n relies more on the sampled trajectory which is less biased but could lead to larger variance.
In TD( λ ), these n -step estimations are further combined to produce a more robust estimate by averaging them exponentially with λ n [Sutton, 1988, Watkins, 1989].
TD errors and the advantage function.
One way to estimate A ( s, a ) is to sample a transition ( s, a, r, s (cid:48) )
and compute A ( s, a ) =
Q ( s, a ) − V ( s ) ≈ r + γV
( s (cid:48) )
− V ( s ) .
The term r + γV ( s (cid:48) )
− V ( s ) is known as TD Error and was introduced in the TD Learning algorithm to update the value estimates.
This method can also be combined with TD( λ ) to produce more accurate estimates of the advantage function [Schulman et al., 2015b].
Policy optimization.
One important application of the advantage function is in policy optimization.
Williams [1992], Sutton et al.
[2000] showed that, for a parameterized policy π θ , the gradient of the
RL objective is given by
[Equation]
where d π ( s ) =
(cid:80) t ≥ 0 γ t p ( s t = s ) and b ( s ) is an arbitrary function of the state known as the baseline function.
The baseline function can be seen as a control variate, which can be used to reduce the variance of the gradient estimator when chosen carefully [Greensmith et al., 2004].
One common choice of b ( s ) is the value function V , which results in the advantage function.
In a different approach, Kakade and Langford [2002] showed that the difference between the expected return of two policies is directly related to the advantage function, which paved the way for some of the most popular policy optimization methods [Schulman et al., 2015a, 2017].
Revisiting the advantage function
We begin by examining the motivating example in Figure 1.
In this example, it is clear that the optimal decision depends solely on the immediate reward at each state, since the transitions are independent of the actions.
However, if we calculate the Q -function, we get
[Equation]
which shows that the Q -values at each state depend on the policy evaluated at all subsequent states despite the actions having only immediate effects on the expected return.
This property can be undesirable as the policy is typically being optimized and varied, which could lead to variations in the
Q -values and make previous estimates unreliable.
This is also known as the problem of distribution shift , where the data generating distribution differs from the distribution we are making predictions on.
Similar problems have also been studied in ofﬂine RL, where distribution shift can cause Q -function estimates to be unreliable Levine et al.
[2020], Kumar et al.
[2020], Fujimoto et al.
One source of this problem is that the expected return by nature associates an action with all its subsequent rewards, whether they are causally related or not.
Alternatively, we would like to focus on the effect of an action.
But how should we deﬁne it?
Questions like this are central to the ﬁeld of causality and have been studied extensively [Holland, 1986, Pearl, 2009].
We draw our inspiration from Splawa-Neyman et al. and Rubin [1974] in which the authors developed a framework, now known as the Neyman-Rubin model, to estimate causal effects of treatments.
It deﬁnes the causal effect of a treatment to be the difference between (1) the outcome of a trial resulting from one treatment, and (2) what would have happened if we had chosen another treatment.
The work focused mainly on the problem of counterfactuals, which we shall not delve into.
Instead, we focus simply on its deﬁnition of causal effects.
In RL, we concern ourselves with the effect of each action on the expected return.
However, the action space usually consists of more than two actions, so it is ambiguous as to which action we should compare to.
One way to resolve this is to compare each action with what would have happened normally [Halpern and Hitchcock, 2015].
This way, we can deﬁne the causal effect of an action a on some quantity X at state s by
[Equation]
If we choose X to be the return, then
[Equation]
is precisely the advantage function, suggesting that it can be interpreted as the causal effect on the expected return.
Readers with more exposure to the causality literature may wonder why machinery such as intervention or counterfactual are not required to deﬁne such quantities.
This is possible mainly because we focus on fully observable environments where confounders are not present, and we expect more sophisticated treatments would be required if one wishes to extend similar ideas to more general problems (e.g., partially observable environments [Bareinboim et al., 2015, Kumor et al., 2021]).
If we calculate the advantage function for the example environment in Figure 1, we get
[Equation]
which now depends solely on the policy evaluated at the state s .
This is because, in this example, the future state distribution is independent of the actions, so the conditional expectation cancels out, leaving only the immediate term.
More generally, we only need the expected return for the future state to be independent of the current action.
Proposition 1.
Given t (cid:48) >
t such that E [ V ( s t (cid:48) )
| s t , a t ] =
E [ V ( s t (cid:48) )
| s t ] , we have
[Equation]
[Figure]
 
[Figure]
[Figure]
See Appendix A for proofs.
This proposition shows that the advantage function is local , in the sense that the advantage of an action is only dependent on the policy of subsequent states to the point where the action stops being relevant to future rewards.
This property is strikingly similar to the Sparse Mechanism Shift hypothesis [Schölkopf et al., 2021], which states that small distribution changes should only lead to sparse or local changes if the representation is causal/disentangled.
The condition of Proposition 1 may be too strict to hold in most cases.
As a relaxation, we hypothesize that, for a large class of problems of interest, the advantage function is more stable under policy variations compared to the Qfunction, making it an easier target to learn.
This may partially explain performance improvements reported by Baird [1993], Wang et al. [2016b].
In Figure 2, we verify this empirically using the Breakout environment from the MinAtar suite [Young and Tian, 2019] by tracking the variations of Q and A throughout the training of a PPO agent.
We can see that the advantage function remains more stable over the course of training compared to the Q -function.
From a more theoretical point of view, this softer condition is satisﬁed when p ( s t (cid:48) | s t , a t )
≈ p ( s t (cid:48) | s t ) for some t (cid:48) > t .
This can happen when, for example, the Markov chain induced by the policy approaches a stationary distribution; as another example, if the environment contains certain bottleneck states (e.g., doorways), that are traversed with high probability, then the state distribution after passing those states would depend weakly on the actions before passing them.
On the other hand, if the environment is "tree-like", where actions lead to completely different branches, then the advantage function would also suffer from instability.
The question we are now interested in is how to learn the advantage function.
In previous work, estimating the advantage function typically relied on using TD errors.
This method also faces the problem of relying on an estimate of the value function that suffers from strong policy dependency.
Instead, we seek methods which could directly model the advantage function and estimate it from data.
Since both Q and V can be learned with bootstrapping via TD learning, one might wonder if there is an equivalent for the advantage function.
Here we give a counterexample showing that the advantage function cannot be learned in the same way.
[Equation]
begins state s actions u d which transition the agent into states s u and s d respectively no immediate reward there is one
1 and
Consider the environment in Figure 3.
In this example, the undiscounted advantage function for s is also equal to Equation 5, while the advantages for both s u and s d are zero because there is only a single action in their action spaces.
Since the immediate rewards and the advantages for the next states are all zero for the state s , we conclude that it is not possible to learn the advantages of s based solely on the immediate reward and the advantages of the next state as in TD learning.
An intuitive explanation for this failure is that, unlike the value function, which contains information about all future rewards, the advantage function only models the causal effect of each action; therefore, if the next action is not causal to a distant reward, then the information about that reward is lost.
Direct Advantage Estimation
In this section, we present the main contribution of this paper.
We develop a simple method to model and estimate the advantage function directly from sampled trajectories.
Before we begin, we ﬁrst observe that the advantage function satisﬁes the equality (cid:80) a ∈A π ( a | s )
A π ( s, a ) = 0 .
Functions with this property, as it turns out, are useful for the development of our method, so we provide a formal deﬁnition.
Deﬁnition 1.
Given a policy π , we say a function f : S ×
A → R is π - centered if for all states s ∈ S
[Equation]
As discussed in Section 2, MC and n -step methods often suffer from excessive variance due to their sampling nature.
It is therefore desirable to ﬁnd unbiased low-variance estimates of the return.
One interesting property of π -centered functions is that, if we transform the reward by r (cid:48) t
= r t − ˆ A ( s t , a t ) with any π -centered function ˆ A , then the expected return remains unchanged, that is E
[ G ( τ )] =
E [ G (cid:48) ( τ )] , where G (cid:48) ( τ ) = (cid:80) t ≥ [-
-]0 γ t r (cid:48) t .
This can be seen as a form of reward shaping [Ng et al., 1999], where the expected return is now invariant to the choice of ˆ A for a ﬁxed policy.
Now, consider the variance of the transformed return G (cid:48)
( τ ) : we have Var [ G (cid:48)
( τ )]
=
E (cid:2) G (cid:48)
( τ ) 2 (cid:3) − E [ G ( τ )] 2 .
In other words, minimizing its variance with respect to ˆ A is equivalent to minimizing the following quantity,
[Equation]
 . Minimizing the variance of the return this way has the additional beneﬁt of simultaneously obtaining the advantage function without having to learn the (action-)value function ﬁrst, as we show in our main theorem.
Theorem 1.
Given a policy π and time step t ≥ 0 , we say ( s, a ) ∈ S ×
A or s ∈ S is reachable within t if there exists 0 ≤ t (cid:48) ≤ t such that p ( s t (cid:48)
= s, a t (cid:48) =
a ) > 0 or p ( s t (cid:48)
= s ) > 0 , respectively.
Let
F π be the set of all π -centered functions, ˆ A t
=
ˆ
A ( s t , a t ) , and
[Equation]
A
[Equation]
F π
E   (cid:32)
G ( τ ) − t (cid:88) t (cid:48) = γ t (cid:48) ˆ A t (cid:48) (cid:33) 2   , then for all ( s, a ) that are reachable within t , we have ˆ A ∗ ( s, a ) =
A π ( s, a ) .
See Appendix A for proofs.
Note that the ﬁrst term inside the expectation is equal to the return, while the summation in the second term is only up to step t .
We observe two interesting properties of the advantage function from this theorem.
(1) Disentanglement : If ( s, a ) is reachable within t , then ˆ A ∗ ( s, a ) =
A π ( s, a ) remains invariant when we apply the theorem to any ˆ t > t , suggesting that additional ˆ A t (cid:48) ’s do not interfere with existing ones.
(2) Additivity : One intuitive way to understand the term (cid:80) t (cid:48)
γ t (cid:48)
ˆ A t (cid:48) is to view it as the total effect from the sequence of actions, which is simply the sum of individual effects ˆ A t (cid:48) .
A similar property was also pointed out by Kakade and Langford [2002], where they showed how the consecutive sum of the advantage function is related to the expected return.
This theorem also suggests a natural way to estimate the advantage function directly from sampled trajectories τ 1 , τ 2 , . .
, τ N by minimizing the following constrained loss function
[Equation]
In practice, the constraint can be forced upon a general function approximator f θ , such as a neural network, by choosing ˆ A
θ ( s, a )
=
f θ ( s, a ) − (cid:80) a ∈A π ( a | s ) f θ ( s, a ) .
In our method, neither V nor
Q is required a priori.
Instead, the advantage function is directly modeled and estimated from data, and hence we name our method Direct Advantage Estimation (DAE).
Unlike learning V or Q where one aims to ﬁnd a mapping between states (state-action pairs) and its expected return, DAE tries to regress the sum of A ( s, a ) along the trajectories towards the associated returns.
This can be interpreted as learning to distribute the return rather than to predict it.
One limitation of this method, similar to the MC methods, is that we have to sample trajectories until they terminate.
To overcome this, we show that the value function can be seamlessly integrated into our method, which enables us to perform updates similarly to the n -step bootstrapping method.
Theorem 2.
Following the notation in Theorem 1, let V target ( s ) be the bootstrapping target and ˆ V ( s ) be the value function to be learned, and
[Equation]
If ˆ V ∗ = arg min ˆ V L ( ˆ A, ˆ V ) and p ( s 0 = s ) > 0 , then for any ˆ A ∈ F π , we have
[Equation]
Additionally, let w t ( s ) =
γ 2 t p ( s t = s ) and W n ( s ) =
(cid:80) nt =
w t ( s ) .
If ˆ A ∗ = arg min ˆ A ∈ F π L ( ˆ A, ˆ V ) and ( s, a ) is reachable within t = 0 to t = n − 1 , then for any ˆ V : S → R , we have ˆ A ∗ ( s, a ) = n − 1 (cid:88) t
=
[Equation]
t ( s )
[Equation]
[Equation]
Note that in Equation 12, the minimizer ˆ A ∗ depends only on the bootstrapping target V target but not the estimated value function ˆ V , and vice versa.
One might notice that Equation 11 is actually the multi-step update for the value function.
This suggests that we can update ˆ V iteratively by letting
V target =
ˆ V k − 1 and ˆ V k ( s )
=
ˆ V ∗ ( s ) =
E (cid:104)(cid:80) n − 1 t {+=
+}
γ t r t + γ n ˆ V k − 1 ( s n ) (cid:12)(cid:12)(cid:12) s
= s (cid:105) , which replicates multi-step TD learning.
Similarly, Equation 12 can be understood as a multi-step estimation for the advantage function with bootstrapping.
At ﬁrst sight, this may seem to suffer from the same problem of spurious policy dependence due to the use of value functions (see discussion in Section 3).
We argue, however, that this can be mitigated by using long backup horizons (i.e. n (cid:29) 1 ) and the assumption that actions have local effects (i.e. p ( s t (cid:48) | s t , a t )
≈ p ( s t (cid:48) | s t ) for some t (cid:48) > t ).
If the backup horizon is large enough such that p ( s t + n | s t , a t )
≈ p ( s t + n | s t )) , then Equation 12 suggests that the dependence on V target can be reduced since E [ V target ( s t + n ) | s t , a t ] − E
[ V target ( s t + n ) | s t ] ≈ 0 . In practice, we can sample n -step trajectories and use Equation 10 as the loss function to estimate both the advantage function and the value function simultaneously.
However, naively applying the theorem this way can cause slow learning because the value function is learned only for the ﬁrst state in each trajectory and distant state-action pairs are heavily discounted.
To mitigate this, we further utilize each trajectory by treating each sub-trajectory as a trajectory sampled from the same policy.
For example, let ( s 0 , a 0 , . . .
, s n ) be an n -step trajectory.
Then we treat ( s i , a i , . . .
, s n ) for all i ∈ { 0 , . .
n − 1 } also as a trajectory sampled from the same policy.
This can be seen as repeated applications of Theorem 2, which leads us to the following loss function
[Equation]
In practice we may use a single function approximator with parameters θ to model both ˆ A and ˆ V simultaneously as done by Wang et al.
[2016b], which reduces the loss function to L A ( θ ) .
Experiments
We evaluate our method empirically by comparing the performance of policy optimization between
DAE and GAE on three different sets of discrete control tasks.
(1) A synthetic environment based on Figure 1 where the true advantage function and the expected return are known, (2) the MinAtar [Young and Tian, 2019] suite, a set of environments inspired by Atari games with similar dynamics but simpler observation space, and (3) the Arcade Learning Environment (ALE)
[Bellemare et al., 2013], a set of tasks with diverse mechanisms and high-dimensional observation spaces.
5.1 Synthetic environment
We consider a ﬁnite variant of the environment shown in Figure 1 with state space S
= { s 1 , ..., s 128 } , and train agents using a simple actor-critic algorithm where in each iteration we (1) sample trajectories, (2) estimate the advantages of the sampled state-action pairs, and (3) perform a single policy gradient step weighted by the estimated advantages.
For both methods, we use the same network architectures and hyperparameters to train the agents, see Appendix C for more details.
We assess the quality of the advantage estimation through mean squared error (MSE) between the estimated and the true advantage function (Equation 5) on the sampled state-action pairs in each iteration.
Additionally, we evaluate the learned policy by computing the true (undiscounted) expected return by E [ (cid:80)
i r i ] = (cid:80)
i π ( u | s i ) .
Note that the optimal policy in this case is π ( u |· ) = 1 with expected return (cid:80)
i π ( u | s i )
= 128 .
The results in Figure 4 demonstrate that DAE can approximate the true advantage function accurately, while GAE struggles.
The inability of GAE to accurately estimate the advantage function is caused by the variance of the n -step returns used in the estimates, which slowly decreases as the policy converges to the deterministic greedy policy.
For policy optimization, we see that GAE performs better in the early phase of training, but is later surpassed by DAE.
This is likely because estimates from GAE include unbiased estimates of n -step returns that are useful to policy optimization, while
DAE relies entirely on the approximated function, which is heavily biased at the beginning.
As DAE becomes more accurate later in training, we see substantial gains in performance, which demonstrates a drawback of using high variance estimates from GAE. [Figure] [Figure] [Figure]
[Figure]
M e a n
Sq u a r e d E rr o r [Figure] Algorithm 1 PPO with DAE (shared network) 1: Initialize ˆ A , ˆ V and the policy π with θ 2: while not converged do 3:
D = {} 4: for i = 1 , . . .
, N actors do 5:
Sample n -step trajectory τ i =
( s 0 , a 0 , r 0 , . . .
, s n ) with policy π 6:
D ← D ∪ { τ i } 7: end for 8: φ ← copy ( θ ) 9:
V target , µ ← V φ , π φ 10: for i = 1 , . .
. , N updates do 11:
Sample a mini-batch of n -step trajectories B from D 12:
Compute ˆ A ( s, a ) for all ( s, a ) ∈ B (centered according to µ ) 13:
Compute PPO clipping loss L π (Equation 49) with stop _ gradient ( ˆ A ( s, a )) , µ , and π 14:
Compute DAE loss L A (Equation 13)with ˆ A , ˆ V and V target 15:
Optimize loss L =
( L π + β V L A ) 16: end for 17: end while
5.2 MinAtar & Arcade Learning Environment (ALE)
We now test DAE’s performance in the more challenging domains of the MinAtar and the ALE environments using PPO as the base algorithm.
We use the PPO implementation and the tuned hyperparameters from Rafﬁn et al.
[2021], Rafﬁn [2020].
A practical implementation of PPO with
DAE can be found in Algorithm 1 (see Appendix B for more details) 3 . For DAE, we tune two of the hyperparameters, namely the scaling coefﬁcient of the value function loss and the number of epochs in each PPO iteration using the MinAtar environments, which are then ﬁxed for the ALE experiments.
Note that DAE has one less hyperparameter ( λ ) to tune because it does not rely on TD( λ ).
Additionally, we increase the number of parallel actors for both methods as it substantially speeds up PPO training in terms of wall-clock time and has become the predominant way to train PPO agents in complex environments [Rudin et al., 2022, Freeman et al., 2021].
We also test how network capacities impact the performance of both methods by increasing the width of each layer in the baseline network (denoted “Wide”).
For ALE, we further test a deep residual network architecture (denoted “Deep”) based on Espeholt et al.
See Appendix C for a more detailed description of the network architectures and hyperparameters.
Each agent is trained independently for 10 million and 40 million in-game frames for the MinAtar and the ALE environments, respectively.
We compare both methods on two metrics, (1) Overall : average undiscounted score of all training episodes, and (2) Last : average undiscounted score of the [Figure] last 100 training episodes, as in Schulman et al.
The metrics reﬂect how fast an agent learns and the ﬁnal performance of the learned policy, respectively.
From Table 1, we observe that, given the same network architecture, DAE outperforms GAE in all the MinAtar environments and most of the ALE environments on both metrics.
In Figure 5, we present the GAE-normalized learning curves for each domain and architecture, where we also see a similar trend that DAE is much more sample efﬁcient while achieving better ﬁnal performance.
Results for individual environments can be found in Appendix E. Finally, we found that increasing the network capacities consistently improves DAE’s performance while having less impact on GAE.
This is likely because the advantage function is directly represented in DAE using the network, such that network capacities should have a larger impact on the performance.
Interestingly, we see signiﬁcant improvements in the MinAtar experiments by switching to the Wide network for DAE but not GAE.
Note that this gap cannot be explained by the increased representation power in the policy network alone, or we should also see signiﬁcant improvements with GAE.
Therefore, it suggests that the improvements are likely originated from the ability to more accurately approximate the advantage function.
In the ALE experiments, we also observe such jumps in performance for some environments (see Appendix C), but for most environments, both DAE and GAE improved signiﬁcantly.
This suggests that the improved performance can be attributed to the increased capacities in the policy networks, which beneﬁt both methods.
Related work
The advantage function was ﬁrst introduced by Baird [1993] to tackle problems with ﬁne-grained time steps, where Q -Learning could be slow.
In policy optimization methods, the advantage function has also been widely used as an alternative to Q -functions to lower the variance of policy updates [Greensmith et al., 2004, Kakade and Langford, 2002, Sutton et al., 2000].
More recently, Wang et al.
[2016b] developed a neural network architecture which can jointly estimate the value function and the advantage function to achieve better performance with DQN [Mnih et al., 2015].
Our work complements theirs by justifying the use of policy centered functions in estimating the advantage function, and showing how it can be combined with multi-step learning.
Schulman et al.
[2015b] proposed GAE, a method which combines TD( λ ) and TD error to estimate the advantage function, and demonstrated strong empirical results in policy optimization.
The link between causality and credit assignment in RL has also been explored in Mesnard et al. [2020], where the authors used the idea of counterfactuals to construct a future-conditional value function to reduce the variance of policy gradient updates.
Connections between causal effects and
the
[Table]
advantage function were also studied by Corcoll and Vicente [2020], where they model the dynamics of environments in a hierarchical manner with a differently deﬁned advantage function.
Conclusions and future directions
In this work, we tackled the problem of credit assignment in the RL setting.
Inspired by previous studies in causality, we found that the advantage function can be interpreted as the causal effect of an action on the expected return.
We then proposed DAE, a novel method that can directly model the advantage function and estimate it from data.
Our experiments in three different domains demonstrated that DAE can converge to the true advantage function, improve the performance of policy optimization compared to GAE in most environments, while also enjoying better scalability with increased network capacities.
Here we note some possible extensions to DAE.
(1) Environments with confounders (e.g., partially observable MDPs) are common in real-world applications, and causal inference is known to be effective in such problems [Bareinboim et al., 2015, Tennenholtz et al., 2020, Kumor et al., 2021].
We expect similar ideas based on the causal effect as proposed in this work may also be useful in those settings.
Due to the difﬁculty of enforcing the centering constraint (Def. 1) in continuous domains, we currently restrict ourselves to discrete action space domains.
One potential solution to this problem was outlined in Wang et al.
[2016a] where they used a sampling method to approximate the centering step for continuous action spaces.
In this work we assumed the data is on-policy, an interesting extension would be to use importance sampling techniques such as Retrace [Munos et al., 2016] or V-Trace [Espeholt et al., 2018] to extend DAE to off-policy settings
Acknowledgments
The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Hsiao-Ru Pan and Alexander Neitz.
This work was supported by the German Federal
Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039B and by the
Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645.
A Proofs
Throughout this section, we use p ( s t = s, a t = a ) to denote the probability of the state-action pair at time step t being equal to ( s, a ) , and the probability of a trajectory by p ( τ ) =
p ( s 0 , a 0 , s 1 , a 1 , ... ) .
A.1 Proof of Proposition
Proof.
By deﬁnition
[Equation]
Similarly, we have
[Equation] [Equation]
If E [ V ( s t (cid:48) ) | s t , a t ] =
E [ V ( s t [Equation]
) | s t ] , then
[Equation]
A.2 Proof of Theorem
Proof.
For simplicity, we denote ˆ A by f .
Notice that the condition f ∈ F π is equivalent to the equality constraint (cid:88)
[Equation]
Using the method of Lagrange multipliers, we have the Lagrangian
[Equation]
(21) with the shorthand f t (cid:48)
=
f ( s t (cid:48) , a t (cid:48) ) .
If we differentiate it with respect to f ( s (cid:48) , a (cid:48) ) , then
[Equation]
= − 2 E τ ∼ π (cid:34)(cid:32)
G ( τ ) − t (cid:88) t (cid:48)
=
γ t (cid:48)
f t (cid:48) (cid:33) (cid:32) t (cid:88) k
= γ k ∂f k ∂f ( s (cid:48) , a (cid:48) )
(cid:33)(cid:35) + λ s (cid:48) π
( a (cid:48) | s (cid:48) )
(22) =
−t (cid:88) k
= γ k E τ ∼ π (cid:34)(cid:32)
G ( τ ) − t (cid:88) t (cid:48)
=
γ t (cid:48) f t (cid:48) (cid:33)
∂f k ∂f ( s (cid:48) , a (cid:48) ) (cid:35) + λ s (cid:48)
π ( a (cid:48) | s (cid:48) )
Note that ∂f k ∂f ( s (cid:48) , a (cid:48) )
=
I ( s k = s, a k = a ) =
(cid:26) 1 , if ( s k , a k ) =
( s (cid:48) , a (cid:48) ) 0 , otherwise (24)
Hence,
[Equation]
[Equation]
If we sum over all a (cid:48) ∈ A , since p ( s k = s (cid:48) , a k = s
[Equation]
= p ( s
[Equation]
= s
[Equation]
s
[Equation]
, then (cid:88)
a
[Equation] [Equation]
which implies
[Equation]
Substituting this back in Equation 22, we get, assuming π ( a (cid:48) | s (cid:48) )
> 0 , t (cid:88)
[Equation]
From the Markov property [Equation] we have
[Equation]
If s (cid:48) is reachable within t (i.e., p ( s k = s (cid:48) ) > 0 for some k ), then
[Equation]
A.3 Proof of Theorem
Proof.
We denote ˆ A by f and V target by U , and use the method of Lagrange multipliers.
Consider the
Lagrangian
[Equation]
Let’s ﬁrst consider the minimum for ˆ V ,
[Equation]
since E (cid:104)(cid:80) [Equation] [Equation] =
[Equation] γ
[Equation] f
[Equation]
(cid:12)(cid:12)(cid:12) s
= s (cid:48) (cid:105)
= 0 .
If p ( s 0 = s
[Equation]
> 0 , then
[Equation]
which completes the ﬁrst part of the proof.
Next, we prove the second part of the theorem regarding f .
Similar to the proof of Theorem 1, we consider
[Equation]
Following the proof of Theorem [Equation] we know that
[Equation]
We substitute this back into Equation 39 and sum over the action space
[Equation]
which results in (cid:88) a
[Equation]
In other words
[Equation]
Substituting this back again into Equation 39, we have
[Equation]
Using the Markov property
[Equation]
the term in the parentheses can be simpliﬁed to
[Equation]
Finally, if ( s (cid:48) , a (cid:48) ) is reachable within t −
[-= (cid:80) t − 1 k =
-]
[Equation]
then
[Equation]
, (48) with the shorthand w k ( s (cid:48) )
=
γ 2 k p ( s k = s (cid:48) ) .
B Integrating PPO with DAE
The PPO clipping loss is given by
[Equation]
where π θ is the policy that is being optimized, µ is the sampling policy and ˆ A ( s, a ) is the estimated advantage function.
In practice, an entropy loss H π
= (cid:80) a ∈A π ( a | s ) log π ( a | s ) is usually added to L π to encourage exploration.
Note that, unlike the original PPO which samples mini-batches of frames, we sample on a trajectory-by-trajectory basis.
For example, assume the batch size is 256 and n = 128 for the backup horizon, then each batch contains 2 128-step trajectories.
C Experiment details C.1 Computational resources
All the experiments were performed on an internal cluster of NVIDIA A100 GPUs.
Training a
MinAtar agent in a single environment takes less than 30 minutes (wall-clock time).
Training an ALE agent on a single environment may take up to 9 hours (wall-clock time) depending on the network capacity.
C.2 Variation Study
We train a DAE baseline PPO agent with the hyperparameters described below and save the states encountered by each actor along with a checkpoint of the policy at the start of each PPO iteration for 50 iterations.
Afterwards, we sample uniformly 1024 states from the set of collected states.
For these states, we perform 128 MC rollouts (up to 512 steps) for every action (the MinAtar Breakout environment has 3 actions) to estimate the true Q ( s, a ) , V ( s ) and A ( s, a ) for every policy.
C.3 Synthetic environment
To avoid giving DAE an unfair advantage because it could learn a function that ignores the input (since the advantage function is the same for each state), we randomly swap the rewards of u and d at each state at the beginning of the experiment.
We summarize the hyperparameters in Table 2.
The policy and the value/advantage network are modeled separately.
We parameterize the policy using θ ij ∈ R 128 × 2 , representing the logit of each action.
θ ij are initialized to 0 at the beginning of each experiment.
The value/advantage network consists of MLPs of 2 hidden layers of size 256, and the input states are represented using one-hot encoding.
[Table]
C
MinAtar.
We turn off sticky action and difﬁculty ramping to further simplify the environments.
ALE.
We follow the standard preprocessing procedures from Mnih et al.
[2015], except for the max frames per episode where we used the default value from the gym implementation [Brockman et al., 2016].
See Table 3 for a summary of the parameters.
[Table]
C.4.2 Network architecture
For the MinAtar and the ALE experiments, we summarize the baseline network architectures in
Figure 6.
For the wide network experiments, we simply multiply the numbers of channels and widths of each hidden layer by 4 (MinAtar) or 2 (ALE).
We follow Espeholt et al.
[2018] for the Deep network architecture used in the ALE experiments, except we multiply the number of channels of each layer by 4 and increase the width of the last fully connected layer to 512.
Additionally, we use SkipInit De and Smith [2020] to initialize the residual connections, which we found to stabilize learning.
[Figure]
Linear (Size: 256)
Linear (Size: |A| ) Linear (Size: |A| )
Linear (Size: 1 )
Input ( 84 × 84 × 4 ) 8×8 conv.
(channels: 32, stride: 4) 4×4 conv.
(channels: 32, stride: 2) 3×3 conv.
(channels: 32, stride: 1)
Linear (Size: 512)
Linear (Size: |A| ) Linear (Size: |A| )
Linear (Size: 1 ) is followed by a ReLU activation.
The three output streams correspond to A ( s, a ) , π ( a | s ) , and V ( s ) , respectively.
C.4.3 PPO hyperparameters
We use the tuned PPO hyperparameters from Rafﬁn [2020] for GAE.
As for DAE, we additionally tune the number of epochs per iteration {4, 6, 8} and the β V coefﬁcient {0.5, 1, 1.5} using the
MinAtar environments.
We also increase the number of parallel actors from 8 to 1024 to speed up training, the number 1024 was chosen to maximize GPU memory usage in the ALE experiments.
Among the tuned hyperparameters, we found that having a large number of parallel actors is the most important one.
This is likely because DAE relies entirely on the network to approximate the advantage function, so having a huge batch of data at each PPO iteration is critical to having reliable estimates.
Aside from the above-mentioned hyperparameters, we have also tried using separate networks for the policy and the advantage function, or replacing PPO clipping with KL-divergence penalties, but found them less effective than the original PPO algorithm.
We summarize the ﬁnal set of hyperparameters in Table 4.
The learning rate and the PPO clipping (cid:15) are linearly annealed towards 0 throughout training
[Table]
[Table]
D Ablation Study
In addition to the GAE baseline, here we consider two additional baselines to demonstrate the effectiveness of DAE.
Indirect: We learn both Q and V separately by minimizing the n -step bootstrapping losses.
[Equation]
where ˆ Q and ˆ V are the learned Q -function and the value function.
ˆ Q and ˆ V can then be used to estimate the advantage function via ˆ A ( s, a ) =
ˆ Q ( s, a ) − ˆ V ( s ) .
This baseline demonstrates the effectiveness of learning the advantage function directly.
Duel:
Based on Wang et al.
[2016b], we slightly modify the original loss function to extend to the n -step setting.
The new loss function is now
[Equation]
where ˆ V and ˆ A are being learned.
Essentially, this differs from the DAE loss in whether we sum over the advantage function over time steps.
Unlike the original dueling architecture, we use the learned policy instead of the uniform policy to enforce the π -centered constraint.
This baseline demonstrates the effectiveness of the DAE loss (sum over advantage functions).
We compare their performance in policy optimization using the MinAtar suite with the settings and
DAE hyperparameters (baseline network) described in Appendix C. We show the learning curves for individual environments and the normalized curves in Figure 7 and Figure 8.
Our results show that the Indirect method performs worse than both DAE and Duel, suggesting that explicitly modelling the advantage function can be beneﬁcial.
Furthermore, by comparing the results from DAE and Duel, we see substantial gains by switching to the DAE loss, even surpassing the performance of GAE.
This suggests that jointly estimating the advantage function across time steps is crucial.
[Figure] [Figure]
E Additional results
[Figure] [Figure]
S c
o r e
[Figure]
e
[Table]
[Table] [Table]
[Table] [Table]