code: def qlambda(
      q_tm1,
      a_tm1,
      r_t,
      pcont_t,
      q_t,
      lambda_,
      name="GeneralizedQLambda",
  ):
      """Implements Peng's and Watkins' Q(lambda) loss as a TensorFlow op.
  
      This function is general enough to implement both Peng's and Watkins'
      Q-lambda algorithms.
  
      See "Reinforcement Learning: An Introduction" by Sutton and Barto.
      (http://incompleteideas.net/book/ebook/node78.html).
  
      Args:
        q_tm1: `Tensor` holding a sequence of Q-values starting at the first
          timestep; shape `[T, B, num_actions]`
        a_tm1: `Tensor` holding a sequence of action indices, shape `[T, B]`
        r_t: Tensor holding a sequence of rewards, shape `[T, B]`
        pcont_t: `Tensor` holding a sequence of pcontinue values, shape `[T, B]`
        q_t: `Tensor` holding a sequence of Q-values for second timestep;
          shape `[T, B, num_actions]`. In a target network setting,
          this quantity is often supplied by the target network.
        lambda_: a scalar or `Tensor` of shape `[T, B]`
          specifying the ratio of mixing between bootstrapped and MC returns;
          if lambda_ is the same for all time steps then the function implements
          Peng's Q-learning algorithm; if lambda_ = 0 at every sub-optimal action
          and a constant otherwise, then the function implements Watkins'
          Q-learning algorithm. Generally lambda_ can be a Tensor of any values
          in the range [0, 1] supplied by the user.
        name: a name of the op.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: a tensor containing the batch of losses, shape `[T, B]`.
        * `extra`: a namedtuple with fields:
            * `target`: batch of target values for `q_tm1[a_tm1]`, shape `[T, B]`.
            * `td_error`: batch of temporal difference errors, shape `[T, B]`.
      """
      # Rank and compatibility checks.
      base_ops.wrap_rank_shape_assert([[q_tm1, q_t]], [3], name)
      if (
          isinstance(lambda_, tf.Tensor)
          and lambda_.get_shape().ndims is not None
          and lambda_.get_shape().ndims > 0
      ):
          base_ops.wrap_rank_shape_assert(
              [[a_tm1, r_t, pcont_t, lambda_]], [2], name
          )
      else:
          base_ops.wrap_rank_shape_assert(
              [[a_tm1, r_t, pcont_t]], [2], name
          )
  
      # QLambda op.
      with tf.name_scope(
          name, values=[q_tm1, a_tm1, r_t, pcont_t, q_t]
      ):
          # Build target and select head to update.
          with tf.name_scope("target"):
              state_values = tf.reduce_max(q_t, axis=2)
              target = sequence_ops.multistep_forward_view(
                  r_t, pcont_t, state_values, lambda_, back_prop=False
              )
              target = tf.stop_gradient(target)
          qa_tm1 = indexing_ops.batched_index(q_tm1, a_tm1)
  
          # Temporal difference error and loss.
          # Loss is MSE scaled by 0.5, so the gradient is equal to the TD error.
          td_error = target - qa_tm1
          loss = 0.5 * tf.square(td_error)
          return base_ops.LossOutput(loss, QExtra(target, td_error))
  
file_path: trfl/trfl/action_value_ops.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  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 6 for calling a global function with example usage: qlambda(q_tm1, a_tm1, r_t, pcont_t, q_t, lambda_, name) and returns: base_ops