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))
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]`.
  
trfl/discrete_policy_gradient_ops.py