class_docstr: Task for value training.
docstr: 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`.
docstr_tok: ["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","`","."]
code: 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
  
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 5 for calling the ValueTrainTask obj's underlying member method code with example usage: obj.value(trajectory_batch, shape_only) and returns: unknown, np