def test_max_steps_doesnt_terminate_epsiodes(self):
    """Test that episodes are not terminated upon reaching `max_steps` steps."""
    task = rl_task.RLTask(
        DummyEnv(), initial_trajectories=2, max_steps=5, time_limit=10
    )
    trajectories = task.trajectories[0]  # Get trajectories from epoch 0.
    self.assertLen(trajectories, 2)
    # The trajectory should be cut in half. The first half should not be "done".
    self.assertFalse(trajectories[0].done)
    self.assertLen(trajectories[0], 6)  # max_steps + 1
    # The second half should be "done".
    self.assertTrue(trajectories[1].done)
    self.assertLen(trajectories[1], 6)  # max_steps + 1
Test that episodes are not terminated upon reaching `max_steps` steps.
trax/rl/task_test.py