code: def transformed_n_step_loss(
      qs: tf.Tensor,
      targnet_qs: tf.Tensor,
      actions: tf.Tensor,
      rewards: tf.Tensor,
      pcontinues: tf.Tensor,
      target_policy_probs: tf.Tensor,
      bootstrap_n: int,
      stop_targnet_gradients: bool = True,
      name: str = "transformed_n_step_loss",
  ) -> trfl.base_ops.LossOutput:
      """Helper function for computing transformed loss on sequences.
  
      Args:
        qs: 3-D tensor corresponding to the Q-values to be learned. Shape is [T+1,
          B, A].
        targnet_qs: Like `qs`, but in the target network setting, these values
          should be computed by the target network. Shape is [T+1, B, A].
        actions: 2-D tensor holding the indices of actions executed during the
          transition that corresponds to each major index. Shape is [T+1, B].
        rewards: 2-D tensor holding rewards received during the transition that
          corresponds to each major index. Shape is [T, B].
        pcontinues: 2-D tensor holding pcontinue values received during the
          transition that corresponds to each major index. Shape is [T, B].
        target_policy_probs: 3-D tensor holding per-action policy probabilities for
          the states encountered just before taking the transitions that correspond
          to each major index, according to the target policy (i.e. the policy we
          wish to learn). For standard Q-learning the probabilities should form a
          one-hot vector over actions where the nonzero index corresponds to the max
          Q. Shape is [T+1, B, A].
        bootstrap_n: Transition length for N-step bootstrapping.
        stop_targnet_gradients: `bool` indicating whether to apply tf.stop_gradients
          to the target values. This should usually be True.
        name: name to prefix ops created by this function.
  
      Returns:
        a tuple of:
        * `loss`: the transformed Q-learning loss summed over `T`.
        * `LossCoreExtra`: namedtuple containing the fields `targets` and `errors`.
      """
  
      with tf.name_scope(name):
          # Require correct tensor ranks---as long as we have shape information
          # available to check. If there isn't any, we print a warning.
          def check_rank(
              tensors: Iterable[tf.Tensor], ranks: Sequence[int]
          ):
              for i, (tensor, rank) in enumerate(zip(tensors, ranks)):
                  if tensor.get_shape():
                      trfl.assert_rank_and_shape_compatibility(
                          [tensor], rank
                      )
                  else:
                      raise ValueError(
                          f'Tensor "{tensor.name}", which was offered as transformed_n_step_loss'
                          f"parameter {i+1}, has no rank at construction time, so cannot verify"
                          f"that it has the necessary rank of {rank}"
                      )
  
          check_rank(
              [
                  qs,
                  targnet_qs,
                  actions,
                  rewards,
                  pcontinues,
                  target_policy_probs,
              ],
              [3, 3, 2, 2, 2, 3],
          )
  
          # Construct arguments to compute bootstrap target.
          a_tm1 = actions[:-1]  # (0:T) x B
          r_t, pcont_t = rewards, pcontinues  # (1:T+1) x B
          q_tm1 = qs[:-1]  # (0:T) x B x A
          target_policy_t = target_policy_probs[1:]  # (1:T+1) x B x A
          targnet_q_t = targnet_qs[1:]  # (1:T+1) x B x A
  
          bootstrap_value = tf.reduce_sum(
              target_policy_t * _signed_parabolic_tx(targnet_q_t), -1
          )
          target = _compute_n_step_sequence_targets(
              r_t=r_t,
              pcont_t=pcont_t,
              bootstrap_value=bootstrap_value,
              n=bootstrap_n,
          )
  
          if stop_targnet_gradients:
              target = tf.stop_gradient(target)
  
          # tx/inv_tx may result in numerical instabilities so mask any NaNs.
          finite_mask = tf.math.is_finite(target)
          target = tf.where(finite_mask, target, tf.zeros_like(target))
  
          qa_tm1 = trfl.batched_index(q_tm1, a_tm1)
          errors = qa_tm1 - _signed_hyperbolic_tx(target)
  
          # Only compute n-step errors w.r.t. finite targets.
          errors = tf.where(finite_mask, errors, tf.zeros_like(errors))
  
          # Sum over time dimension.
          loss = 0.5 * tf.reduce_sum(tf.square(errors), axis=0)
  
          return trfl.base_ops.LossOutput(
              loss, LossCoreExtra(targets=target, errors=errors)
          )
  
file_path: acme/acme/tf/losses/r2d2.py
imports: from typing import Iterable, NamedTuple, Sequence
  import tensorflow
  import trfl
label_desc: Use image node_id 1 for calling a global function with example usage: transformed_n_step_loss(qs, targnet_qs, actions, rewards, pcontinues, target_policy_probs, bootstrap_n, stop_targnet_gradients, name) and returns: trfl