code: def td_learning(v_tm1, r_t, pcont_t, v_t, name="TDLearning"):
      """Implements the TD(0)-learning loss as a TensorFlow op.
  
      The TD loss is `0.5` times the squared difference between `v_tm1` and
      the target `r_t + pcont_t * v_t`.
  
      See "Learning to Predict by the Methods of Temporal Differences" by Sutton.
      (https://link.springer.com/article/10.1023/A:1022633531479).
  
      Args:
        v_tm1: Tensor holding values at previous timestep, shape `[B]`.
        r_t: Tensor holding rewards, shape `[B]`.
        pcont_t: Tensor holding pcontinue values, shape `[B]`.
        v_t: Tensor holding values at current timestep, shape `[B]`.
        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 `v_tm1`, shape `[B]`.
            * `td_error`: batch of temporal difference errors, shape `[B]`.
      """
      # Rank and compatibility checks.
      base_ops.wrap_rank_shape_assert(
          [[v_tm1, v_t, r_t, pcont_t]], [1], name
      )
  
      # TD(0)-learning op.
      with tf.name_scope(name, values=[v_tm1, r_t, pcont_t, v_t]):
          # Build target.
          target = tf.stop_gradient(r_t + pcont_t * v_t)
  
          # Temporal difference error and loss.
          # Loss is MSE scaled by 0.5, so the gradient is equal to the TD error.
          td_error = target - v_tm1
          loss = 0.5 * tf.square(td_error)
          return base_ops.LossOutput(loss, TDExtra(target, td_error))
  
file_path: trfl/trfl/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 sequence_ops
label_desc: Use image node_id 1 for calling a global function with example usage: td_learning(v_tm1, r_t, pcont_t, v_t, name) and returns: base_ops