def test_trajectory_slice_stream_shape(self):
    """Test the shape yielded by trajectory stream."""
    obs = np.zeros((12, 13))
    tr1 = rl_task.Trajectory(obs)
    self._extend(tr1, new_observation=obs, done=True)
    task = rl_task.RLTask(DummyEnv(), initial_trajectories=[tr1], max_steps=9)
    stream = task.trajectory_slice_stream(max_slice_length=1)
    next_slice = next(stream)
    self.assertEqual(next_slice.observation.shape, (1, 12, 13))
Test the shape yielded by trajectory stream.
trax/rl/task_test.py