code: def td_learning(
      v_tm1: Numeric,
      r_t: Numeric,
      discount_t: Numeric,
      v_t: Numeric,
      stop_target_gradients: bool = True,
  ) -> Numeric:
      """Calculates the TD-learning temporal difference error.
  
      See "Learning to Predict by the Methods of Temporal Differences" by Sutton.
      (https://link.springer.com/article/10.1023/A:1022633531479).
  
      Args:
        v_tm1: state values at time t-1.
        r_t: reward at time t.
        discount_t: discount at time t.
        v_t: state values at time t.
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
      Returns:
        TD-learning temporal difference error.
      """
  
      chex.assert_rank([v_tm1, r_t, discount_t, v_t], 0)
      chex.assert_type([v_tm1, r_t, discount_t, v_t], float)
  
      target_tm1 = r_t + discount_t * v_t
      target_tm1 = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(target_tm1),
          target_tm1,
      )
      return target_tm1 - v_tm1
  
imports: from typing import Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
  from rlax._src import clipping
  from rlax._src import distributions
  from rlax._src import multistep
label_desc: Use image node_id 1 for calling a global function with example usage: td_learning(v_tm1, r_t, discount_t, v_t, stop_target_gradients) and returns: unknown