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