def trajectory_batch_stream(
      self,
      batch_size,
      epochs=None,
      max_slice_length=None,
      min_slice_length=None,
      margin=0,
      sample_trajectories_uniformly=False,
      trajectory_stream_preprocessing_fn=None,
  ):
    """Return a stream of trajectory batches from the specified epochs.

    This function returns a stream of tuples of numpy arrays (tensors).
    If tensors have different lengths, they will be padded by 0.

    Args:
      batch_size: the size of the batches to return
      epochs: a list of epochs to use; we use all epochs if None
      max_slice_length: maximum length of the slices of trajectories to return
      min_slice_length: minimum length of the slices of trajectories to return
      margin: number of extra steps after "done" that should be included in
        slices, so that networks see the terminal states in the training data
      sample_trajectories_uniformly: whether to sample trajectories uniformly,
        or proportionally to the number of slices in each trajectory (default)
      trajectory_stream_preprocessing_fn: function to apply to the trajectory
        stream before batching; can be used e.g. to filter trajectories

    Yields:
      batches of trajectory slices sampled uniformly from all slices of length
      at least min_slice_length and up to max_slice_length in all specified
      epochs
    """
    def pad(tensor_list):
      # Replace Nones with valid tensors.
      not_none_tensors = [t for t in tensor_list if t is not None]
      assert not_none_tensors, 'All tensors to pad are None.'
      prototype = np.zeros_like(not_none_tensors[0])
      tensor_list = [t if t is not None else prototype for t in tensor_list]

      max_len = max([t.shape[0] for t in tensor_list])
      if min_slice_length is not None:
        max_len = max(max_len, min_slice_length)
      min_len = min([t.shape[0] for t in tensor_list])
      if max_len == min_len:  # No padding needed.
        return np.array(tensor_list)

      pad_len = 2**int(np.ceil(np.log2(max_len)))
      return np.array([_zero_pad(t, (0, pad_len - t.shape[0]), axis=0)
                       for t in tensor_list])

    trajectory_slice_stream = self.trajectory_slice_stream(
        epochs=epochs,
        max_slice_length=max_slice_length,
        sample_trajectories_uniformly=sample_trajectories_uniformly,
        margin=margin,
        trajectory_stream_preprocessing_fn=trajectory_stream_preprocessing_fn,
    )

    cur_batch = []
    for t in trajectory_slice_stream:
      cur_batch.append(t)
      if len(cur_batch) == batch_size:
        # Make a nested TimeStepBatch of lists out of a list of TimeStepBatches.
        timestep_batch = fastmath.nested_zip(cur_batch)
        # Actions, rewards and returns in the trajectory slice have shape
        # [batch_size, trajectory_length], which we denote as [B, L].
        # Observations are more complex: [B, L] + S, where S is the shape of the
        # observation space (self.observation_space.shape).
        # We stop the recursion at level 1, so we pass lists of arrays into
        # pad().
        yield fastmath.nested_map(pad, timestep_batch, level=1)
        cur_batch = []
Return a stream of trajectory batches from the specified epochs.

    This function returns a stream of tuples of numpy arrays (tensors).
    If tensors have different lengths, they will be padded by 0.

    Args:
      batch_size: the size of the batches to return
      epochs: a list of epochs to use; we use all epochs if None
      max_slice_length: maximum length of the slices of trajectories to return
      min_slice_length: minimum length of the slices of trajectories to return
      margin: number of extra steps after "done" that should be included in
        slices, so that networks see the terminal states in the training data
      sample_trajectories_uniformly: whether to sample trajectories uniformly,
        or proportionally to the number of slices in each trajectory (default)
      trajectory_stream_preprocessing_fn: function to apply to the trajectory
        stream before batching; can be used e.g. to filter trajectories

    Yields:
      batches of trajectory slices sampled uniformly from all slices of length
      at least min_slice_length and up to max_slice_length in all specified
      epochs
    
trax/rl/task.py