def SamplingAWRLoss(beta, w_max, thresholds,  # pylint: disable=invalid-name
                    reweight=False, sampled_all_discrete=False):
  """Definition of the Advantage Weighted Regression (AWR) loss."""
  def f(log_probs, advantages, old_log_probs, mask):
    if reweight:  # Use new policy weights for sampled actions instead.
      mask *= jnp.exp(fastmath.stop_gradient(log_probs) - old_log_probs)
    if sampled_all_discrete:  # Actions were sampled uniformly; weight them.
      mask *= jnp.exp(old_log_probs)
    weights = jnp.minimum(awr_weights(advantages, beta, thresholds), w_max)
    return -jnp.sum(log_probs * weights * mask) / jnp.sum(mask)
  return tl.Fn('SamplingAWRLoss', f)
Definition of the Advantage Weighted Regression (AWR) loss.
trax/rl/actor_critic.py