docs: def _policy_loss(
        self, old_policy, policy, action, advantage, length):
      """Compute the policy loss composed of multiple components.
  
      1. The policy gradient loss is importance sampled from the data-collecting
         policy at the beginning of training.
      2. The second term is a KL penalty between the policy at the beginning of
         training and the current policy.
      3. Additionally, if this KL already changed more than twice the target
         amount, we activate a strong penalty discouraging further divergence.
  
      Args:
        old_policy: Action distribution of the behavioral policy.
        policy: Sequences of distribution params of the current policy.
        action: Sequences of actions.
        advantage: Sequences of advantages.
        length: Batch of sequence lengths.
  
      Returns:
        Tuple of loss tensor and summary tensor.
      """
      with tf.name_scope('policy_loss'):
        kl = tf.contrib.distributions.kl_divergence(old_policy, policy)
        # Infinite values in the KL, even for padding frames that we mask out,
        # cause NaN gradients since TensorFlow computes gradients with respect to
        # the whole input tensor.
        kl = tf.check_numerics(kl, 'kl')
        kl = tf.reduce_mean(self._mask(kl, length), 1)
        policy_gradient = tf.exp(
            policy.log_prob(action) - old_policy.log_prob(action))
        surrogate_loss = -tf.reduce_mean(self._mask(
            policy_gradient * tf.stop_gradient(advantage), length), 1)
        surrogate_loss = tf.check_numerics(surrogate_loss, 'surrogate_loss')
        kl_penalty = self._penalty * kl
        cutoff_threshold = self._config.kl_target * self._config.kl_cutoff_factor
        cutoff_count = tf.reduce_sum(
            tf.cast(kl > cutoff_threshold, tf.int32))
        with tf.control_dependencies([tf.cond(
            cutoff_count > 0,
            lambda: tf.Print(0, [cutoff_count], 'kl cutoff! '), int)]):
          kl_cutoff = (
              self._config.kl_cutoff_coef *
              tf.cast(kl > cutoff_threshold, tf.float32) *
              (kl - cutoff_threshold) ** 2)
        policy_loss = surrogate_loss + kl_penalty + kl_cutoff
        entropy = tf.reduce_mean(policy.entropy(), axis=1)
        if self._config.entropy_regularization:
          policy_loss -= self._config.entropy_regularization * entropy
        summary = tf.summary.merge([
            tf.summary.histogram('entropy', entropy),
            tf.summary.histogram('kl', kl),
            tf.summary.histogram('surrogate_loss', surrogate_loss),
            tf.summary.histogram('kl_penalty', kl_penalty),
            tf.summary.histogram('kl_cutoff', kl_cutoff),
            tf.summary.histogram('kl_penalty_combined', kl_penalty + kl_cutoff),
            tf.summary.histogram('policy_loss', policy_loss),
            tf.summary.scalar('avg_surr_loss', tf.reduce_mean(surrogate_loss)),
            tf.summary.scalar('avg_kl_penalty', tf.reduce_mean(kl_penalty)),
            tf.summary.scalar('avg_policy_loss', tf.reduce_mean(policy_loss))])
        policy_loss = tf.reduce_mean(policy_loss, 0)
        return tf.check_numerics(policy_loss, 'policy_loss'), summary