Fix type hints errors in gymnasium/spaces

gymnasium/spaces/box.py

@@ -69,8 +69,8 @@ def __init__(
         this value across all dimensions.
 
         Args:
-            low (Union[SupportsFloat, np.ndarray]): Lower bounds of the intervals.
-            high (Union[SupportsFloat, np.ndarray]): Upper bounds of the intervals.
+            low (SupportsFloat | np.ndarray): Lower bounds of the intervals.
+            high (SupportsFloat | np.ndarray]): Upper bounds of the intervals.
             shape (Optional[Sequence[int]]): The shape is inferred from the shape of `low` or `high` `np.ndarray`s with
                 `low` and `high` scalars defaulting to a shape of (1,)
             dtype: The dtype of the elements of the space. If this is an integer type, the :class:`Box` is essentially a discrete space.
@@ -104,12 +104,13 @@ def __init__(
 
         # Capture the boundedness information before replacing np.inf with get_inf
         _low = np.full(shape, low, dtype=float) if is_float_integer(low) else low
-        self.bounded_below: bool = -np.inf < _low
+        self.bounded_below: NDArray[np.bool_] = -np.inf < _low
+
         _high = np.full(shape, high, dtype=float) if is_float_integer(high) else high
-        self.bounded_above: bool = np.inf > _high
+        self.bounded_above: NDArray[np.bool_] = np.inf > _high
 
-        low: NDArray[Any] = _broadcast(low, dtype, shape, inf_sign="-")
-        high: NDArray[Any] = _broadcast(high, dtype, shape, inf_sign="+")
+        low = _broadcast(low, self.dtype, shape, inf_sign="-")
+        high = _broadcast(high, self.dtype, shape, inf_sign="+")
 
         assert isinstance(low, np.ndarray)
         assert (
@@ -280,7 +281,7 @@ def __setstate__(self, state: Iterable[tuple[str, Any]] | Mapping[str, Any]):
             self.high_repr = _short_repr(self.high)
 
 
-def get_inf(dtype: np.dtype, sign: str) -> SupportsFloat:
+def get_inf(dtype: np.dtype, sign: str) -> int | float:
     """Returns an infinite that doesn't break things.
 
     Args:

gymnasium/spaces/dict.py

@@ -110,7 +110,7 @@ def __init__(
             ), f"Dict space element is not an instance of Space: key='{key}', space={space}"
 
         # None for shape and dtype, since it'll require special handling
-        super().__init__(None, None, seed)
+        super().__init__(None, None, seed)  # type: ignore
 
     @property
     def is_np_flattenable(self):
@@ -226,7 +226,9 @@ def to_jsonable(self, sample_n: Sequence[dict[str, Any]]) -> dict[str, list[Any]
             for key, space in self.spaces.items()
         }
 
-    def from_jsonable(self, sample_n: dict[str, list[Any]]) -> list[dict[str, Any]]:
+    def from_jsonable(
+        self, sample_n: dict[str, list[Any]]
+    ) -> list[OrderedDict[str, Any]]:
         """Convert a JSONable data type to a batch of samples from this space."""
         dict_of_list: dict[str, list[Any]] = {
             key: space.from_jsonable(sample_n[key])

gymnasium/spaces/discrete.py

@@ -4,6 +4,7 @@
 from typing import Any, Iterable, Mapping, Sequence
 
 import numpy as np
+from numpy.typing import NDArray
 
 from gymnasium.spaces.space import MaskNDArray, Space
 
@@ -55,7 +56,7 @@ def is_np_flattenable(self):
         """Checks whether this space can be flattened to a :class:`spaces.Box`."""
         return True
 
-    def sample(self, mask: MaskNDArray | None = None) -> int:
+    def sample(self, mask: MaskNDArray | None = None) -> np.int64 | NDArray[np.int64]:
         """Generates a single random sample from this space.
 
         A sample will be chosen uniformly at random with the mask if provided

gymnasium/spaces/graph.py

@@ -230,9 +230,9 @@ def __eq__(self, other: Any) -> bool:
 
     def to_jsonable(
         self, sample_n: Sequence[GraphInstance]
-    ) -> list[dict[str, list[int] | list[float]]]:
+    ) -> list[dict[str, list[int | float]]]:
         """Convert a batch of samples from this space to a JSONable data type."""
-        ret_n: list[dict[str, list[int | float]]] = []
+        ret_n = []
         for sample in sample_n:
             ret = {"nodes": sample.nodes.tolist()}
             if sample.edges is not None and sample.edge_links is not None:

gymnasium/spaces/multi_binary.py

@@ -4,12 +4,12 @@
 from typing import Any, Sequence
 
 import numpy as np
-import numpy.typing as npt
+from numpy.typing import NDArray
 
 from gymnasium.spaces.space import MaskNDArray, Space
 
 
-class MultiBinary(Space[npt.NDArray[np.int8]]):
+class MultiBinary(Space[NDArray[np.int8]]):
     """An n-shape binary space.
 
     Elements of this space are binary arrays of a shape that is fixed during construction.
@@ -28,7 +28,7 @@ class MultiBinary(Space[npt.NDArray[np.int8]]):
 
     def __init__(
         self,
-        n: npt.NDArray[np.integer[Any]] | Sequence[int] | int,
+        n: NDArray[np.integer[Any]] | Sequence[int] | int,
         seed: int | np.random.Generator | None = None,
     ):
         """Constructor of :class:`MultiBinary` space.
@@ -58,7 +58,7 @@ def is_np_flattenable(self):
         """Checks whether this space can be flattened to a :class:`spaces.Box`."""
         return True
 
-    def sample(self, mask: MaskNDArray | None = None) -> npt.NDArray[np.int8]:
+    def sample(self, mask: MaskNDArray | None = None) -> NDArray[np.bool_ | np.int8]:
         """Generates a single random sample from this space.
 
         A sample is drawn by independent, fair coin tosses (one toss per binary variable of the space).
@@ -104,15 +104,11 @@ def contains(self, x: Any) -> bool:
             and np.all(np.logical_or(x == 0, x == 1))
         )
 
-    def to_jsonable(
-        self, sample_n: Sequence[npt.NDArray[np.int8]]
-    ) -> list[Sequence[int]]:
+    def to_jsonable(self, sample_n: Sequence[NDArray[np.int8]]) -> list[Sequence[int]]:
         """Convert a batch of samples from this space to a JSONable data type."""
         return np.array(sample_n).tolist()
 
-    def from_jsonable(
-        self, sample_n: list[Sequence[int]]
-    ) -> list[npt.NDArray[np.int8]]:
+    def from_jsonable(self, sample_n: list[Sequence[int]]) -> list[NDArray[np.int8]]:
         """Convert a JSONable data type to a batch of samples from this space."""
         return [np.asarray(sample, self.dtype) for sample in sample_n]
 

gymnasium/spaces/multi_discrete.py

@@ -4,14 +4,14 @@
 from typing import Any, Sequence
 
 import numpy as np
-import numpy.typing as npt
+from numpy.typing import NDArray
 
 import gymnasium as gym
 from gymnasium.spaces.discrete import Discrete
 from gymnasium.spaces.space import MaskNDArray, Space
 
 
-class MultiDiscrete(Space[npt.NDArray[np.integer]]):
+class MultiDiscrete(Space[NDArray[np.integer]]):
     """This represents the cartesian product of arbitrary :class:`Discrete` spaces.
 
     It is useful to represent game controllers or keyboards where each key can be represented as a discrete action space.
@@ -41,7 +41,7 @@ class MultiDiscrete(Space[npt.NDArray[np.integer]]):
 
     def __init__(
         self,
-        nvec: npt.NDArray[np.integer[Any]] | list[int],
+        nvec: NDArray[np.integer[Any]] | list[int],
         dtype: str | type[np.integer[Any]] = np.int64,
         seed: int | np.random.Generator | None = None,
     ):
@@ -72,7 +72,7 @@ def is_np_flattenable(self):
 
     def sample(
         self, mask: tuple[MaskNDArray, ...] | None = None
-    ) -> npt.NDArray[np.integer[Any]]:
+    ) -> NDArray[np.integer[Any]]:
         """Generates a single random sample this space.
 
         Args:
@@ -88,7 +88,7 @@ def sample(
             def _apply_mask(
                 sub_mask: MaskNDArray | tuple[MaskNDArray, ...],
                 sub_nvec: MaskNDArray | np.integer[Any],
-            ) -> int | Sequence[int]:
+            ) -> int | list[Any]:
                 if isinstance(sub_nvec, np.ndarray):
                     assert isinstance(
                         sub_mask, tuple
@@ -144,14 +144,14 @@ def contains(self, x: Any) -> bool:
         )
 
     def to_jsonable(
-        self, sample_n: Sequence[npt.NDArray[np.integer[Any]]]
+        self, sample_n: Sequence[NDArray[np.integer[Any]]]
     ) -> list[Sequence[int]]:
         """Convert a batch of samples from this space to a JSONable data type."""
         return [sample.tolist() for sample in sample_n]
 
     def from_jsonable(
         self, sample_n: list[Sequence[int]]
-    ) -> list[npt.NDArray[np.integer[Any]]]:
+    ) -> list[NDArray[np.integer[Any]]]:
         """Convert a JSONable data type to a batch of samples from this space."""
         return [np.array(sample) for sample in sample_n]
 

gymnasium/spaces/sequence.py

@@ -5,7 +5,7 @@
 from typing import Any, Union
 
 import numpy as np
-import numpy.typing as npt
+from numpy.typing import NDArray
 
 import gymnasium as gym
 from gymnasium.spaces.space import Space
@@ -69,11 +69,11 @@ def sample(
         mask: None
         | (
             tuple[
-                None | np.integer | npt.NDArray[np.integer],
+                None | np.integer | NDArray[np.integer],
                 Any,
             ]
         ) = None,
-    ) -> tuple[Any]:
+    ) -> tuple[Any] | Any:
         """Generates a single random sample from this space.
 
         Args:

