class_docstr: A PPO Agent.
class_docstr_tok: ["A","PPO","Agent","."]
docstr: Performs advantage calculation for the collected experience.
  
  Args:
    experience: A (batch of) experience in the form of a `Trajectory`. The
      structure of `experience` must match that of `self.collect_data_spec`.
      All tensors in `experience` must be shaped `[batch, time + 1, ...]` or
      [time + 1, ...]. The "+1" is needed as the last action from the set of
      trajectories cannot be used for training, as its advantage and returns
      are unknown.
  
  Returns:
    The processed experience which has normalized_advantages and returns
    filled in its policy info. The advantages and returns for the last
    transition are filled with 0s as they cannot be calculated.
docstr_tok: ["Performs","advantage","calculation","for","the","collected","experience",".","Args",":","experience",":","A","(","batch","of",")","experience","in","the","form","of","a","`","Trajectory","`",".","The","structure","of","`","experience","`","must","match","that","of","`","self.collect_data_spec","`",".","All","tensors","in","`","experience","`","must","be","shaped","`","[","batch",",","time","+","1",",","...","]","`","or","[","time","+","1",",","...","]",".","The","``","+1","''","is","needed","as","the","last","action","from","the","set","of","trajectories","can","not","be","used","for","training",",","as","its","advantage","and","returns","are","unknown",".","Returns",":","The","processed","experience","which","has","normalized_advantages","and","returns","filled","in","its","policy","info",".","The","advantages","and","returns","for","the","last","transition","are","filled","with","0s","as","they","can","not","be","calculated","."]
code: def _preprocess(self, experience):
      """Performs advantage calculation for the collected experience.
  
      Args:
        experience: A (batch of) experience in the form of a `Trajectory`. The
          structure of `experience` must match that of `self.collect_data_spec`.
          All tensors in `experience` must be shaped `[batch, time + 1, ...]` or
          [time + 1, ...]. The "+1" is needed as the last action from the set of
          trajectories cannot be used for training, as its advantage and returns
          are unknown.
  
      Returns:
        The processed experience which has normalized_advantages and returns
        filled in its policy info. The advantages and returns for the last
        transition are filled with 0s as they cannot be calculated.
      """
      # Try to be agnostic about the input type of experience before we call
      # to_transition() below.
      outer_rank = nest_utils.get_outer_rank(
          _get_discount(experience), self.collect_data_spec.discount
      )
  
      # Add 1 as the batch dimension for inputs that just have the time dimension,
      # as all utility functions below require the batch dimension.
      if outer_rank == 1:
          batched_experience = nest_utils.batch_nested_tensors(
              experience
          )
      else:
          batched_experience = experience
  
      # Get individual tensors from experience.
      num_steps = _get_discount(batched_experience).shape[1]
      if num_steps and num_steps <= 1:
          raise ValueError(
              "Experience used for advantage calculation must have >1 num_steps."
          )
  
      transition = self._collected_as_transition(batched_experience)
      time_steps, _, next_time_steps = transition
  
      # TODO(b/170680358): Decide if we will require Trajectory for preprocess,
      # or if we want to handle transitions here too; if so, what do we return?
      # Below we use batched_experience assuming it's a Trajectory, and return
      # a trajectory.
  
      # Compute the value predictions for states using the current value function.
      # To be used for return & advantage computation.
      batch_size = nest_utils.get_outer_shape(
          time_steps, self._time_step_spec
      )[0]
      if self._compute_value_and_advantage_in_train:
          value_state = self._collect_policy.get_initial_value_state(
              batch_size
          )
          value_preds, _ = self._collect_policy.apply_value_network(
              batched_experience.observation,
              batched_experience.step_type,
              value_state=value_state,
              training=False,
          )
          value_preds = tf.stop_gradient(value_preds)
      else:
          value_preds = batched_experience.policy_info[
              "value_prediction"
          ]
  
      new_policy_info = {
          "dist_params": batched_experience.policy_info["dist_params"],
          "value_prediction": value_preds,
      }
  
      # Add the calculated advantage and return into the input experience.
      returns, advantages = self.compute_return_and_advantage(
          next_time_steps, value_preds
      )
  
      # Pad returns and normalized_advantages in the time dimension so that the
      # time dimensions are aligned with the input experience's time dimension.
      # When the output trajectory gets sliced by trajectory.to_transition during
      # training, the padded last timesteps will be automatically dropped.
      last_transition_padding = tf.zeros(
          (batch_size, 1), dtype=tf.float32
      )
      new_policy_info["return"] = tf.concat(
          [returns, last_transition_padding], axis=1
      )
      new_policy_info["advantage"] = tf.concat(
          [advantages, last_transition_padding], axis=1
      )
  
      # Remove the batch dimension iff the input experience does not have it.
      if outer_rank == 1:
          new_policy_info = nest_utils.unbatch_nested_tensors(
              new_policy_info
          )
      # The input experience with its policy info filled with the calculated
      # advantages and returns for each action.
      return experience.replace(policy_info=new_policy_info)
  
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 7 for calling the PPOAgent obj's underlying member method code with example usage: obj._preprocess(experience) and returns: experience