I am done

rl/Lib/site-packages/gymnasium/spaces/__init__.py

@@ -0,0 +1,46 @@
+"""This module implements various spaces.
+
+Spaces describe mathematical sets and are used in Gym to specify valid actions and observations.
+Every Gym environment must have the attributes ``action_space`` and ``observation_space``.
+If, for instance, three possible actions (0,1,2) can be performed in your environment and observations
+are vectors in the two-dimensional unit cube, the environment code may contain the following two lines::
+
+    self.action_space = spaces.Discrete(3)
+    self.observation_space = spaces.Box(0, 1, shape=(2,))
+
+All spaces inherit from the :class:`Space` superclass.
+"""
+from gymnasium.spaces.box import Box
+from gymnasium.spaces.dict import Dict
+from gymnasium.spaces.discrete import Discrete
+from gymnasium.spaces.graph import Graph, GraphInstance
+from gymnasium.spaces.multi_binary import MultiBinary
+from gymnasium.spaces.multi_discrete import MultiDiscrete
+from gymnasium.spaces.sequence import Sequence
+from gymnasium.spaces.space import Space
+from gymnasium.spaces.text import Text
+from gymnasium.spaces.tuple import Tuple
+from gymnasium.spaces.utils import flatdim, flatten, flatten_space, unflatten
+
+
+__all__ = [
+    # base space
+    "Space",
+    # fundamental spaces
+    "Box",
+    "Discrete",
+    "Text",
+    "MultiDiscrete",
+    "MultiBinary",
+    # composite spaces
+    "Graph",
+    "GraphInstance",
+    "Tuple",
+    "Sequence",
+    "Dict",
+    # util functions (there are more utility functions in vector/utils/spaces.py)
+    "flatdim",
+    "flatten_space",
+    "flatten",
+    "unflatten",
+]

rl/Lib/site-packages/gymnasium/spaces/box.py

@@ -0,0 +1,337 @@
+"""Implementation of a space that represents closed boxes in euclidean space."""
+from __future__ import annotations
+
+from typing import Any, Iterable, Mapping, Sequence, SupportsFloat
+
+import numpy as np
+from numpy.typing import NDArray
+
+import gymnasium as gym
+from gymnasium.spaces.space import Space
+
+
+def _short_repr(arr: NDArray[Any]) -> str:
+    """Create a shortened string representation of a numpy array.
+
+    If arr is a multiple of the all-ones vector, return a string representation of the multiplier.
+    Otherwise, return a string representation of the entire array.
+
+    Args:
+        arr: The array to represent
+
+    Returns:
+        A short representation of the array
+    """
+    if arr.size != 0 and np.min(arr) == np.max(arr):
+        return str(np.min(arr))
+    return str(arr)
+
+
+def is_float_integer(var: Any) -> bool:
+    """Checks if a variable is an integer or float."""
+    return np.issubdtype(type(var), np.integer) or np.issubdtype(type(var), np.floating)
+
+
+class Box(Space[NDArray[Any]]):
+    r"""A (possibly unbounded) box in :math:`\mathbb{R}^n`.
+
+    Specifically, a Box represents the Cartesian product of n closed intervals.
+    Each interval has the form of one of :math:`[a, b]`, :math:`(-\infty, b]`,
+    :math:`[a, \infty)`, or :math:`(-\infty, \infty)`.
+
+    There are two common use cases:
+
+    * Identical bound for each dimension::
+
+        >>> Box(low=-1.0, high=2.0, shape=(3, 4), dtype=np.float32)
+        Box(-1.0, 2.0, (3, 4), float32)
+
+    * Independent bound for each dimension::
+
+        >>> Box(low=np.array([-1.0, -2.0]), high=np.array([2.0, 4.0]), dtype=np.float32)
+        Box([-1. -2.], [2. 4.], (2,), float32)
+    """
+
+    def __init__(
+        self,
+        low: SupportsFloat | NDArray[Any],
+        high: SupportsFloat | NDArray[Any],
+        shape: Sequence[int] | None = None,
+        dtype: type[np.floating[Any]] | type[np.integer[Any]] = np.float32,
+        seed: int | np.random.Generator | None = None,
+    ):
+        r"""Constructor of :class:`Box`.
+
+        The argument ``low`` specifies the lower bound of each dimension and ``high`` specifies the upper bounds.
+        I.e., the space that is constructed will be the product of the intervals :math:`[\text{low}[i], \text{high}[i]]`.
+
+        If ``low`` (or ``high``) is a scalar, the lower bound (or upper bound, respectively) will be assumed to be
+        this value across all dimensions.
+
+        Args:
+            low (SupportsFloat | np.ndarray): Lower bounds of the intervals. If integer, must be at least ``-2**63``.
+            high (SupportsFloat | np.ndarray]): Upper bounds of the intervals. If integer, must be at most ``2**63 - 2``.
+            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.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
+
+        Raises:
+            ValueError: If no shape information is provided (shape is None, low is None and high is None) then a
+                value error is raised.
+        """
+        assert (
+            dtype is not None
+        ), "Box dtype must be explicitly provided, cannot be None."
+        self.dtype = np.dtype(dtype)
+
+        # determine shape if it isn't provided directly
+        if shape is not None:
+            assert all(
+                np.issubdtype(type(dim), np.integer) for dim in shape
+            ), f"Expected all shape elements to be an integer, actual type: {tuple(type(dim) for dim in shape)}"
+            shape = tuple(int(dim) for dim in shape)  # This changes any np types to int
+        elif isinstance(low, np.ndarray):
+            shape = low.shape
+        elif isinstance(high, np.ndarray):
+            shape = high.shape
+        elif is_float_integer(low) and is_float_integer(high):
+            shape = (1,)
+        else:
+            raise ValueError(
+                f"Box shape is inferred from low and high, expected their types to be np.ndarray, an integer or a float, actual type low: {type(low)}, high: {type(high)}"
+            )
+
+        # 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: NDArray[np.bool_] = -np.inf < _low
+
+        _high = np.full(shape, high, dtype=float) if is_float_integer(high) else high
+        self.bounded_above: NDArray[np.bool_] = np.inf > _high
+
+        low = _broadcast(low, self.dtype, shape)
+        high = _broadcast(high, self.dtype, shape)
+
+        assert isinstance(low, np.ndarray)
+        assert (
+            low.shape == shape
+        ), f"low.shape doesn't match provided shape, low.shape: {low.shape}, shape: {shape}"
+        assert isinstance(high, np.ndarray)
+        assert (
+            high.shape == shape
+        ), f"high.shape doesn't match provided shape, high.shape: {high.shape}, shape: {shape}"
+
+        # check that we don't have invalid low or high
+        if np.any(low > high):
+            raise ValueError(
+                f"Some low values are greater than high, low={low}, high={high}"
+            )
+        if np.any(np.isposinf(low)):
+            raise ValueError(f"No low value can be equal to `np.inf`, low={low}")
+        if np.any(np.isneginf(high)):
+            raise ValueError(f"No high value can be equal to `-np.inf`, high={high}")
+
+        self._shape: tuple[int, ...] = shape
+
+        low_precision = get_precision(low.dtype)
+        high_precision = get_precision(high.dtype)
+        dtype_precision = get_precision(self.dtype)
+        if min(low_precision, high_precision) > dtype_precision:
+            gym.logger.warn(f"Box bound precision lowered by casting to {self.dtype}")
+        self.low = low.astype(self.dtype)
+        self.high = high.astype(self.dtype)
+
+        self.low_repr = _short_repr(self.low)
+        self.high_repr = _short_repr(self.high)
+
+        super().__init__(self.shape, self.dtype, seed)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Has stricter type than gym.Space - never None."""
+        return self._shape
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return True
+
+    def is_bounded(self, manner: str = "both") -> bool:
+        """Checks whether the box is bounded in some sense.
+
+        Args:
+            manner (str): One of ``"both"``, ``"below"``, ``"above"``.
+
+        Returns:
+            If the space is bounded
+
+        Raises:
+            ValueError: If `manner` is neither ``"both"`` nor ``"below"`` or ``"above"``
+        """
+        below = bool(np.all(self.bounded_below))
+        above = bool(np.all(self.bounded_above))
+        if manner == "both":
+            return below and above
+        elif manner == "below":
+            return below
+        elif manner == "above":
+            return above
+        else:
+            raise ValueError(
+                f"manner is not in {{'below', 'above', 'both'}}, actual value: {manner}"
+            )
+
+    def sample(self, mask: None = None) -> NDArray[Any]:
+        r"""Generates a single random sample inside the Box.
+
+        In creating a sample of the box, each coordinate is sampled (independently) from a distribution
+        that is chosen according to the form of the interval:
+
+        * :math:`[a, b]` : uniform distribution
+        * :math:`[a, \infty)` : shifted exponential distribution
+        * :math:`(-\infty, b]` : shifted negative exponential distribution
+        * :math:`(-\infty, \infty)` : normal distribution
+
+        Args:
+            mask: A mask for sampling values from the Box space, currently unsupported.
+
+        Returns:
+            A sampled value from the Box
+        """
+        if mask is not None:
+            raise gym.error.Error(
+                f"Box.sample cannot be provided a mask, actual value: {mask}"
+            )
+
+        high = self.high if self.dtype.kind == "f" else self.high.astype("int64") + 1
+        sample = np.empty(self.shape)
+
+        # Masking arrays which classify the coordinates according to interval type
+        unbounded = ~self.bounded_below & ~self.bounded_above
+        upp_bounded = ~self.bounded_below & self.bounded_above
+        low_bounded = self.bounded_below & ~self.bounded_above
+        bounded = self.bounded_below & self.bounded_above
+
+        # Vectorized sampling by interval type
+        sample[unbounded] = self.np_random.normal(size=unbounded[unbounded].shape)
+
+        sample[low_bounded] = (
+            self.np_random.exponential(size=low_bounded[low_bounded].shape)
+            + self.low[low_bounded]
+        )
+
+        sample[upp_bounded] = (
+            -self.np_random.exponential(size=upp_bounded[upp_bounded].shape)
+            + high[upp_bounded]
+        )
+
+        sample[bounded] = self.np_random.uniform(
+            low=self.low[bounded], high=high[bounded], size=bounded[bounded].shape
+        )
+
+        if self.dtype.kind in ["i", "u", "b"]:
+            sample = np.floor(sample)
+
+        return sample.astype(self.dtype)
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if not isinstance(x, np.ndarray):
+            gym.logger.warn("Casting input x to numpy array.")
+            try:
+                x = np.asarray(x, dtype=self.dtype)
+            except (ValueError, TypeError):
+                return False
+
+        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)
+        )
+
+    def to_jsonable(self, sample_n: Sequence[NDArray[Any]]) -> list[list]:
+        """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: Sequence[float | int]) -> list[NDArray[Any]]:
+        """Convert a JSONable data type to a batch of samples from this space."""
+        return [np.asarray(sample, dtype=self.dtype) for sample in sample_n]
+
+    def __repr__(self) -> str:
+        """A string representation of this space.
+
+        The representation will include bounds, shape and dtype.
+        If a bound is uniform, only the corresponding scalar will be given to avoid redundant and ugly strings.
+
+        Returns:
+            A representation of the space
+        """
+        return f"Box({self.low_repr}, {self.high_repr}, {self.shape}, {self.dtype})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether `other` is equivalent to this instance. Doesn't check dtype equivalence."""
+        return (
+            isinstance(other, Box)
+            and (self.shape == other.shape)
+            # and (self.dtype == other.dtype)
+            and np.allclose(self.low, other.low)
+            and np.allclose(self.high, other.high)
+        )
+
+    def __setstate__(self, state: Iterable[tuple[str, Any]] | Mapping[str, Any]):
+        """Sets the state of the box for unpickling a box with legacy support."""
+        super().__setstate__(state)
+
+        # legacy support through re-adding "low_repr" and "high_repr" if missing from pickled state
+        if not hasattr(self, "low_repr"):
+            self.low_repr = _short_repr(self.low)
+
+        if not hasattr(self, "high_repr"):
+            self.high_repr = _short_repr(self.high)
+
+
+def get_precision(dtype: np.dtype) -> SupportsFloat:
+    """Get precision of a data type."""
+    if np.issubdtype(dtype, np.floating):
+        return np.finfo(dtype).precision
+    else:
+        return np.inf
+
+
+def _broadcast(
+    value: SupportsFloat | NDArray[Any],
+    dtype: np.dtype,
+    shape: tuple[int, ...],
+) -> NDArray[Any]:
+    """Handle infinite bounds and broadcast at the same time if needed.
+
+    This is needed primarily because:
+        >>> import numpy as np
+        >>> np.full((2,), np.inf, dtype=np.int32)
+        array([-2147483648, -2147483648], dtype=int32)
+    """
+    if is_float_integer(value):
+        if np.isneginf(value) and np.dtype(dtype).kind == "i":
+            value = np.iinfo(dtype).min + 2
+        elif np.isposinf(value) and np.dtype(dtype).kind == "i":
+            value = np.iinfo(dtype).max - 2
+
+        return np.full(shape, value, dtype=dtype)
+
+    elif isinstance(value, np.ndarray):
+        # this is needed because we can't stuff np.iinfo(int).min into an array of dtype float
+        casted_value = value.astype(dtype)
+
+        # change bounds only if values are negative or positive infinite
+        if np.dtype(dtype).kind == "i":
+            casted_value[np.isneginf(value)] = np.iinfo(dtype).min + 2
+            casted_value[np.isposinf(value)] = np.iinfo(dtype).max - 2
+
+        return casted_value
+
+    else:
+        # only np.ndarray allowed beyond this point
+        raise TypeError(
+            f"Unknown dtype for `value`, expected `np.ndarray` or float/integer, got {type(value)}"
+        )

