def _trajectory_stream(self, epochs=None, max_slice_length=None,
                         sample_trajectories_uniformly=False, margin=0):
    """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

    Yields:
      random trajectory slices sampled uniformly from all slices of length
      up to max_slice_length in all specified epochs
    """
    # {int: array[int]};
    # epoch_to_ns_slices[epoch][i] = n_slices(self._trajectories[epoch][i])
    # It stores arrays for faster sampling.
    epoch_to_ns_slices = {}
    # {int: int};
    # epoch_to_total_n_slices[epoch] = sum(epoch_to_ns_slices[epoch])
    epoch_to_total_n_slices = {}
    # [int]: list of epoch indices to sample from.
    epoch_indices = []
    # epoch_to_total_n_slices filtered using epoch_indices. It's an array for
    # faster sampling.
    sampling_epoch_weights = None

    def new_epoch(epoch_id):
      """Updates the lists defined above to include the new epoch."""
      all_epochs = list(self._trajectories.keys())
      max_epoch = max(all_epochs) + 1

      # Calculate the numbers of slices for the new epoch.
      epoch_to_ns_slices[epoch_id] = np.array([
          _n_slices(trajectory, max_slice_length, margin)
          for trajectory in self._trajectories[epoch_id]
      ])
      epoch_to_total_n_slices[epoch_id] = np.sum(
          epoch_to_ns_slices[epoch_id]
      )

      # Update the indices of epochs to sample from.
      new_epoch_indices = epochs or all_epochs
      new_epoch_indices = [
          # So -1 means "last".
          ep % max_epoch for ep in new_epoch_indices
      ]
      # Remove duplicates and consider only epochs where some trajectories
      # were recorded and that we have processed in new_epoch.
      new_epoch_indices = [
          epoch_id for epoch_id in set(new_epoch_indices)
          if self._trajectories[epoch_id] and epoch_id in epoch_to_ns_slices
      ]
      epoch_indices[:] = new_epoch_indices

      nonlocal sampling_epoch_weights
      sampling_epoch_weights = np.array(list(
          epoch_to_total_n_slices[ep] for ep in epoch_indices
      ))

    while True:
      # If we haven't collected any trajectories yet, yield an example
      # trajectory. It's needed to determine the input/output shapes of
      # networks.
      if not self._trajectories:
        yield self._example_trajectory
        continue

      # Catch up if we have a new epoch or we've restarted the experiment.
      for epoch_id in self._trajectories.keys() - epoch_to_ns_slices.keys():  # pylint:disable=g-builtin-op
        new_epoch(epoch_id)

      # Sample an epoch proportionally to number of slices in each epoch.
      epoch_id = _sample_proportionally(epoch_indices, sampling_epoch_weights)
      epoch = self._trajectories[epoch_id]

      # Sample a trajectory proportionally to number of slices in each one.
      if sample_trajectories_uniformly:
        slices_per_trajectory = np.ones((len(epoch),))
      else:
        slices_per_trajectory = epoch_to_ns_slices[epoch_id]
      trajectory = _sample_proportionally(epoch, slices_per_trajectory)

      yield trajectory
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

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