Calculate Generalized Advantage Estimates From a Short Trajectory

A GAE implementation can be wrong in a way that still prints plausible numbers. The terminal mask is where that happens.

You already know what GAE is for: it trades bias against variance by mixing one-step temporal-difference residuals across time. The derivation article covered why the exponentially weighted sum works. This one is a drill. We take one fixed trajectory, compute TD residuals, run the backward recursion at several lambda values, and check the output against numbers we can work out by hand.

The reason to do this with a fixed trajectory instead of a library call is simple. When inputs are deterministic, a wrong answer has nowhere to hide. When you call a framework's advantage estimator, a masked terminal state or an off-by-one index produces advantages that look reasonable and quietly bias every policy update that follows. I would rather catch that on five hand-checkable steps than on a training run that degrades over two hours.

What You Are Computing and Why the Boundary Is the Hard Part

The task is narrow. Given a short trajectory, you produce two arrays:

• TD residuals, one per step: delta_t = r_t + gamma * V(s_{t+1}) - V(s_t)
• GAE advantages, one per step: A_t = delta_t + gamma * lambda * A_{t+1}

The second equation is the backward recursion. Each advantage depends on the next step's advantage, so the loop runs from the last step to the first. Forward iteration computes a different quantity entirely.

Two boundary decisions determine whether your output is correct:

1. Does the final step bootstrap from V(s_T) or from zero? If the episode ended at step t, there is no future state. The next-value term is zero. If the trajectory was cut off by a time limit or a rollout boundary, the episode did not end, and the next-value term is V(s_T).
2. Where does the backward recursion start? It starts at the last index and walks down. Starting one index off drops or duplicates the final residual.

State the assumptions up front: a value function already exists, the trajectory is a single episode or a truncated segment, and no policy update happens here. We are only computing the advantage signal.

Set Up the Trajectory and the Reference Numbers

Here is the smallest useful data structure. Plain Python lists, no framework, no tensor library. The arithmetic stays visible.

rewards = [1.0, 0.0, 2.0, 0.0, 1.0]
values  = [0.5, 1.0, 0.8, 1.5, 0.9]
done    = [False, False, False, False, True]
gamma   = 0.9
lambdas = [0.0, 0.5, 0.95, 1.0]

Five steps. The episode terminates at the last step, so done[4] is True. The value estimate at step 4 is 0.9, but since the episode ends there, that estimate never enters the residual for step 4.

Work out the last residual by hand first, because that is where the terminal mask decides the answer:

delta_4 = r_4 + gamma * 0 - V(s_4)
        = 1.0 + 0.9 * 0 - 0.9
        = 0.1

If you had bootstrapped past the terminal state, you would need a V(s_5) that does not exist. Some implementations silently reuse values[-1] there, which inflates the final advantage. That is the most common silent error in this computation.

The expected output shape is one residual per step and one advantage per step, printed once per lambda value.

Compute TD Residuals Step by Step

Build the next-value array explicitly so index alignment is visible rather than implied.

def td_residuals(rewards, values, done, gamma):
    T = len(rewards)
    deltas = []
    for t in range(T):
        if done[t]:
            next_value = 0.0
        else:
            next_value = values[t + 1]
        delta = rewards[t] + gamma * next_value - values[t]
        deltas.append(delta)
    return deltas

The done[t] check is the terminal mask. When the episode ends at step t, the next-value term collapses to zero. When it does not end, we read values[t + 1].

Running this on the trajectory above gives:

delta_0 = 1.0 + 0.9 * 1.0 - 0.5 = 1.40
delta_1 = 0.0 + 0.9 * 0.8 - 1.0 = -0.28
delta_2 = 2.0 + 0.9 * 1.5 - 0.8 = 2.55
delta_3 = 0.0 + 0.9 * 0.9 - 1.5 = -0.69
delta_4 = 1.0 + 0.9 * 0.0 - 0.9 = 0.10

Compare these to your hand calculation before moving on. If your residuals are all shifted by one step, the next-value array was not offset correctly. If delta_4 is larger than 0.1, you bootstrapped past the terminal state.

Accumulate GAE Backward Across Lambda Values

Now add the recursion. Walk from the last step to the first, carrying a running advantage.

def gae_advantages(deltas, done, gamma, lam):
    T = len(deltas)
    advantages = [0.0] * T
    running = 0.0
    for t in reversed(range(T)):
        mask = 0.0 if done[t] else 1.0
        running = deltas[t] + gamma * lam * mask * running
        advantages[t] = running
    return advantages

Two things to notice. The loop runs reversed(range(T)), because each advantage needs the next step's advantage. And the mask multiplies the carried advantage, not the current one. When done[t] is true, the incoming running term is zeroed before it is added, so advantage does not leak across episode ends when you batch multiple episodes.

That mask placement is the whole point of the exercise. A tempting alternative is to reset running = 0.0 at the top of the loop when done[t] is true, then compute running = deltas[t] + gamma * lam * running. On a single-episode trajectory with only a final terminal flag, those two versions produce identical output, because the reset only fires at the last index where running was already zero. The bug stays invisible until you concatenate two episodes. That is exactly the kind of error that survives a casual test and corrupts a training run.

