class_docstr: A PPO Agent.
class_docstr_tok: ["A","PPO","Agent","."]
docstr: Compute a loss that penalizes policy steps with high KL.
  
  Based on KL divergence from old (data-collection) policy to new (updated)
  policy.
  
  All tensors should have a single batch dimension.
  
  Args:
    time_steps: TimeStep tuples with observations for each timestep. Used for
      computing new action distributions.
    action_distribution_parameters: Action distribution params of the data
      collection policy, used for reconstruction old action distributions.
    current_policy_distribution: The policy distribution, evaluated on all
      time_steps.
    weights: Optional scalar or element-wise (per-batch-entry) importance
      weights.  Inlcudes a mask for invalid timesteps.
    debug_summaries: True if debug summaries should be created.
  
  Returns:
    kl_penalty_loss: The sum of a squared penalty for KL over a constant
      threshold, plus an adaptive penalty that encourages updates toward a
      target KL divergence.
docstr_tok: ["Compute","a","loss","that","penalizes","policy","steps","with","high","KL",".","Based","on","KL","divergence","from","old","(","data-collection",")","policy","to","new","(","updated",")","policy",".","All","tensors","should","have","a","single","batch","dimension",".","Args",":","time_steps",":","TimeStep","tuples","with","observations","for","each","timestep",".","Used","for","computing","new","action","distributions",".","action_distribution_parameters",":","Action","distribution","params","of","the","data","collection","policy",",","used","for","reconstruction","old","action","distributions",".","current_policy_distribution",":","The","policy","distribution",",","evaluated","on","all","time_steps",".","weights",":","Optional","scalar","or","element-wise","(","per-batch-entry",")","importance","weights",".","Inlcudes","a","mask","for","invalid","timesteps",".","debug_summaries",":","True","if","debug","summaries","should","be","created",".","Returns",":","kl_penalty_loss",":","The","sum","of","a","squared","penalty","for","KL","over","a","constant","threshold",",","plus","an","adaptive","penalty","that","encourages","updates","toward","a","target","KL","divergence","."]
code: def kl_penalty_loss(
      self,
      time_steps: ts.TimeStep,
      action_distribution_parameters: types.NestedTensor,
      current_policy_distribution: types.NestedDistribution,
      weights: types.Tensor,
      debug_summaries: bool = False,
  ) -> types.Tensor:
      """Compute a loss that penalizes policy steps with high KL.
  
      Based on KL divergence from old (data-collection) policy to new (updated)
      policy.
  
      All tensors should have a single batch dimension.
  
      Args:
        time_steps: TimeStep tuples with observations for each timestep. Used for
          computing new action distributions.
        action_distribution_parameters: Action distribution params of the data
          collection policy, used for reconstruction old action distributions.
        current_policy_distribution: The policy distribution, evaluated on all
          time_steps.
        weights: Optional scalar or element-wise (per-batch-entry) importance
          weights.  Inlcudes a mask for invalid timesteps.
        debug_summaries: True if debug summaries should be created.
  
      Returns:
        kl_penalty_loss: The sum of a squared penalty for KL over a constant
          threshold, plus an adaptive penalty that encourages updates toward a
          target KL divergence.
      """
      kl_divergence = self._kl_divergence(
          time_steps,
          action_distribution_parameters,
          current_policy_distribution,
      )
      kl_divergence *= weights
  
      # TODO(b/171573175): remove the condition once histograms are supported
      # on TPUs.
      if debug_summaries and not tf.config.list_logical_devices("TPU"):
          tf.compat.v2.summary.histogram(
              name="kl_divergence",
              data=kl_divergence,
              step=self.train_step_counter,
          )
  
      kl_cutoff_loss = self.kl_cutoff_loss(
          kl_divergence, debug_summaries
      )
      adaptive_kl_loss = self.adaptive_kl_loss(
          kl_divergence, debug_summaries
      )
      return tf.add(
          kl_cutoff_loss, adaptive_kl_loss, name="kl_penalty_loss"
      )
  
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 19 for calling the PPOAgent obj's underlying member method code with example usage: obj.kl_penalty_loss(time_steps, action_distribution_parameters, current_policy_distribution, weights, debug_summaries) and returns: tf