rl/Lib/site-packages/gymnasium/spaces/dict.py

@@ -0,0 +1,238 @@
+"""Implementation of a space that represents the cartesian product of other spaces as a dictionary."""
+from __future__ import annotations
+
+import collections.abc
+import typing
+from collections import OrderedDict
+from typing import Any, KeysView, Sequence
+
+import numpy as np
+
+from gymnasium.spaces.space import Space
+
+
+class Dict(Space[typing.Dict[str, Any]], typing.Mapping[str, Space[Any]]):
+    """A dictionary of :class:`Space` instances.
+
+    Elements of this space are (ordered) dictionaries of elements from the constituent spaces.
+
+    Example:
+        >>> from gymnasium.spaces import Dict, Box, Discrete
+        >>> observation_space = Dict({"position": Box(-1, 1, shape=(2,)), "color": Discrete(3)}, seed=42)
+        >>> observation_space.sample()
+        OrderedDict([('color', 0), ('position', array([-0.3991573 ,  0.21649833], dtype=float32))])
+
+        With a nested dict:
+
+        >>> from gymnasium.spaces import Box, Dict, Discrete, MultiBinary, MultiDiscrete
+        >>> Dict(  # doctest: +SKIP
+        ...     {
+        ...         "ext_controller": MultiDiscrete([5, 2, 2]),
+        ...         "inner_state": Dict(
+        ...             {
+        ...                 "charge": Discrete(100),
+        ...                 "system_checks": MultiBinary(10),
+        ...                 "job_status": Dict(
+        ...                     {
+        ...                         "task": Discrete(5),
+        ...                         "progress": Box(low=0, high=100, shape=()),
+        ...                     }
+        ...                 ),
+        ...             }
+        ...         ),
+        ...     }
+        ... )
+
+    It can be convenient to use :class:`Dict` spaces if you want to make complex observations or actions more human-readable.
+    Usually, it will not be possible to use elements of this space directly in learning code. However, you can easily
+    convert `Dict` observations to flat arrays by using a :class:`gymnasium.wrappers.FlattenObservation` wrapper.
+    Similar wrappers can be implemented to deal with :class:`Dict` actions.
+    """
+
+    def __init__(
+        self,
+        spaces: None | dict[str, Space] | Sequence[tuple[str, Space]] = None,
+        seed: dict | int | np.random.Generator | None = None,
+        **spaces_kwargs: Space,
+    ):
+        """Constructor of :class:`Dict` space.
+
+        This space can be instantiated in one of two ways: Either you pass a dictionary
+        of spaces to :meth:`__init__` via the ``spaces`` argument, or you pass the spaces as separate
+        keyword arguments (where you will need to avoid the keys ``spaces`` and ``seed``)
+
+        Args:
+            spaces: A dictionary of spaces. This specifies the structure of the :class:`Dict` space
+            seed: Optionally, you can use this argument to seed the RNGs of the spaces that make up the :class:`Dict` space.
+            **spaces_kwargs: If ``spaces`` is ``None``, you need to pass the constituent spaces as keyword arguments, as described above.
+        """
+        # Convert the spaces into an OrderedDict
+        if isinstance(spaces, collections.abc.Mapping) and not isinstance(
+            spaces, OrderedDict
+        ):
+            try:
+                spaces = OrderedDict(sorted(spaces.items()))
+            except TypeError:
+                # Incomparable types (e.g. `int` vs. `str`, or user-defined types) found.
+                # The keys remain in the insertion order.
+                spaces = OrderedDict(spaces.items())
+        elif isinstance(spaces, Sequence):
+            spaces = OrderedDict(spaces)
+        elif spaces is None:
+            spaces = OrderedDict()
+        else:
+            assert isinstance(
+                spaces, OrderedDict
+            ), f"Unexpected Dict space input, expecting dict, OrderedDict or Sequence, actual type: {type(spaces)}"
+
+        # Add kwargs to spaces to allow both dictionary and keywords to be used
+        for key, space in spaces_kwargs.items():
+            if key not in spaces:
+                spaces[key] = space
+            else:
+                raise ValueError(
+                    f"Dict space keyword '{key}' already exists in the spaces dictionary."
+                )
+
+        self.spaces: dict[str, Space[Any]] = spaces
+        for key, space in self.spaces.items():
+            assert isinstance(
+                space, Space
+            ), 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)  # type: ignore
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return all(space.is_np_flattenable for space in self.spaces.values())
+
+    def seed(self, seed: dict[str, Any] | int | None = None) -> list[int]:
+        """Seed the PRNG of this space and all subspaces.
+
+        Depending on the type of seed, the subspaces will be seeded differently
+
+        * ``None`` - All the subspaces will use a random initial seed
+        * ``Int`` - The integer is used to seed the :class:`Dict` space that is used to generate seed values for each of the subspaces. Warning, this does not guarantee unique seeds for all of the subspaces.
+        * ``Dict`` - Using all the keys in the seed dictionary, the values are used to seed the subspaces. This allows the seeding of multiple composite subspaces (``Dict["space": Dict[...], ...]`` with ``{"space": {...}, ...}``).
+
+        Args:
+            seed: An optional list of ints or int to seed the (sub-)spaces.
+        """
+        seeds: list[int] = []
+
+        if isinstance(seed, dict):
+            assert (
+                seed.keys() == self.spaces.keys()
+            ), f"The seed keys: {seed.keys()} are not identical to space keys: {self.spaces.keys()}"
+            for key in seed.keys():
+                seeds += self.spaces[key].seed(seed[key])
+        elif isinstance(seed, int):
+            seeds = super().seed(seed)
+            # Using `np.int32` will mean that the same key occurring is extremely low, even for large subspaces
+            subseeds = self.np_random.integers(
+                np.iinfo(np.int32).max, size=len(self.spaces)
+            )
+            for subspace, subseed in zip(self.spaces.values(), subseeds):
+                seeds += subspace.seed(int(subseed))
+        elif seed is None:
+            for space in self.spaces.values():
+                seeds += space.seed(None)
+        else:
+            raise TypeError(
+                f"Expected seed type: dict, int or None, actual type: {type(seed)}"
+            )
+
+        return seeds
+
+    def sample(self, mask: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Generates a single random sample from this space.
+
+        The sample is an ordered dictionary of independent samples from the constituent spaces.
+
+        Args:
+            mask: An optional mask for each of the subspaces, expects the same keys as the space
+
+        Returns:
+            A dictionary with the same key and sampled values from :attr:`self.spaces`
+        """
+        if mask is not None:
+            assert isinstance(
+                mask, dict
+            ), f"Expects mask to be a dict, actual type: {type(mask)}"
+            assert (
+                mask.keys() == self.spaces.keys()
+            ), f"Expect mask keys to be same as space keys, mask keys: {mask.keys()}, space keys: {self.spaces.keys()}"
+            return OrderedDict(
+                [(k, space.sample(mask[k])) for k, space in self.spaces.items()]
+            )
+
+        return OrderedDict([(k, space.sample()) for k, space in self.spaces.items()])
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, dict) and x.keys() == self.spaces.keys():
+            return all(x[key] in self.spaces[key] for key in self.spaces.keys())
+        return False
+
+    def __getitem__(self, key: str) -> Space[Any]:
+        """Get the space that is associated to `key`."""
+        return self.spaces[key]
+
+    def keys(self) -> KeysView:
+        """Returns the keys of the Dict."""
+        return KeysView(self.spaces)
+
+    def __setitem__(self, key: str, value: Space[Any]):
+        """Set the space that is associated to `key`."""
+        assert isinstance(
+            value, Space
+        ), f"Trying to set {key} to Dict space with value that is not a gymnasium space, actual type: {type(value)}"
+        self.spaces[key] = value
+
+    def __iter__(self):
+        """Iterator through the keys of the subspaces."""
+        yield from self.spaces
+
+    def __len__(self) -> int:
+        """Gives the number of simpler spaces that make up the `Dict` space."""
+        return len(self.spaces)
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        return (
+            "Dict(" + ", ".join([f"{k!r}: {s}" for k, s in self.spaces.items()]) + ")"
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether `other` is equivalent to this instance."""
+        return (
+            isinstance(other, Dict)
+            # Comparison of `OrderedDict`s is order-sensitive
+            and self.spaces == other.spaces  # OrderedDict.__eq__
+        )
+
+    def to_jsonable(self, sample_n: Sequence[dict[str, Any]]) -> dict[str, list[Any]]:
+        """Convert a batch of samples from this space to a JSONable data type."""
+        # serialize as dict-repr of vectors
+        return {
+            key: space.to_jsonable([sample[key] for sample in sample_n])
+            for key, space in self.spaces.items()
+        }
+
+    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])
+            for key, space in self.spaces.items()
+        }
+
+        n_elements = len(next(iter(dict_of_list.values())))
+        result = [
+            OrderedDict({key: value[n] for key, value in dict_of_list.items()})
+            for n in range(n_elements)
+        ]
+        return result

rl/Lib/site-packages/gymnasium/spaces/discrete.py

