def td_k(gamma, margin):
  """Calculate TD-k advantage.

  The k parameter is assumed to be the same as margin.

  We calculate advantage(s_i) as:

    gamma^n_steps * value(s_{i + n_steps}) - value(s_i) + discounted_rewards

  where discounted_rewards is the sum of rewards in these steps with
  discounting by powers of gamma.

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

  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)
    # Here we calculate advantage with TD-k, where k=margin.
    k = margin
    assert k > 0
    advantages = np.zeros_like(values[:, k:])
    discount = 1.0
    for i in range(margin):
      advantages += discount * rewards[:, i:-(margin - i)]
      discount *= gammas[:, i:-(margin - i)]
    advantages += discount * values[:, k:]
    # Zero out the future returns at "done" states.
    dones = dones[:, :-k]
    # TPU friendly version of the formula
    # advantages[dones] = rewards[:, :-k][dones]
    advantages = fastmath.index_update(advantages,
                                       dones,
                                       rewards[:, :-k][dones])
    # Subtract the baseline (value).
    advantages -= values[:, :-k]
    return advantages
  return estimator
Calculate TD-k advantage.

  The k parameter is assumed to be the same as margin.

  We calculate advantage(s_i) as:

    gamma^n_steps * value(s_{i + n_steps}) - value(s_i) + discounted_rewards

  where discounted_rewards is the sum of rewards in these steps with
  discounting by powers of gamma.

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

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