code: def policy_entropy_loss(
      policies,
      policy_vars=None,
      scale_op=None,
      name="policy_entropy_loss",
  ):
      """Calculates entropy 'loss' for policies represented by a distributions.
  
      Given a (possible nested structure of) batch(es) of policies, this
      calculates the total entropy and corrects the sign so that minimizing the
      resulting loss op is equivalent to increasing entropy in the batch.
  
      This function accepts a nested structure of `policies` in order to allow for
      multiple distribution types or for multiple action dimensions in the case
      where there is no corresponding mutivariate form for available for a given
      univariate distribution. In this case, the loss is `sum_i(H(p_i, p_i))`
      where `p_i` are members of the `policies` nest. It can be shown that this is
      equivalent to calculating the entropy loss on the Cartesian product space
      over all the action dimensions, if the sampled actions are independent.
  
      The entropy loss is optionally scaled by some function of the policies.
      E.g. for Categorical distributions there exists such a scaling which maps
      the entropy loss into the range `[-1, 0]` in order to make it invariant to
      the size of the action space - specifically one can divide the loss by
      `sum_i(log(A_i))` where `A_i` is the number of categories in the i'th
      Categorical distribution in the `policies` nest).
  
      Args:
        policies: A (possibly nested structure of) batch distribution(s)
            supporting an `entropy` method that returns an N-D Tensor with shape
            equal to the `batch_shape` of the distribution, e.g. an instance of
            `tfp.distributions.Distribution`.
        policy_vars: An optional (possibly nested structure of) iterable(s) of
            Tensors used by `policies`. If provided is used in scope checks.
        scale_op: An optional op that takes `policies` as its only argument and
            returns a scalar Tensor that is used to scale the entropy loss.
            E.g. for Diag(sigma) Gaussian policies dividing by the number of
            dimensions makes entropy loss invariant to the action space dimension.
        name: Optional, name of this op.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: a tensor containing the batch of losses, shape `[B1, B2, ...]`.
        * `extra`: a namedtuple with fields:
            * `entropy`: entropy of the policy, shape `[B1, B2, ...]`.
        where [B1, B2, ... ] == policy.batch_shape
      """
      flat_policy_vars = (
          nest.flatten(policy_vars) if policy_vars else list()
      )
      with tf.name_scope(name, values=flat_policy_vars):
          # We want a value that we can minimize along with other losses, and where
          # minimizing means driving the policy towards a uniform distribution over
          # the actions. We thus scale it by negative one so that it can be simply
          # added to other losses.
          scale = tf.constant(-1.0, dtype=tf.float32)
          if scale_op:
              scale *= scale_op(policies)
  
          policies = nest.flatten(policies)
          entropy = tf.add_n(
              [policy.entropy() for policy in policies], name="entropy"
          )
          loss = tf.multiply(scale, entropy, name="entropy_loss")
          return base_ops.LossOutput(loss, PolicyEntropyExtra(entropy))
  
file_path: trfl/trfl/policy_gradient_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from six.moves import zip
  import tensorflow.compat.v1
  import tree
  from trfl import base_ops
  from trfl import value_ops
label_desc: Use image node_id 3 for calling a global function with example usage: policy_entropy_loss(policies, policy_vars, scale_op, name) and returns: base_ops