class_docstr: Task for policy evaluation.
class_docstr_tok: ["Task","for","policy","evaluation","."]
docstr: Initializes PolicyEvalTask.
  
  Args:
    train_task: PolicyTrainTask used to train the policy network.
    n_eval_batches: Number of batches per evaluation.
    head_selector: Layer to apply to the network output to select the value
      head. Only needed in multitask training.
docstr_tok: ["Initializes","PolicyEvalTask",".","Args",":","train_task",":","PolicyTrainTask","used","to","train","the","policy","network",".","n_eval_batches",":","Number","of","batches","per","evaluation",".","head_selector",":","Layer","to","apply","to","the","network","output","to","select","the","value","head",".","Only","needed","in","multitask","training","."]
code: def __init__(self, train_task, n_eval_batches=1, head_selector=()):
      """Initializes PolicyEvalTask.
  
      Args:
        train_task: PolicyTrainTask used to train the policy network.
        n_eval_batches: Number of batches per evaluation.
        head_selector: Layer to apply to the network output to select the value
          head. Only needed in multitask training.
      """
      self._train_task = train_task
      self._policy_dist = train_task.policy_distribution
      labeled_data = map(
          self._eval_batch, train_task.trajectory_batch_stream
      )
      sample_batch = self._eval_batch(
          next(train_task.trajectory_batch_stream), shape_only=True
      )
      # TODO(pkozakowski): Implement more metrics.
      metrics = {
          "policy_entropy": self.entropy_metric,
      }
      metrics.update(self.advantage_metrics)
      metrics.update(self.weight_metrics)
      metrics = {
          name: tl.Serial(head_selector, metric)
          for (name, metric) in metrics.items()
      }
      (metric_names, metric_layers) = zip(*metrics.items())
      # Select the appropriate head for evaluation.
      super().__init__(
          labeled_data,
          metric_layers,
          sample_batch=sample_batch,
          metric_names=metric_names,
          n_eval_batches=n_eval_batches,
      )
  
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 1 to create a new PolicyEvalTask object from inherited base classes: training with example: obj = PolicyEvalTask(train_task, n_eval_batches, head_selector)