Generalised Advantage Estimation (GAE)

Background

Generalised Advantage Estimation (GAE) interpolates between the high-bias TD(0) advantage estimator and the high-variance Monte-Carlo return. The $\lambda$ knob trades bias for variance. GAE is the standard advantage estimator inside PPO and most modern RLHF stacks.

Problem statement

Implement gae(rewards, values, gamma=0.99, lam=0.95). Walk backwards:

$$\delta_t = r_t + \gamma V_{t+1} - V_t,\qquad A_t = \delta_t + \gamma\lambda A_{t+1}$$

Assume $V_{T+1} = 0$ (episode boundary) and $A_{T+1} = 0$.

Input

• rewards — array-like length $T$.
• values — array-like length $T$, the critic's value estimates at each step.
• gamma — float in $(0, 1]$, discount factor (default 0.99).
• lam — float in $[0, 1]$, GAE mixing parameter (default 0.95).

Output

Returns np.ndarray length $T$ of per-step advantages.

Examples

Example 1 — single step

Input:  rewards=[1.0], values=[0.2]
Output: A_0 = r_0 - V_0 = 0.8

Example 2 — hand-computed two-step

Input:  rewards=[0, 1.0], values=[0, 0]; gamma=0.5, lam=1.0
Output: A_1 = 1.0; A_0 = 0 + 0.5 * 1.0 * 1.0 = 0.5

Example 3 — output length matches input

Input:  zeros of length 5
Output: shape (5,)

Constraints

• rewards and values must be the same length.
• Returns np.ndarray of floats.
• Final advantage uses $V_{T+1} = 0$ (assume terminal state).

Notes

• $\lambda = 0$ is TD(0). $A_t = r_t + \gamma V_{t+1} - V_t$. Lowest variance, highest bias.
• $\lambda = 1$ is Monte Carlo. $A_t = \sum_{k \ge 0} \gamma^k r_{t+k} - V_t$. Lowest bias, highest variance.
• PPO defaults. Most PPO implementations use $\gamma = 0.99$, $\lambda = 0.95$.