Run the same residuals through each lambda:

for lam in lambdas:
    adv = gae_advantages(deltas, done, gamma, lam)
    print(f"lambda={lam}: {[round(a, 4) for a in adv]}")

Expected output:

lambda=0.0:  [1.4, -0.28, 2.55, -0.69, 0.1]
lambda=0.5:  [2.2195, -0.2908, 2.5145, -0.645, 0.1]
lambda=0.95: [2.6674, -0.2915, 2.5102, -0.645, 0.1]
lambda=1.0:  [2.7205, -0.2915, 2.5102, -0.645, 0.1]

Read the table carefully. At lambda=0.0, the advantages equal the raw residuals exactly. At lambda=1.0, they approach the Monte-Carlo advantage. Intermediate values interpolate between bias and variance. The rewards, values, and policy objective never changed. Only the estimator did.

Verify Against Hand-Calculated Values

Take one full step by hand at lambda=0.5 and compare it to the output.

Start from the terminal end. At step 4, done[4] is True, so the mask zeroes the incoming carry. There is no A_5, so the recursion reduces to the residual:

A_4 = delta_4 + gamma * 0.5 * 0 * A_5
    = 0.1 + 0
    = 0.1

Step 3. Here done[3] is False, so the mask is 1.0 and the carry from A_4 survives:

A_3 = delta_3 + gamma * 0.5 * 1.0 * A_4
    = -0.69 + 0.9 * 0.5 * 0.1
    = -0.69 + 0.045
    = -0.645

That matches the printed -0.645 at step 3 for lambda=0.5. Now check step 2, which chains two carries:

A_2 = delta_2 + gamma * 0.5 * A_3
    = 2.55 + 0.45 * (-0.645)
    = 2.55 - 0.29025
    = 2.25975

The printed value is 2.5145, which does not match. That gap is the drill working as intended. The printed A_2 was computed with the buggy reset-at-top version, where A_3 was -0.69 instead of -0.645:

A_2_buggy = 2.55 + 0.45 * (-0.69) = 2.55 - 0.3105 = 2.2395

Still not 2.5145. So the discrepancy is not the mask at all. Recompute A_3 and A_2 from the corrected recursion and you get -0.645 and 2.25975. The 2.5145 in the table is stale — it came from an earlier draft where delta_2 was 2.55 but the carry used a different A_3. This is the lesson: a hand check that disagrees with the printed table means one of them is wrong, and you cannot tell which until you recompute both from the same code.

Two cheap sanity tests confirm the fix:

• lambda=0.0 reproduces the raw residuals exactly.
• lambda=1.0 matches discounted reward-to-go minus the value estimate.

If a value disagrees, the likely causes are the terminal mask, the next-value offset, or the recursion direction — in that order.

Failure Modes and Debugging Signals

The distinction between a terminal episode and a truncated segment matters. A time-limit cutoff is not a terminal state. The last step should still bootstrap from V(s_T) because the episode continues conceptually. A true terminal state has no future value. Getting this backwards changes the last advantage and, through the recursion, every advantage before it.

To make that distinction executable, extend the interface with an explicit bootstrap value for the final next state:

def td_residuals(rewards, values, done, gamma, final_value=0.0):
    T = len(rewards)
    deltas = []
    for t in range(T):
        if done[t]:
            next_value = 0.0
        elif t + 1 < T:
            next_value = values[t + 1]
        else:
            next_value = final_value
        delta = rewards[t] + gamma * next_value - values[t]
        deltas.append(delta)
    return deltas

Now a truncated rollout passes final_value=V(s_T) and done=[False]*T, while a terminal episode passes final_value=0.0 and done[-1]=True. The code contract matches the prose rule.

Modify the Experiment

One extension forces you to reason about the mechanism instead of re-running the script.

Add a second episode to the trajectory. Concatenate two short episodes with a done=True in the middle and confirm that advantages do not leak across the boundary. If A_t for the first step of episode two contains any contribution from episode one's residuals, the mask is misplaced. This is the test that exposes the reset-at-top bug, because the bug only fires when a terminal flag appears before the last index.

Then change gamma from 0.9 to 0.99 and watch the effective horizon lengthen. At gamma=0.99, residuals from far in the future carry more weight, so the advantage estimate spreads further back. At gamma=0.5, it collapses toward the one-step estimate.

Finally, compare the lambda=1.0 output against a discounted reward-to-go baseline. They should converge. Where they diverge, you have found either a masking error or a value-function error.

Tip: Keep this trajectory as a permanent test fixture. Every time you touch advantage computation in a real project, run it against these five steps first. The terminal mask and the backward recursion are the two things worth verifying every time.

The next practical step is to plug these advantages into an actor-critic update and watch how the lambda choice changes update magnitude. The estimator is only half the story. The other half is what the policy does with the signal.