code: 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
  
file_path: trax/trax/rl/advantages.py
imports: import gin
  import numpy
  from trax import fastmath
label_desc: Use image node_id 6 for calling a global function with example usage: gae(gamma, margin, lambda_) and returns: estimator, advantages