def test_collects_specified_number_of_interactions(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_interactions=10)
    trajectories = task.trajectories[1]  # Get trajectories from epoch 1.
    n_interactions = 0
    for trajectory in trajectories:
      n_interactions += len(trajectory) - 1
    self.assertEqual(n_interactions, 10)
Test that the specified number of interactions are collected.
trax/rl/task_test.py