class_docstr: Task for value training.
docstr: 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.
docstr_tok: ["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","."]
code: 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",
      )
  
file_path: trax/trax/rl/value_tasks.py
imports: import copy
  import numpy
  from trax import layers
  from trax.fastmath import numpy
  from trax.supervised import training
label_desc: Use image node_id 1 to create a new ValueTrainTask object from inherited base classes: training with example: obj = ValueTrainTask(trajectory_batch_stream, optimizer, lr_schedule, advantage_estimator, model, target_model, target_scale, sync_at, loss_layer, head_selector)