@@ -0,0 +1,145 @@
+"""Implementation of a space consisting of finitely many elements."""
+from __future__ import annotations
+
+from typing import Any, Iterable, Mapping, Sequence
+
+import numpy as np
+
+from gymnasium.spaces.space import MaskNDArray, Space
+
+
+class Discrete(Space[np.int64]):
+    r"""A space consisting of finitely many elements.
+
+    This class represents a finite subset of integers, more specifically a set of the form :math:`\{ a, a+1, \dots, a+n-1 \}`.
+
+    Example:
+        >>> from gymnasium.spaces import Discrete
+        >>> observation_space = Discrete(2, seed=42) # {0, 1}
+        >>> observation_space.sample()
+        0
+        >>> observation_space = Discrete(3, start=-1, seed=42)  # {-1, 0, 1}
+        >>> observation_space.sample()
+        -1
+    """
+
+    def __init__(
+        self,
+        n: int | np.integer[Any],
+        seed: int | np.random.Generator | None = None,
+        start: int | np.integer[Any] = 0,
+    ):
+        r"""Constructor of :class:`Discrete` space.
+
+        This will construct the space :math:`\{\text{start}, ..., \text{start} + n - 1\}`.
+
+        Args:
+            n (int): The number of elements of this space.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the ``Dict`` space.
+            start (int): The smallest element of this space.
+        """
+        assert np.issubdtype(
+            type(n), np.integer
+        ), f"Expects `n` to be an integer, actual dtype: {type(n)}"
+        assert n > 0, "n (counts) have to be positive"
+        assert np.issubdtype(
+            type(start), np.integer
+        ), f"Expects `start` to be an integer, actual type: {type(start)}"
+
+        self.n = np.int64(n)
+        self.start = np.int64(start)
+        super().__init__((), np.int64, seed)
+
+    @property
+    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) -> np.int64:
+        """Generates a single random sample from this space.
+
+        A sample will be chosen uniformly at random with the mask if provided
+
+        Args:
+            mask: An optional mask for if an action can be selected.
+                Expected `np.ndarray` of shape `(n,)` and dtype `np.int8` where `1` represents valid actions and `0` invalid / infeasible actions.
+                If there are no possible actions (i.e. `np.all(mask == 0)`) then `space.start` will be returned.
+
+        Returns:
+            A sampled integer from the space
+        """
+        if mask is not None:
+            assert isinstance(
+                mask, np.ndarray
+            ), f"The expected type of the mask is np.ndarray, actual type: {type(mask)}"
+            assert (
+                mask.dtype == np.int8
+            ), f"The expected dtype of the mask is np.int8, actual dtype: {mask.dtype}"
+            assert mask.shape == (
+                self.n,
+            ), f"The expected shape of the mask is {(self.n,)}, actual shape: {mask.shape}"
+            valid_action_mask = mask == 1
+            assert np.all(
+                np.logical_or(mask == 0, valid_action_mask)
+            ), f"All values of a mask should be 0 or 1, actual values: {mask}"
+            if np.any(valid_action_mask):
+                return self.start + self.np_random.choice(
+                    np.where(valid_action_mask)[0]
+                )
+            else:
+                return self.start
+
+        return self.start + self.np_random.integers(self.n)
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, int):
+            as_int64 = np.int64(x)
+        elif isinstance(x, (np.generic, np.ndarray)) and (
+            np.issubdtype(x.dtype, np.integer) and x.shape == ()
+        ):
+            as_int64 = np.int64(x)
+        else:
+            return False
+
+        return bool(self.start <= as_int64 < self.start + self.n)
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        if self.start != 0:
+            return f"Discrete({self.n}, start={self.start})"
+        return f"Discrete({self.n})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether ``other`` is equivalent to this instance."""
+        return (
+            isinstance(other, Discrete)
+            and self.n == other.n
+            and self.start == other.start
+        )
+
+    def __setstate__(self, state: Iterable[tuple[str, Any]] | Mapping[str, Any]):
+        """Used when loading a pickled space.
+
+        This method has to be implemented explicitly to allow for loading of legacy states.
+
+        Args:
+            state: The new state
+        """
+        # Don't mutate the original state
+        state = dict(state)
+
+        # Allow for loading of legacy states.
+        # See https://github.com/openai/gym/pull/2470
+        if "start" not in state:
+            state["start"] = np.int64(0)
+
+        super().__setstate__(state)
+
+    def to_jsonable(self, sample_n: Sequence[np.int64]) -> list[int]:
+        """Converts a list of samples to a list of ints."""
+        return [int(x) for x in sample_n]
+
+    def from_jsonable(self, sample_n: list[int]) -> list[np.int64]:
+        """Converts a list of json samples to a list of np.int64."""
+        return [np.int64(x) for x in sample_n]

rl/Lib/site-packages/gymnasium/spaces/graph.py

@@ -0,0 +1,263 @@
+"""Implementation of a space that represents graph information where nodes and edges can be represented with euclidean space."""
+from __future__ import annotations
+
+from typing import Any, NamedTuple, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+import gymnasium as gym
+from gymnasium.spaces.box import Box
+from gymnasium.spaces.discrete import Discrete
+from gymnasium.spaces.multi_discrete import MultiDiscrete
+from gymnasium.spaces.space import Space
+
+
+class GraphInstance(NamedTuple):
+    """A Graph space instance.
+
+    * nodes (np.ndarray): an (n x ...) sized array representing the features for n nodes, (...) must adhere to the shape of the node space.
+    * edges (Optional[np.ndarray]): an (m x ...) sized array representing the features for m edges, (...) must adhere to the shape of the edge space.
+    * edge_links (Optional[np.ndarray]): an (m x 2) sized array of ints representing the indices of the two nodes that each edge connects.
+    """
+
+    nodes: NDArray[Any]
+    edges: NDArray[Any] | None
+    edge_links: NDArray[Any] | None
+
+
+class Graph(Space[GraphInstance]):
+    r"""A space representing graph information as a series of `nodes` connected with `edges` according to an adjacency matrix represented as a series of `edge_links`.
+
+    Example:
+        >>> from gymnasium.spaces import Graph, Box, Discrete
+        >>> observation_space = Graph(node_space=Box(low=-100, high=100, shape=(3,)), edge_space=Discrete(3), seed=42)
+        >>> observation_space.sample()
+        GraphInstance(nodes=array([[-12.224312 ,  71.71958  ,  39.473606 ],
+               [-81.16453  ,  95.12447  ,  52.22794  ],
+               [ 57.21286  , -74.37727  ,  -9.922812 ],
+               [-25.840395 ,  85.353    ,  28.773024 ],
+               [ 64.55232  , -11.317161 , -54.552258 ],
+               [ 10.916958 , -87.23655  ,  65.52624  ],
+               [ 26.33288  ,  51.61755  , -29.094807 ],
+               [ 94.1396   ,  78.62422  ,  55.6767   ],
+               [-61.072258 ,  -6.6557994, -91.23925  ],
+               [-69.142105 ,  36.60979  ,  48.95243  ]], dtype=float32), edges=array([2, 0, 1, 1, 0, 0, 1, 0]), edge_links=array([[7, 5],
+               [6, 9],
+               [4, 1],
+               [8, 6],
+               [7, 0],
+               [3, 7],
+               [8, 4],
+               [8, 8]], dtype=int32))
+    """
+
+    def __init__(
+        self,
+        node_space: Box | Discrete,
+        edge_space: None | Box | Discrete,
+        seed: int | np.random.Generator | None = None,
+    ):
+        r"""Constructor of :class:`Graph`.
+
+        The argument ``node_space`` specifies the base space that each node feature will use.
+        This argument must be either a Box or Discrete instance.
+
+        The argument ``edge_space`` specifies the base space that each edge feature will use.
+        This argument must be either a None, Box or Discrete instance.
+
+        Args:
+            node_space (Union[Box, Discrete]): space of the node features.
+            edge_space (Union[None, Box, Discrete]): space of the edge features.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
+        """
+        assert isinstance(
+            node_space, (Box, Discrete)
+        ), f"Values of the node_space should be instances of Box or Discrete, got {type(node_space)}"
+        if edge_space is not None:
+            assert isinstance(
+                edge_space, (Box, Discrete)
+            ), f"Values of the edge_space should be instances of None Box or Discrete, got {type(node_space)}"
+
+        self.node_space = node_space
+        self.edge_space = edge_space
+
+        super().__init__(None, None, seed)
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return False
+
+    def _generate_sample_space(
+        self, base_space: None | Box | Discrete, num: int
+    ) -> Box | MultiDiscrete | None:
+        if num == 0 or base_space is None:
+            return None
+
+        if isinstance(base_space, Box):
+            return Box(
+                low=np.array(max(1, num) * [base_space.low]),
+                high=np.array(max(1, num) * [base_space.high]),
+                shape=(num,) + base_space.shape,
+                dtype=base_space.dtype,
+                seed=self.np_random,
+            )
+        elif isinstance(base_space, Discrete):
+            return MultiDiscrete(nvec=[base_space.n] * num, seed=self.np_random)
+        else:
+            raise TypeError(
+                f"Expects base space to be Box and Discrete, actual space: {type(base_space)}."
+            )
+
+    def sample(
+        self,
+        mask: None
+        | (
+            tuple[
+                NDArray[Any] | tuple[Any, ...] | None,
+                NDArray[Any] | tuple[Any, ...] | None,
+            ]
+        ) = None,
+        num_nodes: int = 10,
+        num_edges: int | None = None,
+    ) -> GraphInstance:
+        """Generates a single sample graph with num_nodes between 1 and 10 sampled from the Graph.
+
+        Args:
+            mask: An optional tuple of optional node and edge mask that is only possible with Discrete spaces
+                (Box spaces don't support sample masks).
+                If no `num_edges` is provided then the `edge_mask` is multiplied by the number of edges
+            num_nodes: The number of nodes that will be sampled, the default is 10 nodes
+            num_edges: An optional number of edges, otherwise, a random number between 0 and `num_nodes` ^ 2
+
+        Returns:
+            A :class:`GraphInstance` with attributes `.nodes`, `.edges`, and `.edge_links`.
+        """
+        assert (
+            num_nodes > 0
+        ), f"The number of nodes is expected to be greater than 0, actual value: {num_nodes}"
+
+        if mask is not None:
+            node_space_mask, edge_space_mask = mask
+        else:
+            node_space_mask, edge_space_mask = None, None
+
+        # we only have edges when we have at least 2 nodes
+        if num_edges is None:
+            if num_nodes > 1:
+                # maximal number of edges is `n*(n-1)` allowing self connections and two-way is allowed
+                num_edges = self.np_random.integers(num_nodes * (num_nodes - 1))
+            else:
+                num_edges = 0
+
+            if edge_space_mask is not None:
+                edge_space_mask = tuple(edge_space_mask for _ in range(num_edges))
+        else:
+            if self.edge_space is None:
+                gym.logger.warn(
+                    f"The number of edges is set ({num_edges}) but the edge space is None."
+                )
+            assert (
+                num_edges >= 0
+            ), f"Expects the number of edges to be greater than 0, actual value: {num_edges}"
+        assert num_edges is not None
+
+        sampled_node_space = self._generate_sample_space(self.node_space, num_nodes)
+        sampled_edge_space = self._generate_sample_space(self.edge_space, num_edges)
+
+        assert sampled_node_space is not None
+        sampled_nodes = sampled_node_space.sample(node_space_mask)
+        sampled_edges = (
+            sampled_edge_space.sample(edge_space_mask)
+            if sampled_edge_space is not None
+            else None
+        )
+
+        sampled_edge_links = None
+        if sampled_edges is not None and num_edges > 0:
+            sampled_edge_links = self.np_random.integers(
+                low=0, high=num_nodes, size=(num_edges, 2), dtype=np.int32
+            )
+
+        return GraphInstance(sampled_nodes, sampled_edges, sampled_edge_links)
+
+    def contains(self, x: GraphInstance) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, GraphInstance):
+            # Checks the nodes
+            if isinstance(x.nodes, np.ndarray):
+                if all(node in self.node_space for node in x.nodes):
+                    # Check the edges and edge links which are optional
+                    if isinstance(x.edges, np.ndarray) and isinstance(
+                        x.edge_links, np.ndarray
+                    ):
+                        assert x.edges is not None
+                        assert x.edge_links is not None
+                        if self.edge_space is not None:
+                            if all(edge in self.edge_space for edge in x.edges):
+                                if np.issubdtype(x.edge_links.dtype, np.integer):
+                                    if x.edge_links.shape == (len(x.edges), 2):
+                                        if np.all(
+                                            np.logical_and(
+                                                x.edge_links >= 0,
+                                                x.edge_links < len(x.nodes),
+                                            )
+                                        ):
+                                            return True
+                    else:
+                        return x.edges is None and x.edge_links is None
+        return False
+
+    def __repr__(self) -> str:
+        """A string representation of this space.
+
+        The representation will include node_space and edge_space
+
+        Returns:
+            A representation of the space
+        """
+        return f"Graph({self.node_space}, {self.edge_space})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether `other` is equivalent to this instance."""
+        return (
+            isinstance(other, Graph)
+            and (self.node_space == other.node_space)
+            and (self.edge_space == other.edge_space)
+        )
+
+    def to_jsonable(
+        self, sample_n: Sequence[GraphInstance]
+    ) -> list[dict[str, list[int | float]]]:
+        """Convert a batch of samples from this space to a JSONable data type."""
+        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:
+                ret["edges"] = sample.edges.tolist()
+                ret["edge_links"] = sample.edge_links.tolist()
+            ret_n.append(ret)
+        return ret_n
+
+    def from_jsonable(
+        self, sample_n: Sequence[dict[str, list[list[int] | list[float]]]]
+    ) -> list[GraphInstance]:
+        """Convert a JSONable data type to a batch of samples from this space."""
+        ret: list[GraphInstance] = []
+        for sample in sample_n:
+            if "edges" in sample:
+                assert self.edge_space is not None
+                ret_n = GraphInstance(
+                    np.asarray(sample["nodes"], dtype=self.node_space.dtype),
+                    np.asarray(sample["edges"], dtype=self.edge_space.dtype),
+                    np.asarray(sample["edge_links"], dtype=np.int32),
+                )
+            else:
+                ret_n = GraphInstance(
+                    np.asarray(sample["nodes"], dtype=self.node_space.dtype),
+                    None,
+                    None,
+                )
+            ret.append(ret_n)
+        return ret

