def test_trajectory_len(self):
    """Test that trajectory length is equal to the number of observations."""
    tr = rl_task.Trajectory(observation=0)
    for _ in range(5):
      self._extend(tr)
    self.assertLen(tr, 6)
Test that trajectory length is equal to the number of observations.
trax/rl/task_test.py