def test_flatten_n_tuple(self): box = Box(0.0, 1.0, (3, 4)) obs = ((([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]), ([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]), ([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]))) arr = box.flatten_n(obs) assert arr.shape == (3, 12)
def test_flatten_n(self): box = Box(0.0, 1.0, (3, 4)) obs = np.asarray((([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]), ([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]), ([1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]))) arr = box.flatten_n(obs) assert obs.shape == (3, 4, 3) assert arr.shape == (3, 12)