class_docstr: A PPO Agent.
class_docstr_tok: ["A","PPO","Agent","."]
docstr: Compute the loss and create optimization op for one training epoch.
  
  All tensors should have a single batch dimension.
  
  Args:
    time_steps: A minibatch of TimeStep tuples.
    actions: A minibatch of actions.
    act_log_probs: A minibatch of action probabilities (probability under the
      sampling policy).
    returns: A minibatch of per-timestep returns.
    normalized_advantages: A minibatch of normalized per-timestep advantages.
    action_distribution_parameters: Parameters of data-collecting action
      distribution. Needed for KL computation.
    weights: Optional scalar or element-wise (per-batch-entry) importance
      weights.  Includes a mask for invalid timesteps.
    train_step: A train_step variable to increment for each train step.
      Typically the global_step.
    debug_summaries: True if debug summaries should be created.
    old_value_predictions: (Optional) The saved value predictions, used for
      calculating the value estimation loss when value clipping is performed.
    training: Whether this loss is being used for training.
  
  Returns:
    A tf_agent.LossInfo named tuple with the total_loss and all intermediate
      losses in the extra field contained in a PPOLossInfo named tuple.
docstr_tok: ["Compute","the","loss","and","create","optimization","op","for","one","training","epoch",".","All","tensors","should","have","a","single","batch","dimension",".","Args",":","time_steps",":","A","minibatch","of","TimeStep","tuples",".","actions",":","A","minibatch","of","actions",".","act_log_probs",":","A","minibatch","of","action","probabilities","(","probability","under","the","sampling","policy",")",".","returns",":","A","minibatch","of","per-timestep","returns",".","normalized_advantages",":","A","minibatch","of","normalized","per-timestep","advantages",".","action_distribution_parameters",":","Parameters","of","data-collecting","action","distribution",".","Needed","for","KL","computation",".","weights",":","Optional","scalar","or","element-wise","(","per-batch-entry",")","importance","weights",".","Includes","a","mask","for","invalid","timesteps",".","train_step",":","A","train_step","variable","to","increment","for","each","train","step",".","Typically","the","global_step",".","debug_summaries",":","True","if","debug","summaries","should","be","created",".","old_value_predictions",":","(","Optional",")","The","saved","value","predictions",",","used","for","calculating","the","value","estimation","loss","when","value","clipping","is","performed",".","training",":","Whether","this","loss","is","being","used","for","training",".","Returns",":","A","tf_agent.LossInfo","named","tuple","with","the","total_loss","and","all","intermediate","losses","in","the","extra","field","contained","in","a","PPOLossInfo","named","tuple","."]
code: def get_loss(
      self,
      time_steps: ts.TimeStep,
      actions: types.NestedTensorSpec,
      act_log_probs: types.Tensor,
      returns: types.Tensor,
      normalized_advantages: types.Tensor,
      action_distribution_parameters: types.NestedTensor,
      weights: types.Tensor,
      train_step: tf.Variable,
      debug_summaries: bool,
      old_value_predictions: Optional[types.Tensor] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Compute the loss and create optimization op for one training epoch.
  
      All tensors should have a single batch dimension.
  
      Args:
        time_steps: A minibatch of TimeStep tuples.
        actions: A minibatch of actions.
        act_log_probs: A minibatch of action probabilities (probability under the
          sampling policy).
        returns: A minibatch of per-timestep returns.
        normalized_advantages: A minibatch of normalized per-timestep advantages.
        action_distribution_parameters: Parameters of data-collecting action
          distribution. Needed for KL computation.
        weights: Optional scalar or element-wise (per-batch-entry) importance
          weights.  Includes a mask for invalid timesteps.
        train_step: A train_step variable to increment for each train step.
          Typically the global_step.
        debug_summaries: True if debug summaries should be created.
        old_value_predictions: (Optional) The saved value predictions, used for
          calculating the value estimation loss when value clipping is performed.
        training: Whether this loss is being used for training.
  
      Returns:
        A tf_agent.LossInfo named tuple with the total_loss and all intermediate
          losses in the extra field contained in a PPOLossInfo named tuple.
      """
      # Evaluate the current policy on timesteps.
  
      # batch_size from time_steps
      batch_size = nest_utils.get_outer_shape(
          time_steps, self._time_step_spec
      )[0]
      policy_state = self._collect_policy.get_initial_state(batch_size)
      # We must use _distribution because the distribution API doesn't pass down
      # the training= kwarg.
      distribution_step = self._collect_policy._distribution(  # pylint: disable=protected-access
          time_steps, policy_state, training=training
      )
      # TODO(eholly): Rename policy distributions to something clear and uniform.
      current_policy_distribution = distribution_step.action
  
      # Call all loss functions and add all loss values.
      value_estimation_loss = self.value_estimation_loss(
          time_steps=time_steps,
          returns=returns,
          old_value_predictions=old_value_predictions,
          weights=weights,
          debug_summaries=debug_summaries,
          training=training,
      )
      policy_gradient_loss = self.policy_gradient_loss(
          time_steps,
          actions,
          tf.stop_gradient(act_log_probs),
          tf.stop_gradient(normalized_advantages),
          current_policy_distribution,
          weights,
          debug_summaries=debug_summaries,
      )
  
      if (
          self._policy_l2_reg > 0.0
          or self._value_function_l2_reg > 0.0
          or self._shared_vars_l2_reg > 0.0
      ):
          l2_regularization_loss = self.l2_regularization_loss(
              debug_summaries
          )
      else:
          l2_regularization_loss = tf.zeros_like(policy_gradient_loss)
  
      with tf.name_scope("entropy_regularization"):
          outer_rank = nest_utils.get_outer_rank(
              time_steps, self.time_step_spec
          )
          entropy = common.entropy(
              current_policy_distribution, self.action_spec, outer_rank
          )
          if entropy is None:
              entropy = tf.zeros_like(weights, tf.float32)
          else:
              entropy = tf.cast(entropy, tf.float32)
      if self._entropy_regularization > 0.0:
          entropy_regularization_loss = (
              self.entropy_regularization_loss(
                  time_steps, entropy, weights, debug_summaries
              )
          )
      else:
          entropy_regularization_loss = tf.zeros_like(
              policy_gradient_loss
          )
  
      with tf.name_scope("Losses/"):
          tf.compat.v2.summary.scalar(
              name="entropy",
              data=tf.reduce_mean(entropy * weights),
              step=self.train_step_counter,
          )
  
      # TODO(b/161365079): Move this logic to PPOKLPenaltyAgent.
      if (
          self._initial_adaptive_kl_beta == 0
          and self._kl_cutoff_factor == 0
      ):
          kl_penalty_loss = tf.zeros_like(policy_gradient_loss)
      else:
          kl_penalty_loss = self.kl_penalty_loss(
              time_steps,
              action_distribution_parameters,
              current_policy_distribution,
              weights,
              debug_summaries,
          )
  
      total_loss = (
          policy_gradient_loss
          + value_estimation_loss
          + l2_regularization_loss
          + entropy_regularization_loss
          + kl_penalty_loss
      )
  
      return tf_agent.LossInfo(
          total_loss,
          PPOLossInfo(
              policy_gradient_loss=policy_gradient_loss,
              value_estimation_loss=value_estimation_loss,
              l2_regularization_loss=l2_regularization_loss,
              entropy_regularization_loss=entropy_regularization_loss,
              kl_penalty_loss=kl_penalty_loss,
              clip_fraction=self._clip_fraction,
          ),
      )
  
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 5 for calling the PPOAgent obj's underlying member method code with example usage: obj.get_loss(time_steps, actions, act_log_probs, returns, normalized_advantages, action_distribution_parameters, weights, train_step, debug_summaries, old_value_predictions, training) and returns: tf_agent