def extend(self, new_observation, mask=1, **kwargs):
    """Take action in the last state, getting reward and going to new state."""
    self._timesteps.append(TimeStepBatch(
        observation=self._last_observation, mask=mask, **kwargs
    ))
    self._last_observation = new_observation
Take action in the last state, getting reward and going to new state.
trax/rl/task.py