def calculate_weights(self, advantages):
    """Calculates advantage-based weights for log loss in policy training."""
    if self._advantage_normalization:
      # Normalize advantages.
      advantages -= jnp.mean(advantages)
      advantage_std = jnp.std(advantages)
      advantages /= advantage_std + self._advantage_normalization_epsilon
    weights = self._weight_fn(advantages)
    assert weights.shape == advantages.shape
    return weights
Calculates advantage-based weights for log loss in policy training.
trax/rl/policy_tasks.py