def gae(gamma, margin, lambda_=0.95):
  """Calculate Generalized Advantage Estimation.

  Calculate state values bootstrapping off the following state values -
  Generalized Advantage Estimation https://arxiv.org/abs/1506.02438

  Args:
    gamma: float, gamma parameter for TD from the underlying task
    margin: number of extra steps in the sequence
    lambda_: float, the lambda parameter of GAE

  Returns:
    Function (rewards, returns, values, dones) -> advantages, where advantages
    advantages is an array of shape [batch_size, length - margin].
  """
  def estimator(rewards, returns, values, dones, discount_mask):
    del returns
    gammas = mask_discount(gamma, discount_mask)
    lambdas = mask_discount(lambda_, discount_mask)
    advantages = np.zeros_like(rewards)
    (_, length) = rewards.shape

    for i in reversed(range(length - 1)):
      bellman_delta = rewards[:, i] - values[:, i] + (1 - dones[:, i]) * (
          gammas[:, i] * values[:, i + 1]
      )
      advantages[:, i] = bellman_delta + (1 - dones[:, i]) * (
          gammas[:, i] * lambdas[:, i] * advantages[:, i + 1]
      )

    return advantages[:, :(rewards.shape[1] - margin)]
  return estimator
Calculate Generalized Advantage Estimation.

  Calculate state values bootstrapping off the following state values -
  Generalized Advantage Estimation https://arxiv.org/abs/1506.02438

  Args:
    gamma: float, gamma parameter for TD from the underlying task
    margin: number of extra steps in the sequence
    lambda_: float, the lambda parameter of GAE

  Returns:
    Function (rewards, returns, values, dones) -> advantages, where advantages
    advantages is an array of shape [batch_size, length - margin].
  
trax/rl/advantages.py