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 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
docstr_tok: ["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"]
code: 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 = []
  
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 23 for calling the RLTask obj's underlying member method code with example usage: obj.trajectory_batch_stream(batch_size, epochs, max_slice_length, min_slice_length, margin, sample_trajectories_uniformly, trajectory_stream_preprocessing_fn) and returns: np, np