def sequence_advantage_actor_critic_loss(
    policy_logits, baseline_values, actions, rewards,
    pcontinues, bootstrap_value, lambda_=1, entropy_cost=None,
    baseline_cost=1, normalise_entropy=False,
    name="SequenceAdvantageActorCriticLoss"):
  """Calculates the loss for an A2C update along a batch of trajectories.

  Technically A2C is the special case where lambda=1; for general lambda
  this is the loss for Generalized Advantage Estimation (GAE), modulo chunking
  behaviour if passing chunks of episodes (see `generalized_lambda_returns` for
  more detail).

  Note: This function takes policy _logits_ as input, not the log-policy like
  `learning.deepmind.lua.rl.learners.Reinforce` does.

  This loss jointly learns the policy and the baseline. Therefore, gradients
  for this loss flow through each tensor in `policy_logits` and
  `baseline_values`, but no other input tensors. The policy is learnt with the
  advantage actor-critic loss, plus an optional entropy term. The baseline is
  regressed towards the n-step bootstrapped returns given by the
  reward/pcontinue sequence.  The `baseline_cost` parameter scales the
  gradients w.r.t the baseline relative to the policy gradient. i.e:
  `d(loss) / d(baseline) = baseline_cost * (n_step_return - baseline)`.

  `rewards` and `pcontinues` are the sequences of data taken directly from the
  environment, possibly modulated by a discount. `baseline_values` are the
  sequences of (typically learnt) estimates of the values of the states
  visited along a batch of trajectories as observed by the agent given the
  sequences of one or more actions sampled from the `policy_logits`.

  The sequences in the tensors should be aligned such that an agent in a state
  with value `V` that takes an action `a` transitions into another state
  with value `V'`, receiving reward `r` and pcontinue `p`. Then `V`, `a`, `r`
  and `p` are all at the same index `i` in the corresponding tensors. `V'` is
  at index `i+1`, or in the `bootstrap_value` tensor if `i == T`.

  This function accepts a nested array of `policy_logits` and `actions` in order
  to allow for multidimensional discrete action spaces. In this case, the loss
  is given by `sum_i(loss(p_i, a_i))` where `p_i` are members of the
  `policy_logits` nest, and `a_i` are members of the `actions` nest.
  We assume that a single baseline is used across all action dimensions for
  each timestep.

  Args:
    policy_logits: A (possibly nested structure of) 3-D Tensor(s) with shape
        `[T, B, num_actions]` and possibly different dimension `num_actions`.
    baseline_values: 2-D Tensor containing an estimate of state values `[T, B]`.
    actions: A (possibly nested structure of) 2-D Tensor(s) with shape
        `[T, B]` and integer type.
    rewards: 2-D Tensor with shape `[T, B]`.
    pcontinues: 2-D Tensor with shape `[T, B]`.
    bootstrap_value: 1-D Tensor with shape `[B]`.
    lambda_: an optional scalar or 2-D Tensor with shape `[T, B]` for
        Generalised Advantage Estimation as per
        https://arxiv.org/abs/1506.02438.
    entropy_cost: optional scalar cost that pushes the policy to have high
        entropy, larger values cause higher entropies.
    baseline_cost: scalar cost that scales the derivatives of the baseline
        relative to the policy gradient.
    normalise_entropy: if True, the entropy loss is normalised to the range
        `[-1, 0]` by dividing by the log number of actions. This makes it more
        invariant to the size of the action space. Default is False.
    name: Customises the name_scope for this op.

  Returns:
    A namedtuple with fields:

    * `loss`: a tensor containing the total loss, shape `[B]`.
    * `extra`: a namedtuple with fields:
        * `entropy`: total loss per sequence, shape `[B]`.
        * `entropy_loss`: scaled entropy loss per sequence, shape `[B]`.
        * `baseline_loss`: scaled baseline loss per sequence, shape `[B]`.
        * `policy_gradient_loss`: policy gradient loss per sequence,
            shape `[B]`.
        * `advantages`: advantange estimates per timestep, shape `[T, B]`.
        * `discounted_returns`: discounted returns per timestep,
            shape `[T, B]`.
  """
  scoped_values = (nest.flatten(policy_logits) + nest.flatten(actions) +
                   [baseline_values, rewards, pcontinues, bootstrap_value])
  with tf.name_scope(name, values=scoped_values):
    # Loss for the baseline, summed over the time dimension.
    baseline_loss_td, td_lambda = value_ops.td_lambda(
        baseline_values, rewards, pcontinues, bootstrap_value, lambda_)

    # The TD error provides an estimate of the advantages of the actions.
    advantages = td_lambda.temporal_differences
    baseline_loss = tf.multiply(
        tf.convert_to_tensor(baseline_cost, dtype=tf.float32),
        baseline_loss_td,
        name="baseline_loss")

    # Loss for the policy. Doesn't push additional gradients through
    # the advantages.
    policy_gradient_loss = discrete_policy_gradient_loss(
        policy_logits, actions, advantages, name="policy_gradient_loss")

    total_loss = tf.add(policy_gradient_loss, baseline_loss, name="total_loss")

    if entropy_cost is not None:
      entropy_loss_op, policy_entropy = discrete_policy_entropy_loss(
          policy_logits, normalise=normalise_entropy)  # [T,B].
      entropy = tf.reduce_sum(
          policy_entropy.entropy, axis=0, name="entropy")  # [B].
      entropy_loss = tf.multiply(
          tf.convert_to_tensor(entropy_cost, dtype=tf.float32),
          tf.reduce_sum(entropy_loss_op, axis=0),
          name="scaled_entropy_loss")  # [B].
      total_loss = tf.add(total_loss, entropy_loss,
                          name="total_loss_with_entropy")
    else:
      entropy = None
      entropy_loss = None

    extra = SequenceAdvantageActorCriticExtra(
        entropy=entropy, entropy_loss=entropy_loss,
        baseline_loss=baseline_loss,
        policy_gradient_loss=policy_gradient_loss,
        advantages=advantages,
        discounted_returns=td_lambda.discounted_returns)

    return base_ops.LossOutput(total_loss, extra)
