code: 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)
  
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 2 for calling a global function with example usage: sequence_advantage_actor_critic_loss(policy_logits, baseline_values, actions, rewards, pcontinues, bootstrap_value, lambda_, entropy_cost, baseline_cost, normalise_entropy, name) and returns: base_ops