def test_trajectory_to_np_shape_after_extend(self):
    """Test that the shape of a to_np result grows after calling extend()."""
    tr = rl_task.Trajectory(observation=0)
    for _ in range(5):
      self._extend(tr)
    len_before = tr.to_np().observation.shape[0]
    self._extend(tr)
    len_after = tr.to_np().observation.shape[0]
    self.assertEqual(len_after, len_before + 1)
Test that the shape of a to_np result grows after calling extend().
trax/rl/task_test.py