rl/Lib/site-packages/gymnasium/spaces/multi_binary.py

@@ -0,0 +1,121 @@
+"""Implementation of a space that consists of binary np.ndarrays of a fixed shape."""
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from gymnasium.spaces.space import MaskNDArray, Space
+
+
+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.
+
+    Example:
+        >>> from gymnasium.spaces import MultiBinary
+        >>> observation_space = MultiBinary(5, seed=42)
+        >>> observation_space.sample()
+        array([1, 0, 1, 0, 1], dtype=int8)
+        >>> observation_space = MultiBinary([3, 2], seed=42)
+        >>> observation_space.sample()
+        array([[1, 0],
+               [1, 0],
+               [1, 1]], dtype=int8)
+    """
+
+    def __init__(
+        self,
+        n: NDArray[np.integer[Any]] | Sequence[int] | int,
+        seed: int | np.random.Generator | None = None,
+    ):
+        """Constructor of :class:`MultiBinary` space.
+
+        Args:
+            n: This will fix the shape of elements of the space. It can either be an integer (if the space is flat)
+                or some sort of sequence (tuple, list or np.ndarray) if there are multiple axes.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
+        """
+        if isinstance(n, (Sequence, np.ndarray)):
+            self.n = input_n = tuple(int(i) for i in n)
+            assert (np.asarray(input_n) > 0).all()  # n (counts) have to be positive
+        else:
+            self.n = n = int(n)
+            input_n = (n,)
+            assert (np.asarray(input_n) > 0).all()  # n (counts) have to be positive
+
+        super().__init__(input_n, np.int8, seed)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Has stricter type than gym.Space - never None."""
+        return self._shape  # type: ignore
+
+    @property
+    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) -> NDArray[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).
+
+        Args:
+            mask: An optional np.ndarray to mask samples with expected shape of ``space.shape``.
+                For mask == 0 then the samples will be 0 and mask == 1 then random samples will be generated.
+                The expected mask shape is the space shape and mask dtype is `np.int8`.
+
+        Returns:
+            Sampled values from space
+        """
+        if mask is not None:
+            assert isinstance(
+                mask, np.ndarray
+            ), f"The expected type of the mask is np.ndarray, actual type: {type(mask)}"
+            assert (
+                mask.dtype == np.int8
+            ), f"The expected dtype of the mask is np.int8, actual dtype: {mask.dtype}"
+            assert (
+                mask.shape == self.shape
+            ), f"The expected shape of the mask is {self.shape}, actual shape: {mask.shape}"
+            assert np.all(
+                (mask == 0) | (mask == 1) | (mask == 2)
+            ), f"All values of a mask should be 0, 1 or 2, actual values: {mask}"
+
+            return np.where(
+                mask == 2,
+                self.np_random.integers(low=0, high=2, size=self.n, dtype=self.dtype),
+                mask.astype(self.dtype),
+            )
+
+        return self.np_random.integers(low=0, high=2, size=self.n, dtype=self.dtype)
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, Sequence):
+            x = np.array(x)  # Promote list to array for contains check
+
+        return bool(
+            isinstance(x, np.ndarray)
+            and self.shape == x.shape
+            and np.all(np.logical_or(x == 0, x == 1))
+        )
+
+    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[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]
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        return f"MultiBinary({self.n})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether `other` is equivalent to this instance."""
+        return isinstance(other, MultiBinary) and self.n == other.n

rl/Lib/site-packages/gymnasium/spaces/multi_discrete.py

