Reinforcement Learning Part 16: Generalized Advantage Estimation

While we looked at the surrogate objective in the previous post, we’re actually going to return to the discussion of advantage functions from part 14. In this post, we’ll find a middle ground between full Monte Carlo (MC) returns and one-step bootstrapping in order to estimate the advantage function. This becomes a useful parameter in modern reinforcement learning (RL) algorithms, as it allows us to strike a balance between bias and variance.

Advantage Function Review

In part 14, we provided two definitions of the advantage function, which (if you recall) attempts to quantify the answer to “how much better was a specific action than the policy would do on average from this state?” The first mathematical definition used the return-to-go, G_t, as a sampled estimate of the action-value, q(s,a). 

\(\hat{A}_t = G_t – V(S_t; \boldsymbol{\phi})\) 

We subtract the value of the approximator, V(s;φ), from the sampled G_t in order to get the estimated advantage value.

Later in that same post, we demonstrated how the TD error can be used as a valid estimate of the same advantage. We swapped in the bootstrapped estimate (R_t+1 + γ · V(S_t+1;φ)) for G_t to get the following version of the advantage function:

\(\hat{A}_t = \delta_t = R_{t+1} + \gamma \cdot V(S_{t+1}; \boldsymbol{\phi}) – V(S_t; \boldsymbol{\phi})\)

Both equations are valid estimates of the advantage function. The first requires us to wait until the end of an episode in order to get the full Monte Carlo returns to calculate G_t at each timestep. The second allows us to compute an estimated advantage after each step (by bootstrapping the estimated remaining state values using V(S_t+1;φ)).

Waiting for the full episode to complete eliminates bias (as we can calculate the true G_t at each step), but it introduces variance, as G_t is the sum of many random rewards across an episode. More variance can slow down training by requiring more samples. Using the bootstrapped TD error to calculate advantage, on the other hand, lowers variance but introduces bias (as we’re only using one real reward at a time).

We’ve seen this conundrum before: is there a way to balance the high bias of a one-step bootstrapped estimate with the high variance of waiting for a full MC episode to complete? We discussed this exact tradeoff in part 9, where we introduced the TD(λ) algorithm. 

The n-Step Advantage

As we saw in part 9, we can generalize the idea of n-step discounted return, followed by a bootstrapped state value estimate, as follows:

\(G_t^{(n)} = R_{t+1} + \gamma R_{t+2} + \cdots + \gamma^{n-1} R_{t+n} + \gamma^n V(S_{t+n}; \boldsymbol{\phi})\)

Here, we take a rollout of n steps and compute the discounted rewards for those steps, from t+1 to t+n. We then bootstrap the rest of the return-to-go with the discounted V(S_t+n;φ). The only difference between the version from post 9 and this one is that we’re using the approximator for the state value function (often a neural network), as given by the parameter set φ.

We can simply plug this n-step version of G_t into our advantage function to get a generalized version, where we can choose some n to vary between 1 (single step) and ∞ (full MC).

\(\hat{A}_t^{(n)} = G_t^{(n)} – V(S_t; \boldsymbol{\phi})\)

Let’s look at this equation in more detail. First, we’ll examine the n=1 case:

\(\hat{A}_t^{(1)} = G_t^{(1)} – V(S_t; \boldsymbol{\phi}) = R_{t+1} + \gamma V(S_{t+1}; \boldsymbol{\phi}) – V(S_t; \boldsymbol{\phi}) = \delta_t\)

Here, we look at only the most recent real reward and bootstrap the rest. If you recall from post 8, this is the definition of the TD error. Now, we can return to the general form of the equation for some arbitrary n. This time, we factor the discount factor (γ) out of every term after R_t+1 in order to arrive at the recursive definition for G_t⁽ⁿ⁾ (note that we assume n>2 for this to work):

\(G_t^{(n)} = R_{t+1} + \gamma \left[ R_{t+2} + \gamma R_{t+3} + \cdots + \gamma^{n-2} R_{t+n} + \gamma^{n-1} V(S_{t+n}; \boldsymbol{\phi}) \right] = R_{t+1} + \gamma G_{t+1}^{(n-1)}\)