Calculates the loss for an A2C update along a batch of trajectories.

  Technically A2C is the special case where lambda=1; for general lambda
  this is the loss for Generalized Advantage Estimation (GAE), modulo chunking
  behaviour if passing chunks of episodes (see `generalized_lambda_returns` for
  more detail).

  Note: This function takes policy _logits_ as input, not the log-policy like
  `learning.deepmind.lua.rl.learners.Reinforce` does.

  This loss jointly learns the policy and the baseline. Therefore, gradients
  for this loss flow through each tensor in `policy_logits` and
  `baseline_values`, but no other input tensors. The policy is learnt with the
  advantage actor-critic loss, plus an optional entropy term. The baseline is
  regressed towards the n-step bootstrapped returns given by the
  reward/pcontinue sequence.  The `baseline_cost` parameter scales the
  gradients w.r.t the baseline relative to the policy gradient. i.e:
  `d(loss) / d(baseline) = baseline_cost * (n_step_return - baseline)`.

  `rewards` and `pcontinues` are the sequences of data taken directly from the
  environment, possibly modulated by a discount. `baseline_values` are the
  sequences of (typically learnt) estimates of the values of the states
  visited along a batch of trajectories as observed by the agent given the
  sequences of one or more actions sampled from the `policy_logits`.

  The sequences in the tensors should be aligned such that an agent in a state
  with value `V` that takes an action `a` transitions into another state
  with value `V'`, receiving reward `r` and pcontinue `p`. Then `V`, `a`, `r`
  and `p` are all at the same index `i` in the corresponding tensors. `V'` is
  at index `i+1`, or in the `bootstrap_value` tensor if `i == T`.

  This function accepts a nested array of `policy_logits` and `actions` in order
  to allow for multidimensional discrete action spaces. In this case, the loss
  is given by `sum_i(loss(p_i, a_i))` where `p_i` are members of the
  `policy_logits` nest, and `a_i` are members of the `actions` nest.
  We assume that a single baseline is used across all action dimensions for
  each timestep.

  Args:
    policy_logits: A (possibly nested structure of) 3-D Tensor(s) with shape
        `[T, B, num_actions]` and possibly different dimension `num_actions`.
    baseline_values: 2-D Tensor containing an estimate of state values `[T, B]`.
    actions: A (possibly nested structure of) 2-D Tensor(s) with shape
        `[T, B]` and integer type.
    rewards: 2-D Tensor with shape `[T, B]`.
    pcontinues: 2-D Tensor with shape `[T, B]`.
    bootstrap_value: 1-D Tensor with shape `[B]`.
    lambda_: an optional scalar or 2-D Tensor with shape `[T, B]` for
        Generalised Advantage Estimation as per
        https://arxiv.org/abs/1506.02438.
    entropy_cost: optional scalar cost that pushes the policy to have high
        entropy, larger values cause higher entropies.
    baseline_cost: scalar cost that scales the derivatives of the baseline
        relative to the policy gradient.
    normalise_entropy: if True, the entropy loss is normalised to the range
        `[-1, 0]` by dividing by the log number of actions. This makes it more
        invariant to the size of the action space. Default is False.
    name: Customises the name_scope for this op.

  Returns:
    A namedtuple with fields:

    * `loss`: a tensor containing the total loss, shape `[B]`.
    * `extra`: a namedtuple with fields:
        * `entropy`: total loss per sequence, shape `[B]`.
        * `entropy_loss`: scaled entropy loss per sequence, shape `[B]`.
        * `baseline_loss`: scaled baseline loss per sequence, shape `[B]`.
        * `policy_gradient_loss`: policy gradient loss per sequence,
            shape `[B]`.
        * `advantages`: advantange estimates per timestep, shape `[T, B]`.
        * `discounted_returns`: discounted returns per timestep,
            shape `[T, B]`.
  
trfl/discrete_policy_gradient_ops.py