code: def transformed_n_step_q_learning(
      q_tm1: Array,
      a_tm1: Array,
      target_q_t: Array,
      a_t: Array,
      r_t: Array,
      discount_t: Array,
      n: int,
      stop_target_gradients: bool = True,
      tx_pair: TxPair = IDENTITY_PAIR,
  ) -> Array:
      """Calculates transformed n-step TD errors.
  
      See "Recurrent Experience Replay in Distributed Reinforcement Learning" by
      Kapturowski et al. (https://openreview.net/pdf?id=r1lyTjAqYX).
  
      Args:
        q_tm1: Q-values at times [0, ..., T - 1].
        a_tm1: action index at times [0, ..., T - 1].
        target_q_t: target Q-values at time [1, ..., T].
        a_t: action index at times [[1, ... , T]] used to select target q-values to
          bootstrap from; max(target_q_t) for normal Q-learning, max(q_t) for double
          Q-learning.
        r_t: reward at times [1, ..., T].
        discount_t: discount at times [1, ..., T].
        n: number of steps over which to accumulate reward before bootstrapping.
        stop_target_gradients: bool indicating whether or not to apply stop gradient
          to targets.
        tx_pair: TxPair of value function transformation and its inverse.
  
      Returns:
        Transformed N-step TD error.
      """
      chex.assert_rank(
          [q_tm1, target_q_t, a_tm1, a_t, r_t, discount_t],
          [2, 2, 1, 1, 1, 1],
      )
      chex.assert_type(
          [q_tm1, target_q_t, a_tm1, a_t, r_t, discount_t],
          [float, float, int, int, float, float],
      )
  
      v_t = base.batched_index(target_q_t, a_t)
      target_tm1 = transformed_n_step_returns(
          tx_pair,
          r_t,
          discount_t,
          v_t,
          n,
          stop_target_gradients=stop_target_gradients,
      )
      q_a_tm1 = base.batched_index(q_tm1, a_tm1)
      return target_tm1 - q_a_tm1
  
file_path: rlax/rlax/_src/nonlinear_bellman.py
imports: import collections
  import functools
  import chex
  import jax.numpy
  from rlax._src import base
  from rlax._src import multistep
  from rlax._src import transforms
label_desc: Use image node_id 8 for calling a global function with example usage: transformed_n_step_q_learning(q_tm1, a_tm1, target_q_t, a_t, r_t, discount_t, n, stop_target_gradients, tx_pair) and returns: unknown