Repository: openai/gym

Issue:
[Bug Report] Space utils break for Discrete with non-zero start
**Describe the bug**
Many of the utility functions for working with spaces are not working as intended with `Discrete` spaces with the `start` argument.

**Code example**
 - Flattening a sample from a `Discrete` space can raise an error:
 ```python
from gym.spaces import Discrete
from gym.spaces.utils import flatten

space = Discrete(3, start=2)
assert 4 in space
flattened = flatten(space, 4)
# IndexError: index 4 is out of bounds for axis 0 with size 3
```
 - `flatten` and `unflatten` are no longer inverse of one another:
 ```python
from gym.spaces import Discrete
from gym.spaces.utils import flatten, unflatten

space = Discrete(5, start=-2)
flattened = flatten(space, -2)  # [0 0 0 1 0]
sample = unflatten(space, flattened)  # 3
```
 - `batch_space` of `Discrete` returns a `MultiDiscrete` space, which has no `start` counterpart:
```python
from gym.spaces import Discrete
from gym.vector.utils import batch_space

space = Discrete(3, start=2)
batched_space = batch_space(space, n=4)
# MultiDiscrete([3 3 3 3])
assert (4, 4, 4, 4) in batched_space
# AssertionError
```

**System Info**
master c8321e68bbd7e452ef65fc237525669979bb45af

- [x] I have checked that there is no similar [issue](https://github.com/openai/gym/issues) in the repo (**required**)



Discussion:

Resolving patch:
diff --git a/gym/spaces/utils.py b/gym/spaces/utils.py
--- a/gym/spaces/utils.py
+++ b/gym/spaces/utils.py
@@ -78,7 +78,7 @@ def _flatten_box_multibinary(space, x) -> np.ndarray:
 @flatten.register(Discrete)
 def _flatten_discrete(space, x) -> np.ndarray:
     onehot = np.zeros(space.n, dtype=space.dtype)
-    onehot[x] = 1
+    onehot[x - space.start] = 1
     return onehot
 
 
@@ -124,7 +124,7 @@ def _unflatten_box_multibinary(space: Box | MultiBinary, x: np.ndarray) -> np.nd
 
 @unflatten.register(Discrete)
 def _unflatten_discrete(space: Discrete, x: np.ndarray) -> int:
-    return int(np.nonzero(x)[0][0])
+    return int(space.start + np.nonzero(x)[0][0])
 
 
 @unflatten.register(MultiDiscrete)
diff --git a/gym/vector/utils/spaces.py b/gym/vector/utils/spaces.py
--- a/gym/vector/utils/spaces.py
+++ b/gym/vector/utils/spaces.py
@@ -43,37 +43,44 @@ def batch_space(space, n=1):
 
 
 @batch_space.register(Box)
-@batch_space.register(Discrete)
-@batch_space.register(MultiDiscrete)
-@batch_space.register(MultiBinary)
-def batch_space_base(space, n=1):
-    if isinstance(space, Box):
-        repeats = tuple([n] + [1] * space.low.ndim)
-        low, high = np.tile(space.low, repeats), np.tile(space.high, repeats)
-        return Box(low=low, high=high, dtype=space.dtype)
+def _batch_space_box(space, n=1):
+    repeats = tuple([n] + [1] * space.low.ndim)
+    low, high = np.tile(space.low, repeats), np.tile(space.high, repeats)
+    return Box(low=low, high=high, dtype=space.dtype)
 
-    elif isinstance(space, Discrete):
+
+@batch_space.register(Discrete)
+def _batch_space_discrete(space, n=1):
+    if space.start == 0:
         return MultiDiscrete(np.full((n,), space.n, dtype=space.dtype))
+    else:
+        return Box(
+            low=space.start,
+            high=space.start + space.n - 1,
+            shape=(n,),
+            dtype=space.dtype,
+        )
 
-    elif isinstance(space, MultiDiscrete):
-        repeats = tuple([n] + [1] * space.nvec.ndim)
-        high = np.tile(space.nvec, repeats) - 1
-        return Box(low=np.zeros_like(high), high=high, dtype=space.dtype)
 
-    elif isinstance(space, MultiBinary):
-        return Box(low=0, high=1, shape=(n,) + space.shape, dtype=space.dtype)
+@batch_space.register(MultiDiscrete)
+def _batch_space_multidiscrete(space, n=1):
+    repeats = tuple([n] + [1] * space.nvec.ndim)
+    high = np.tile(space.nvec, repeats) - 1
+    return Box(low=np.zeros_like(high), high=high, dtype=space.dtype)
 
-    else:
-        raise ValueError(f"Space type `{type(space)}` is not supported.")
+
+@batch_space.register(MultiBinary)
+def _batch_space_multibinary(space, n=1):
+    return Box(low=0, high=1, shape=(n,) + space.shape, dtype=space.dtype)
 
 
 @batch_space.register(Tuple)
-def batch_space_tuple(space, n=1):
+def _batch_space_tuple(space, n=1):
     return Tuple(tuple(batch_space(subspace, n=n) for subspace in space.spaces))
 
 
 @batch_space.register(Dict)
-def batch_space_dict(space, n=1):
+def _batch_space_dict(space, n=1):
     return Dict(
         OrderedDict(
             [
@@ -85,7 +92,7 @@ def batch_space_dict(space, n=1):
 
 
 @batch_space.register(Space)
-def batch_space_custom(space, n=1):
+def _batch_space_custom(space, n=1):
     return Tuple(tuple(space for _ in range(n)))
 
 
@@ -130,14 +137,14 @@ def iterate(space, items):
 
 
 @iterate.register(Discrete)
-def iterate_discrete(space, items):
+def _iterate_discrete(space, items):
     raise TypeError("Unable to iterate over a space of type `Discrete`.")
 
 
 @iterate.register(Box)
 @iterate.register(MultiDiscrete)
 @iterate.register(MultiBinary)
-def iterate_base(space, items):
+def _iterate_base(space, items):
     try:
         return iter(items)
     except TypeError:
@@ -145,7 +152,7 @@ def iterate_base(space, items):
 
 
 @iterate.register(Tuple)
-def iterate_tuple(space, items):
+def _iterate_tuple(space, items):
     # If this is a tuple of custom subspaces only, then simply iterate over items
     if all(
         isinstance(subspace, Space)
@@ -160,7 +167,7 @@ def iterate_tuple(space, items):
 
 
 @iterate.register(Dict)
-def iterate_dict(space, items):
+def _iterate_dict(space, items):
     keys, values = zip(
         *[
             (key, iterate(subspace, items[key]))
@@ -172,7 +179,7 @@ def iterate_dict(space, items):
 
 
 @iterate.register(Space)
-def iterate_custom(space, items):
+def _iterate_custom(space, items):
     raise CustomSpaceError(
         f"Unable to iterate over {items}, since {space} "
         "is a custom `gym.Space` instance (i.e. not one of "


Test patch:
diff --git a/tests/spaces/test_utils.py b/tests/spaces/test_utils.py
--- a/tests/spaces/test_utils.py
+++ b/tests/spaces/test_utils.py
@@ -28,9 +28,11 @@
             ),
         }
     ),
+    Discrete(3, start=2),
+    Discrete(8, start=-5),
 ]
 
-flatdims = [3, 4, 4, 15, 7, 9, 14, 10, 7]
+flatdims = [3, 4, 4, 15, 7, 9, 14, 10, 7, 3, 8]
 
 
 @pytest.mark.parametrize(["space", "flatdim"], zip(spaces, flatdims))
@@ -123,6 +125,8 @@ def compare_nested(left, right):
     np.int64,
     np.int8,
     np.float64,
+    np.int64,
+    np.int64,
 ]
 
 
@@ -187,6 +191,8 @@ def compare_sample_types(original_space, original_sample, unflattened_sample):
     OrderedDict(
         [("position", 3), ("velocity", np.array([0.5, 3.5], dtype=np.float32))]
     ),
+    3,
+    -2,
 ]
 
 
@@ -200,6 +206,8 @@ def compare_sample_types(original_space, original_sample, unflattened_sample):
     np.array([1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], dtype=np.int64),
     np.array([0, 1, 1, 0, 0, 0, 1, 1, 1, 1], dtype=np.int8),
     np.array([0, 0, 0, 1, 0, 0.5, 3.5], dtype=np.float64),
+    np.array([0, 1, 0], dtype=np.int64),
+    np.array([0, 0, 0, 1, 0, 0, 0, 0], dtype=np.int64),
 ]
 
 
@@ -243,6 +251,8 @@ def test_unflatten(space, flattened_sample, expected_sample):
         high=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 5.0], dtype=np.float64),
         dtype=np.float64,
     ),
