code: def clipped_surrogate_pg_loss(
      prob_ratios_t: Array,
      adv_t: Array,
      epsilon: Scalar,
      use_stop_gradient=True,
  ) -> Array:
      """Computes the clipped surrogate policy gradient loss.
  
      L_clipₜ(θ) = - min(rₜ(θ)Âₜ, clip(rₜ(θ), 1-ε, 1+ε)Âₜ)
  
      Where rₜ(θ) = π_θ(aₜ| sₜ) / π_θ_old(aₜ| sₜ) and Âₜ are the advantages.
  
      See Proximal Policy Optimization Algorithms, Schulman et al.:
      https://arxiv.org/abs/1707.06347
  
      Args:
        prob_ratios_t: Ratio of action probabilities for actions a_t:
            rₜ(θ) = π_θ(aₜ| sₜ) / π_θ_old(aₜ| sₜ)
        adv_t: the observed or estimated advantages from executing actions a_t.
        epsilon: Scalar value corresponding to how much to clip the objecctive.
        use_stop_gradient: bool indicating whether or not to apply stop gradient to
          advantages.
  
      Returns:
        Loss whose gradient corresponds to a clipped surrogate policy gradient
            update.
      """
      chex.assert_rank([prob_ratios_t, adv_t], [1, 1])
      chex.assert_type([prob_ratios_t, adv_t], [float, float])
  
      adv_t = jax.lax.select(
          use_stop_gradient, jax.lax.stop_gradient(adv_t), adv_t
      )
      clipped_ratios_t = jnp.clip(
          prob_ratios_t, 1.0 - epsilon, 1.0 + epsilon
      )
      clipped_objective = jnp.fmin(
          prob_ratios_t * adv_t, clipped_ratios_t * adv_t
      )
      return -jnp.mean(clipped_objective)
  
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 9 for calling a global function with example usage: clipped_surrogate_pg_loss(prob_ratios_t, adv_t, epsilon, use_stop_gradient) and returns: unknown