code: def discounted_rewards(reward, done, gae_gamma, end_values):
      """Discounted rewards."""
      not_done = tf.expand_dims(1 - tf.cast(done, tf.float32), axis=2)
      end_values = end_values * not_done[-1, :, :]
      return_ = tf.scan(
          lambda agg, cur: cur + gae_gamma * agg,
          tf.expand_dims(reward, axis=2) * not_done,
          initializer=end_values,
          reverse=True,
          back_prop=False,
          parallel_iterations=2,
      )
      return tf.check_numerics(return_, "return")
  
file_path: tensor2tensor/tensor2tensor/rl/ppo.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models.research.rl import get_policy
  from tensor2tensor.utils import learning_rate
  from tensor2tensor.utils import optimize
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 5 for calling a global function with example usage: discounted_rewards(reward, done, gae_gamma, end_values) and returns: tf