Generalised Advantage Estimation

Compute advantages with generalised advantage estimation (GAE), as used in PPO.

An advantage says how much better an action turned out than the critic expected. The one-step version is the TD error:

delta_t = r_t + gamma * V(s_{t+1}) * (1 - done_t) - V(s_t)

It has low variance but is biased when the critic is wrong. Summing many future TD errors removes the bias but adds variance. GAE takes an exponentially weighted sum, with lam choosing the trade-off:

A_t = delta_t + gamma * lam * (1 - done_t) * A_{t+1}

Compute it backwards from the last step, with A_T = 0. values has one more entry than rewards: the last entry is the critic's value for the state after the final step, which bootstraps the return when the rollout is cut off.

Function signature

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

Example

Input

rewards = np.array([1.0, 1.0])
values  = np.array([0.5, 0.5, 0.5])
dones   = np.array([0.0, 0.0])
gae(rewards, values, dones, gamma=0.9, lam=0.8)

Output

array([1.634, 0.95 ])

Reasoning

Both TD errors are 1 + 0.9 * 0.5 - 0.5 = 0.95. The last advantage is just 0.95. The first is 0.95 + 0.9 * 0.8 * 0.95 = 1.634.

Constraints and notes

• rewards and dones have shape (T,), and values has shape (T + 1,).
• done_t = 1 means the episode ended after step t. Then step t must not use V(s_{t+1}) or A_{t+1}, since those belong to a new episode.
• With lam = 1 and no done flags, A_t + V(s_t) is the discounted return bootstrapped from the last value.