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