def discounted_returns(rewards, gammas):
  """Computes discounted returns for a trajectory or a batch of them."""
  returns = np.zeros_like(rewards)
  ret = 0.0
  for i in reversed(range(rewards.shape[-1])):
    ret = rewards[..., i] + gammas[..., i] * ret
    returns[..., i] = ret
  return returns
Computes discounted returns for a trajectory or a batch of them.
trax/rl/advantages.py