code: def discrete_policy_gradient(
      policy_logits,
      actions,
      action_values,
      name="discrete_policy_gradient",
  ):
      """Computes a batch of discrete-action policy gradient losses.
  
      See notes by Silver et al here:
      http://www0.cs.ucl.ac.uk/staff/D.Silver/web/Teaching_files/pg.pdf
  
      From slide 41, denoting by `policy` the probability distribution with
      log-probabilities `policy_logit`:
      ```
      *   `action` should have been sampled according to `policy`.
      *   `action_value` can be any estimate of `Q^{policy}(s, a)`, potentially
          minus a baseline that doesn't depend on the action. This admits
          many possible algorithms:
          * `v_t` (Monte-Carlo return for time t) : REINFORCE
          * `Q^w(s, a)` : Q Actor-Critic
          * `v_t - V(s)` : Monte-Carlo Advantage Actor-Critic
          * `A^{GAE(gamma, lambda)}` : Generalized Avantage Actor Critic
          * + many more.
      ```
  
      Gradients for this op are only defined with respect to the `policy_logits`,
      not `actions` or `action_values`.
  
      This op supports multiple batch dimensions. The first N >= 1 dimensions of
      each input/output tensor index into independent values. All tensors must
      having matching sizes for each batch dimension.
  
      Args:
        policy_logits: (N+1)-D Tensor of shape
            `[batch_size_1, ..., batch_size_N, num_actions]` containing uncentered
            log-probabilities.
        actions: N-D Tensor of shape `[batch_size_1, ..., batch_size_N]` and integer
            type, containing indices for the selected actions.
        action_values: N-D Tensor of shape `[batch_size_1, ..., batch_size_N]`
            containing an estimate of the value of the selected `actions`.
        name: Customises the name_scope for this op.
  
      Returns:
        loss: N-D Tensor of shape `[batch_size_1, ..., batch_size_N]` containing the
            loss. Differentiable w.r.t `policy_logits` only.
  
      Raises:
        ValueError: If the batch dimensions of `policy_logits` and `action_values`
            do not match.
      """
      with tf.name_scope(
          name, values=[policy_logits, actions, action_values]
      ):
          cross_entropy = (
              tf.nn.sparse_softmax_cross_entropy_with_logits(
                  labels=actions, logits=policy_logits
              )
          )
          action_values = tf.stop_gradient(action_values)
          # Prevent accidental broadcasting if possible at construction time.
          action_values.get_shape().assert_is_compatible_with(
              cross_entropy.get_shape()
          )
          return tf.multiply(cross_entropy, action_values)
  
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 3 for calling a global function with example usage: discrete_policy_gradient(policy_logits, actions, action_values, name) and returns: tf