@@ -0,0 +1,226 @@
+"""Implementation of a space that represents the cartesian product of `Discrete` spaces."""
+from __future__ import annotations
+
+from typing import Any, Iterable, Mapping, Sequence
+
+import numpy as np
+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[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.
+
+    Note:
+        Some environment wrappers assume a value of 0 always represents the NOOP action.
+
+    e.g. Nintendo Game Controller - Can be conceptualized as 3 discrete action spaces:
+
+    1. Arrow Keys: Discrete 5  - NOOP[0], UP[1], RIGHT[2], DOWN[3], LEFT[4]  - params: min: 0, max: 4
+    2. Button A:   Discrete 2  - NOOP[0], Pressed[1] - params: min: 0, max: 1
+    3. Button B:   Discrete 2  - NOOP[0], Pressed[1] - params: min: 0, max: 1
+
+    It can be initialized as ``MultiDiscrete([ 5, 2, 2 ])`` such that a sample might be ``array([3, 1, 0])``.
+
+    Although this feature is rarely used, :class:`MultiDiscrete` spaces may also have several axes
+    if ``nvec`` has several axes:
+
+    Example:
+        >>> from gymnasium.spaces import MultiDiscrete
+        >>> import numpy as np
+        >>> observation_space = MultiDiscrete(np.array([[1, 2], [3, 4]]), seed=42)
+        >>> observation_space.sample()
+        array([[0, 0],
+               [2, 2]])
+    """
+
+    def __init__(
+        self,
+        nvec: NDArray[np.integer[Any]] | list[int],
+        dtype: str | type[np.integer[Any]] = np.int64,
+        seed: int | np.random.Generator | None = None,
+        start: NDArray[np.integer[Any]] | list[int] | None = None,
+    ):
+        """Constructor of :class:`MultiDiscrete` space.
+
+        The argument ``nvec`` will determine the number of values each categorical variable can take. If
+        ``start`` is provided, it will define the minimal values corresponding to each categorical variable.
+
+        Args:
+            nvec: vector of counts of each categorical variable. This will usually be a list of integers. However,
+                you may also pass a more complicated numpy array if you'd like the space to have several axes.
+            dtype: This should be some kind of integer type.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
+            start: Optionally, the starting value the element of each class will take (defaults to 0).
+        """
+        self.nvec = np.array(nvec, dtype=dtype, copy=True)
+        if start is not None:
+            self.start = np.array(start, dtype=dtype, copy=True)
+        else:
+            self.start = np.zeros(self.nvec.shape, dtype=dtype)
+
+        assert (
+            self.start.shape == self.nvec.shape
+        ), "start and nvec (counts) should have the same shape"
+        assert (self.nvec > 0).all(), "nvec (counts) have to be positive"
+
+        super().__init__(self.nvec.shape, dtype, seed)
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Has stricter type than :class:`gym.Space` - never None."""
+        return self._shape  # type: ignore
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return True
+
+    def sample(
+        self, mask: tuple[MaskNDArray, ...] | None = None
+    ) -> NDArray[np.integer[Any]]:
+        """Generates a single random sample this space.
+
+        Args:
+            mask: An optional mask for multi-discrete, expects tuples with a `np.ndarray` mask in the position of each
+                action with shape `(n,)` where `n` is the number of actions and `dtype=np.int8`.
+                Only mask values == 1 are possible to sample unless all mask values for an action are 0 then the default action `self.start` (the smallest element) is sampled.
+
+        Returns:
+            An `np.ndarray` of shape `space.shape`
+        """
+        if mask is not None:
+
+            def _apply_mask(
+                sub_mask: MaskNDArray | tuple[MaskNDArray, ...],
+                sub_nvec: MaskNDArray | np.integer[Any],
+                sub_start: MaskNDArray | np.integer[Any],
+            ) -> int | list[Any]:
+                if isinstance(sub_nvec, np.ndarray):
+                    assert isinstance(
+                        sub_mask, tuple
+                    ), f"Expects the mask to be a tuple for sub_nvec ({sub_nvec}), actual type: {type(sub_mask)}"
+                    assert len(sub_mask) == len(
+                        sub_nvec
+                    ), f"Expects the mask length to be equal to the number of actions, mask length: {len(sub_mask)}, nvec length: {len(sub_nvec)}"
+                    return [
+                        _apply_mask(new_mask, new_nvec, new_start)
+                        for new_mask, new_nvec, new_start in zip(
+                            sub_mask, sub_nvec, sub_start
+                        )
+                    ]
+                else:
+                    assert np.issubdtype(
+                        type(sub_nvec), np.integer
+                    ), f"Expects the sub_nvec to be an action, actually: {sub_nvec}, {type(sub_nvec)}"
+                    assert isinstance(
+                        sub_mask, np.ndarray
+                    ), f"Expects the sub mask to be np.ndarray, actual type: {type(sub_mask)}"
+                    assert (
+                        len(sub_mask) == sub_nvec
+                    ), f"Expects the mask length to be equal to the number of actions, mask length: {len(sub_mask)}, action: {sub_nvec}"
+                    assert (
+                        sub_mask.dtype == np.int8
+                    ), f"Expects the mask dtype to be np.int8, actual dtype: {sub_mask.dtype}"
+
+                    valid_action_mask = sub_mask == 1
+                    assert np.all(
+                        np.logical_or(sub_mask == 0, valid_action_mask)
+                    ), f"Expects all masks values to 0 or 1, actual values: {sub_mask}"
+
+                    if np.any(valid_action_mask):
+                        return (
+                            self.np_random.choice(np.where(valid_action_mask)[0])
+                            + sub_start
+                        )
+                    else:
+                        return sub_start
+
+            return np.array(_apply_mask(mask, self.nvec, self.start), dtype=self.dtype)
+
+        return (self.np_random.random(self.nvec.shape) * self.nvec).astype(
+            self.dtype
+        ) + self.start
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, Sequence):
+            x = np.array(x)  # Promote list to array for contains check
+
+        # if nvec is uint32 and space dtype is uint32, then 0 <= x < self.nvec guarantees that x
+        # is within correct bounds for space dtype (even though x does not have to be unsigned)
+        return bool(
+            isinstance(x, np.ndarray)
+            and x.shape == self.shape
+            and x.dtype != object
+            and np.all(self.start <= x)
+            and np.all(x - self.start < self.nvec)
+        )
+
+    def to_jsonable(
+        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[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]
+
+    def __repr__(self):
+        """Gives a string representation of this space."""
+        if np.any(self.start != 0):
+            return f"MultiDiscrete({self.nvec}, start={self.start})"
+        return f"MultiDiscrete({self.nvec})"
+
+    def __getitem__(self, index: int | tuple[int, ...]):
+        """Extract a subspace from this ``MultiDiscrete`` space."""
+        nvec = self.nvec[index]
+        start = self.start[index]
+        if nvec.ndim == 0:
+            subspace = Discrete(nvec, start=start)
+        else:
+            subspace = MultiDiscrete(nvec, self.dtype, start=start)
+
+        # you don't need to deepcopy as np random generator call replaces the state not the data
+        subspace.np_random.bit_generator.state = self.np_random.bit_generator.state
+
+        return subspace
+
+    def __len__(self):
+        """Gives the ``len`` of samples from this space."""
+        if self.nvec.ndim >= 2:
+            gym.logger.warn(
+                "Getting the length of a multi-dimensional MultiDiscrete space."
+            )
+        return len(self.nvec)
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether ``other`` is equivalent to this instance."""
+        return bool(
+            isinstance(other, MultiDiscrete)
+            and np.all(self.nvec == other.nvec)
+            and np.all(self.start == other.start)
+        )
+
+    def __setstate__(self, state: Iterable[tuple[str, Any]] | Mapping[str, Any]):
+        """Used when loading a pickled space.
+
+        This method has to be implemented explicitly to allow for loading of legacy states.
+
+        Args:
+            state: The new state
+        """
+        state = dict(state)
+
+        if "start" not in state:
+            state["start"] = np.zeros(state["_shape"], dtype=state["dtype"])
+
+        super().__setstate__(state)

rl/Lib/site-packages/gymnasium/spaces/sequence.py

@@ -0,0 +1,179 @@
+"""Implementation of a space that represents finite-length sequences."""
+from __future__ import annotations
+
+import typing
+from typing import Any, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+import gymnasium as gym
+from gymnasium.spaces.space import Space
+
+
+class Sequence(Space[Union[typing.Tuple[Any, ...], Any]]):
+    r"""This space represent sets of finite-length sequences.
+
+    This space represents the set of tuples of the form :math:`(a_0, \dots, a_n)` where the :math:`a_i` belong
+    to some space that is specified during initialization and the integer :math:`n` is not fixed
+
+    Example:
+        >>> from gymnasium.spaces import Sequence, Box
+        >>> observation_space = Sequence(Box(0, 1), seed=2)
+        >>> observation_space.sample()
+        (array([0.26161215], dtype=float32),)
+        >>> observation_space = Sequence(Box(0, 1), seed=0)
+        >>> observation_space.sample()
+        (array([0.6369617], dtype=float32), array([0.26978672], dtype=float32), array([0.04097353], dtype=float32))
+    """
+
+    def __init__(
+        self,
+        space: Space[Any],
+        seed: int | np.random.Generator | None = None,
+        stack: bool = False,
+    ):
+        """Constructor of the :class:`Sequence` space.
+
+        Args:
+            space: Elements in the sequences this space represent must belong to this space.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space.
+            stack: If `True` then the resulting samples would be stacked.
+        """
+        assert isinstance(
+            space, Space
+        ), f"Expects the feature space to be instance of a gym Space, actual type: {type(space)}"
+        self.feature_space = space
+        self.stack = stack
+        if self.stack:
+            self.stacked_feature_space: Space = gym.vector.utils.batch_space(
+                self.feature_space, 1
+            )
+
+        # None for shape and dtype, since it'll require special handling
+        super().__init__(None, None, seed)
+
+    def seed(self, seed: int | None = None) -> list[int]:
+        """Seed the PRNG of this space and the feature space."""
+        seeds = super().seed(seed)
+        seeds += self.feature_space.seed(seed)
+        return seeds
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return False
+
+    def sample(
+        self,
+        mask: None
+        | (
+            tuple[
+                None | np.integer | NDArray[np.integer],
+                Any,
+            ]
+        ) = None,
+    ) -> tuple[Any] | Any:
+        """Generates a single random sample from this space.
+
+        Args:
+            mask: An optional mask for (optionally) the length of the sequence and (optionally) the values in the sequence.
+                If you specify `mask`, it is expected to be a tuple of the form `(length_mask, sample_mask)` where `length_mask`
+                is
+
+                * ``None`` The length will be randomly drawn from a geometric distribution
+                * ``np.ndarray`` of integers, in which case the length of the sampled sequence is randomly drawn from this array.
+                * ``int`` for a fixed length sample
+
+                The second element of the mask tuple `sample` mask specifies a mask that is applied when
+                sampling elements from the base space. The mask is applied for each feature space sample.
+
+        Returns:
+            A tuple of random length with random samples of elements from the :attr:`feature_space`.
+        """
+        if mask is not None:
+            length_mask, feature_mask = mask
+        else:
+            length_mask, feature_mask = None, None
+
+        if length_mask is not None:
+            if np.issubdtype(type(length_mask), np.integer):
+                assert (
+                    0 <= length_mask
+                ), f"Expects the length mask to be greater than or equal to zero, actual value: {length_mask}"
+                length = length_mask
+            elif isinstance(length_mask, np.ndarray):
+                assert (
+                    len(length_mask.shape) == 1
+                ), f"Expects the shape of the length mask to be 1-dimensional, actual shape: {length_mask.shape}"
+                assert np.all(
+                    0 <= length_mask
+                ), f"Expects all values in the length_mask to be greater than or equal to zero, actual values: {length_mask}"
+                assert np.issubdtype(
+                    length_mask.dtype, np.integer
+                ), f"Expects the length mask array to have dtype to be an numpy integer, actual type: {length_mask.dtype}"
+                length = self.np_random.choice(length_mask)
+            else:
+                raise TypeError(
+                    f"Expects the type of length_mask to an integer or a np.ndarray, actual type: {type(length_mask)}"
+                )
+        else:
+            # The choice of 0.25 is arbitrary
+            length = self.np_random.geometric(0.25)
+
+        # Generate sample values from feature_space.
+        sampled_values = tuple(
+            self.feature_space.sample(mask=feature_mask) for _ in range(length)
+        )
+
+        if self.stack:
+            # Concatenate values if stacked.
+            out = gym.vector.utils.create_empty_array(
+                self.feature_space, len(sampled_values)
+            )
+            return gym.vector.utils.concatenate(self.feature_space, sampled_values, out)
+
+        return sampled_values
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        # by definition, any sequence is an iterable
+        if self.stack:
+            return all(
+                item in self.feature_space
+                for item in gym.vector.utils.iterate(self.stacked_feature_space, x)
+            )
+        else:
+            return isinstance(x, tuple) and all(
+                self.feature_space.contains(item) for item in x
+            )
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        return f"Sequence({self.feature_space}, stack={self.stack})"
+
+    def to_jsonable(
+        self, sample_n: typing.Sequence[tuple[Any, ...] | Any]
+    ) -> list[list[Any]]:
+        """Convert a batch of samples from this space to a JSONable data type."""
+        if self.stack:
+            return self.stacked_feature_space.to_jsonable(sample_n)
+        else:
+            return [self.feature_space.to_jsonable(sample) for sample in sample_n]
+
+    def from_jsonable(self, sample_n: list[list[Any]]) -> list[tuple[Any, ...] | Any]:
+        """Convert a JSONable data type to a batch of samples from this space."""
+        if self.stack:
+            return self.stacked_feature_space.from_jsonable(sample_n)
+        else:
+            return [
+                tuple(self.feature_space.from_jsonable(sample)) for sample in sample_n
+            ]
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether ``other`` is equivalent to this instance."""
+        return (
+            isinstance(other, Sequence)
+            and self.feature_space == other.feature_space
+            and self.stack == other.stack
+        )

rl/Lib/site-packages/gymnasium/spaces/space.py

@@ -0,0 +1,152 @@
+"""Implementation of the `Space` metaclass."""
+from __future__ import annotations
+
+from typing import Any, Generic, Iterable, Mapping, Sequence, TypeVar
+
+import numpy as np
+import numpy.typing as npt
+
+from gymnasium.utils import seeding
+
+
+T_cov = TypeVar("T_cov", covariant=True)
+
+
+MaskNDArray = npt.NDArray[np.int8]
+
+
+class Space(Generic[T_cov]):
+    """Superclass that is used to define observation and action spaces.
+
+    Spaces are crucially used in Gym to define the format of valid actions and observations.
+    They serve various purposes:
+
+    * They clearly define how to interact with environments, i.e. they specify what actions need to look like
+      and what observations will look like
+    * They allow us to work with highly structured data (e.g. in the form of elements of :class:`Dict` spaces)
+      and painlessly transform them into flat arrays that can be used in learning code
+    * They provide a method to sample random elements. This is especially useful for exploration and debugging.
+
+    Different spaces can be combined hierarchically via container spaces (:class:`Tuple` and :class:`Dict`) to build a
+    more expressive space
+
+    Warning:
+        Custom observation & action spaces can inherit from the ``Space``
+        class. However, most use-cases should be covered by the existing space
+        classes (e.g. :class:`Box`, :class:`Discrete`, etc...), and container classes (:class`Tuple` &
+        :class:`Dict`). Note that parametrized probability distributions (through the
+        :meth:`Space.sample()` method), and batching functions (in :class:`gym.vector.VectorEnv`), are
+        only well-defined for instances of spaces provided in gym by default.
+        Moreover, some implementations of Reinforcement Learning algorithms might
+        not handle custom spaces properly. Use custom spaces with care.
+    """
+
+    def __init__(
+        self,
+        shape: Sequence[int] | None = None,
+        dtype: npt.DTypeLike | None = None,
+        seed: int | np.random.Generator | None = None,
+    ):
+        """Constructor of :class:`Space`.
+
+        Args:
+            shape (Optional[Sequence[int]]): If elements of the space are numpy arrays, this should specify their shape.
+            dtype (Optional[Type | str]): If elements of the space are numpy arrays, this should specify their dtype.
+            seed: Optionally, you can use this argument to seed the RNG that is used to sample from the space
+        """
+        self._shape = None if shape is None else tuple(shape)
+        self.dtype = None if dtype is None else np.dtype(dtype)
+        self._np_random = None
+        if seed is not None:
+            if isinstance(seed, np.random.Generator):
+                self._np_random = seed
+            else:
+                self.seed(seed)
+
+    @property
+    def np_random(self) -> np.random.Generator:
+        """Lazily seed the PRNG since this is expensive and only needed if sampling from this space.
+
+        As :meth:`seed` is not guaranteed to set the `_np_random` for particular seeds. We add a
+        check after :meth:`seed` to set a new random number generator.
+        """
+        if self._np_random is None:
+            self.seed()
+
+        # As `seed` is not guaranteed (in particular for composite spaces) to set the `_np_random` then we set it randomly.
+        if self._np_random is None:
+            self._np_random, _ = seeding.np_random()
+
+        return self._np_random
+
+    @property
+    def shape(self) -> tuple[int, ...] | None:
+        """Return the shape of the space as an immutable property."""
+        return self._shape
+
+    @property
+    def is_np_flattenable(self) -> bool:
+        """Checks whether this space can be flattened to a :class:`gymnasium.spaces.Box`."""
+        raise NotImplementedError
+
+    def sample(self, mask: Any | None = None) -> T_cov:
+        """Randomly sample an element of this space.
+
+        Can be uniform or non-uniform sampling based on boundedness of space.
+
+        Args:
+            mask: A mask used for sampling, expected ``dtype=np.int8`` and see sample implementation for expected shape.
+
+        Returns:
+            A sampled actions from the space
+        """
+        raise NotImplementedError
+
+    def seed(self, seed: int | None = None) -> list[int]:
+        """Seed the PRNG of this space and possibly the PRNGs of subspaces."""
+        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."""
+        raise NotImplementedError
+
+    def __contains__(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        return self.contains(x)
+
+    def __setstate__(self, state: Iterable[tuple[str, Any]] | Mapping[str, Any]):
+        """Used when loading a pickled space.
+
+        This method was implemented explicitly to allow for loading of legacy states.
+
+        Args:
+            state: The updated state value
+        """
+        # Don't mutate the original state
+        state = dict(state)
+
+        # Allow for loading of legacy states.
+        # See:
+        #   https://github.com/openai/gym/pull/2397 -- shape
+        #   https://github.com/openai/gym/pull/1913 -- np_random
+        #
+        if "shape" in state:
+            state["_shape"] = state.get("shape")
+            del state["shape"]
+        if "np_random" in state:
+            state["_np_random"] = state["np_random"]
+            del state["np_random"]
+
+        # Update our state
+        self.__dict__.update(state)
+
+    def to_jsonable(self, sample_n: Sequence[T_cov]) -> list[Any]:
+        """Convert a batch of samples from this space to a JSONable data type."""
+        # By default, assume identity is JSONable
+        return list(sample_n)
+
+    def from_jsonable(self, sample_n: list[Any]) -> list[T_cov]:
+        """Convert a JSONable data type to a batch of samples from this space."""
+        # By default, assume identity is JSONable
+        return sample_n

rl/Lib/site-packages/gymnasium/spaces/text.py

@@ -0,0 +1,191 @@
+"""Implementation of a space that represents textual strings."""
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from gymnasium.spaces.space import Space
+
+
+alphanumeric: frozenset[str] = frozenset(
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+)
+
+
+class Text(Space[str]):
+    r"""A space representing a string comprised of characters from a given charset.
+
+    Example:
+        >>> from gymnasium.spaces import Text
+        >>> # {"", "B5", "hello", ...}
+        >>> Text(5)
+        Text(1, 5, charset=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz)
+        >>> # {"0", "42", "0123456789", ...}
+        >>> import string
+        >>> Text(min_length = 1,
+        ...      max_length = 10,
+        ...      charset = string.digits)
+        Text(1, 10, charset=0123456789)
+    """
+
+    def __init__(
+        self,
+        max_length: int,
+        *,
+        min_length: int = 1,
+        charset: frozenset[str] | str = alphanumeric,
+        seed: int | np.random.Generator | None = None,
+    ):
+        r"""Constructor of :class:`Text` space.
+
+        Both bounds for text length are inclusive.
+
+        Args:
+            min_length (int): Minimum text length (in characters). Defaults to 1 to prevent empty strings.
+            max_length (int): Maximum text length (in characters).
+            charset (Union[set], str): Character set, defaults to the lower and upper english alphabet plus latin digits.
+            seed: The seed for sampling from the space.
+        """
+        assert np.issubdtype(
+            type(min_length), np.integer
+        ), f"Expects the min_length to be an integer, actual type: {type(min_length)}"
+        assert np.issubdtype(
+            type(max_length), np.integer
+        ), f"Expects the max_length to be an integer, actual type: {type(max_length)}"
+        assert (
+            0 <= min_length
+        ), f"Minimum text length must be non-negative, actual value: {min_length}"
+        assert (
+            min_length <= max_length
+        ), f"The min_length must be less than or equal to the max_length, min_length: {min_length}, max_length: {max_length}"
+
+        self.min_length: int = int(min_length)
+        self.max_length: int = int(max_length)
+
+        self._char_set: frozenset[str] = frozenset(charset)
+        self._char_list: tuple[str, ...] = tuple(charset)
+        self._char_index: dict[str, np.int32] = {
+            val: np.int32(i) for i, val in enumerate(tuple(charset))
+        }
+        self._char_str: str = "".join(sorted(tuple(charset)))
+
+        # As the shape is dynamic (between min_length and max_length) then None
+        super().__init__(dtype=str, seed=seed)
+
+    def sample(
+        self,
+        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`.
+
+        Args:
+            mask: An optional tuples of length and mask for the text.
+                The length is expected to be between the `min_length` and `max_length` otherwise a random integer between `min_length` and `max_length` is selected.
+                For the mask, we expect a numpy array of length of the charset passed with `dtype == np.int8`.
+                If the charlist mask is all zero then an empty string is returned no matter the `min_length`
+
+        Returns:
+            A sampled string from the space
+        """
+        if mask is not None:
+            assert isinstance(
+                mask, tuple
+            ), f"Expects the mask type to be a tuple, actual type: {type(mask)}"
+            assert (
+                len(mask) == 2
+            ), f"Expects the mask length to be two, actual length: {len(mask)}"
+            length, charlist_mask = mask
+
+            if length is not None:
+                assert np.issubdtype(
+                    type(length), np.integer
+                ), f"Expects the Text sample length to be an integer, actual type: {type(length)}"
+                assert (
+                    self.min_length <= length <= self.max_length
+                ), f"Expects the Text sample length be between {self.min_length} and {self.max_length}, actual length: {length}"
+
+            if charlist_mask is not None:
+                assert isinstance(
+                    charlist_mask, np.ndarray
+                ), f"Expects the Text sample mask to be an np.ndarray, actual type: {type(charlist_mask)}"
+                assert (
+                    charlist_mask.dtype == np.int8
+                ), f"Expects the Text sample mask to be an np.ndarray, actual dtype: {charlist_mask.dtype}"
+                assert charlist_mask.shape == (
+                    len(self.character_set),
+                ), f"expects the Text sample mask to be {(len(self.character_set),)}, actual shape: {charlist_mask.shape}"
+                assert np.all(
+                    np.logical_or(charlist_mask == 0, charlist_mask == 1)
+                ), f"Expects all masks values to 0 or 1, actual values: {charlist_mask}"
+        else:
+            length, charlist_mask = None, None
+
+        if length is None:
+            length = self.np_random.integers(self.min_length, self.max_length + 1)
+
+        if charlist_mask is None:
+            string = self.np_random.choice(self.character_list, size=length)
+        else:
+            valid_mask = charlist_mask == 1
+            valid_indexes = np.where(valid_mask)[0]
+            if len(valid_indexes) == 0:
+                if self.min_length == 0:
+                    string = ""
+                else:
+                    # Otherwise the string will not be contained in the space
+                    raise ValueError(
+                        f"Trying to sample with a minimum length > 0 ({self.min_length}) but the character mask is all zero meaning that no character could be sampled."
+                    )
+            else:
+                string = "".join(
+                    self.character_list[index]
+                    for index in self.np_random.choice(valid_indexes, size=length)
+                )
+
+        return "".join(string)
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, str):
+            if self.min_length <= len(x) <= self.max_length:
+                return all(c in self.character_set for c in x)
+        return False
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        return f"Text({self.min_length}, {self.max_length}, charset={self.characters})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether ``other`` is equivalent to this instance."""
+        return (
+            isinstance(other, Text)
+            and self.min_length == other.min_length
+            and self.max_length == other.max_length
+            and self.character_set == other.character_set
+        )
+
+    @property
+    def character_set(self) -> frozenset[str]:
+        """Returns the character set for the space."""
+        return self._char_set
+
+    @property
+    def character_list(self) -> tuple[str, ...]:
+        """Returns a tuple of characters in the space."""
+        return self._char_list
+
+    def character_index(self, char: str) -> np.int32:
+        """Returns a unique index for each character in the space's character set."""
+        return self._char_index[char]
+
+    @property
+    def characters(self) -> str:
+        """Returns a string with all Text characters."""
+        return self._char_str
+
+    @property
+    def is_np_flattenable(self) -> bool:
+        """The flattened version is an integer array for each character, padded to the max character length."""
+        return True

