code: def vtrace(
      v_tm1: Array,
      v_t: Array,
      r_t: Array,
      discount_t: Array,
      rho_tm1: Array,
      lambda_: Numeric = 1.0,
      clip_rho_threshold: float = 1.0,
      stop_target_gradients: bool = True,
  ) -> Array:
      """Calculates V-Trace errors from importance weights.
  
      V-trace computes TD-errors from multistep trajectories by applying
      off-policy corrections based on clipped importance sampling ratios.
  
      See "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor
      Learner Architectures" by Espeholt et al. (https://arxiv.org/abs/1802.01561).
  
      Args:
        v_tm1: values at time t-1.
        v_t: values at time t.
        r_t: reward at time t.
        discount_t: discount at time t.
        rho_tm1: importance sampling ratios at time t-1.
        lambda_: mixing parameter; a scalar or a vector for timesteps t.
        clip_rho_threshold: clip threshold for importance weights.
        stop_target_gradients: whether or not to apply stop gradient to targets.
  
      Returns:
        V-Trace error.
      """
      chex.assert_rank(
          [v_tm1, v_t, r_t, discount_t, rho_tm1, lambda_],
          [1, 1, 1, 1, 1, {0, 1}],
      )
      chex.assert_type(
          [v_tm1, v_t, r_t, discount_t, rho_tm1, lambda_],
          [float, float, float, float, float, float],
      )
      chex.assert_equal_shape([v_tm1, v_t, r_t, discount_t, rho_tm1])
  
      # Clip importance sampling ratios.
      c_tm1 = jnp.minimum(1.0, rho_tm1) * lambda_
      clipped_rhos_tm1 = jnp.minimum(clip_rho_threshold, rho_tm1)
  
      # Compute the temporal difference errors.
      td_errors = clipped_rhos_tm1 * (r_t + discount_t * v_t - v_tm1)
  
      # Work backwards computing the td-errors.
      def _body(acc, xs):
          td_error, discount, c = xs
          acc = td_error + discount * c * acc
          return acc, acc
  
      _, errors = jax.lax.scan(
          _body, 0.0, (td_errors, discount_t, c_tm1), reverse=True
      )
  
      # Return errors, maybe disabling gradient flow through bootstrap targets.
      return jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(errors + v_tm1) - v_tm1,
          errors,
      )
  
file_path: rlax/rlax/_src/vtrace.py
imports: import collections
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: vtrace(v_tm1, v_t, r_t, discount_t, rho_tm1, lambda_, clip_rho_threshold, stop_target_gradients) and returns: jax, acc,  acc