+    Box(low=0, high=1, shape=(3,), dtype=np.int64),
+    Box(low=0, high=1, shape=(8,), dtype=np.int64),
 ]
 
 
diff --git a/tests/vector/test_shared_memory.py b/tests/vector/test_shared_memory.py
--- a/tests/vector/test_shared_memory.py
+++ b/tests/vector/test_shared_memory.py
@@ -26,6 +26,7 @@
     Array("B", 1),
     Array("B", 32 * 32 * 3),
     Array("i", 1),
+    Array("i", 1),
     (Array("i", 1), Array("i", 1)),
     (Array("i", 1), Array("f", 2)),
     Array("B", 3),
diff --git a/tests/vector/test_spaces.py b/tests/vector/test_spaces.py
--- a/tests/vector/test_spaces.py
+++ b/tests/vector/test_spaces.py
@@ -33,6 +33,7 @@
     Box(low=0, high=255, shape=(4,), dtype=np.uint8),
     Box(low=0, high=255, shape=(4, 32, 32, 3), dtype=np.uint8),
     MultiDiscrete([2, 2, 2, 2]),
+    Box(low=-2, high=2, shape=(4,), dtype=np.int64),
     Tuple((MultiDiscrete([3, 3, 3, 3]), MultiDiscrete([5, 5, 5, 5]))),
     Tuple(
         (
diff --git a/tests/vector/utils.py b/tests/vector/utils.py
--- a/tests/vector/utils.py
+++ b/tests/vector/utils.py
@@ -18,6 +18,7 @@
     Box(low=0, high=255, shape=(), dtype=np.uint8),
     Box(low=0, high=255, shape=(32, 32, 3), dtype=np.uint8),
     Discrete(2),
+    Discrete(5, start=-2),
     Tuple((Discrete(3), Discrete(5))),
     Tuple(
         (