def test_task_random_initial_trajectories_and_max_steps(self):
    """Test generating initial random trajectories, stop at max steps."""
    task = rl_task.RLTask(DummyEnv(), initial_trajectories=1, max_steps=9)
    stream = task.trajectory_slice_stream(max_slice_length=1)
    next_slice = next(stream)
    self.assertEqual(next_slice.observation.shape, (1, 2))
Test generating initial random trajectories, stop at max steps.
trax/rl/task_test.py