Generalized Advantage Estimation (Pairs With Policy Gradient)
The problem, from first principles
rl-alignment-policy-gradient-loss's raw discounted returns work, but they're noisy: a single episode's return depends on every random action taken throughout the WHOLE episode, so two runs of the exact same good policy can produce wildly different returns just from randomness elsewhere. GAE (Generalized Advantage Estimation) fixes this by blending many different "look-ahead lengths" of the TD residual into one smoothed advantage estimate, controlled by a single interpolation parameter, lambda.
From theory to code
Implement td_residuals (the one-step prediction error a value function makes) and generalized_advantage_estimation, computing the full GAE via one efficient backward recursive pass.
Constraints
td_residuals(rewards, values, next_values, dones, gamma)[t] = rewards[t] + gamma * next_values[t] * (1 - dones[t]) - values[t].generalized_advantage_estimation(..., lam)[t] = td_residuals[t] + gamma * lam * (1 - dones[t]) * generalized_advantage_estimation[t+1], computed backward, with the value past the last index treated as0.lam=0must reduce EXACTLY totd_residualsitself.lam=1must reduce EXACTLY to the full Monte-Carlo advantage:discounted_returns(rewards, gamma, dones) - values.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
generalized_advantage_estimation is structurally identical to rl-alignment-policy-gradient-loss's discounted_returns backward loop, same reset-at-dones pattern, except the thing being accumulated is td_residuals scaled by gamma * lam each step, instead of raw rewards scaled by gamma alone.
Hint 2
Verify your implementation against a completely different (much slower) brute-force approach: directly sum sum_k (gamma*lam)^k * delta[t+k] for every t, stopping at the next episode boundary: if the fast recursive version and this direct definition don't agree, something in the recursion is wrong.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.