code: def discrete_policy_gradient_loss(
      policy_logits,
      actions,
      action_values,
      name="discrete_policy_gradient_loss",
  ):
      """Computes discrete policy gradient losses for a batch of trajectories.
  
      This wraps `discrete_policy_gradient` to accept a possibly nested array of
      `policy_logits` and `actions` in order to allow for multiple discrete actions.
      It also sums up losses along the time dimension, and is more restrictive about
      shapes, assuming a [T, B] layout.
  
      Args:
        policy_logits: A (possibly nested structure of) Tensor(s) of shape
            `[T, B, num_actions]` containing uncentered log-probabilities.
        actions: A (possibly nested structure of) Tensor(s) of shape
            `[T, B]` and integer type, containing indices for the selected actions.
        action_values: Tensor of shape `[T, B]`
            containing an estimate of the value of the selected `actions`, see
            `discrete_policy_gradient`.
        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.
      """
      policy_logits = nest.flatten(policy_logits)
      actions = nest.flatten(actions)
  
      # Check happens after flatten so that we can be more flexible on
      # nest structures. This is equivalent to asserting that
      # `len(policy_logits) == len(actions)`, which is sufficient for what we're
      # doing here. In particular, it means that we can allow one argument to be
      # a tensor, while the other one to be a single-element tensor iterable.
      nest.assert_same_structure(policy_logits, actions)
      for scalar_policy_logits in policy_logits:
          scalar_policy_logits.get_shape().assert_has_rank(3)
      for scalar_actions in actions:
          scalar_actions.get_shape().assert_has_rank(2)
  
      scoped_values = policy_logits + 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(
              [
                  discrete_policy_gradient(
                      scalar_policy_logits,
                      scalar_actions,
                      action_values,
                  )
                  for scalar_policy_logits, scalar_actions in zip(
                      policy_logits, actions
                  )
              ]
          )
  
          return tf.reduce_sum(
              policy_gradient_loss_sequence,
              axis=[0],
              name="policy_gradient_loss",
          )
  
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 4 for calling a global function with example usage: discrete_policy_gradient_loss(policy_logits, actions, action_values, name) and returns: tf