We then substitute this recursive definition of G_t⁽ⁿ⁾ into our advantage formula:

\(\hat{A}_t^{(n)} = G_t^{(n)} – V(S_t; \boldsymbol{\phi}) = R_{t+1} + \gamma G_{t+1}^{(n-1)} – V(S_t; \boldsymbol{\phi})\)

Now, we add and subtract γV(S_t+1;φ) in order to help us group terms for a future step. Note that this is adding a net 0 to the right side of the equation, but it assumes that the φ are the same for both state value estimates (i.e. the neural network is frozen for this advantage calculation),.

\(\hat{A}_t^{(n)} = R_{t+1} + \gamma G_{t+1}^{(n-1)} – V(S_t; \boldsymbol{\phi}) + \gamma V(S_{t+1}; \boldsymbol{\phi}) – \gamma V(S_{t+1}; \boldsymbol{\phi})\)

We then group the terms as follows:

\(\hat{A}_t^{(n)} = \left[ R_{t+1} + \gamma V(S_{t+1}; \boldsymbol{\phi}) – V(S_t; \boldsymbol{\phi}) \right] + \gamma \left[ G_{t+1}^{(n-1)} – V(S_{t+1}; \boldsymbol{\phi}) \right]\)

Everything in the first bracket is the TD error (δ_t) by definition, and everything in the second bracket is the advantage function (one timestep later, one step shorter). As a result, we can rewrite our advantage function in the recursive form as follows:

\(\hat{A}_t^{(n)} = \delta_t + \gamma \hat{A}_{t+1}^{(n-1)}\)

Here, we’re saying that we can build an n-step advantage estimate one step at a time. We start with the TD error at time t (how much the critic’s forecast moved once it saw R_t+1 and S_t+1). We then add the advantage estimate of the action taken at t+1, discounted by γ, and with one fewer step in the lookahead process (n-1). 

We can then unroll the advantages to arrive at how the TD error relates to the n-step advantage function.

\(\begin{align*}\hat{A}_t^{(n)} &= \delta_t + \gamma \hat{A}_{t+1}^{(n-1)} \\ &= \delta_t + \gamma \delta_{t+1} + \gamma^2 \hat{A}_{t+2}^{(n-2)} \\ &= \delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2} + \gamma^3 \hat{A}_{t+3}^{(n-3)} \\ &= \delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2} + \cdots + \gamma^{n-1} \delta_{t+n-1} \\ &= \sum_{l=0}^{n-1} \gamma^l \delta_{t+l} \end{align*}\)

As a result, we can see that the estimated advantage function can be found by summing the discounted TD errors from t to t+n-1. The n-step advantage is a running total of discounted “surprises” (how much the critic’s forecast of G_t moved over n steps of real experience).

Rather than having to choose a specific n value, we can blend the MC and boostrapped advantages.

λ-Weighted Advantage

Once again, we return to post 9 to review the definition of the generalized λ-weighted return:

\(G_t^{\lambda} = (1-\lambda) \sum\limits_{n=1}^{T-t-1} \lambda^{n-1} G_t^{(n)} + \lambda^{T-t-1}G_{t}\)

We choose a value of λ (between 0 and 1) that determines how we mix the MC versus bootstrapped estimates. With λ=0, we end up with a single-step, fully bootstrapped estimate of G_t. With λ=1, we’d calculate the full MC return.

Remember that we also need to wait for the full episode to end in order to blend with the MC returns (i.e. λ>0). As a result, this does not work with infinitely long episodes. We’ll discuss how this works with finite rollouts later.

Instead of returns, what if we applied this bootstrap versus MC mixing to the advantage function? This is known as Generalized Advantage Estimation (GAE), and it was introduced and analyzed in a 2015 paper by Schulman et al. We end up with a formula that looks similar to our λ-weighted return but uses the advantage instead:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = (1-\lambda) \left( \hat{A}_t^{(1)} + \lambda \hat{A}_t^{(2)} + \lambda^2 \hat{A}_t^{(3)} + \cdots \right)\)

