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