class_docstr: Task for value training.
docstr: Computes a value training batch based on a trajectory batch.
  
  Args:
    trajectory_batch: trax.rl.task.TimeStepBatch with a batch of trajectory
      slices. Elements should have shape (batch_size, seq_len, ...).
    shape_only: Whether to return dummy zero arrays of correct shape. Useful
      for initializing models.
  
  Returns:
    Triple (observations, targets, weights), where targets are the target
    values for network training and weights are used for masking in loss
    computation. Shapes:
    - observations: (batch_size, seq_len) + observation_shape
    - targets: (batch_size, seq_len)
    - weights: (batch_size, seq_len)
docstr_tok: ["Computes","a","value","training","batch","based","on","a","trajectory","batch",".","Args",":","trajectory_batch",":","trax.rl.task.TimeStepBatch","with","a","batch","of","trajectory","slices",".","Elements","should","have","shape","(","batch_size",",","seq_len",",","...",")",".","shape_only",":","Whether","to","return","dummy","zero","arrays","of","correct","shape",".","Useful","for","initializing","models",".","Returns",":","Triple","(","observations",",","targets",",","weights",")",",","where","targets","are","the","target","values","for","network","training","and","weights","are","used","for","masking","in","loss","computation",".","Shapes",":","-","observations",":","(","batch_size",",","seq_len",")","+","observation_shape","-","targets",":","(","batch_size",",","seq_len",")","-","weights",":","(","batch_size",",","seq_len",")"]
code: def value_batch(self, trajectory_batch, shape_only=False):
      """Computes a value training batch based on a trajectory batch.
  
      Args:
        trajectory_batch: trax.rl.task.TimeStepBatch with a batch of trajectory
          slices. Elements should have shape (batch_size, seq_len, ...).
        shape_only: Whether to return dummy zero arrays of correct shape. Useful
          for initializing models.
  
      Returns:
        Triple (observations, targets, weights), where targets are the target
        values for network training and weights are used for masking in loss
        computation. Shapes:
        - observations: (batch_size, seq_len) + observation_shape
        - targets: (batch_size, seq_len)
        - weights: (batch_size, seq_len)
      """
      if self._sync_at(self._step) and not shape_only:
          self._sync_target_model()
  
      (batch_size, seq_len) = trajectory_batch.observation.shape[:2]
      assert trajectory_batch.action.shape[:2] == (batch_size, seq_len)
      assert trajectory_batch.mask.shape == (batch_size, seq_len)
      # Compute the value from the target network.
      values = np.array(
          self.value(trajectory_batch, shape_only=shape_only)
      )
      assert values.shape == (batch_size, seq_len)
      # Compute the advantages - the TD errors of the target network.
      advantages = self._advantage_estimator(
          rewards=trajectory_batch.reward,
          returns=trajectory_batch.return_,
          dones=trajectory_batch.done,
          values=values,
          discount_mask=trajectory_batch.env_info.discount_mask,
      )
      adv_seq_len = advantages.shape[1]
      # The advantage sequence should be shorter by the margin. For more details,
      # see the comment in policy_tasks.PolicyTrainTask.policy_batch.
      assert adv_seq_len <= seq_len
      assert advantages.shape == (batch_size, adv_seq_len)
      # Compute the targets based on the target values and their TD errors. The
      # network gives perfect predictions when targets == values, so the
      # advantages are zero.
      targets = (
          values[:, :adv_seq_len] + advantages
      ) * self._target_scale
      # Trim observations and the mask to match the target length.
      observations = trajectory_batch.observation[:, :adv_seq_len]
      mask = trajectory_batch.mask[:, :adv_seq_len]
      # Add a singleton depth dimension to the targets and the mask.
      targets = targets[:, :, None]
      mask = mask[:, :, None]
      return (observations, targets, mask)
  
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 4 for calling the ValueTrainTask obj's underlying member method code with example usage: obj.value_batch(trajectory_batch, shape_only) and returns: observations,  targets,  mask