Generalized advantage estimation (GAE) estimates policy-gradient advantages by blending multi-step temporal-difference information. It addresses the trade-off between biased single-step estimates and high-variance complete-trajectory returns. The decay parameter λ moves the estimator from one-step TD at λ=0 toward full-trajectory returns at λ=1; λ=0.95 is a common default for LLM fine-tuning. section 6.4.1
The estimator is built from three operational pieces:
- Compute the value-function TD residual δᵗ = rᵗ + γV(sᵗ⁺¹) − V(sᵗ), measuring the prediction error at each step. section 6.4.1
- Combine residuals backward with exponentially decreasing weights, producing Âᵗ = Σₗ(γλ)ˡδᵗ⁺ˡ. section 6.4.1
- Use terminal masks to stop bootstrapping and reset accumulation, allowing independently computed advantages for packed or multiple sequences. section 6.4.1
Equation 6.56
\[\begin{aligned} \hat{A}_t^{GAE(\gamma,\lambda)} &= (1-\lambda)(\hat{A}_t^{(1)} + \lambda\hat{A}_t^{(2)} + \lambda^2\hat{A}_t^{(3)} + \cdots) \notag \\ &= (1-\lambda)(\delta_t^V + \lambda(\delta_t^V + \gamma\delta_{t+1}^V) + \lambda^2(\delta_t^V + \gamma\delta_{t+1}^V + \gamma^2\delta_{t+2}^V) + \cdots) \notag \\ &= (1-\lambda)(\delta_t^V(1 + \lambda + \lambda^2 + \cdots) + \gamma\delta_{t+1}^V(\lambda + \lambda^2 + \cdots) + \cdots) \notag \\ &= (1-\lambda)(\delta_t^V\frac{1}{1-\lambda} + \gamma\delta_{t+1}^V\frac{\lambda}{1-\lambda} + \cdots) \notag \\ &= \sum_{l=0}^{\infty}(\gamma\lambda)^l\delta_{t+l}^V. \label{eq:GAE_DFN} \end{aligned}\]Follow the derivation from weighted n-step advantages to the compact exponentially weighted sum of TD residuals. section 6.4.1
Token-level PPO implementation
For language-model RLHF, process post-KL per-token rewards, current value predictions, terminal masks, γ, and λ in a reversed sequence loop. At each token, calculate the residual, update the accumulator with γλ and the terminal mask, store the result as the advantage, and form value targets by adding advantages to values. Detach the resulting advantages before applying the policy loss. section 6.4.1 section 6.2.7
B, L = rewards.shape
advantages = torch.zeros_like(rewards)
next_v = torch.zeros(
B,
device=rewards.device,
dtype=rewards.dtype,
)
gae = torch.zeros(
B,
device=rewards.device,
dtype=rewards.dtype,
)
for t in reversed(range(L)):
not_done = 1.0 - done_mask[:, t]1
delta = (
rewards[:, t] + gamma * not_done * next_v - values[:, t]
) 2
gae = delta + gamma * lam * not_done * gae
advantages[:, t] = gae
next_v = values[:, t]
targets = advantages + values # y_t for value regression3
advantages = advantages.detach() # for policy loss4
Inspect the reversed PyTorch-style loop, including terminal handling, residual accumulation, value targets, and advantage detachment. section 6.4.1