gymnasium/spaces/space.py

@@ -104,8 +104,8 @@ def sample(self, mask: Any | None = None) -> T_cov:
 
     def seed(self, seed: int | None = None) -> list[int]:
         """Seed the PRNG of this space and possibly the PRNGs of subspaces."""
-        self._np_random, seed = seeding.np_random(seed)
-        return [seed]
+        self._np_random, np_random_seed = seeding.np_random(seed)
+        return [np_random_seed]
 
     def contains(self, x: Any) -> bool:
         """Return boolean specifying if x is a valid member of this space."""

gymnasium/spaces/text.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 import numpy as np
-import numpy.typing as npt
+from numpy.typing import NDArray
 
 from gymnasium.spaces.space import Space
 
@@ -35,7 +35,7 @@ def __init__(
         max_length: int,
         *,
         min_length: int = 1,
-        charset: set[str] | str = alphanumeric,
+        charset: frozenset[str] | str = alphanumeric,
         seed: int | np.random.Generator | None = None,
     ):
         r"""Constructor of :class:`Text` space.
@@ -76,7 +76,7 @@ def __init__(
 
     def sample(
         self,
-        mask: None | (tuple[int | None, npt.NDArray[np.int8] | None]) = None,
+        mask: None | (tuple[int | None, NDArray[np.int8] | None]) = None,
     ) -> str:
         """Generates a single random sample from this space with by default a random length between `min_length` and `max_length` and sampled from the `charset`.
 

