code: def sarse(
      q_tm1,
      a_tm1,
      r_t,
      pcont_t,
      q_t,
      probs_a_t,
      debug=False,
      name="Sarse",
  ):
      """Implements the SARSE (Expected SARSA) loss as a TensorFlow op.
  
      The loss is `0.5` times the squared difference between `q_tm1[a_tm1]` and
      the target `r_t + pcont_t * (sum_a probs_a_t[a] * q_t[a])`.
  
      See "A Theoretical and Empirical Analysis of Expected Sarsa" by Seijen,
      van Hasselt, Whiteson et al.
      (http://www.cs.ox.ac.uk/people/shimon.whiteson/pubs/vanseijenadprl09.pdf).
  
      Args:
        q_tm1: Tensor holding Q-values for first timestep in a batch of
          transitions, shape `[B x num_actions]`.
        a_tm1: Tensor holding action indices, shape `[B]`.
        r_t: Tensor holding rewards, shape `[B]`.
        pcont_t: Tensor holding pcontinue values, shape `[B]`.
        q_t: Tensor holding Q-values for second timestep in a batch of
          transitions, shape `[B x num_actions]`.
        probs_a_t: Tensor holding action probabilities for second timestep,
          shape `[B x num_actions]`.
        debug: Boolean flag, when set to True adds ops to check whether probs_a_t
          is a batch of (approximately) valid probability distributions.
        name: name to prefix ops created by this function.
  
      Returns:
        A namedtuple with fields:
  
        * `loss`: a tensor containing the batch of losses, shape `[B]`.
        * `extra`: a namedtuple with fields:
            * `target`: batch of target values for `q_tm1[a_tm1]`, shape `[B]`.
            * `td_error`: batch of temporal difference errors, shape `[B]`.
      """
      # Rank and compatibility checks.
      base_ops.wrap_rank_shape_assert(
          [[q_tm1, q_t, probs_a_t], [a_tm1, r_t, pcont_t]], [2, 1], name
      )
  
      # SARSE (Expected SARSA) op.
      with tf.name_scope(
          name, values=[q_tm1, a_tm1, r_t, pcont_t, q_t, probs_a_t]
      ):
          # Debug ops.
          deps = []
          if debug:
              cumulative_prob = tf.reduce_sum(probs_a_t, axis=1)
              almost_prob = tf.less(
                  tf.abs(tf.subtract(cumulative_prob, 1.0)), 1e-6
              )
              deps.append(
                  tf.Assert(
                      tf.reduce_all(almost_prob),
                      ["probs_a_t tensor does not sum to 1", probs_a_t],
                  )
              )
  
          # With dependency on possible debug ops.
          with tf.control_dependencies(deps):
              # Select head to update and build target.
              qa_tm1 = indexing_ops.batched_index(q_tm1, a_tm1)
              target = tf.stop_gradient(
                  r_t
                  + pcont_t
                  * tf.reduce_sum(tf.multiply(q_t, probs_a_t), axis=1)
              )
  
              # 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 5 for calling a global function with example usage: sarse(q_tm1, a_tm1, r_t, pcont_t, q_t, probs_a_t, debug, name) and returns: base_ops