rl/Lib/site-packages/gymnasium/spaces/tuple.py

@@ -0,0 +1,161 @@
+"""Implementation of a space that represents the cartesian product of other spaces."""
+from __future__ import annotations
+
+import collections.abc
+import typing
+from typing import Any, Iterable
+
+import numpy as np
+
+from gymnasium.spaces.space import Space
+
+
+class Tuple(Space[typing.Tuple[Any, ...]], typing.Sequence[Any]):
+    """A tuple (more precisely: the cartesian product) of :class:`Space` instances.
+
+    Elements of this space are tuples of elements of the constituent spaces.
+
+    Example:
+        >>> from gymnasium.spaces import Tuple, Box, Discrete
+        >>> observation_space = Tuple((Discrete(2), Box(-1, 1, shape=(2,))), seed=42)
+        >>> observation_space.sample()
+        (0, array([-0.3991573 ,  0.21649833], dtype=float32))
+    """
+
+    def __init__(
+        self,
+        spaces: Iterable[Space[Any]],
+        seed: int | typing.Sequence[int] | np.random.Generator | None = None,
+    ):
+        r"""Constructor of :class:`Tuple` space.
+
+        The generated instance will represent the cartesian product :math:`\text{spaces}[0] \times ... \times \text{spaces}[-1]`.
+
+        Args:
+            spaces (Iterable[Space]): The spaces that are involved in the cartesian product.
+            seed: Optionally, you can use this argument to seed the RNGs of the ``spaces`` to ensure reproducible sampling.
+        """
+        self.spaces = tuple(spaces)
+        for space in self.spaces:
+            assert isinstance(
+                space, Space
+            ), f"{space} does not inherit from `gymnasium.Space`. Actual Type: {type(space)}"
+        super().__init__(None, None, seed)  # type: ignore
+
+    @property
+    def is_np_flattenable(self):
+        """Checks whether this space can be flattened to a :class:`spaces.Box`."""
+        return all(space.is_np_flattenable for space in self.spaces)
+
+    def seed(self, seed: int | typing.Sequence[int] | None = None) -> list[int]:
+        """Seed the PRNG of this space and all subspaces.
+
+        Depending on the type of seed, the subspaces will be seeded differently
+
+        * ``None`` - All the subspaces will use a random initial seed
+        * ``Int`` - The integer is used to seed the `Tuple` space that is used to generate seed values for each of the subspaces. Warning, this does not guarantee unique seeds for all of the subspaces.
+        * ``List`` - Values used to seed the subspaces. This allows the seeding of multiple composite subspaces (``List(42, 54, ...``).
+
+        Args:
+            seed: An optional list of ints or int to seed the (sub-)spaces.
+        """
+        seeds: list[int] = []
+
+        if isinstance(seed, collections.abc.Sequence):
+            assert len(seed) == len(
+                self.spaces
+            ), f"Expects that the subspaces of seeds equals the number of subspaces. Actual length of seeds: {len(seeds)}, length of subspaces: {len(self.spaces)}"
+            for subseed, space in zip(seed, self.spaces):
+                seeds += space.seed(subseed)
+        elif isinstance(seed, int):
+            seeds = super().seed(seed)
+            subseeds = self.np_random.integers(
+                np.iinfo(np.int32).max, size=len(self.spaces)
+            )
+            for subspace, subseed in zip(self.spaces, subseeds):
+                seeds += subspace.seed(int(subseed))
+        elif seed is None:
+            for space in self.spaces:
+                seeds += space.seed(seed)
+        else:
+            raise TypeError(
+                f"Expected seed type: list, tuple, int or None, actual type: {type(seed)}"
+            )
+
+        return seeds
+
+    def sample(self, mask: tuple[Any | None, ...] | None = None) -> tuple[Any, ...]:
+        """Generates a single random sample inside this space.
+
+        This method draws independent samples from the subspaces.
+
+        Args:
+            mask: An optional tuple of optional masks for each of the subspace's samples,
+                expects the same number of masks as spaces
+
+        Returns:
+            Tuple of the subspace's samples
+        """
+        if mask is not None:
+            assert isinstance(
+                mask, tuple
+            ), f"Expected type of mask is tuple, actual type: {type(mask)}"
+            assert len(mask) == len(
+                self.spaces
+            ), f"Expected length of mask is {len(self.spaces)}, actual length: {len(mask)}"
+
+            return tuple(
+                space.sample(mask=sub_mask)
+                for space, sub_mask in zip(self.spaces, mask)
+            )
+
+        return tuple(space.sample() for space in self.spaces)
+
+    def contains(self, x: Any) -> bool:
+        """Return boolean specifying if x is a valid member of this space."""
+        if isinstance(x, (list, np.ndarray)):
+            x = tuple(x)  # Promote list and ndarray to tuple for contains check
+
+        return (
+            isinstance(x, tuple)
+            and len(x) == len(self.spaces)
+            and all(space.contains(part) for (space, part) in zip(self.spaces, x))
+        )
+
+    def __repr__(self) -> str:
+        """Gives a string representation of this space."""
+        return "Tuple(" + ", ".join([str(s) for s in self.spaces]) + ")"
+
+    def to_jsonable(
+        self, sample_n: typing.Sequence[tuple[Any, ...]]
+    ) -> list[list[Any]]:
+        """Convert a batch of samples from this space to a JSONable data type."""
+        # serialize as list-repr of tuple of vectors
+        return [
+            space.to_jsonable([sample[i] for sample in sample_n])
+            for i, space in enumerate(self.spaces)
+        ]
+
+    def from_jsonable(self, sample_n: list[list[Any]]) -> list[tuple[Any, ...]]:
+        """Convert a JSONable data type to a batch of samples from this space."""
+        return [
+            sample
+            for sample in zip(
+                *[
+                    space.from_jsonable(sample_n[i])
+                    for i, space in enumerate(self.spaces)
+                ]
+            )
+        ]
+
+    def __getitem__(self, index: int) -> Space[Any]:
+        """Get the subspace at specific `index`."""
+        return self.spaces[index]
+
+    def __len__(self) -> int:
+        """Get the number of subspaces that are involved in the cartesian product."""
+        return len(self.spaces)
+
+    def __eq__(self, other: Any) -> bool:
+        """Check whether ``other`` is equivalent to this instance."""
+        return isinstance(other, Tuple) and self.spaces == other.spaces

