def done(self, done):
    """Sets the `done` flag in the last timestep."""
    if not self._timesteps:
      raise ValueError('No interactions yet in the trajectory.')
    self._timesteps[-1] = self._timesteps[-1]._replace(done=done)
Sets the `done` flag in the last timestep.
trax/rl/task.py