gymnasium/spaces/utils.py

@@ -184,15 +184,17 @@ def _flatten_multidiscrete(
 def _flatten_tuple(space: Tuple, x: tuple[Any, ...]) -> tuple[Any, ...] | NDArray[Any]:
     if space.is_np_flattenable:
         return np.concatenate(
-            [flatten(s, x_part) for x_part, s in zip(x, space.spaces)]
+            [np.array(flatten(s, x_part)) for x_part, s in zip(x, space.spaces)]
         )
     return tuple(flatten(s, x_part) for x_part, s in zip(x, space.spaces))
 
 
 @flatten.register(Dict)
 def _flatten_dict(space: Dict, x: dict[str, Any]) -> dict[str, Any] | NDArray[Any]:
     if space.is_np_flattenable:
-        return np.concatenate([flatten(s, x[key]) for key, s in space.spaces.items()])
+        return np.concatenate(
+            [np.array(flatten(s, x[key])) for key, s in space.spaces.items()]
+        )
     return OrderedDict((key, flatten(s, x[key])) for key, s in space.spaces.items())
 
 
@@ -354,7 +356,17 @@ def _unflatten_graph(space: Graph, x: GraphInstance) -> GraphInstance:
     nodes and edges in the graph.
     """
 
-    def _graph_unflatten(unflatten_space, unflatten_x):
+    def _graph_unflatten_nodes(
+        unflatten_space: Box | Discrete, unflatten_x: NDArray[Any]
+    ) -> NDArray[Any]:
+        if isinstance(unflatten_space, Box):
+            return unflatten_x.reshape(-1, *unflatten_space.shape)
+        elif isinstance(unflatten_space, Discrete):
+            return np.asarray(np.nonzero(unflatten_x))[-1, :]
+
+    def _graph_unflatten_edges(
+        unflatten_space: Box | Discrete | None, unflatten_x: NDArray[Any] | None
+    ) -> NDArray[Any] | None:
         result = None
         if unflatten_space is not None and unflatten_x is not None:
             if isinstance(unflatten_space, Box):
@@ -363,8 +375,8 @@ def _graph_unflatten(unflatten_space, unflatten_x):
                 result = np.asarray(np.nonzero(unflatten_x))[-1, :]
         return result
 
-    nodes = _graph_unflatten(space.node_space, x.nodes)
-    edges = _graph_unflatten(space.edge_space, x.edges)
+    nodes = _graph_unflatten_nodes(space.node_space, x.nodes)
+    edges = _graph_unflatten_edges(space.edge_space, x.edges)
 
     return GraphInstance(nodes, edges, x.edge_links)