Because these λ weights sum to 1, they form a probability distribution over n with mean 1/(1−λ). Practitioners often use a relatively high value, such as λ=0.95, that works out to an average lookahead of 20 steps. This means that short-term (e.g. 1 or 2 step) estimates carry relatively more weight than ones from a longer horizon (e.g. 20 steps), but the blend includes a good mix of both.

With this definition, we mix the n-step advantages, with more emphasis on lower n values (i.e. 1- and 2-step advantage values have more weight than larger n-step advantages). We can rewrite this equation as a sum from 1 to infinite steps:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = (1-\lambda) \sum_{n=1}^{\infty} \lambda^{n-1} \hat{A}_t^{(n)}\)

We can then substitute in our definition of the advantage function (in terms of TD error) as follows:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = (1-\lambda) \sum_{n=1}^{\infty} \lambda^{n-1} \sum_{l=0}^{n-1} \gamma^l \delta_{t+l}\)

Next, we swap the order of summation:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = (1-\lambda) \sum_{l=0}^{\infty} \gamma^l \delta_{t+l} \sum_{n=l+1}^{\infty} \lambda^{n-1}\)

Note that the term γ^lδ_t+l appears in every n-step advantage that looks at least l+1 steps ahead, which is why n now starts at l+1.

By choosing a value of λ such that 0 < λ < 1, the inner summation is a geometric series that converges to the following:

\(\sum_{n=l+1}^{\infty} \lambda^{n-1} = \lambda^l + \lambda^{l+1} + \lambda^{l+2} + \cdots = \frac{\lambda^l}{1-\lambda}\)

As a result, we can factor out a 1-λ denominator term from the summation, which cancels with the 1-λ term in front of the summation. That leaves us with the paper’s official definition of the generalized advantage estimator:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l}\)

Note that the paper calls out two special cases: when λ=0 and when λ=1. With λ=0, every term where l ≥ 1 contains a factor of (γ·0)ˡ = 0, which zeroes out every term past l=0. Only the l=0 term survives, leaving us with the single-step TD error:

\(GAE(\gamma,0): \; \hat{A}_{t} = R_{t+1} + \gamma V(S_{t+1}; \boldsymbol{\phi}) – V(S_t; \boldsymbol{\phi}) = \delta_t\)

For the λ=1 case, the derivation breaks, as the geometric series no longer converges. So, while the derivation assumed λ<1, we can see that the final formula still converges for the λ=1 case. We see that the advantage is found by subtracting the approximated value of the current state from the sum of the discounted real rewards. In other words, what is the difference between the critic’s state value approximation and the MC return.

\(GAE(\gamma,1): \; \hat{A}_{t} = \sum_{l=0}^{\infty} \gamma^l \delta_{t+l} = \sum_{l=0}^{\infty} \gamma^l R_{t+l+1} – V(S_t; \boldsymbol{\phi}) = G_t – V(S_t; \boldsymbol{\phi})\)

Another way to think about GAE: it is the λ-weighted return (that we developed in post 9) minus the baseline (that we developed in post 14).

This definition of GAE is incredibly useful for deep RL algorithms going forward, including PPO, as it helps us balance the bias vs. variance tradeoff found in 1-step vs. full MC advantage calculations (just like it did for the TD(λ) algorithm). Before we move on, let’s take a moment to talk about how this works in episodes or rollouts with a finite number of steps. After all, how can we assume we’ll have enough data to compute the advantage using l=0..∞ steps?

Finite Rollouts

Let’s look at two different scenarios. In the first, we assume that the episode ends (termination). In the second, we assume that we’ve collected a limited set of samples in a rollout batch (as we discussed in post 15).

If the episode ends, we can still assume that we’re looking at an infinitely long rollout in order to calculate the GAE. At a terminal state, no further rewards arrive, and the estimated value from that state is zero by definition (V(S_T;φ) = 0). As a result, the TD errors for any timestamp after the terminating state (t > T) evaluate to 0. So, we can simplify the GAE to the following:

