(self, margin=1, timestep_to_np=None)
return timestep_batch
Create a tuple of numpy arrays from a given trajectory.

    Args:
        margin (int): Number of dummy timesteps past the trajectory end to
            include. By default we include 1, which contains the last
            observation.
        timestep_to_np (callable or None): Optional function
            TimeStepBatch[Any] -> TimeStepBatch[np.array], converting the
            timestep data into numpy arrays.

    Returns:
        TimeStepBatch, where all fields have shape
        (len(self) + margin - 1, ...).
Create a tuple of numpy arrays from a given trajectory.
Create a tuple of numpy arrays from a given trajectory .
def to_np(self, margin=1, timestep_to_np=None):
    """Create a tuple of numpy arrays from a given trajectory.

    Args:
        margin (int): Number of dummy timesteps past the trajectory end to
            include. By default we include 1, which contains the last
            observation.
        timestep_to_np (callable or None): Optional function
            TimeStepBatch[Any] -> TimeStepBatch[np.array], converting the
            timestep data into numpy arrays.

    Returns:
        TimeStepBatch, where all fields have shape
        (len(self) + margin - 1, ...).
    """
    timestep_to_np = timestep_to_np or self._default_timestep_to_np
    args = (margin, timestep_to_np)

    # Return the cached result if the arguments agree and the trajectory has not
    # grown.
    if self._timestep_batch:
      result_length = len(self) + margin - 1
      length_ok = self._timestep_batch.observation.shape[0] == result_length
      if args == self._cached_to_np_args and length_ok:
        return self._timestep_batch

    # observation, action, reward, etc.
    fields = TimeStepBatch._fields
    # List of timestep data for each field.
    data_lists = TimeStepBatch(**{field: [] for field in fields})
    for timestep in self._timesteps:
      timestep_np = timestep_to_np(timestep)
      # Append each field of timestep_np to the appropriate list.
      for field in fields:
        getattr(data_lists, field).append(getattr(timestep_np, field))
    # Append the last observation.
    data_lists.observation.append(self._last_observation)

    # TODO(pkozakowski): The case len(obs) == 1 is for handling
    # "dummy trajectories", that are only there to determine data shapes. Check
    # if they're still required.
    if len(data_lists.observation) > 1:
      # Extend the trajectory with a given margin - this is to make sure that
      # the networks always "see" the "done" states in the training data, even
      # when a suffix is added to the trajectory slice for better estimation of
      # returns.
      # We set `mask` to 0, so the added timesteps don't influence the loss. We
      # set `done` to True for easier implementation of advantage estimators.
      # The rest of the fields don't matter, so we set them to 0 for easy
      # debugging (unless they're None). The list of observations is longer, so
      # we pad it with margin - 1.
      data_lists.mask.extend([0] * margin)
      data_lists.done.extend([True] * margin)
      data_lists.observation.extend(
          [np.zeros_like(data_lists.observation[-1])] * (margin - 1)
      )
      for field in set(fields) - {'mask', 'done', 'observation'}:
        l = getattr(data_lists, field)
        filler = None if l[-1] is None else np.zeros_like(l[-1])
        l.extend([filler] * margin)

      # Trim the observations to have the same length as the rest of the fields.
      # This is not be the case when margin=0.
      if margin == 0:
        data_lists.observation.pop()

    def stack(x):
      if not x:
        return None
      return fastmath.nested_stack(x)

    # Stack the data_lists into numpy arrays.
    timestep_batch = TimeStepBatch(*map(stack, data_lists))

    self._timestep_batch = timestep_batch
    self._cached_to_np_args = args

    return timestep_batch
def to_np ( self , margin = 1 , timestep_to_np = None ) : timestep_to_np = timestep_to_np or self . _default_timestep_to_np args = ( margin , timestep_to_np ) # Return the cached result if the arguments agree and the trajectory has not # grown. if self . _timestep_batch : result_length = len ( self ) + margin - 1 length_ok = self . _timestep_batch . observation . shape [ 0 ] == result_length if args == self . _cached_to_np_args and length_ok : return self . _timestep_batch # observation, action, reward, etc. fields = TimeStepBatch . _fields # List of timestep data for each field. data_lists = TimeStepBatch ( * * { field : [ ] for field in fields } ) for timestep in self . _timesteps : timestep_np = timestep_to_np ( timestep ) # Append each field of timestep_np to the appropriate list. for field in fields : getattr ( data_lists , field ) . append ( getattr ( timestep_np , field ) ) # Append the last observation. data_lists . observation . append ( self . _last_observation ) # TODO(pkozakowski): The case len(obs) == 1 is for handling # "dummy trajectories", that are only there to determine data shapes. Check # if they're still required. if len ( data_lists . observation ) > 1 : # Extend the trajectory with a given margin - this is to make sure that # the networks always "see" the "done" states in the training data, even # when a suffix is added to the trajectory slice for better estimation of # returns. # We set `mask` to 0, so the added timesteps don't influence the loss. We # set `done` to True for easier implementation of advantage estimators. # The rest of the fields don't matter, so we set them to 0 for easy # debugging (unless they're None). The list of observations is longer, so # we pad it with margin - 1. data_lists . mask . extend ( [ 0 ] * margin ) data_lists . done . extend ( [ True ] * margin ) data_lists . observation . extend ( [ np . zeros_like ( data_lists . observation [ - 1 ] ) ] * ( margin - 1 ) ) for field in set ( fields ) - { 'mask' , 'done' , 'observation' } : l = getattr ( data_lists , field ) filler = None if l [ - 1 ] is None else np . zeros_like ( l [ - 1 ] ) l . extend ( [ filler ] * margin ) # Trim the observations to have the same length as the rest of the fields. # This is not be the case when margin=0. if margin == 0 : data_lists . observation . pop ( ) def stack ( x ) : if not x : return None return fastmath . nested_stack ( x ) # Stack the data_lists into numpy arrays. timestep_batch = TimeStepBatch ( * map ( stack , data_lists ) ) self . _timestep_batch = timestep_batch self . _cached_to_np_args = args return timestep_batch