Add proper handling of nullable boxes

Farama-Foundation · Dec 4, 2023 · 27bf126 · 27bf126
1 parent 35f1805
commit 27bf126
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/gymnasium/spaces/box.py b/gymnasium/spaces/box.py
@@ -246,11 +246,17 @@ def contains(self, x: Any) -> bool:
             except (ValueError, TypeError):
                 return False
 
+        bounded_below = x >= self.low
+        bounded_above = x <= self.high
+
+        bounded = bounded_below & bounded_above
+        if self.nullable:
+            bounded |= np.isnan(x)
+
         return bool(
             np.can_cast(x.dtype, self.dtype)
             and x.shape == self.shape
-            and np.all(x >= self.low)
-            and np.all(x <= self.high)
+            and np.all(bounded)
         )
 
     def to_jsonable(self, sample_n: Sequence[NDArray[Any]]) -> list[list]:

diff --git a/gymnasium/spaces/utils.py b/gymnasium/spaces/utils.py
@@ -575,4 +575,4 @@ def _flatten_space_oneof(space: OneOf) -> Box:
     high = np.concatenate([[num_subspaces - 1], np.full(max_flatdim - 1, overall_high)])
 
     dtype = np.result_type(*[s.dtype for s in space.spaces if hasattr(s, "dtype")])
-    return Box(low=low, high=high, shape=(max_flatdim,), dtype=dtype)
+    return Box(low=low, high=high, shape=(max_flatdim,), dtype=dtype, nullable=True)