def test_trajectory_slice_stream_sampling_uniform(self):
    """Test if the trajectory stream samples uniformly."""
    # Long trajectory of 0s.
    tr1 = rl_task.Trajectory(0)
    for _ in range(100):
      self._extend(tr1)
    self._extend(tr1, new_observation=200, done=True)
    # Short trajectory of 101.
    tr2 = rl_task.Trajectory(101)
    self._extend(tr2, new_observation=200, done=True)
    task = rl_task.RLTask(
        DummyEnv(), initial_trajectories=[tr1, tr2], max_steps=9)

    # Stream of both. Check that we're sampling by slice, not by trajectory.
    stream = task.trajectory_slice_stream(max_slice_length=1)
    slices = []
    for _ in range(10):
      next_slice = next(stream)
      assert next_slice.observation.shape[0] == 1
      slices.append(next_slice.observation[-1])
    mean_obs = sum(slices) / float(len(slices))
    # Average should be around 1 sampling from 0x100, 101 uniformly.
    self.assertLess(mean_obs, 31)  # Sampling 101 even 3 times is unlikely.
    self.assertLen(slices, 10)
Test if the trajectory stream samples uniformly.
trax/rl/task_test.py