code: def policy_gradient_loss(
      policies,
      actions,
      action_values,
      policy_vars=None,
      name="policy_gradient_loss",
  ):
      """Computes policy gradient losses for a batch of trajectories.
  
      This wraps `policy_gradient` to accept a possibly nested array of `policies`
      and `actions` in order to allow for multiple action distribution types or
      independent multivariate distributions if not directly available. It also sums
      up losses along the time dimension, and is more restrictive about shapes,
      assuming a [T, B] layout for the `batch_shape` of the policies and a
      concatenate(`[T, B]`, `event_shape` of the policies) shape for the actions.
  
      Args:
        policies: A (possibly nested structure of) distribution(s) supporting
            `batch_shape` and `event_shape` properties along with a `log_prob`
            method (e.g. an instance of `tfp.distributions.Distribution`),
            with `batch_shape` equal to `[T, B]`.
        actions: A (possibly nested structure of) N-D Tensor(s) with shape
            `[T, B, ...]` where the final dimensions are the `event_shape` of the
            corresponding distribution in the nested structure (the shape can be
            just `[T, B]` if the `event_shape` is scalar).
        action_values: Tensor of shape `[T, B]` containing an estimate of the value
            of the selected `actions`.
        policy_vars: An optional (possibly nested structure of) iterable(s) of
            Tensors used by `policies`. If provided is used in scope checks.
        name: Customises the name_scope for this op.
  
      Returns:
        loss: Tensor of shape `[B]` containing the total loss for each sequence
        in the batch. Differentiable w.r.t `policy_logits` only.
      """
      actions = nest.flatten(actions)
      if policy_vars:
          policy_vars = nest.flatten_up_to(policies, policy_vars)
      else:
          policy_vars = [list()] * len(actions)
      policies = nest.flatten(policies)
  
      # Check happens after flatten so that we can be more flexible on nest
      # structures. This is equivalent to asserting that `len(policies) ==
      # len(actions)`, which is sufficient for what we're doing here.
      nest.assert_same_structure(policies, actions)
  
      for policies_, actions_ in zip(policies, actions):
          policies_.batch_shape.assert_has_rank(2)
          actions_.get_shape().assert_is_compatible_with(
              policies_.batch_shape.concatenate(policies_.event_shape)
          )
  
      scoped_values = policy_vars + actions + [action_values]
      with tf.name_scope(name, values=scoped_values):
          # Loss for the policy gradient. Doesn't push additional gradients through
          # the action_values.
          policy_gradient_loss_sequence = tf.add_n(
              [
                  policy_gradient(
                      policies_, actions_, action_values, pvars
                  )
                  for policies_, actions_, pvars in zip(
                      policies, actions, policy_vars
                  )
              ]
          )
  
          return tf.reduce_sum(
              policy_gradient_loss_sequence,
              axis=[0],
              name="policy_gradient_loss",
          )
  
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 2 for calling a global function with example usage: policy_gradient_loss(policies, actions, action_values, policy_vars, name) and returns: tf