code: def vtrace_td_error_and_advantage(
      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,
      clip_pg_rho_threshold: float = 1.0,
      stop_target_gradients: bool = True,
  ) -> VTraceOutput:
      """Calculates V-Trace errors and PG advantage from importance weights.
  
      This functions computes the TD-errors and policy gradient Advantage terms
      as used by the IMPALA distributed actor-critic agent.
  
      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 weights at time t-1.
        lambda_: mixing parameter; a scalar or a vector for timesteps t.
        clip_rho_threshold: clip threshold for importance ratios.
        clip_pg_rho_threshold: clip threshold for policy gradient importance ratios.
        stop_target_gradients: whether or not to apply stop gradient to targets.
  
      Returns:
        a tuple of V-Trace error, policy gradient advantage, and estimated Q-values.
      """
      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])
  
      # If scalar make into vector.
      lambda_ = jnp.ones_like(discount_t) * lambda_
  
      errors = vtrace(
          v_tm1,
          v_t,
          r_t,
          discount_t,
          rho_tm1,
          lambda_,
          clip_rho_threshold,
          stop_target_gradients,
      )
      targets_tm1 = errors + v_tm1
      q_bootstrap = jnp.concatenate(
          [
              lambda_[:-1] * targets_tm1[1:]
              + (1 - lambda_[:-1]) * v_tm1[1:],
              v_t[-1:],
          ],
          axis=0,
      )
      q_estimate = r_t + discount_t * q_bootstrap
      clipped_pg_rho_tm1 = jnp.minimum(clip_pg_rho_threshold, rho_tm1)
      pg_advantages = clipped_pg_rho_tm1 * (q_estimate - v_tm1)
      return VTraceOutput(
          errors=errors,
          pg_advantage=pg_advantages,
          q_estimate=q_estimate,
      )
  
file_path: rlax/rlax/_src/vtrace.py
imports: import collections
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: vtrace_td_error_and_advantage(v_tm1, v_t, r_t, discount_t, rho_tm1, lambda_, clip_rho_threshold, clip_pg_rho_threshold, stop_target_gradients) and returns: VTraceOutput