def __init__(
      self,
      trajectory_batch_stream,
      optimizer,
      lr_schedule,
      advantage_estimator,
      model,
      target_model=None,
      target_scale=1.0,
      sync_at=(lambda step: step % 100 == 0),
      loss_layer=None,
      head_selector=(),
  ):
    """Initializes ValueTrainTask.

    Args:
      trajectory_batch_stream: Generator of trax.rl.task.TimeStepBatch.
      optimizer: Optimizer for network training.
      lr_schedule: Learning rate schedule for network training.
      advantage_estimator: Function
        (rewards, returns, values, dones) -> advantages, created by one of the
        functions from trax.rl.advantages.
      model: Model being trained, used to synchronize weights of the target
        model.
      target_model: Model for calculating TD targets. If `None`, use `model`.
      target_scale: Multiplier for the targets. Useful for rescaling the targets
        to a more reasonable range for model training.
      sync_at: Function step -> bool, indicating when to synchronize the target
        network with the trained network. This is necessary for training the
        network on bootstrapped targets, e.g. using TD-k.
      loss_layer: The value loss layer. The default is L2 loss.
      head_selector: Layer to apply to the network output to select the value
        head. Only needed in multitask training.
    """
    self._trajectory_batch_stream = trajectory_batch_stream
    self._advantage_estimator = advantage_estimator
    self._target_scale = target_scale

    self._synced = False
    def sync_also_on_initial_batch(step):
      return sync_at(step) or not self._synced
    self._sync_at = sync_also_on_initial_batch

    self._head_selector = head_selector

    def attach_head(model):
      return tl.Serial(model, self._head_selector)
    self._train_model = attach_head(model)
    if target_model is None:
      target_model = model
    # TODO(pkozakowski): Use target_model.clone() once it's implemented.
    self._target_model = attach_head(copy.deepcopy(target_model))

    # Count the steps, so we know when to synchronize the target network.
    self._step = 0
    def labeled_data():
      for trajectory_batch in self._trajectory_batch_stream:
        self._step += 1
        yield self.value_batch(trajectory_batch)
    sample_batch = self.value_batch(
        next(trajectory_batch_stream), shape_only=True
    )
    if loss_layer is None:
      loss_layer = tl.L2Loss()
    loss_layer = tl.Serial(head_selector, loss_layer)
    super().__init__(
        labeled_data(), loss_layer, optimizer,
        sample_batch=sample_batch,
        lr_schedule=lr_schedule,
        loss_name='value_loss',
    )
Initializes ValueTrainTask.

    Args:
      trajectory_batch_stream: Generator of trax.rl.task.TimeStepBatch.
      optimizer: Optimizer for network training.
      lr_schedule: Learning rate schedule for network training.
      advantage_estimator: Function
        (rewards, returns, values, dones) -> advantages, created by one of the
        functions from trax.rl.advantages.
      model: Model being trained, used to synchronize weights of the target
        model.
      target_model: Model for calculating TD targets. If `None`, use `model`.
      target_scale: Multiplier for the targets. Useful for rescaling the targets
        to a more reasonable range for model training.
      sync_at: Function step -> bool, indicating when to synchronize the target
        network with the trained network. This is necessary for training the
        network on bootstrapped targets, e.g. using TD-k.
      loss_layer: The value loss layer. The default is L2 loss.
      head_selector: Layer to apply to the network output to select the value
        head. Only needed in multitask training.
    
trax/rl/value_tasks.py