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))
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
  
trfl/policy_gradient_ops.py