Direct Advantage Estimation
Anonymous Author(s)
AfﬁliationAddress email
Abstract
The predominant approach in reinforcement learning is to assign credit to actionsbased on the expected return.
However, we show that the return may depend on thepolicy in a way which could lead to excessive variance in value estimation and slowdown learning.
Instead, we show that the advantage function can be interpreted ascausal effects and shares similar properties with causal representations.
Based onthis insight, we propose Direct Advantage Estimation (DAE), a novel method thatcan model the advantage function and estimate it directly from on-policy data whilesimultaneously minimizing the variance of the return without requiring the (action-)value function.
We also relate our method to Temporal Difference methods byshowing how value functions can be seamlessly integrated into DAE.
The proposedmethod is easy to implement and can be readily adapted by modern actor-criticmethods.
We evaluate DAE empirically on three discrete control domains and showthat it can outperform generalized advantage estimation (GAE), a strong baselinefor advantage estimation, on a majority of the environments when applied to policyoptimization.
Introduction
Reinforcement learning (RL) methods aim to maximize cumulative rewards in sequential decisionmaking problems [Sutton and Barto, 2018].
Through interactions with the environment, agents learnto 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 itcomplex to identify which decisions are crucial to the outcomes.
This is also known as the creditassignment problem [Minsky, 1961].
A straightforward approach is to assign credit according tothe expected return, which is aligned with the RL objective[Sutton and Barto, 2018].
However, theexpected 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 stateare independent of the actions chosen, the effect of any action should be immediate.
The expectedreturn conditioned on any action, however, depends on the policy of all subsequent states.
Thisproperty could be undesirable because the expected return would vary signiﬁcantly whenever thepolicy 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 takeinspiration from the causality literature, where the notion of causal effects is deﬁned.
We show thatthe 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 expectedto improve generalization and transfer learning [Schölkopf et al., 2021].
We prove that the advantage function minimizes the variance of the return under a normalizationconstraint, and propose Direct Advantage Estimation (DAE), a method that can model and estimatethe advantage function directly from sampled trajectories unlike previous methods that rely onlearning either the value function or the Q -function beforehand.
Furthermore, we show that DAE sss... r =u
r =d r =u r =d r =u r =d 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 deepactor-critic methods using the example of Proximal Policy Optimization (PPO)
[Schulman et al.
Finally, we test our method empirically on three discrete domains including (1) a syntheticenvironment, (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 tocausality.•
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 , ﬁniteaction space A , transition probability P ( s (cid:48) | s, a ) , expected reward function r : S ×
A → R , anddiscount factor γ ∈ [0 , 1) .
A policy is a function π : S × A → R with π ( a | s ) denoting theprobability of action a given state s .
A trajectory τ = ( s 0 , a 0 , s 1 , a 1 , ... ) is said to be sampled froma policy π if a t ∼ π ( ·| s t ) and s t +1 ∼ P ( ·| s t , a t ) .
We denote the return of a trajectory by G ( τ ) = 56 (cid:80) t ≥ [-
-]0 γ t r t , where r t = r ( s t , a t ) .
The goal of an RL agent is to ﬁnd a policy π ∗ that maximizes theexpected return, i.e., π ∗ = arg max π E τ ∼ π
[ G ( τ )] .
Given a policy π , we deﬁne the value functionby 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 usuallyunknown to the agent, and have to be learned from its interactions with the environment.
Below, webrieﬂy summarize how they are typically learned.
Monte Carlo (MC) method.
By deﬁnition, an unbiased way to estimate V ( s ) is by samplingtrajectories starting from state s and averaging the sample returns V ( s ) ≈ 1 n (cid:80) ni
=1 G ( τ i ) .
Whilebeing conceptually simple and unbiased, MC methods are known to suffer from high variance, whichcan 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 improveour estimates iteratively by sampling transitions ( s, a, r, s (cid:48) ) and updating our current estimates tomatch 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 dueto their use of bootstrapping targets.
Methods based on bootstrapping have been shown to convergefaster than MC methods empirically and have been widely adapted in recent deep RL methods [Mnihet 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 ton -steps ahead and update value estimates by r t + γr t +1 + ... + γ n V ( s t + n ) .
This allows the reward information to propagate faster to relevant state-action pairs.
Additionally, this view uniﬁes bothmethods 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 lowern relies more on bootstrapping which could be biased while higher n relies more on the sampledtrajectory which is less biased but could lead to larger variance.
In TD( λ ), these n -step estimationsare 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 advantagefunction [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 baselinefunction.
The baseline function can be seen as a control variate, which can be used to reduce thevariance of the gradient estimator when chosen carefully [Greensmith et al., 2004].
One commonchoice 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 expectedreturn of two policies is directly related to the advantage function, which paved the way for some ofthe 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 theoptimal decision depends solely on the immediate reward at each state, since the transitions areindependent 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 statesdespite the actions having only immediate effects on the expected return.
This property can beundesirable 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 distributionshift , 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 -functionestimates to be unreliable Levine et al.
[2020], Kumar et al.
[2020], Fujimoto et al.
Onesource of this problem is that the expected return by nature associates an action with all its subsequentrewards, whether they are causally related or not.
Alternatively, we would like to focus on the effectof 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 Rubin [1974] in which the author developed aframework, now known as the Neyman-Rubin model, to estimate causal effects of treatments.
Itdeﬁnes the causal effect of a treatment to be the difference between (1) the outcome of a trial resultingfrom one treatment, and (2) what would have happened if we had chosen another treatment.
Thework focused mainly on the problem of counterfactuals, which we shall not delve into.
Instead, wefocus simply on its deﬁnition of causal effects.
In RL, we concern ourselves with the effect of eachaction 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, wecan 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 theexpected return.
If we calculate the advantage function for the example environment in Figure 1, weget
[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 futurestate 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]
See Appendix A for proofs.
propositionshows that the advantage function is local , inthe sense that the advantage of an action is onlydependent on the policy of subsequent states tothe point where the action stops being relevantto future rewards.
This property is strikinglysimilar to the Sparse Mechanism Shift hypothesis [Schölkopf et al., 2021] which states thatsmall distribution changes should only lead tosparse or local changes if the representation iscausal/disentangled.
The condition of Proposition 1 may be too strictto hold in most cases.
As a relaxation, we hypothesize that, for a large class of problemsof 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 empiricallyusing the Breakout environment from the MinAtar suite [Young and Tian, 2019] by trackingthe variations of Q and A throughout the training of a PPO agent.
We can see that the advantagefunction remains more stable over the course of training compared to the Q -function.
From a moretheoretical 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 somet (cid:48) > t . This can happen when, for example, the Markov chain induced by the policy approachesa 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 passingthose states would depend weakly on the actions before passing them.
On the other hand, if theenvironment is "tree-like", where actions lead to completely different branches, then the advantagefunction 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 theproblem 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 ifthere is an equivalent for the advantage function.
Here we give a counterexample showing that theadvantage function cannot be learned in the same
way.s s u s d r =r =
u d
An agent begins in state s with actions u and d which transition the agent into states s u and s d , respectively, with no immediate reward.
In both s u and s d , there is only one action with immediate rewards 1 and 0, respectively.
Consider the environment in Figure 3.
In this example, the undiscounted advantage function fors is also equal to Equation 5, while the advantages for both s u and s d are zero because thereis only a single action in their action spaces.
Since the immediate rewards and the advantagesfor the next states are all zero for the state s ,we conclude that it is not possible to learn theadvantages of s based solely on the immediatereward and the advantages of the next state asin TD learning.
An intuitive explanation for thisfailure is that, unlike the value function, whichcontains information about all future rewards,the advantage function only models the causaleffect of each action; therefore, if the next action is not causal to a distant reward, then theinformation 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 modeland estimate the advantage function directly from sampled trajectories.
Before we begin, we ﬁrst observe that the advantage function satisﬁes the equality 194 (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 statess ∈
[Equation] [Equation]
As discussed in Section 2, MC and n -step methods often suffer from excessive variance due to theirsampling nature.
It is therefore desirable to ﬁnd unbiased low-variance estimates of the return.
Oneinteresting 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 ( τ )]
= 202 E [ G (cid:48) ( τ )] , where G (cid:48) ( τ ) = (cid:80) t ≥
γ 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 thevariance 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 otherwords, 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 obtainingthe advantage function without having to learn the (action-)value function ﬁrst, as we show in ourmain theorem.
Theorem 1.
Given a policy π and time step t ≥ 0 , we say ( s, a ) ∈ S ×
A or s ∈ S is reachablewithin 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]
then for all ( s, a ) that are reachable within t , we have ˆ A ∗ ( s, a ) =
A π ( s, a ) .
Note that the ﬁrst term inside the expectation is equal to the return, while the summation in the secondterm is only up to step t .
We observe two interesting properties of the advantage function from thistheorem.
(1) Disentanglement : If ( s, a ) is reachable within t , then ˆ A ∗ ( s, a ) =
A π ( s, a ) remainsinvariant when we apply the theorem to any ˆ t > t , suggesting that additional ˆ A t (cid:48) ’s do not interferewith 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 itas the total effect from the sequence of actions, which is simply the sum of individual effects ˆ A t (cid:48) .
Asimilar property was also pointed out by Kakade and Langford [2002], where they showed how theconsecutive 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 sampledtrajectories τ 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 neuralnetwork, 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 Qwhere one aims to ﬁnd a mapping between states (state-action pairs) and its expected return, DAEtries to regress the sum of A ( s, a ) along the trajectories towards the associated returns.
This can beinterpreted 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 untilthey terminate.
To overcome this, we show that the value function can be seamlessly integrated intoour method, which enables us to perform updates similarly to the n -step bootstrapping method.
Theorem 2.
Following the notations 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 t ( s ) =
(cid:80) tt (cid:48)
=
w t (cid:48) ( s ) .
If ˆ A ∗ = arg min ˆ A ∈ F π L ( ˆ A, ˆ V )and ( s, a ) is reachable within t − 1 , then for any ˆ V : S → R , we haveˆ A ∗ ( s, a )
=
[Equation]
[Equation]
Note that in Equation 12, the minimizer ˆ A ∗ depends only on the bootstrapping target V target but notthe estimated value function ˆ V , and vice versa.
One might notice that Equation 11 is actually themulti-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) t − 1 t (cid:48) = γ t (cid:48) r t (cid:48)
{++ γ t ˆ V k − 1 ( s t ) (cid:12)(cid:12)(cid:12) s 
+}
= s (cid:105) , which replicatesmulti-step TD learning.
Similarly, Equation 12 can be understood as a multi-step estimation for theadvantage function with bootstrapping.
In practice, we can sample n -step trajectories and use Equation 10 as the loss function to estimate boththe advantage function and the value function simultaneously.
However, naively applying the theoremthis 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 utilizeeach trajectory by treating each sub-trajectory as a trajectory sampled from the same policy.
Forexample, let ( s 0 , a 0 , . . .
, s t ) be a trajectory.
Then we treat ( s i , a i , . . .
, s t ) for all i ∈ { 0 , . .
, t − 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 loss function
[Equation]
In practice we may use a single approximator with parameters θ to model both ˆ A and ˆ V simultaneouslyas 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 basedon 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 dynamicsbut 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 gradientstep weighted by the estimated advantages.
For both methods, we use the same network architecturesand hyperparameters to train the agents, see Appendix C for more details.
We assess the quality of theadvantage estimation through mean squared error (MSE) between the estimated and the true advantagefunction (Equation 5) on the sampled state-action pairs in each iteration.
Additionally, we evaluatethe 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 causedby the variance of the n -step returns used in the estimates, which slowly decreases as the policyconverges to the deterministic greedy policy.
For policy optimization, we see that GAE performsbetter in the early phase of training, but is later surpassed by DAE.
This is likely because estimatesfrom 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 DAEbecomes more accurate later in training, we see substantial gains in performance, which demonstratesa drawback of using high variance estimates from GAE.
5.2 MinAtar & Arcade Learning Environment (ALE)
We now test DAE’s performance in the more challenging domains of the MinAtar and the ALEenvironments using PPO as the base algorithm.
We use the PPO implementation and the tunedhyperparameters from Rafﬁn et al.
[2021], Rafﬁn [2020].
A practical implementation of PPO with
DAE can be found in Appendix B. For DAE, we tune two of the hyperparameters, namely the scalingcoefﬁ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 lesshyperparameter ( λ ) to tune because it does not rely on TD( λ ).
Additionally, we increase the numberof parallel actors for both methods as it substantially speeds up PPO training in terms of wall-clocktime 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 capacity impacts the performance of bothmethods 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. [Figure]
Each agent is trained independently for 10 million and 40 million in-game frames for the MinAtarand 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 thelast 100 training episodes, as in Schulman et al.
The metrics reﬂect how fast an agent learnsand the ﬁnal performance of the learned policy, respectively.
From Table 1, we observe that, given the same network architecture, DAE outperforms GAE in allthe MinAtar environments and most of the ALE environments on both metrics.
In Figure 5, wepresent the GAE-normalized learning curves for each domain and architecture, where we also seea 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 increasingthe network capacity consistently improves DAE’s performance while having less impact on GAE.
This is likely due to the way 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 seesigniﬁcant improvements in the MinAtar experiments by switching to the Wide network for DAE butnot GAE. Note that this gap cannot be explained by the increased representation power in the policynetwork alone, or we should also see signiﬁcant improvements with GAE. Therefore, it suggests thatthe improvements are likely originated from the ability to more accurately approximate the advantagefunction.
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 canbe attributed to the increased capacities in the policy networks, which beneﬁt both methods.
The advantage function was ﬁrst introduced by Baird [1993] to tackle problems with ﬁne-grainedtime steps, where Q -Learning could be slow.
In policy optimization methods, the advantage functionhas 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, Wanget al.
[2016b] developed a neural network architecture which can jointly estimate the value functionand the advantage function to achieve better performance with DQN [Mnih et al., 2015].
Our workcomplements theirs by justifying the use of policy centered functions in estimating the advantagefunction, 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 valuefunction to reduce the variance of policy gradient updates.
Connections between causal effects and
the
[Figure]
[Table]
advantage function were also studied by Corcoll and Vicente [2020], where they model the dynamicsof environments in a hierarchical manner with a differently deﬁned advantage function
[Table]
Conclusions and future directions
In this work, we tackled the problem of credit assignment in the RL setting.
Inspired by previousstudies in causality, we found that the advantage function can be interpreted as the causal effectof an action on the expected return.
We then proposed DAE, a novel method that can directlymodel the advantage function and estimate it from data.
Our experiments in three different domainsdemonstrated that DAE can converge to the true advantage function, improve the performance ofpolicy optimization compared to GAE in most environments while also enjoying better scalabilitywith network capacity.
Here we note some possible extensions to DAE.
Due to the difﬁculty of enforcing the normalizingconstraint (Def. 1) in continuous domains, we currently restrict ourselves to discrete action spacedomains.
One potential solution to this problem was outlined in Wang et al.
[2016a] where they useda sampling method to approximate the normalization step for continuous action spaces.
In thiswork we assumed the data is on-policy, an interesting extension would be to use importance samplingtechniques such as Retrace [Munos et al., 2016] or V-Trace [Espeholt et al., 2018] to extend DAE tooff-policy settings.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] See Section 4 for the theory and Section 5 for theempirical results.
(b) Did you describe the limitations of your work?
[Yes] We address the limitations inSection 7(c)
Did you discuss any potential negative societal impacts of your work?
[No] Reinforcement learning is a powerful framework for task automation which may lead to malicious uses.
However, our work focuses on a theoretical aspect of the framework,and hence we expect minimal societal impact.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 4(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Section 5 and Appendix C(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix C
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[No
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Proofs
Throughout this section, we use p ( s t = s, a t = a ) to denote the probability of the state-action pair attime 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 1
Proof.
By deﬁnition
[Equation]
(cid:48)
[Equation] [Equation]
Similarly, we have
[Equation] [Equation]
If E [ V ( s t (cid:48) ) | s t , a t ] =
E [ V ( s t (cid:48) )
| 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 theequality 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∂ L ∂f ( s (cid:48) , a (cid:48)
[Equation]
E τ ∼ [Equation] (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 =
[Equation] k ∂f k ∂f ( s (cid:48) , a (cid:48) )
(cid:33)(cid:35) [Equation] (cid:48)
[Equation]
( a (cid:48) [Equation] (cid:48)
[Equation] (22) =
[Equation] (cid:88)
[Equation] [Equation] (cid:34)(cid:32)
G ( τ ) [Equation] t (cid:88) t (cid:48)
[Equation] [Equation]
t (cid:48) f t (cid:48) (cid:33) ∂f
[Equation]
∂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) , [Equation] (cid:48) )
=
I ( [Equation] s, a k [Equation]
a ) = (cid:26) 1 , if ( s k , a k ) =
( s (cid:48) , a (cid:48) ) 0 , otherwise (24)
Hence,
[Equation] [Equation]
[Equation]
If we sum over all a (cid:48) ∈ A , since p ( s k = s [Equation] a k = s
[Equation]
= p
[Equation]
s
[Equation]
= s
[Equation]
, then
[Equation]
a
[Equation] [Equation]
which implies
[Equation]
Substituting this back in Equation 22, we get, assuming π ( a (cid:48) | s (cid:48) )
> 0 ,t
[Equation]
From the Markov property, 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) t − 1 t (cid:48) =
γ t (cid:48)
f t (cid:48)
(cid:12)(cid:12)(cid:12) s
= s (cid:48) (cid:105)
= 0 .
If p ( s 0 = s (cid:48) ) > 0 , then
[Equation] [Equation]
which completes the ﬁrst part of the proof.
Next, we prove the second part of the theorem regardingf .
Similar to the proof of Theorem 1, we consider
[Equation]
Following the proof of Theorem 1, we know that
[Equation]
We substitute this back into Equation 39 and sum over the action space
[Equation]
which results in(cid:88)
[Equation]
In other words
[Equation] [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]
[Equation]
f
[Equation]
s
[Equation]
a [Equation] (cid:33) (46) Finally, if ( s (cid:48) , a (cid:48) ) is reachable within t − 1 , thenf ( s (cid:48) , a (cid:48) )
= (cid:80) t − 1 k =
w k ( s (cid:48) )
t − 1 (cid:88) k
= w k ( s (cid:48) ) (cid:32) E τ ∼ π (cid:34)
[Equation]
− k −
[Equation]
(cid:88)
[Equation]
[Equation]
(cid:12) s
[Equation]
(cid:35) (47) − E
[Equation]
π (cid:34)
[Equation]
48) with the shorthand w k ( s (cid:48) )
=
γ 2 k p ( s k = s (cid:48) ) .
B Integrating PPO with DAE
Algorithm 1 shows a practical implementation of PPO with DAE.
The PPO clipping loss is
[Equation]
where π θ is the policy that is being optimized, µ is the sampling policy and ˆ A ( s, a ) is the estimatedadvantage function.
In practice, an entropy loss H π
= (cid:80) a ∈A π ( a | s ) log π ( a | s ) is also 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 = 128for
the backup horizon, then each batch would contain 2 128-step trajectories.
← θ − α ∇ ( L π + β V L A + β ent H π ) 15: end for 16: end while
C Experiment detailsC.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 ALEagent on a single environment may take up to 9 hours (wall-clock time) depending on the networkcapacity.
C.2 Variation Study
We train a DAE baseline PPO agent with the hyperparameters described below and save the statesencountered by each actor along with a checkpoint of the policy at the start of each PPO iterationfor 50 iterations.
Afterwards, we sample uniformly 1024 states from the set of collected states.
Forthese states, we perform 128 MC rollouts (up to 512 steps) for every action (the MinAtar Breakoutenvironment 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 ateach state at the beginning of the experiment.
We summarize the hyperparameters in Table 2.
The policy and the value/advantage network aremodeled separately.
We parameterize the policy using θ ij ∈ R 128 × 2 , representing the logit of eachaction.
θ ij are initialized to 0 at the beginning of each experiment.
The value/advantage networkconsists of MLPs of 2 hidden layers of size 256, and the input states are represented using one-hotencoding.
[Table]
Whole dataset
[Equation]
C.4 MinAtar & ALE C.4.1 Preprocessing
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 maxframes 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.
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
[Table]
True
Max
[Table]
per
[Table]
400K of each hidden layer by 4 (MinAtar) or 2 (ALE).
We follow Espeholt et al.
[2018] for the Deepnetwork architecture used in the ALE experiments, except we multiply the number of channels ofeach layer by 4 and increase the width of the last fully connected layer to 512.
Additionally, weuse SkipInit De and Smith [2020] to initialize the residual connections, which we found to stabilizelearning.
Input ( 10 × 10 × n ) 3×3 conv.
(channels: 32 stride: 1 ) 3×3 conv.
(channels: 32 stride: 1 )
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 )
C.4.3 PPO hyperparameters
We use the tuned PPO hyperparameters from Rafﬁn [2020] for GAE.
As for DAE, we additionallytune 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 uptraining, 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 themost important one.
This is likely because DAE relies entirely on the network to approximate theadvantage function, so having a huge batch of data at each PPO iteration is critical to having reliableestimates.
Aside from the above-mentioned hyperparameters, we have also tried using separatenetworks for the policy and the advantage function, or replacing PPO clipping with KL-divergencepenalties, 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]
D Ablation Study
In addition to the GAE baseline, here we consider two additional baselines to demonstrate theeffectiveness of DAE.1.
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 thenbe used to estimate the advantage function via ˆ A ( s, a )
=
ˆ Q ( s, a ) − ˆ V ( s ) .
This baselinedemonstrates the effectiveness of learning the advantage function directly.2.
Duel:
Based on Wang et al.
[2016b], we slightly modify the original loss function to extendto the n -step setting.
The new loss function is now
[Equation] [Equation]
where ˆ V and ˆ A are being learned.
Essentially, this differs from the DAE loss in whether wesum 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 forindividual environments and the normalized curves in Figure 7 and Figure 8.
Our results show thatthe Indirect method performs worse than both DAE and Duel, suggesting that explicitly modellingthe 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.
Thissuggests that jointly estimating the advantage function across time steps is crucial to the performanceof DAE.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
e
[Table]
[Table] [Table]
[Table] [Table]