(self, trajectory_batch, shape_only=False)
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)
Computes a value training batch based on a trajectory batch.
Computes a value training batch based on a trajectory batch .
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)
def value_batch ( self , trajectory_batch , shape_only = False ) : 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 )