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