code: def entropy_loss(
      logits_t: Array,
      w_t: Array,
  ) -> Array:
      """Calculates the entropy regularization loss.
  
      See "Function Optimization using Connectionist RL Algorithms" by Williams.
      (https://www.tandfonline.com/doi/abs/10.1080/09540099108946587)
  
      Args:
        logits_t: a sequence of unnormalized action preferences.
        w_t: a per timestep weighting for the loss.
  
      Returns:
        Entropy loss.
      """
      chex.assert_rank([logits_t, w_t], [2, 1])
      chex.assert_type([logits_t, w_t], float)
  
      entropy_per_timestep = distributions.softmax().entropy(logits_t)
      return -jnp.mean(entropy_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 4 for calling a global function with example usage: entropy_loss(logits_t, w_t) and returns: unknown