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.
Methods
compute
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).
compute_from_precomputed_bootstrap
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.
next_non_terminal
next_non_terminal(dones: np.ndarray, last_dones: np.ndarray) -> np.ndarrayBuild 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.
next_values
next_values(values: np.ndarray, last_values: np.ndarray) -> np.ndarrayBuild the per-step next-value table used by GAE bootstrapping.