code: def discrete_policy_entropy_loss(
      policy_logits,
      normalise=False,
      name="discrete_policy_entropy_loss",
  ):
      """Computes the entropy 'loss' for a batch of policy logits.
  
      Given a batch of policy logits, calculates the entropy and corrects the sign
      so that minimizing the resulting loss op is equivalent to increasing entropy
      in the batch. This loss is optionally normalised to the range `[-1, 0]` by
      dividing by the log number of actions. This makes it more invariant to the
      size of the action space.
  
      This function accepts a nested array of `policy_logits` in order
      to allow for multiple discrete actions. In this case, the loss is given by
      `-sum_i(H(p_i))` where `p_i` are members of the `policy_logits` nest and
      H is the Shannon entropy.
  
      Args:
        policy_logits: A (possibly nested structure of) (N+1)-D Tensor(s) with
            shape `[..., A]`,  representing the log-probabilities of a set of
            Categorical distributions, where `...` represents at least one
            dimension (e.g., batch, sequence), and `A` is the number of discrete
            actions (which need not be identical across all tensors).
            Does not need to be centered.
        normalise: If True, divide the loss by the `sum_i(log(A_i))` where `A_i`
            is the number of actions for the i'th tensor in the `policy_logits`
            nest. Default is False.
        name: Optional, name of this op.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: Entropy 'loss', shape `[B]`.
        * `extra`: a namedtuple with fields:
            * `entropy`: Entropy of the policy, shape `[B]`.
      """
      policy_logits = nest.flatten(policy_logits)
  
      with tf.name_scope(name, values=policy_logits):
          entropy = tf.add_n(
              [
                  tf.reduce_sum(
                      -tf.nn.softmax(scalar_policy_logits)
                      * tf.nn.log_softmax(scalar_policy_logits),
                      axis=-1,
                  )
                  for scalar_policy_logits in policy_logits
              ],
              name="entropy",
          )
          # 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 normalise:
              num_actions = [
                  tf.to_float(tf.shape(scalar_policy_logits)[-1])
                  for scalar_policy_logits in policy_logits
              ]
              scale /= tf.reduce_sum(tf.log(tf.stack(num_actions)))
          loss = tf.multiply(scale, entropy, name="entropy_loss")
  
      return base_ops.LossOutput(
          loss, DiscretePolicyEntropyExtra(entropy)
      )
  
file_path: trfl/trfl/discrete_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 1 for calling a global function with example usage: discrete_policy_entropy_loss(policy_logits, normalise, name) and returns: base_ops