code: def td_lambda(gamma, margin, lambda_=0.95):
      """Calculate TD-lambda advantage.
  
      The estimated return is an exponentially-weighted average of different TD-k
      returns.
  
      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 TD-lambda
  
      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):
          gammas = mask_discount(gamma, discount_mask)
          lambdas = mask_discount(lambda_, discount_mask)
          td_returns = np.zeros_like(returns)
          (_, length) = returns.shape
          td_returns[:, -1] = values[:, -1]
          for i in reversed(range(length - 1)):
              lambda_i = lambdas[:, i]
              td_returns[:, i] = rewards[:, i] + (
                  1 - dones[:, i]
              ) * gammas[:, i] * (
                  (1 - lambda_i) * values[:, i + 1]
                  + lambda_i * td_returns[:, i + 1]
              )
          return (td_returns - values)[:, : (returns.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 5 for calling a global function with example usage: td_lambda(gamma, margin, lambda_) and returns: estimator, unknown