code: def policy_gradient(
      policies,
      actions,
      action_values,
      policy_vars=None,
      name="policy_gradient",
  ):
      """Computes policy gradient losses for a batch of trajectories.
  
      See `policy_gradient_loss` for more information on expected inputs and usage.
  
      Args:
        policies: A distribution over a batch supporting a `log_prob` method, e.g.
            an instance of `tfp.distributions.Distribution`. For example, for
            a diagonal gaussian policy:
            `policies = tfp.distributions.MultivariateNormalDiag(mus, sigmas)`
        actions: An action batch Tensor used as the argument for `log_prob`. Has
            shape equal to the batch shape of the policies concatenated with the
            event shape of the policies (which may be scalar, in which case
            concatenation leaves shape just equal to batch shape).
        action_values: A Tensor containing estimates of the values of the `actions`.
            Has shape equal to the batch shape of the policies.
        policy_vars: An optional iterable of Tensors used by `policies`. If provided
            is used in scope checks. For the multivariate normal example above this
            would be `[mus, sigmas]`.
        name: Customises the name_scope for this op.
  
      Returns:
        loss: Tensor with same shape as `actions` containing the total loss for each
            element in the batch. Differentiable w.r.t the variables in `policies`
            only.
      """
      policy_vars = list(policy_vars) if policy_vars else list()
      with tf.name_scope(
          values=policy_vars + [actions, action_values], name=name
      ):
          actions = tf.stop_gradient(actions)
          action_values = tf.stop_gradient(action_values)
          log_prob_actions = policies.log_prob(actions)
          # Prevent accidental broadcasting if possible at construction time.
          action_values.get_shape().assert_is_compatible_with(
              log_prob_actions.get_shape()
          )
          return -tf.multiply(log_prob_actions, action_values)
  
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 1 for calling a global function with example usage: policy_gradient(policies, actions, action_values, policy_vars, name) and returns: unknown