class_docstr: A PPO Agent.
class_docstr_tok: ["A","PPO","Agent","."]
docstr: Compute the Monte Carlo return and advantage.
  
  Args:
    next_time_steps: batched tensor of TimeStep tuples after action is taken.
    value_preds: Batched value prediction tensor. Should have one more entry
      in time index than time_steps, with the final value corresponding to the
      value prediction of the final state.
  
  Returns:
    tuple of (return, advantage), both are batched tensors.
docstr_tok: ["Compute","the","Monte","Carlo","return","and","advantage",".","Args",":","next_time_steps",":","batched","tensor","of","TimeStep","tuples","after","action","is","taken",".","value_preds",":","Batched","value","prediction","tensor",".","Should","have","one","more","entry","in","time","index","than","time_steps",",","with","the","final","value","corresponding","to","the","value","prediction","of","the","final","state",".","Returns",":","tuple","of","(","return",",","advantage",")",",","both","are","batched","tensors","."]
code: def compute_return_and_advantage(
      self, next_time_steps: ts.TimeStep, value_preds: types.Tensor
  ) -> Tuple[types.Tensor, types.Tensor]:
      """Compute the Monte Carlo return and advantage.
  
      Args:
        next_time_steps: batched tensor of TimeStep tuples after action is taken.
        value_preds: Batched value prediction tensor. Should have one more entry
          in time index than time_steps, with the final value corresponding to the
          value prediction of the final state.
  
      Returns:
        tuple of (return, advantage), both are batched tensors.
      """
      discounts = next_time_steps.discount * tf.constant(
          self._discount_factor, dtype=tf.float32
      )
  
      rewards = next_time_steps.reward
      if self._debug_summaries:
          # Summarize rewards before they get normalized below.
          # TODO(b/171573175): remove the condition once histograms are
          # supported on TPUs.
          if not tf.config.list_logical_devices("TPU"):
              tf.compat.v2.summary.histogram(
                  name="rewards",
                  data=rewards,
                  step=self.train_step_counter,
              )
          tf.compat.v2.summary.scalar(
              name="rewards_mean",
              data=tf.reduce_mean(rewards),
              step=self.train_step_counter,
          )
  
      # Normalize rewards if self._reward_normalizer is defined.
      if self._reward_normalizer:
          rewards = self._reward_normalizer.normalize(
              rewards,
              center_mean=False,
              clip_value=self._reward_norm_clipping,
          )
          if self._debug_summaries:
              # TODO(b/171573175): remove the condition once histograms are
              # supported on TPUs.
              if not tf.config.list_logical_devices("TPU"):
                  tf.compat.v2.summary.histogram(
                      name="rewards_normalized",
                      data=rewards,
                      step=self.train_step_counter,
                  )
              tf.compat.v2.summary.scalar(
                  name="rewards_normalized_mean",
                  data=tf.reduce_mean(rewards),
                  step=self.train_step_counter,
              )
  
      # Make discount 0.0 at end of each episode to restart cumulative sum
      #   end of each episode.
      episode_mask = common.get_episode_mask(next_time_steps)
      discounts *= episode_mask
  
      # Compute Monte Carlo returns. Data from incomplete trajectories, not
      #   containing the end of an episode will also be used, with a bootstrapped
      #   estimation from the last value.
      # Note that when a trajectory driver is used, then the final step is
      #   terminal, the bootstrapped estimation will not be used, as it will be
      #   multiplied by zero (the discount on the last step).
      final_value_bootstrapped = value_preds[:, -1]
      returns = value_ops.discounted_return(
          rewards,
          discounts,
          time_major=False,
          final_value=final_value_bootstrapped,
      )
      # TODO(b/171573175): remove the condition once histograms are
      # supported on TPUs.
      if self._debug_summaries and not tf.config.list_logical_devices(
          "TPU"
      ):
          tf.compat.v2.summary.histogram(
              name="returns", data=returns, step=self.train_step_counter
          )
  
      # Compute advantages.
      advantages = self.compute_advantages(
          rewards, returns, discounts, value_preds
      )
  
      # TODO(b/171573175): remove the condition once historgrams are
      # supported on TPUs.
      if self._debug_summaries and not tf.config.list_logical_devices(
          "TPU"
      ):
          tf.compat.v2.summary.histogram(
              name="advantages",
              data=advantages,
              step=self.train_step_counter,
          )
  
      # Return TD-Lambda returns if both use_td_lambda_return and use_gae.
      if self._use_td_lambda_return:
          if not self._use_gae:
              logging.warning(
                  "use_td_lambda_return was True, but use_gae was "
                  "False. Using Monte Carlo return."
              )
          else:
              returns = tf.add(
                  advantages,
                  value_preds[:, :-1],
                  name="td_lambda_returns",
              )
  
      return returns, advantages
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Optional, Text, Tuple
  from absl import logging
  import gin
  from six.moves import range
  from six.moves import zip
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.agents.ppo import ppo_policy
  from tf_agents.agents.ppo import ppo_utils
  from tf_agents.networks import network
  from tf_agents.policies import greedy_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
  from tf_agents.utils import nest_utils
  from tf_agents.utils import object_identity
  from tf_agents.utils import tensor_normalizer
  from tf_agents.utils import value_ops
label_desc: Use image node_id 6 for calling the PPOAgent obj's underlying member method code with example usage: obj.compute_return_and_advantage(next_time_steps, value_preds) and returns: returns,  advantages