code: def policy_gradient_loss(
      logits_t: Array,
      a_t: Array,
      adv_t: Array,
      w_t: Array,
      use_stop_gradient: bool = True,
  ) -> Array:
      """Calculates the policy gradient loss.
  
      See "Simple Gradient-Following Algorithms for Connectionist RL" by Williams.
      (http://www-anw.cs.umass.edu/~barto/courses/cs687/williams92simple.pdf)
  
      Args:
        logits_t: a sequence of unnormalized action preferences.
        a_t: a sequence of actions sampled from the preferences `logits_t`.
        adv_t: the observed or estimated advantages from executing actions `a_t`.
        w_t: a per timestep weighting for the loss.
        use_stop_gradient: bool indicating whether or not to apply stop gradient to
          advantages.
  
      Returns:
        Loss whose gradient corresponds to a policy gradient update.
      """
      chex.assert_rank([logits_t, a_t, adv_t, w_t], [2, 1, 1, 1])
      chex.assert_type(
          [logits_t, a_t, adv_t, w_t], [float, int, float, float]
      )
  
      log_pi_a_t = distributions.softmax().logprob(a_t, logits_t)
      adv_t = jax.lax.select(
          use_stop_gradient, jax.lax.stop_gradient(adv_t), adv_t
      )
      loss_per_timestep = -log_pi_a_t * adv_t
      return jnp.mean(loss_per_timestep * w_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 3 for calling a global function with example usage: policy_gradient_loss(logits_t, a_t, adv_t, w_t, use_stop_gradient) and returns: jnp