\(\hat{A}_t^{GAE(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l} = \sum_{l=0}^{T-t-1} (\gamma\lambda)^l \delta_{t+l}\)

In the other scenario, we collect a finite number of steps in a rollout in order to perform a policy (and/or critic) update. The GAE is calculated similarly to how we handled the terminated episode case (we use T here to denote the end of the rollout, not the terminating step):

\(\hat{A}_t = \sum_{l=0}^{T-t-1} (\gamma\lambda)^l \delta_{t+l}\)

The difference is that the predicted state value at T (given by V(S_T;φ)) is not 0 (unlike for the terminating case). As a result, we truncate the GAE calculation to T steps (instead of the usual ∞), and we bootstrap the rest, relying on the critic’s estimation of the final state in the rollout. We calculate that final TD error as follows:

\(\delta_{T-1} = R_T + \gamma V(S_T; \boldsymbol{\phi}) – V(S_{T-1}; \boldsymbol{\phi})\)

By establishing rules for handling the terminating and truncating cases, we can formally introduce the algorithm used to calculate GAE.

The GAE Algorithm

While not formally cited in academic papers, the GAE calculation is often implemented as a separate algorithm inside a larger RL algorithm, like PPO. You can find a discussion of the GAE implementation details in this blog post by S. Huang et al.

The calculation relies on a recursive form of the advantage estimation, where we mix together the bootstrapped and MC advantages within the rollout. Let’s look again at the definition of the GAE we established, but this time, we’ll reindex the summation (originally starting from l=1) to m=0. That allows us to pull out a γ·λ term and rewrite the remaining summation as GAE from time t+1.

\(\begin{align*} \hat{A}_t^{GAE(\gamma,\lambda)} &= \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l} \\ &= \delta_t + \sum_{l=1}^{\infty} (\gamma\lambda)^l \delta_{t+l} \\ &= \delta_t + \gamma\lambda \sum_{m=0}^{\infty} (\gamma\lambda)^{m} \delta_{t+1+m} \\ &= \delta_t + \gamma\lambda \hat{A}_{t+1}^{GAE(\gamma,\lambda)} \end{align*}\)

Now, we have a recursive version of the GAE that we can easily implement in a computer program. The last remaining part is to handle the termination versus truncation issue. To do that, we introduce a d_t+1 term, which is set to 1 if S_t+1 is terminal and set to 0 otherwise.

We can use this flag as a mask to set the bootstrapped state-value estimate to 0, as we discussed in the finite rollout section: γ · (1 – d_t+1)  · V(S_t+1; φ).

With these details in place, we can formally write the GAE calculation as follows:

Input: rollout S₀..S_T, R₁..R_T, terminal flags d₁..d_T
Input: critic parameters φ, discount γ, trace parameter λ

ÂT ← 0
Loop for t = T-1, T-2, ..., 0:
    m ← 1 - dt+1
    δ ← Rt+1 + γ · m · V(St+1; φ) - V(St; φ)
    Ât ← δ + γ · λ · m · Ât+1
return Â₀, ..., ÂT-1

This recursive calculation allows us to consider the forward view (from t looking forward to truncation/termination at T) as well as the backward view (from T looking backward to t). The γλ decay here is the same one we saw from post 9’s eligibility traces. In TD(λ), a trace decayed by γλ as the agent moved forward in time, spreading the TD error back over the recently visited states. Here, we compute a γλ-weighted sum of future TD errors in a single backward pass over the stored rollout. This is a similar decay and credit assignment as TD(λ), but it’s used for a batched rollout algorithm instead of an online one.

Even if we’re not using the full GAE, this λ-weighted mix of bootstrapped and MC advantages (in the finite rollout) affords us the ability to balance bias and variance. As a result, λ becomes a knob (hyperparameter) that we can choose, and it means that we also need to choose an appropriate rollout buffer size (another hyperparameter) to appropriately accommodate our chosen λ.

Choosing λ and Rollout Buffer Size

The mixing factor, λ, plays a role in determining how much weight should be placed on samples later in the rollout. With a shorter rollout (or for samples near the end of the rollout), we have to rely on the bootstrapped estimate given by V(S_T;φ), which means we’re not able to mix in real MC rewards.

