class_docstr: A trajectory of interactions with a RL environment.
  
  Trajectories are created when interacting with an RL environment. They can
  be prolonged and sliced and when completed, allow to re-calculate returns.
class_docstr_tok: ["A","trajectory","of","interactions","with","a","RL","environment",".","Trajectories","are","created","when","interacting","with","an","RL","environment",".","They","can","be","prolonged","and","sliced","and","when","completed",",","allow","to","re-calculate","returns","."]
docstr: Calculate discounted returns.
docstr_tok: ["Calculate","discounted","returns","."]
code: def calculate_returns(self, gamma):
      """Calculate discounted returns."""
      rewards = np.array([ts.reward for ts in self._timesteps])
      discount_mask = np.array(
          [ts.env_info.discount_mask for ts in self._timesteps]
      )
      gammas = advantages.mask_discount(gamma, discount_mask)
      returns = advantages.discounted_returns(rewards, gammas)
      for i, return_ in enumerate(returns):
          self._timesteps[i] = self._timesteps[i]._replace(
              return_=return_
          )
  
imports: import collections
  import os
  import gin
  import gym
  import numpy
  from trax import fastmath
  from trax.rl import advantages
  from trax.supervised import training
label_desc: Use image node_id 11 for calling the Trajectory obj's underlying member method code with example usage: obj.calculate_returns(gamma) without return types