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,
    )
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.
    
trax/rl/policy_tasks.py