rl/Lib/site-packages/gymnasium/spaces/utils.py

@@ -0,0 +1,528 @@
+"""Implementation of utility functions that can be applied to spaces.
+
+These functions mostly take care of flattening and unflattening elements of spaces
+ to facilitate their usage in learning code.
+"""
+from __future__ import annotations
+
+import operator as op
+import typing
+from collections import OrderedDict
+from functools import reduce, singledispatch
+from typing import Any, TypeVar, Union, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+import gymnasium as gym
+from gymnasium.spaces import (
+    Box,
+    Dict,
+    Discrete,
+    Graph,
+    GraphInstance,
+    MultiBinary,
+    MultiDiscrete,
+    Sequence,
+    Space,
+    Text,
+    Tuple,
+)
+
+
+@singledispatch
+def flatdim(space: Space[Any]) -> int:
+    """Return the number of dimensions a flattened equivalent of this space would have.
+
+    Args:
+        space: The space to return the number of dimensions of the flattened spaces
+
+    Returns:
+        The number of dimensions for the flattened spaces
+
+    Raises:
+         NotImplementedError: if the space is not defined in :mod:`gym.spaces`.
+         ValueError: if the space cannot be flattened into a :class:`gymnasium.spaces.Box`
+
+    Example:
+        >>> from gymnasium.spaces import Dict, Discrete
+        >>> space = Dict({"position": Discrete(2), "velocity": Discrete(3)})
+        >>> flatdim(space)
+        5
+    """
+    if space.is_np_flattenable is False:
+        raise ValueError(
+            f"{space} cannot be flattened to a numpy array, probably because it contains a `Graph` or `Sequence` subspace"
+        )
+
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@flatdim.register(Box)
+@flatdim.register(MultiBinary)
+def _flatdim_box_multibinary(space: Box | MultiBinary) -> int:
+    return reduce(op.mul, space.shape, 1)
+
+
+@flatdim.register(Discrete)
+def _flatdim_discrete(space: Discrete) -> int:
+    return int(space.n)
+
+
+@flatdim.register(MultiDiscrete)
+def _flatdim_multidiscrete(space: MultiDiscrete) -> int:
+    return int(np.sum(space.nvec))
+
+
+@flatdim.register(Tuple)
+def _flatdim_tuple(space: Tuple) -> int:
+    if space.is_np_flattenable:
+        return sum(flatdim(s) for s in space.spaces)
+    raise ValueError(
+        f"{space} cannot be flattened to a numpy array, probably because it contains a `Graph` or `Sequence` subspace"
+    )
+
+
+@flatdim.register(Dict)
+def _flatdim_dict(space: Dict) -> int:
+    if space.is_np_flattenable:
+        return sum(flatdim(s) for s in space.spaces.values())
+    raise ValueError(
+        f"{space} cannot be flattened to a numpy array, probably because it contains a `Graph` or `Sequence` subspace"
+    )
+
+
+@flatdim.register(Graph)
+def _flatdim_graph(space: Graph):
+    raise ValueError(
+        "Cannot get flattened size as the Graph Space in Gym has a dynamic size."
+    )
+
+
+@flatdim.register(Text)
+def _flatdim_text(space: Text) -> int:
+    return space.max_length
+
+
+T = TypeVar("T")
+FlatType = Union[
+    NDArray[Any], typing.Dict[str, Any], typing.Tuple[Any, ...], GraphInstance
+]
+
+
+@singledispatch
+def flatten(space: Space[T], x: T) -> FlatType:
+    """Flatten a data point from a space.
+
+    This is useful when e.g. points from spaces must be passed to a neural
+    network, which only understands flat arrays of floats.
+
+    Args:
+        space: The space that ``x`` is flattened by
+        x: The value to flatten
+
+    Returns:
+        The flattened datapoint
+
+            - For :class:`gymnasium.spaces.Box` and :class:`gymnasium.spaces.MultiBinary`, this is a flattened array
+            - For :class:`gymnasium.spaces.Discrete` and :class:`gymnasium.spaces.MultiDiscrete`, this is a flattened one-hot array of the sample
+            - For :class:`gymnasium.spaces.Tuple` and :class:`gymnasium.spaces.Dict`, this is a concatenated array the subspaces (does not support graph subspaces)
+            - For graph spaces, returns :class:`GraphInstance` where:
+                - :attr:`GraphInstance.nodes` are n x k arrays
+                - :attr:`GraphInstance.edges` are either:
+                    - m x k arrays
+                    - None
+                - :attr:`GraphInstance.edge_links` are either:
+                    - m x 2 arrays
+                    - None
+
+    Raises:
+        NotImplementedError: If the space is not defined in :mod:`gymnasium.spaces`.
+
+    Example:
+        >>> from gymnasium.spaces import Box, Discrete, Tuple
+        >>> space = Box(0, 1, shape=(3, 5))
+        >>> flatten(space, space.sample()).shape
+        (15,)
+        >>> space = Discrete(4)
+        >>> flatten(space, 2)
+        array([0, 0, 1, 0])
+        >>> space = Tuple((Box(0, 1, shape=(2,)), Box(0, 1, shape=(3,)), Discrete(3)))
+        >>> example = ((.5, .25), (1., 0., .2), 1)
+        >>> flatten(space, example)
+        array([0.5 , 0.25, 1.  , 0.  , 0.2 , 0.  , 1.  , 0.  ])
+    """
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@flatten.register(Box)
+@flatten.register(MultiBinary)
+def _flatten_box_multibinary(space: Box | MultiBinary, x: NDArray[Any]) -> NDArray[Any]:
+    return np.asarray(x, dtype=space.dtype).flatten()
+
+
+@flatten.register(Discrete)
+def _flatten_discrete(space: Discrete, x: np.int64) -> NDArray[np.int64]:
+    onehot = np.zeros(space.n, dtype=space.dtype)
+    onehot[x - space.start] = 1
+    return onehot
+
+
+@flatten.register(MultiDiscrete)
+def _flatten_multidiscrete(
+    space: MultiDiscrete, x: NDArray[np.int64]
+) -> NDArray[np.int64]:
+    offsets = np.zeros((space.nvec.size + 1,), dtype=np.int32)
+    offsets[1:] = np.cumsum(space.nvec.flatten())
+
+    onehot = np.zeros((offsets[-1],), dtype=space.dtype)
+    onehot[offsets[:-1] + (x - space.start).flatten()] = 1
+    return onehot
+
+
+@flatten.register(Tuple)
+def _flatten_tuple(space: Tuple, x: tuple[Any, ...]) -> tuple[Any, ...] | NDArray[Any]:
+    if space.is_np_flattenable:
+        return np.concatenate(
+            [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(
+            [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())
+
+
+@flatten.register(Graph)
+def _flatten_graph(space: Graph, x: GraphInstance) -> GraphInstance:
+    """We're not using ``.unflatten()`` for :class:`Box` and :class:`Discrete` because a graph is not a homogeneous space, see `.flatten` docstring."""
+
+    def _graph_unflatten(
+        unflatten_space: Discrete | Box | None,
+        unflatten_x: NDArray[Any] | None,
+    ) -> NDArray[Any] | None:
+        ret = None
+        if unflatten_space is not None and unflatten_x is not None:
+            if isinstance(unflatten_space, Box):
+                ret = unflatten_x.reshape(unflatten_x.shape[0], -1)
+            else:
+                assert isinstance(unflatten_space, Discrete)
+                ret = np.zeros(
+                    (unflatten_x.shape[0], unflatten_space.n - unflatten_space.start),
+                    dtype=unflatten_space.dtype,
+                )
+                ret[
+                    np.arange(unflatten_x.shape[0]), unflatten_x - unflatten_space.start
+                ] = 1
+        return ret
+
+    nodes = _graph_unflatten(space.node_space, x.nodes)
+    assert nodes is not None
+    edges = _graph_unflatten(space.edge_space, x.edges)
+
+    return GraphInstance(nodes, edges, x.edge_links)
+
+
+@flatten.register(Text)
+def _flatten_text(space: Text, x: str) -> NDArray[np.int32]:
+    arr = np.full(
+        shape=(space.max_length,), fill_value=len(space.character_set), dtype=np.int32
+    )
+    for i, val in enumerate(x):
+        arr[i] = space.character_index(val)
+    return arr
+
+
+@flatten.register(Sequence)
+def _flatten_sequence(
+    space: Sequence, x: tuple[Any, ...] | Any
+) -> tuple[Any, ...] | Any:
+    if space.stack:
+        samples_iters = gym.vector.utils.iterate(space.stacked_feature_space, x)
+        flattened_samples = [
+            flatten(space.feature_space, sample) for sample in samples_iters
+        ]
+        flattened_space = flatten_space(space.feature_space)
+        out = gym.vector.utils.create_empty_array(
+            flattened_space, n=len(flattened_samples)
+        )
+        return gym.vector.utils.concatenate(flattened_space, flattened_samples, out)
+    else:
+        return tuple(flatten(space.feature_space, item) for item in x)
+
+
+@singledispatch
+def unflatten(space: Space[T], x: FlatType) -> T:
+    """Unflatten a data point from a space.
+
+    This reverses the transformation applied by :func:`flatten`. You must ensure
+    that the ``space`` argument is the same as for the :func:`flatten` call.
+
+    Args:
+        space: The space used to unflatten ``x``
+        x: The array to unflatten
+
+    Returns:
+        A point with a structure that matches the space.
+
+    Raises:
+        NotImplementedError: if the space is not defined in :mod:`gymnasium.spaces`.
+    """
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@unflatten.register(Box)
+@unflatten.register(MultiBinary)
+def _unflatten_box_multibinary(
+    space: Box | MultiBinary, x: NDArray[Any]
+) -> NDArray[Any]:
+    return np.asarray(x, dtype=space.dtype).reshape(space.shape)
+
+
+@unflatten.register(Discrete)
+def _unflatten_discrete(space: Discrete, x: NDArray[np.int64]) -> np.int64:
+    nonzero = np.nonzero(x)
+    if len(nonzero[0]) == 0:
+        raise ValueError(
+            f"{x} is not a valid one-hot encoded vector and can not be unflattened to space {space}. "
+            "Not all valid samples in a flattened space can be unflattened."
+        )
+    return space.start + nonzero[0][0]
+
+
+@unflatten.register(MultiDiscrete)
+def _unflatten_multidiscrete(
+    space: MultiDiscrete, x: NDArray[np.integer[Any]]
+) -> NDArray[np.integer[Any]]:
+    offsets = np.zeros((space.nvec.size + 1,), dtype=space.dtype)
+    offsets[1:] = np.cumsum(space.nvec.flatten())
+    nonzero = np.nonzero(x)
+    if len(nonzero[0]) == 0:
+        raise ValueError(
+            f"{x} is not a concatenation of one-hot encoded vectors and can not be unflattened to space {space}. "
+            "Not all valid samples in a flattened space can be unflattened."
+        )
+    (indices,) = cast(type(offsets[:-1]), nonzero)
+    return (
+        np.asarray(indices - offsets[:-1], dtype=space.dtype).reshape(space.shape)
+        + space.start
+    )
+
+
+@unflatten.register(Tuple)
+def _unflatten_tuple(
+    space: Tuple, x: NDArray[Any] | tuple[Any, ...]
+) -> tuple[Any, ...]:
+    if space.is_np_flattenable:
+        assert isinstance(
+            x, np.ndarray
+        ), f"{space} is numpy-flattenable. Thus, you should only unflatten numpy arrays for this space. Got a {type(x)}"
+        dims = np.asarray([flatdim(s) for s in space.spaces], dtype=np.int_)
+        list_flattened = np.split(x, np.cumsum(dims[:-1]))
+        return tuple(
+            unflatten(s, flattened)
+            for flattened, s in zip(list_flattened, space.spaces)
+        )
+    assert isinstance(
+        x, tuple
+    ), f"{space} is not numpy-flattenable. Thus, you should only unflatten tuples for this space. Got a {type(x)}"
+    return tuple(unflatten(s, flattened) for flattened, s in zip(x, space.spaces))
+
+
+@unflatten.register(Dict)
+def _unflatten_dict(space: Dict, x: NDArray[Any] | dict[str, Any]) -> dict[str, Any]:
+    if space.is_np_flattenable:
+        dims = np.asarray([flatdim(s) for s in space.spaces.values()], dtype=np.int_)
+        list_flattened = np.split(x, np.cumsum(dims[:-1]))
+        return OrderedDict(
+            [
+                (key, unflatten(s, flattened))
+                for flattened, (key, s) in zip(list_flattened, space.spaces.items())
+            ]
+        )
+    assert isinstance(
+        x, dict
+    ), f"{space} is not numpy-flattenable. Thus, you should only unflatten dictionary for this space. Got a {type(x)}"
+    return OrderedDict((key, unflatten(s, x[key])) for key, s in space.spaces.items())
+
+
+@unflatten.register(Graph)
+def _unflatten_graph(space: Graph, x: GraphInstance) -> GraphInstance:
+    """We're not using `.unflatten() for :class:`Box` and :class:`Discrete` because a graph is not a homogeneous space.
+
+    The size of the outcome is actually not fixed, but determined based on the number of
+    nodes and edges in the graph.
+    """
+
+    def _graph_unflatten(unflatten_space, unflatten_x):
+        result = None
+        if unflatten_space is not None and unflatten_x is not None:
+            if isinstance(unflatten_space, Box):
+                result = unflatten_x.reshape(-1, *unflatten_space.shape)
+            elif isinstance(unflatten_space, Discrete):
+                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)
+
+    return GraphInstance(nodes, edges, x.edge_links)
+
+
+@unflatten.register(Text)
+def _unflatten_text(space: Text, x: NDArray[np.int32]) -> str:
+    return "".join(
+        [space.character_list[val] for val in x if val < len(space.character_set)]
+    )
+
+
+@unflatten.register(Sequence)
+def _unflatten_sequence(space: Sequence, x: tuple[Any, ...]) -> tuple[Any, ...] | Any:
+    if space.stack:
+        flattened_space = flatten_space(space.feature_space)
+        flatten_iters = gym.vector.utils.iterate(flattened_space, x)
+        unflattened_samples = [
+            unflatten(space.feature_space, sample) for sample in flatten_iters
+        ]
+        out = gym.vector.utils.create_empty_array(
+            space.feature_space, len(unflattened_samples)
+        )
+        return gym.vector.utils.concatenate(
+            space.feature_space, unflattened_samples, out
+        )
+    else:
+        return tuple(unflatten(space.feature_space, item) for item in x)
+
+
+@singledispatch
+def flatten_space(space: Space[Any]) -> Box | Dict | Sequence | Tuple | Graph:
+    """Flatten a space into a space that is as flat as possible.
+
+    This function will attempt to flatten ``space`` into a single :class:`gymnasium.spaces.Box` space.
+    However, this might not be possible when ``space`` is an instance of :class:`gymnasium.spaces.Graph`,
+    :class:`gymnasium.spaces.Sequence` or a compound space that contains a :class:`gymnasium.spaces.Graph`
+    or :class:`gymnasium.spaces.Sequence` space.
+    This is equivalent to :func:`flatten`, but operates on the space itself. The
+    result for non-graph spaces is always a :class:`gymnasium.spaces.Box` with flat boundaries. While
+    the result for graph spaces is always a :class:`gymnasium.spaces.Graph` with
+    :attr:`Graph.node_space` being a ``Box``
+    with flat boundaries and :attr:`Graph.edge_space` being a ``Box`` with flat boundaries or
+    ``None``. The box has exactly :func:`flatdim` dimensions. Flattening a sample
+    of the original space has the same effect as taking a sample of the flattened
+    space. However, sampling from the flattened space is not necessarily reversible.
+    For example, sampling from a flattened Discrete space is the same as sampling from
+    a Box, and the results may not be integers or one-hot encodings. This may result in
+    errors or non-uniform sampling.
+
+    Args:
+        space: The space to flatten
+
+    Returns:
+        A flattened Box
+
+    Raises:
+        NotImplementedError: if the space is not defined in :mod:`gymnasium.spaces`.
+
+    Example:
+        Flatten spaces.Box:
+        >>> from gymnasium.spaces import Box
+        >>> box = Box(0.0, 1.0, shape=(3, 4, 5))
+        >>> box
+        Box(0.0, 1.0, (3, 4, 5), float32)
+        >>> flatten_space(box)
+        Box(0.0, 1.0, (60,), float32)
+        >>> flatten(box, box.sample()) in flatten_space(box)
+        True
+
+        Flatten spaces.Discrete:
+        >>> from gymnasium.spaces import Discrete
+        >>> discrete = Discrete(5)
+        >>> flatten_space(discrete)
+        Box(0, 1, (5,), int64)
+        >>> flatten(discrete, discrete.sample()) in flatten_space(discrete)
+        True
+
+        Flatten spaces.Dict:
+        >>> from gymnasium.spaces import Dict, Discrete, Box
+        >>> space = Dict({"position": Discrete(2), "velocity": Box(0, 1, shape=(2, 2))})
+        >>> flatten_space(space)
+        Box(0.0, 1.0, (6,), float64)
+        >>> flatten(space, space.sample()) in flatten_space(space)
+        True
+
+        Flatten spaces.Graph:
+        >>> from gymnasium.spaces import Graph, Discrete, Box
+        >>> space = Graph(node_space=Box(low=-100, high=100, shape=(3, 4)), edge_space=Discrete(5))
+        >>> flatten_space(space)
+        Graph(Box(-100.0, 100.0, (12,), float32), Box(0, 1, (5,), int64))
+        >>> flatten(space, space.sample()) in flatten_space(space)
+        True
+    """
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@flatten_space.register(Box)
+def _flatten_space_box(space: Box) -> Box:
+    return Box(space.low.flatten(), space.high.flatten(), dtype=space.dtype)
+
+
+@flatten_space.register(Discrete)
+@flatten_space.register(MultiBinary)
+@flatten_space.register(MultiDiscrete)
+def _flatten_space_binary(space: Discrete | MultiBinary | MultiDiscrete) -> Box:
+    return Box(low=0, high=1, shape=(flatdim(space),), dtype=space.dtype)
+
+
+@flatten_space.register(Tuple)
+def _flatten_space_tuple(space: Tuple) -> Box | Tuple:
+    if space.is_np_flattenable:
+        space_list = [flatten_space(s) for s in space.spaces]
+        return Box(
+            low=np.concatenate([s.low for s in space_list]),
+            high=np.concatenate([s.high for s in space_list]),
+            dtype=np.result_type(*[s.dtype for s in space_list]),
+        )
+    return Tuple(spaces=[flatten_space(s) for s in space.spaces])
+
+
+@flatten_space.register(Dict)
+def _flatten_space_dict(space: Dict) -> Box | Dict:
+    if space.is_np_flattenable:
+        space_list = [flatten_space(s) for s in space.spaces.values()]
+        return Box(
+            low=np.concatenate([s.low for s in space_list]),
+            high=np.concatenate([s.high for s in space_list]),
+            dtype=np.result_type(*[s.dtype for s in space_list]),
+        )
+    return Dict(
+        spaces=OrderedDict(
+            (key, flatten_space(space)) for key, space in space.spaces.items()
+        )
+    )
+
+
+@flatten_space.register(Graph)
+def _flatten_space_graph(space: Graph) -> Graph:
+    return Graph(
+        node_space=flatten_space(space.node_space),
+        edge_space=flatten_space(space.edge_space)
+        if space.edge_space is not None
+        else None,
+    )
+
+
+@flatten_space.register(Text)
+def _flatten_space_text(space: Text) -> Box:
+    return Box(
+        low=0, high=len(space.character_set), shape=(space.max_length,), dtype=np.int32
+    )
+
+
+@flatten_space.register(Sequence)
+def _flatten_space_sequence(space: Sequence) -> Sequence:
+    return Sequence(flatten_space(space.feature_space), stack=space.stack)