code: def rpg_loss(
      logits_t: Array,
      q_t: Array,
      use_stop_gradient: bool = True,
  ) -> Array:
      """Computes the RPG (Regret Policy Gradient) loss.
  
      The gradient of this loss adapts the Regret Matching rule by weighting the
      standard PG update with regret.
  
      See "Actor-Critic Policy Optimization in Partially Observable Multiagent
      Environments" by Srinivasan, Lanctot (https://arxiv.org/abs/1810.09026).
  
      Args:
        logits_t: a sequence of unnormalized action preferences.
        q_t: the observed or estimated action value from executing actions `a_t` at
          time t.
        use_stop_gradient: bool indicating whether or not to apply stop gradient to
          advantages.
  
      Returns:
        RPG Loss.
      """
      chex.assert_rank([logits_t, q_t], 2)
      chex.assert_type([logits_t, q_t], float)
  
      _, adv_t = _compute_advantages(logits_t, q_t, use_stop_gradient)
      regrets_t = jnp.sum(jax.nn.relu(adv_t), axis=1)
  
      total_regret_t = jnp.mean(regrets_t, axis=0)
      return total_regret_t
  
file_path: rlax/rlax/_src/policy_gradients.py
imports: from typing import Optional, Tuple
  import chex
  import jax
  import jax.numpy
  from rlax._src import distributions
  from rlax._src import losses
label_desc: Use image node_id 8 for calling a global function with example usage: rpg_loss(logits_t, q_t, use_stop_gradient) and returns: total_regret_t