code: def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
    r
  
    return rewards_to_go(td_deltas, mask, lambda_ * gamma)
docs: r"""Computes the GAE advantages given the one step TD-residuals.
  
    The formula for a GAE advantage estimator is as follows:
  
    A_{bt} = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{b,t+l}).
  
    Internally we just call rewards_to_go, since it is the same computation.
  
    Args:
      td_deltas: np.ndarray of shape (B, T) of one step TD-residuals.
      mask: np.ndarray of shape (B, T) of mask for the residuals. It maybe the
        case that the `td_deltas` are already masked correctly since they are
        produced by `deltas(...)`
      lambda_: float, lambda parameter for GAE estimators.
      gamma: float, lambda parameter for GAE estimators.
  
    Returns:
      GAE advantage estimates.