def truncate(self, num_to_keep=1):
    """Truncate trajectories, keeping the last `num_to_keep` time-steps."""

    # We return `ts_copy` back to the truncator.
    ts_copy = self._time_steps[:]

    # We keep the last few observations.
    self._time_steps = self._time_steps[-num_to_keep:]

    # NOTE: We will need to set the rewards to 0, to eliminate double counting.
    for i in range(self.num_time_steps):
      self._time_steps[i] = self._time_steps[i].replace(
          raw_reward=0, processed_reward=0)

    return Trajectory(time_steps=ts_copy)
Truncate trajectories, keeping the last `num_to_keep` time-steps.
tensor2tensor/envs/trajectory.py