code: def retrace(
      lambda_,
      qs,
      targnet_qs,
      actions,
      rewards,
      pcontinues,
      target_policy_probs,
      behaviour_policy_probs,
      stop_targnet_gradients=True,
      name=None,
  ):
      """Retrace algorithm loss calculation op.
  
      Given a minibatch of temporally-contiguous sequences of Q values, policy
      probabilities, and various other typical RL algorithm inputs, this
      Op creates a subgraph that computes a loss according to the
      Retrace multi-step off-policy value learning algorithm. This Op supports the
      use of target networks, but does not require them.
  
      For more details of Retrace, refer to
      [the arXiv paper](http://arxiv.org/abs/1606.02647).
  
      In argument descriptions, `T` counts the number of transitions over which
      the Retrace loss is computed, and `B` is the minibatch size. Note that all
      tensor arguments list a first-dimension (time dimension) size of T+1;
      this is because in order to compute the loss over T timesteps, the
      algorithm must be aware of the values of many of its inputs at timesteps
      before and after each transition.
  
      All tensor arguments are indexed first by transition, with specific
      details of this indexing in the argument descriptions.
  
      Args:
        lambda_: Positive scalar value or 0-D `Tensor` controlling the degree to
          which future timesteps contribute to the loss computed at each
          transition.
        qs: 3-D tensor holding per-action Q-values for the states encountered
          just before taking the transitions that correspond to each major index.
          Since these values are the predicted values we wish to update (in other
          words, the values we intend to change as we learn), in a target network
          setting, these nearly always come from the "non-target" network, which
          we usually call the "learning network".
          Shape is `[(T+1), B, num_actions]`.
        targnet_qs: Like `qs`, but in the target network setting, these values
          should be computed by the target network. We use these values to
          compute multi-step error values for timesteps that follow the first
          timesteps in each sequence and sequence fragment we consider.
          Shape is `[(T+1), B, num_actions]`.
        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+1), B]`.
        pcontinues: 2-D tensor holding pcontinue values received during the
          transition that corresponds to each major index.
          Shape is `[(T+1), 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). These probabilities usually derive from
          the learning net.
          Shape is `[(T+1), B, num_actions]`.
        behaviour_policy_probs: 2-D tensor holding the *behaviour* policy's
          probabilities of having taken actions `action` during the transitions
          that correspond to each major index. These probabilities derive from
          whatever policy you used to generate the data.
          Shape is `[(T+1), B]`.
        stop_targnet_gradients: `bool` that enables a sensible default way of
          handling gradients through the Retrace op (essentially, gradients
          are not permitted to involve the `targnet_qs` inputs). Can be disabled
          if you require a different arrangement, but you'll probably want to
          block some gradients somewhere.
        name: name to prefix ops created by this function.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: Tensor containing the batch of losses, shape `[B]`.
        * `extra`: None
      """
      all_args = [
          lambda_,
          qs,
          targnet_qs,
          actions,
          rewards,
          pcontinues,
          target_policy_probs,
          behaviour_policy_probs,
      ]
      with tf.name_scope(name, "Retrace", values=all_args):
          # Mainly to simplify testing:
          (
              lambda_,
              qs,
              targnet_qs,
              actions,
              rewards,
              pcontinues,
              target_policy_probs,
              behaviour_policy_probs,
          ) = (tf.convert_to_tensor(arg) for arg in all_args)
  
          # 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, ranks):
              for i, (tensor, rank) in enumerate(zip(tensors, ranks)):
                  if tensor.get_shape():
                      base_ops.assert_rank_and_shape_compatibility(
                          [tensor], rank
                      )
                  else:
                      tf.logging.error(
                          'Tensor "%s", which was offered as Retrace parameter %d, has '
                          "no rank at construction time, so Retrace can't verify that "
                          "it has the necessary rank of %d",
                          tensor.name,
                          i + 1,
                          rank,
                      )
  
          check_rank(
              [
                  lambda_,
                  qs,
                  targnet_qs,
                  actions,
                  rewards,
                  pcontinues,
                  target_policy_probs,
                  behaviour_policy_probs,
              ],
              [0, 3, 3, 2, 2, 2, 3, 2],
          )
  
          # Deduce the shapes of the arguments we'll create for retrace_core.
          qs_shape = tf.shape(qs)
          timesteps = qs_shape[0]  # Batch size is qs_shape[1].
  
          # Deduce the time indices for the arguments we'll create for retrace_core.
          timestep_indices_tm1 = tf.range(0, timesteps - 1)
          timestep_indices_t = tf.range(1, timesteps)
  
          # Construct arguments for retrace_core and call.
          q_tm1 = tf.gather(qs, timestep_indices_tm1)
          a_tm1 = tf.gather(actions, timestep_indices_tm1)
  
          r_t = tf.gather(rewards, timestep_indices_tm1)
          pcont_t = tf.gather(pcontinues, timestep_indices_tm1)
  
          target_policy_t = tf.gather(
              target_policy_probs, timestep_indices_t
          )
          behaviour_policy_t = tf.gather(
              behaviour_policy_probs, timestep_indices_t
          )
          targnet_q_t = tf.gather(targnet_qs, timestep_indices_t)
          a_t = tf.gather(actions, timestep_indices_t)
  
          core = retrace_core(
              lambda_,
              q_tm1,
              a_tm1,
              r_t,
              pcont_t,
              target_policy_t,
              behaviour_policy_t,
              targnet_q_t,
              a_t,
              stop_targnet_gradients,
          )
  
          return base_ops.LossOutput(core.loss, None)
  
file_path: trfl/trfl/retrace_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
  from trfl import base_ops
  from trfl import indexing_ops
  from trfl import sequence_ops
label_desc: Use image node_id 1 for calling a global function with example usage: retrace(lambda_, qs, targnet_qs, actions, rewards, pcontinues, target_policy_probs, behaviour_policy_probs, stop_targnet_gradients, name) and returns: base_ops