Skip to content

GAE ​

Generalized Advantage Estimation (GAE) calculator.

Computes advantages and returns from rollout data using the GAE-Lambda method. Separated from the rollout buffer to allow reuse and independent testing.

Source

Methods ​

compute ​

python
compute(rewards: np.ndarray, values: np.ndarray, dones: np.ndarray, last_values: np.ndarray, last_dones: np.ndarray, gamma: float, gae_lambda: float) -> tuple[np.ndarray, np.ndarray]

Compute Generalized Advantage Estimation (GAE).

Parameters:

  • rewards (ndarray) – Rewards with shape (rollout_length, num_envs).
  • values (ndarray) – Value estimates with shape (rollout_length, num_envs).
  • dones (ndarray) – Done flags with shape (rollout_length, num_envs).
  • last_values (ndarray) – Value estimates for the state after the last step, shape (num_envs,).
  • last_dones (ndarray) – Done flags for the state after the last step, shape (num_envs,).
  • gamma (float) – Discount factor.
  • gae_lambda (float) – GAE lambda parameter.

Returns:

  • tuple[ndarray, ndarray] – A tuple of (advantages, returns) each with shape (rollout_length, num_envs).

Source

compute_from_precomputed_bootstrap ​

python
compute_from_precomputed_bootstrap(rewards: np.ndarray, values: np.ndarray, next_values: np.ndarray, next_non_terminal: np.ndarray, gamma: float, gae_lambda: float) -> tuple[np.ndarray, np.ndarray]

Compute GAE from precomputed bootstrap tables so shared masks can be reused.

Source

next_non_terminal ​

python
next_non_terminal(dones: np.ndarray, last_dones: np.ndarray) -> np.ndarray

Build the per-step non-terminal bootstrap mask used by GAE.

dones[t] is the done flag of the transition out of step t (the actor records the value V(s_t) and the done resulting from that step's action at the same buffer row). As a result, the bootstrap of V(s_{t+1}) for step t is masked by dones[t] itself -- when the step terminated there is no successor to bootstrap. (Using dones[t+1] here, the CleanRL formula, is only correct when the done flag is recorded against the next state; with this recording convention it is off by one and bootstraps the next episode into terminal transitions.) last_dones is the last step's done.

Source

next_values ​

python
next_values(values: np.ndarray, last_values: np.ndarray) -> np.ndarray

Build the per-step next-value table used by GAE bootstrapping.

Source