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