def trajectory_slice_stream(
      self,
      epochs=None,
      max_slice_length=None,
      sample_trajectories_uniformly=False,
      margin=0,
      trajectory_stream_preprocessing_fn=None,
  ):
    """Return a stream of random trajectory slices from the specified epochs.

    Args:
      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
      sample_trajectories_uniformly: whether to sample trajectories uniformly,
        or proportionally to the number of slices in each trajectory (default)
      margin: number of extra steps after "done" that should be included in
        slices, so that networks see the terminal states in the training data
      trajectory_stream_preprocessing_fn: function to apply to the trajectory
        stream before batching; can be used e.g. to filter trajectories

    Yields:
      random trajectory slices sampled uniformly from all slices of length
      up to max_slice_length in all specified epochs
    """
    trajectory_stream = self._trajectory_stream(
        epochs=epochs,
        max_slice_length=max_slice_length,
        sample_trajectories_uniformly=sample_trajectories_uniformly,
        margin=margin,
    )

    if trajectory_stream_preprocessing_fn is not None:
      trajectory_stream = trajectory_stream_preprocessing_fn(trajectory_stream)

    for trajectory in trajectory_stream:
      yield self._random_slice(trajectory, max_slice_length, margin)
Return a stream of random trajectory slices from the specified epochs.

    Args:
      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
      sample_trajectories_uniformly: whether to sample trajectories uniformly,
        or proportionally to the number of slices in each trajectory (default)
      margin: number of extra steps after "done" that should be included in
        slices, so that networks see the terminal states in the training data
      trajectory_stream_preprocessing_fn: function to apply to the trajectory
        stream before batching; can be used e.g. to filter trajectories

    Yields:
      random trajectory slices sampled uniformly from all slices of length
      up to max_slice_length in all specified epochs
    
trax/rl/task.py