def test_collects_specified_number_of_trajectories(self):
    """Test that the specified number of interactions are collected."""
    task = rl_task.RLTask(
        DummyEnv(), initial_trajectories=0, max_steps=3, time_limit=20
    )
    task.collect_trajectories(policy=(lambda _: (0, 0)), n_trajectories=3)
    trajectories = task.trajectories[1]  # Get trajectories from epoch 1.
    self.assertLen(trajectories, 3)
Test that the specified number of interactions are collected.
trax/rl/task_test.py