(
      self, task, model_fn,
      optimizer=adam.Adam,
      policy_lr_schedule=lr.multifactor,
      policy_n_steps_per_epoch=1000,
      policy_weight_fn=(lambda x: x),
      value_lr_schedule=lr.multifactor,
      value_n_steps_per_epoch=1000,
      value_sync_at=(lambda x: x % 100 == 0),
      advantage_estimator=rl_advantages.monte_carlo,
      batch_size=64,
      network_eval_at=None,
      n_eval_batches=1,
      max_slice_length=1,
      margin=0,
      n_replay_epochs=1,
      **kwargs
  )
Initializes LoopActorCriticAgent.

    Args:
      task: `RLTask` instance to use.
      model_fn: Function mode -> Trax model, building a joint policy and value
        network.
      optimizer: Optimizer for the policy and value networks.
      policy_lr_schedule: Learning rate schedule for the policy network.
      policy_n_steps_per_epoch: Number of steps to train the policy network for
        in each epoch.
      policy_weight_fn: Function advantages -> weights for calculating the
        log probability weights in policy training.
      value_lr_schedule: Learning rate schedule for the value network.
      value_n_steps_per_epoch: Number of steps to train the value network for
        in each epoch.
      value_sync_at: Function step -> bool indicating when to synchronize the
        target network with the trained network in value training.
      advantage_estimator: Advantage estimator to use in policy and value
        training.
      batch_size: Batch size for training the networks.
      network_eval_at: Function step -> bool indicating in when to evaluate the
        networks.
      n_eval_batches: Number of batches to compute the network evaluation
        metrics on.
      max_slice_length: Maximum length of a trajectory slice to train on.
      margin: Number of timesteps to add at the end of each trajectory slice for
        better advantage estimation.
      n_replay_epochs: Number of epochs of trajectories to store in the replay
        buffer.
      **kwargs: Keyword arguments forwarded to Agent.
Initializes LoopActorCriticAgent.
Initializes LoopActorCriticAgent .
def __init__(
      self, task, model_fn,
      optimizer=adam.Adam,
      policy_lr_schedule=lr.multifactor,
      policy_n_steps_per_epoch=1000,
      policy_weight_fn=(lambda x: x),
      value_lr_schedule=lr.multifactor,
      value_n_steps_per_epoch=1000,
      value_sync_at=(lambda x: x % 100 == 0),
      advantage_estimator=rl_advantages.monte_carlo,
      batch_size=64,
      network_eval_at=None,
      n_eval_batches=1,
      max_slice_length=1,
      margin=0,
      n_replay_epochs=1,
      **kwargs
  ):
    """Initializes LoopActorCriticAgent.

    Args:
      task: `RLTask` instance to use.
      model_fn: Function mode -> Trax model, building a joint policy and value
        network.
      optimizer: Optimizer for the policy and value networks.
      policy_lr_schedule: Learning rate schedule for the policy network.
      policy_n_steps_per_epoch: Number of steps to train the policy network for
        in each epoch.
      policy_weight_fn: Function advantages -> weights for calculating the
        log probability weights in policy training.
      value_lr_schedule: Learning rate schedule for the value network.
      value_n_steps_per_epoch: Number of steps to train the value network for
        in each epoch.
      value_sync_at: Function step -> bool indicating when to synchronize the
        target network with the trained network in value training.
      advantage_estimator: Advantage estimator to use in policy and value
        training.
      batch_size: Batch size for training the networks.
      network_eval_at: Function step -> bool indicating in when to evaluate the
        networks.
      n_eval_batches: Number of batches to compute the network evaluation
        metrics on.
      max_slice_length: Maximum length of a trajectory slice to train on.
      margin: Number of timesteps to add at the end of each trajectory slice for
        better advantage estimation.
      n_replay_epochs: Number of epochs of trajectories to store in the replay
        buffer.
      **kwargs: Keyword arguments forwarded to Agent.
    """
    super().__init__(task, **kwargs)

    self._policy_dist = distributions.create_distribution(
        self.task.action_space
    )
    model_fn = functools.partial(
        model_fn,
        policy_distribution=self._policy_dist,
    )
    train_model = model_fn(mode='train')
    eval_model = model_fn(mode='eval')

    trajectory_batch_stream = self._init_trajectory_batch_stream(
        batch_size, max_slice_length, margin, n_replay_epochs
    )
    advantage_estimator = advantage_estimator(task.gamma, margin=margin)
    (value_train_task, value_eval_task) = self._init_value_tasks(
        trajectory_batch_stream,
        optimizer=optimizer(),
        lr_schedule=value_lr_schedule(),
        advantage_estimator=advantage_estimator,
        train_model=train_model,
        eval_model=eval_model,
        sync_at=value_sync_at,
        n_steps_per_epoch=value_n_steps_per_epoch,
        n_eval_batches=n_eval_batches,
    )
    (policy_train_task, policy_eval_task) = self._init_policy_tasks(
        trajectory_batch_stream,
        optimizer=optimizer(),
        lr_schedule=policy_lr_schedule(),
        advantage_estimator=advantage_estimator,
        value_train_task=value_train_task,
        weight_fn=policy_weight_fn,
        n_eval_batches=n_eval_batches,
    )
    self._init_loop(
        train_model=train_model,
        eval_model=eval_model,
        policy_train_and_eval_task=(policy_train_task, policy_eval_task),
        value_train_and_eval_task=(value_train_task, value_eval_task),
        eval_at=network_eval_at,
        policy_n_steps_per_epoch=policy_n_steps_per_epoch,
        value_n_steps_per_epoch=value_n_steps_per_epoch,
    )
    self._init_collection(model_fn, policy_train_task.sample_batch)
