def test_task_epochs_index_minusone(self):
    """Test that the epoch index -1 means last epoch and updates to it."""
    obs = np.zeros((2,))
    tr1 = rl_task.Trajectory(obs)
    self._extend(tr1, new_observation=obs, done=True)
    task = rl_task.RLTask(DummyEnv(), initial_trajectories=[tr1], max_steps=9)
    stream = task.trajectory_slice_stream(epochs=[-1], max_slice_length=1)
    next_slice = next(stream)
    np.testing.assert_equal(next_slice.observation, np.zeros((1, 2)))
    task.collect_trajectories(policy=(lambda _: (0, 0)), n_trajectories=1)
    next_slice = next(stream)
    np.testing.assert_equal(next_slice.observation, np.ones((1, 2)))
Test that the epoch index -1 means last epoch and updates to it.
trax/rl/task_test.py