code: def persistent_qlearning(
      q_tm1,
      a_tm1,
      r_t,
      pcont_t,
      q_t,
      action_gap_scale=0.5,
      name="PersistentQLearning",
  ):
      """Implements the persistent Q-learning loss as a TensorFlow op.
  
      The loss is `0.5` times the squared difference between `q_tm1[a_tm1]` and
      `r_t + pcont_t * [(1-action_gap_scale) max q_t + action_gap_scale qa_t]`
  
      See "Increasing the Action Gap: New Operators for Reinforcement Learning"
      by Bellemare, Ostrovski, Guez et al. (https://arxiv.org/abs/1512.04860).
  
      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]`.
          These values are used for estimating the value of the best action. In
          DQN they come from the target network.
        action_gap_scale: coefficient in [0, 1] for scaling the action gap term.
        name: name to prefix ops created within this op.
  
      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], [a_tm1, r_t, pcont_t]], [2, 1], name
      )
      base_ops.assert_arg_bounded(
          action_gap_scale, 0, 1, name, "action_gap_scale"
      )
  
      # persistent Q-learning 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"):
              max_q_t = tf.reduce_max(q_t, axis=1)
              qa_t = indexing_ops.batched_index(q_t, a_tm1)
              corrected_q_t = (
                  1 - action_gap_scale
              ) * max_q_t + action_gap_scale * qa_t
              target = tf.stop_gradient(r_t + pcont_t * corrected_q_t)
          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 3 for calling a global function with example usage: persistent_qlearning(q_tm1, a_tm1, r_t, pcont_t, q_t, action_gap_scale, name) and returns: base_ops