GAE Over a Full Trajectory

hardPROPolicy Gradients and Advantage EstimationReinforcement LearningRL 1: Foundations

Problem Statement

Compute GAE(gamma, lambda) advantages for a full trajectory. Given rewards R_1..R_T and values values of length T+1, the advantages satisfy the backward recursion:

$\delta_t = R_{t+1} + \gamma V(s_{t+1}) - V(s_t)$ $\hat{A}_t = \delta_t + \gamma\lambda\, \hat{A}_{t+1}, \qquad \hat{A}_T = 0$

Implement gae(rewards, values, gamma, lam) returning the list of T advantages.

Example

Input

gae([1.0, 1.0], [0.0, 0.0, 0.0], 0.9, 1.0)

Output

[1.9, 1.0]

Reasoning

Initialize the accumulated advantage to $0$ and begin the backward pass at the final time step $t=1$ , where the future advantage is zero by definition.
Calculate the temporal difference error $\delta_1$ using the reward and values at step 1: $\delta_1 = R_1 + \gamma V(s_2) - V(s_1) = 1.0 + 0.9(0.0) - 0.0 = 1.0$ .
Update the accumulated advantage for $t=1$ by adding the discounted future advantage: $\hat{A}_1 = \delta_1 + \gamma\lambda \hat{A}_2 = 1.0 + 0.9(1.0)(0) = 1.0$ .
Move to the previous time step $t=0$ and calculate the temporal difference error $\delta_0$ : $\delta_0 = R_0 + \gamma V(s_1) - V(s_0) = 1.0 + 0.9(0.0) - 0.0 = 1.0$ .
Update the accumulated advantage for $t=0$ by incorporating the previously computed advantage: $\hat{A}_0 = \delta_0 + \gamma\lambda \hat{A}_1 = 1.0 + 0.9(1.0)(1.0) = 1.9$ .
The final output is [1.9, 1.0]

Constraints

len(values) == len(rewards) + 1.
Compute in a single backward pass.
Return a list of T floats.

Learn this in

Actor-Critic & PPOReinforcement Learning

More from RL 1: Foundations

PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium Entropy of a Policyeasy

Concept primerPREVIEW

1. Background Knowledge

Generalized Advantage Estimation (GAE) is a technique used in policy gradient reinforcement learning to balance the bias and variance of advantage estimates. The advantage function $\hat{A}_t$ measures how much better an action is compared to the baseline value estimate $V(s_t)$ . A naive Monte Carlo return has low bias but high variance, while a single-step temporal difference (TD) error has high bias but low variance. GAE interpolates between these extremes using a decay parameter $\lambda \in [0, 1]$ .

The core building block is the TD residual (or TD error) at time step $t$ : $\delta_t = R_{t+1} + \gamma V(s_{t+1}) - V(s_t)$ where $\gamma$ is the discount factor, $R_{t+1}$ is the reward received after taking an action in state $s_t$ , and $V(s_t)$ , $V(s_{t+1})$ are the value estimates for the current and next states. The GAE advantage is then a weighted sum of future TD residuals: $\hat{A}_t = \sum_{l=0}^{T-t} (\gamma\lambda)^l \delta_{t+l}$ This can be computed efficiently via the backward recursion $\hat{A}_t = \delta_t + \gamma\lambda\, \hat{A}_{t+1}$ , with the boundary condition $\hat{A}_T = 0$ (assuming the episode terminates at step $T$ ).

In practice, you are given a full trajectory of rewards and pre-computed value estimates. The values array has length $T+1$ because it includes the value of the terminal state $V(s_{T+1})$ , which is typically $0$ for episodic tasks but may be non-zero if the episode is truncated. The rewards array has length $T$ , corresponding to rewards $R_1, \dots, R_T$ .

2. Algorithm Approach

The problem is a classic backward dynamic programming or backward recursion pattern. Since $\hat{A}_t$ depends on $\hat{A}_{t+1}$ , you must compute advantages from the end of the trajectory backward to the beginning.

The algorithm follows these principles:

Initialize the advantage at the final step to $0$ .
Iterate backward from $t = T-1$ down to $t = 0$ .
At each step, compute the TD residual $\delta_t$ using the current reward, the next state's value, and the current state's value.
Update the advantage using the recursion $\hat{A}_t = \delta_t + \gamma\lambda\, \hat{A}_{t+1}$ .
Store the computed advantage in the output list.

This approach avoids the $O(T^2)$ cost of directly summing the geometric series for each $t$ by reusing the previously computed advantage, resulting in an $O(T)$ solution.

3. Step-by-Step Strategy

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results