(
      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.
Initializes ValueTrainTask.
Initializes ValueTrainTask .
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',
    )
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 = ( ) , ) : 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' , )