def _random_slice(self, trajectory, max_slice_length, margin):
    """Returns a random TimeStepBatch slice from a given trajectory."""
    # Sample a slice from the trajectory.
    slice_start = np.random.randint(
        _n_slices(trajectory, max_slice_length, margin)
    )

    # Convert the whole trajectory to Numpy while adding the margin. The
    # result is cached, so we don't have to repeat this for every sample.
    timestep_batch = trajectory.to_np(margin, self._timestep_to_np)

    # Slice and yield the result.
    slice_end = slice_start + (
        max_slice_length or timestep_batch.observation.shape[0]
    )
    return fastmath.nested_map(
        lambda x: x[slice_start:slice_end], timestep_batch
    )
Returns a random TimeStepBatch slice from a given trajectory.
trax/rl/task.py