def monte_carlo(gamma, margin):
  """Calculate Monte Carlo advantage.

  We assume the values are a tensor of shape [batch_size, length] and this
  is the same shape as rewards and returns.

  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].
  """
  del gamma
  def estimator(rewards, returns, values, dones, discount_mask):
    del discount_mask
    (_, length) = returns.shape
    # Make sure that the future returns and values at "done" states are zero.
    returns[dones] = rewards[dones]
    values[dones] = 0
    return (returns - values)[:, :(length - margin)]
  return estimator
Calculate Monte Carlo advantage.

  We assume the values are a tensor of shape [batch_size, length] and this
  is the same shape as rewards and returns.

  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