Generalized Advantage Estimation

Problem Statement

Compute GAE(gamma, lambda) advantages and the corresponding value targets for a rollout.

Background

An actor-critic agent needs an advantage estimate $A_t \approx q(s_t,a_t) - v(s_t)$: how much better this action was than the critic's expectation. The one-step TD error is the cheapest such estimate,

$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$

low variance but biased by whatever the critic gets wrong. The full Monte Carlo advantage is unbiased but high variance. GAE dials between them with exactly the geometric average you already met as the lambda return:

$A_t^{\text{GAE}} = \sum_{l=0}^{\infty} (\gamma\lambda)^l \delta_{t+l}$

which computes in one backward pass:

$A_t = \delta_t + \gamma\lambda\,(1 - d_t)\, A_{t+1}$

Two implementation details do all the damage in practice:

1. The done mask appears twice. It kills the bootstrap inside $\delta_t$ (a terminal state has no successor value) and it stops the accumulator from leaking advantage across an episode boundary into the previous episode.
2. **The value targets are **A_t + V(s_t)****, not the raw discounted returns. This is what keeps the critic's regression target consistent with the advantages the actor is being trained on.

lam = 0 recovers the plain TD error and lam = 1 recovers the Monte Carlo advantage — check your implementation against both.

Your Task

Implement:

def gae(rewards, values, dones, last_value, gamma, lam):
    ...

• rewards[t], values[t] (the critic's $V(s_t)$), dones[t] (bool, True if step t ended an episode).
• last_value — the critic's value of the state after the final step of the rollout, used only when the last step is not terminal.
• Process the rollout backwards.

Return a **tuple **(advantages, returns)****, both lists of T floats, where returns[t] = advantages[t] + values[t].

Input / Output Format

Lists of floats/bools in; a tuple of two lists of floats out, rounded to 4 decimals by the grader.

Sample

rewards = [1.0, 1.0]
values  = [0.0, 0.0]
dones   = [False, True]
adv, ret = gae(rewards, values, dones, 0.0, 0.9, 0.95)
print([round(x, 4) for x in adv])
print([round(x, 4) for x in ret])

Output:

[1.855, 1.0]
[1.855, 1.0]

Step 1 is terminal so delta_1 = 1.0 and A_1 = 1.0. Step 0 has delta_0 = 1.0 + 0.9*0.0 - 0.0 = 1.0 and A_0 = 1.0 + 0.90.951.0 = 1.855. The values are all zero, so the returns equal the advantages.