class_docstr: Task for policy training.
  
  Trains the policy based on action advantages.
class_docstr_tok: ["Task","for","policy","training",".","Trains","the","policy","based","on","action","advantages","."]
docstr: Calculates advantage-based weights for log loss in policy training.
docstr_tok: ["Calculates","advantage-based","weights","for","log","loss","in","policy","training","."]
code: 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
  
file_path: trax/trax/rl/policy_tasks.py
imports: import numpy
  from trax import layers
  from trax.fastmath import numpy
  from trax.rl import distributions
  from trax.supervised import training
label_desc: Use image node_id 3 for calling the PolicyTrainTask obj's underlying member method code with example usage: obj.calculate_weights(advantages) and returns: weights