class_docstr: Task for policy evaluation.
class_docstr_tok: ["Task","for","policy","evaluation","."]
code: def weight_metrics(self):
      def make_metric(aggregate_fn):  # pylint: disable=invalid-name
          def WeightMetric(policy_inputs, actions, advantages, mask):
              del policy_inputs, actions, mask
              weights = self._train_task.calculate_weights(advantages)
              return aggregate_fn(weights)
  
          return tl.Fn("WeightMetric", WeightMetric)
  
      return {  # pylint: disable=g-complex-comprehension
          "weight_" + name: make_metric(fn)
          for (name, fn) in [
              ("mean", jnp.mean),
              ("std", jnp.std),
              ("min", jnp.min),
              ("max", jnp.max),
          ]
      }
  
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 5 for calling the PolicyEvalTask obj's underlying member method code with example usage: obj.weight_metrics() and returns: unknown, tl, aggregate_fn