def __init__(
      self,
      task,
      model_fn,
      value_fn,
      weight_fn,
      n_replay_epochs,
      n_train_steps_per_epoch,
      advantage_normalization,
      optimizer=adam.Adam,
      lr_schedule=lr.multifactor,
      batch_size=64,
      network_eval_at=None,
      n_eval_batches=1,
      max_slice_length=1,
      trajectory_stream_preprocessing_fn=None,
      **kwargs
  ):
    """Initializes LoopPolicyAgent.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      value_fn: Function TimeStepBatch -> array (batch_size, seq_len)
        calculating the baseline for advantage calculation.
      weight_fn: Function float -> float to apply to advantages when calculating
        policy loss.
      n_replay_epochs: Number of last epochs to take into the replay buffer;
        only makes sense for off-policy algorithms.
      n_train_steps_per_epoch: Number of steps to train the policy network for
        in each epoch.
      advantage_normalization: Whether to normalize the advantages before
        passing them to weight_fn.
      optimizer: Optimizer for network training.
      lr_schedule: Learning rate schedule for network training.
      batch_size: Batch size for network training.
      network_eval_at: Function step -> bool indicating the training steps, when
        network evaluation should be performed.
      n_eval_batches: Number of batches to run during network evaluation.
      max_slice_length: The length of trajectory slices to run the network on.
      trajectory_stream_preprocessing_fn: Function to apply to the trajectory
        stream before batching. Can be used e.g. to filter trajectories.
      **kwargs: Keyword arguments passed to the superclass.
    """
    self._n_train_steps_per_epoch = n_train_steps_per_epoch
    super().__init__(task, **kwargs)

    task.set_n_replay_epochs(n_replay_epochs)
    self._max_slice_length = max_slice_length
    trajectory_batch_stream = task.trajectory_batch_stream(
        batch_size,
        epochs=[-(ep + 1) for ep in range(n_replay_epochs)],
        max_slice_length=self._max_slice_length,
        sample_trajectories_uniformly=True,
        trajectory_stream_preprocessing_fn=trajectory_stream_preprocessing_fn,
    )
    self._policy_dist = distributions.create_distribution(task.action_space)
    train_task = policy_tasks.PolicyTrainTask(
        trajectory_batch_stream,
        optimizer(),
        lr_schedule(),
        self._policy_dist,
        # Without a value network it doesn't make a lot of sense to use
        # a better advantage estimator than MC.
        advantage_estimator=advantages.monte_carlo(task.gamma, margin=0),
        advantage_normalization=advantage_normalization,
        value_fn=value_fn,
        weight_fn=weight_fn,
    )
    eval_task = policy_tasks.PolicyEvalTask(train_task, n_eval_batches)
    model_fn = functools.partial(
        model_fn,
        policy_distribution=self._policy_dist,
    )

    if self._output_dir is not None:
      policy_output_dir = os.path.join(self._output_dir, 'policy')
    else:
      policy_output_dir = None
    # Checkpoint every epoch.
    checkpoint_at = lambda step: step % n_train_steps_per_epoch == 0
    self._loop = supervised.training.Loop(
        model=model_fn(mode='train'),
        tasks=[train_task],
        eval_model=model_fn(mode='eval'),
        eval_tasks=[eval_task],
        output_dir=policy_output_dir,
        eval_at=network_eval_at,
        checkpoint_at=checkpoint_at,
    )
    self._collect_model = model_fn(mode='collect')
    self._collect_model.init(shapes.signature(train_task.sample_batch))

    # Validate the restored checkpoints.
    # TODO(pkozakowski): Move this to the base class once all Agents use Loop.
    if self._loop.step != self._epoch * self._n_train_steps_per_epoch:
      raise ValueError(
          'The number of Loop steps must equal the number of Agent epochs '
          'times the number of steps per epoch, got {}, {} and {}.'.format(
              self._loop.step, self._epoch, self._n_train_steps_per_epoch
          )
      )
Initializes LoopPolicyAgent.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      value_fn: Function TimeStepBatch -> array (batch_size, seq_len)
        calculating the baseline for advantage calculation.
      weight_fn: Function float -> float to apply to advantages when calculating
        policy loss.
      n_replay_epochs: Number of last epochs to take into the replay buffer;
        only makes sense for off-policy algorithms.
      n_train_steps_per_epoch: Number of steps to train the policy network for
        in each epoch.
      advantage_normalization: Whether to normalize the advantages before
        passing them to weight_fn.
      optimizer: Optimizer for network training.
      lr_schedule: Learning rate schedule for network training.
      batch_size: Batch size for network training.
      network_eval_at: Function step -> bool indicating the training steps, when
        network evaluation should be performed.
      n_eval_batches: Number of batches to run during network evaluation.
      max_slice_length: The length of trajectory slices to run the network on.
      trajectory_stream_preprocessing_fn: Function to apply to the trajectory
        stream before batching. Can be used e.g. to filter trajectories.
      **kwargs: Keyword arguments passed to the superclass.
    
trax/rl/training.py