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
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, ...).
    
trax/rl/task.py