code: def qpg_loss(
      logits_t: Array,
      q_t: Array,
      use_stop_gradient: bool = True,
  ) -> Array:
      """Computes the QPG (Q-based Policy Gradient) loss.
  
      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:
        QPG Loss.
      """
      chex.assert_rank([logits_t, q_t], 2)
      chex.assert_type([logits_t, q_t], float)
  
      policy_t, advantage_t = _compute_advantages(logits_t, q_t)
      advantage_t = jax.lax.select(
          use_stop_gradient,
          jax.lax.stop_gradient(advantage_t),
          advantage_t,
      )
      policy_advantages = -policy_t * advantage_t
      loss = jnp.mean(jnp.sum(policy_advantages, axis=1), axis=0)
      return loss
  
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 6 for calling a global function with example usage: qpg_loss(logits_t, q_t, use_stop_gradient) and returns: loss