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
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].
  
trax/rl/advantages.py