def __init__ ( self , task , model_fn , optimizer = adam . Adam , policy_lr_schedule = lr . multifactor , policy_n_steps_per_epoch = 1000 , policy_weight_fn = ( lambda x : x ) , value_lr_schedule = lr . multifactor , value_n_steps_per_epoch = 1000 , value_sync_at = ( lambda x : x % 100 == 0 ) , advantage_estimator = rl_advantages . monte_carlo , batch_size = 64 , network_eval_at = None , n_eval_batches = 1 , max_slice_length = 1 , margin = 0 , n_replay_epochs = 1 , * * kwargs ) : super ( ) . __init__ ( task , * * kwargs ) self . _policy_dist = distributions . create_distribution ( self . task . action_space ) model_fn = functools . partial ( model_fn , policy_distribution = self . _policy_dist , ) train_model = model_fn ( mode = 'train' ) eval_model = model_fn ( mode = 'eval' ) trajectory_batch_stream = self . _init_trajectory_batch_stream ( batch_size , max_slice_length , margin , n_replay_epochs ) advantage_estimator = advantage_estimator ( task . gamma , margin = margin ) ( value_train_task , value_eval_task ) = self . _init_value_tasks ( trajectory_batch_stream , optimizer = optimizer ( ) , lr_schedule = value_lr_schedule ( ) , advantage_estimator = advantage_estimator , train_model = train_model , eval_model = eval_model , sync_at = value_sync_at , n_steps_per_epoch = value_n_steps_per_epoch , n_eval_batches = n_eval_batches , ) ( policy_train_task , policy_eval_task ) = self . _init_policy_tasks ( trajectory_batch_stream , optimizer = optimizer ( ) , lr_schedule = policy_lr_schedule ( ) , advantage_estimator = advantage_estimator , value_train_task = value_train_task , weight_fn = policy_weight_fn , n_eval_batches = n_eval_batches , ) self . _init_loop ( train_model = train_model , eval_model = eval_model , policy_train_and_eval_task = ( policy_train_task , policy_eval_task ) , value_train_and_eval_task = ( value_train_task , value_eval_task ) , eval_at = network_eval_at , policy_n_steps_per_epoch = policy_n_steps_per_epoch , value_n_steps_per_epoch = value_n_steps_per_epoch , ) self . _init_collection ( model_fn , policy_train_task . sample_batch )