def test_trajectory_to_np_shape(self):
    """Test that the shape of a to_np result matches the trajectory length."""
    tr = rl_task.Trajectory(observation=np.zeros((2, 3)))
    for _ in range(5):
      self._extend(tr, new_observation=np.zeros((2, 3)))
    tr_np = tr.to_np()
    self.assertEqual(tr_np.observation.shape, (len(tr), 2, 3))
    self.assertEqual(tr_np.action.shape, (len(tr),))
Test that the shape of a to_np result matches the trajectory length.
trax/rl/task_test.py