(
      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.
Initializes LoopPolicyAgent.
Initializes LoopPolicyAgent .
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
          )
      )
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 ) : 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 ) )