[RLlib] Fix assertion in Simplex space. (ray-project#47880)

Signed-off-by: ujjawal-khare <[email protected]>
ujjawal-khare-27 · Oct 15, 2024 · 0224d7b · 0224d7b
1 parent 960e28d
commit 0224d7b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/rllib/utils/spaces/simplex.py b/rllib/utils/spaces/simplex.py
@@ -29,8 +29,8 @@ def __init__(self, shape, concentration=None, dtype=np.float32):
 
         if concentration is not None:
             assert (
-                concentration.shape == shape[:-1]
-            ), f"{concentration.shape} vs {shape[:-1]}"
+                concentration.shape[0] == shape[-1]
+            ), f"{concentration.shape[0]} vs {shape[-1]}"
             self.concentration = concentration
         else:
             self.concentration = np.array([1] * self.dim)

diff --git a/rllib/utils/tests/test_serialization.py b/rllib/utils/tests/test_serialization.py
@@ -141,7 +141,7 @@ def test_dict_space_with_ordered_dict(self):
         self.assertTrue(input_space == deserialized_space)
 
     def test_simplex_space(self):
-        space = Simplex(shape=(3, 4), concentration=np.array((1, 2, 1)))
+        space = Simplex(shape=(3, 4), concentration=np.array((1, 2, 1, 2)))
 
         d = gym_space_to_dict(space)
         sp = gym_space_from_dict(d)