We often want to choose a rollout buffer size that allows for better mixing. Obviously, this is an engineering tradeoff: smaller buffers are easier to capture and compute whereas larger buffers offer better mixing. There is no single right answer here, but we can derive some general guidance.

We can come up with a general measurement for how much a sample (on average) in a buffer is missing the GAE weight it would otherwise have with unlimited lookahead. Remember that a truncated rollout only has TD errors out to the end of the buffer, so some of the (γλ)^l weights never get used. To analyze that, we want to answer the question, “what fraction of the estimator’s total weight is missing?”

The total weight is a geometric series:

\(\text{total} = \sum_{l=0}^{\infty} (\gamma\lambda)^l = \frac{1}{1-\gamma\lambda}\)

The lost weight (from T-t onwards) is also a geometric series:

\(\text{lost} = \sum_{l=T-t}^{\infty} (\gamma\lambda)^l = \frac{(\gamma\lambda)^{T-t}}{1-\gamma\lambda}\)

We can get the fraction of the weight lost via truncation by a sample at time t (which is T-t steps away from truncation).

\(\frac{\text{loss}}{\text{total}} = (\gamma\lambda)^{T-t}\)

We can then take the average weight loss over all the timesteps in the rollout to get the following formula:

\(\text{avg weight loss} = \frac{1}{T}\sum_{t=0}^{T-1} (\gamma\lambda)^{T-t}\)

Note that as t runs from 0 to T-1, the exponent T-t runs from T down to 1. This is a geometric series written in reverse (i.e. γλ + (γλ)² + ⋯ + (γλ)^T), which converges to the following:

\(\text{avg weight loss} = \frac{1}{T} \cdot \frac{\gamma\lambda \left(1 – (\gamma\lambda)^T\right)}{1-\gamma\lambda}\)

Many implementations choose a discount factor of γ=0.99 and a mixing factor of λ=0.95 (which blends estimates over a horizon of about 20 steps). Let’s choose a relatively small buffer size of T=32. Plugging those values into the above formula gives us a lost GAE weight of about 42.5%. That means, on average, samples lose about half their GAE weighted value compared to the infinitely long buffer, which is quite a lot.

Now, let’s increase the buffer size to T=512. Plugging in the values, we get a GAE weight loss of about 3.1%, which is much more acceptable (i.e. the advantage estimate from this rollout is relatively close to the full, infinitely long GAE). Obviously, we could make the buffer infinitely long to reduce this difference to 0%, but that defeats the purpose of using a finite buffer length.

There is no established rule here, but our calculations give us a way to reason about the tradeoff. Getting the loss under about 5% means the buffer can easily accommodate your chosen λ. Above about 20%, a good chunk of the lookahead simply is not there, and you would likely do just as well with a smaller λ. 

As an example, the basic CleanRL PPO implementation defaults to a buffer length of T=128 to achieve a loss of about 12.3%.

Conclusion

As we saw in post 9, we can effectively mix MC and bootstrapped values using the λ-weighted formula. We start with the λ-weighted returns (from post 9) and derive a formula for calculating the n-step advantage function mixed with the bootstrapped advantage, which is known as the generalized advantage estimation. While this calculation assumes that rollouts continue to infinity, we have to carefully look at how early episode termination and finite rollout buffers factor into the equation.

Most implementations rely on the recursive version of the advantage estimation, working backward from the end of the buffer to time t. For each environment, you would keep track of the advantage estimation at each timestep. This means you need a buffer of size (T, num_envs) to store the advantage estimations.

As we’ll see in the next post, the  advantage estimate for each t in the rollout is part of the loss function used to update the actor. Additionally, the value estimate at each t can be added back to the advantage to get the λ-weighted return:

\(\hat{A}_t + V(S_t; \boldsymbol{\phi}) = G_t^{\lambda} – V(S_t; \boldsymbol{\phi}) + V(S_t; \boldsymbol{\phi}) = G_t^{\lambda}\)

We’ll see this trick applied for the PPO implementation, as the returns from each t are needed for the critic update as well.

If you have any questions or feedback, please leave a comment below!