def value(self, trajectory_batch, shape_only=False):
    """Computes values of states in a given batch of trajectory slices.

    Can be passed as value_fn to PolicyTrainTask to implement a critic baseline
    for advantage calculation.

    Args:
      trajectory_batch: Batch of trajectory slices to compute values for.
      shape_only: Whether to return dummy zero arrays of correct shape. Useful
        for initializing models.

    Returns:
      Array of values of all states in `trajectory_batch`.
    """
    if shape_only:
      # The target model hasn't been initialized yet, and we are asked for the
      # initial, sample batch. Only shape matters here, so just return zeros.
      return np.zeros(trajectory_batch.observation.shape[:2])

    if not self._synced:
      self._sync_target_model()

    values = self._target_model(trajectory_batch.observation)
    # Squeeze the singleton depth axis.
    return np.squeeze(values, axis=-1) / self._target_scale
Computes values of states in a given batch of trajectory slices.

    Can be passed as value_fn to PolicyTrainTask to implement a critic baseline
    for advantage calculation.

    Args:
      trajectory_batch: Batch of trajectory slices to compute values for.
      shape_only: Whether to return dummy zero arrays of correct shape. Useful
        for initializing models.

    Returns:
      Array of values of all states in `trajectory_batch`.
    
trax/rl/value_tasks.py