class_docstr: A RL task: environment and a collection of trajectories.
class_docstr_tok: ["A","RL","task",":","environment","and","a","collection","of","trajectories","."]
docstr: 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
docstr_tok: ["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"]
code: 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
  
imports: import collections
  import os
  import gin
  import gym
  import numpy
  from trax import fastmath
  from trax.rl import advantages
  from trax.supervised import training
label_desc: Use image node_id 21 for calling the RLTask obj's underlying member method code with example usage: obj._trajectory_stream(epochs, max_slice_length, sample_trajectories_uniformly, margin) without return types