Repository: openai/gym

Issue:
MultiDiscrete flatten_space inconsistencies
I ran into some seemingly glaring inconsistencies with [flattening](https://github.com/openai/gym/blob/bfd4f6d179c90c56356b7cd3bba323fdabd0d2d2/gym/spaces/utils.py#L103) a [MultiDiscrete space](https://github.com/openai/gym/blob/master/gym/spaces/multi_discrete.py)

1. `flatten_space(MultiDiscrete)` does __NOT__ behave the same as `flatten_space(Tuple([Discrete, Discrete]))`. Flattening a `Discrete` space results in a one-hot vector space. In turn, flattening a `Tuple([Discrete, Discrete])` results in a space of concatenated one-hot vectors. Therefore it would be expected that flattening a `MultiDiscrete` space would also result in a concatenated one-hot vectors space. This is __NOT__ the case. Instead `flatten_space(MultiDiscrete)` results in a space with the same dimensions as the original `MultiDiscrete` space with bounds drawn from the ranges of each variable in the original `MultiDiscrete` space. Therefore the flattening of `MultiDiscrete` has _almost_ no functional purpose; but it does seem to introduce a significant bug (see 2) 
2. `flatten_space(MultiDiscrete)` upper bound is invalid and causes invalid sample draws! The upper bound shouldn't be inclusive. This can be easily demonstrated with a simple `MultiDiscrete([2,2])` space:
```
from gym import spaces
a = spaces.MultiDiscrete([2,2])
valid_samples = [[0,0], [1,0], [0,1], [1,1]]

# this works
for i in range(100):
    sample = a.sample()
    assert list(sample) in valid_samples, "{} is not a valid sample of {}".format(flat_sample, valid_samples)

af = spaces.flatten_space(a)
# this will fail because a sample will be drawn with a value of 2 in one of the elements
for i in range(100):
    flat_sample = af.sample()
    assert list(flat_sample) in valid_samples, "{} is not a valid sample of {}".format(flat_sample, valid_samples)
```

See related thread [here](https://github.com/openai/gym/issues/1830), however inconsistency 2 was not identified nor was inconsistency 1 really addressed.  


Discussion:
Thanks for pointing this out. I believe that part of this is a duplicate of https://github.com/openai/gym/issues/1915, but ya the rest needs to be dealt with.

Resolving patch:
diff --git a/gym/spaces/utils.py b/gym/spaces/utils.py
--- a/gym/spaces/utils.py
+++ b/gym/spaces/utils.py
@@ -1,5 +1,7 @@
 from collections import OrderedDict
+from functools import singledispatch, reduce
 import numpy as np
+import operator as op
 
 from gym.spaces import Box
 from gym.spaces import Discrete
@@ -9,6 +11,7 @@
 from gym.spaces import Dict
 
 
+@singledispatch
 def flatdim(space):
     """Return the number of dimensions a flattened equivalent of this space
     would have.
@@ -16,22 +19,36 @@ def flatdim(space):
     Accepts a space and returns an integer. Raises ``NotImplementedError`` if
     the space is not defined in ``gym.spaces``.
     """
-    if isinstance(space, Box):
-        return int(np.prod(space.shape))
-    elif isinstance(space, Discrete):
-        return int(space.n)
-    elif isinstance(space, Tuple):
-        return int(sum([flatdim(s) for s in space.spaces]))
-    elif isinstance(space, Dict):
-        return int(sum([flatdim(s) for s in space.spaces.values()]))
-    elif isinstance(space, MultiBinary):
-        return int(space.n)
-    elif isinstance(space, MultiDiscrete):
-        return int(np.prod(space.shape))
-    else:
-        raise NotImplementedError
+    raise NotImplementedError(f"Unknown space: `{space}`")
 
 
+@flatdim.register(Box)
+@flatdim.register(MultiBinary)
+def flatdim_box_multibinary(space):
+    return reduce(op.mul, space.shape, 1)
+
+
+@flatdim.register(Discrete)
+def flatdim_discrete(space):
+    return int(space.n)
+
+
+@flatdim.register(MultiDiscrete)
+def flatdim_multidiscrete(space):
+    return int(np.sum(space.nvec))
+
+
+@flatdim.register(Tuple)
+def flatdim_tuple(space):
+    return sum([flatdim(s) for s in space.spaces])
+
+
+@flatdim.register(Dict)
+def flatdim_dict(space):
+    return sum([flatdim(s) for s in space.spaces.values()])
+
+
+@singledispatch
 def flatten(space, x):
     """Flatten a data point from a space.
 
@@ -42,26 +59,43 @@ def flatten(space, x):
     Raises ``NotImplementedError`` if the space is not defined in
     ``gym.spaces``.
     """
-    if isinstance(space, Box):
-        return np.asarray(x, dtype=space.dtype).flatten()
-    elif isinstance(space, Discrete):
-        onehot = np.zeros(space.n, dtype=space.dtype)
-        onehot[x] = 1
-        return onehot
-    elif isinstance(space, Tuple):
-        return np.concatenate(
-            [flatten(s, x_part) for x_part, s in zip(x, space.spaces)]
-        )
-    elif isinstance(space, Dict):
-        return np.concatenate([flatten(s, x[key]) for key, s in space.spaces.items()])
-    elif isinstance(space, MultiBinary):
-        return np.asarray(x, dtype=space.dtype).flatten()
-    elif isinstance(space, MultiDiscrete):
-        return np.asarray(x, dtype=space.dtype).flatten()
-    else:
-        raise NotImplementedError
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@flatten.register(Box)
+@flatten.register(MultiBinary)
+def flatten_box_multibinary(space, x):
+    return np.asarray(x, dtype=space.dtype).flatten()
+
+
+@flatten.register(Discrete)
+def flatten_discrete(space, x):
+    onehot = np.zeros(space.n, dtype=space.dtype)
+    onehot[x] = 1
+    return onehot
+
+
+@flatten.register(MultiDiscrete)
+def flatten_multidiscrete(space, x):
+    offsets = np.zeros((space.nvec.size + 1,), dtype=space.dtype)
+    offsets[1:] = np.cumsum(space.nvec.flatten())
+
+    onehot = np.zeros((offsets[-1],), dtype=space.dtype)
+    onehot[offsets[:-1] + x.flatten()] = 1
+    return onehot
+
+
+@flatten.register(Tuple)
+def flatten_tuple(space, x):
+    return np.concatenate([flatten(s, x_part) for x_part, s in zip(x, space.spaces)])
+
+
+@flatten.register(Dict)
+def flatten_dict(space, x):
+    return np.concatenate([flatten(s, x[key]) for key, s in space.spaces.items()])
 
 
+@singledispatch
 def unflatten(space, x):
     """Unflatten a data point from a space.
 
@@ -72,34 +106,51 @@ def unflatten(space, x):
     that matches the space. Raises ``NotImplementedError`` if the space is not
     defined in ``gym.spaces``.
     """
-    if isinstance(space, Box):
-        return np.asarray(x, dtype=space.dtype).reshape(space.shape)
-    elif isinstance(space, Discrete):
-        return int(np.nonzero(x)[0][0])
-    elif isinstance(space, Tuple):
-        dims = [flatdim(s) for s in space.spaces]
-        list_flattened = np.split(x, np.cumsum(dims)[:-1])
-        list_unflattened = [
-            unflatten(s, flattened)
-            for flattened, s in zip(list_flattened, space.spaces)
-        ]
-        return tuple(list_unflattened)
-    elif isinstance(space, Dict):
-        dims = [flatdim(s) for s in space.spaces.values()]
-        list_flattened = np.split(x, np.cumsum(dims)[:-1])
-        list_unflattened = [
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@unflatten.register(Box)
+@unflatten.register(MultiBinary)
+def unflatten_box_multibinary(space, x):
+    return np.asarray(x, dtype=space.dtype).reshape(space.shape)
+
+
+@unflatten.register(Discrete)
+def unflatten_discrete(space, x):
+    return int(np.nonzero(x)[0][0])
+
+
+@unflatten.register(MultiDiscrete)
+def unflatten_multidiscrete(space, x):
+    offsets = np.zeros((space.nvec.size + 1,), dtype=space.dtype)
+    offsets[1:] = np.cumsum(space.nvec.flatten())
+
+    (indices,) = np.nonzero(x)
+    return np.asarray(indices - offsets[:-1], dtype=space.dtype).reshape(space.shape)
+
+
+@unflatten.register(Tuple)
+def unflatten_tuple(space, 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)
+    )
+
+
+@unflatten.register(Dict)
+def unflatten_dict(space, x):
+    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())
         ]
-        return OrderedDict(list_unflattened)
-    elif isinstance(space, MultiBinary):
-        return np.asarray(x, dtype=space.dtype).reshape(space.shape)
-    elif isinstance(space, MultiDiscrete):
-        return np.asarray(x, dtype=space.dtype).reshape(space.shape)
-    else:
-        raise NotImplementedError
+    )
 
 
+@singledispatch
 def flatten_space(space):
     """Flatten a space into a single ``Box``.
 
@@ -138,26 +189,36 @@ def flatten_space(space):
         >>> flatten(space, space.sample()) in flatten_space(space)
         True
     """
-    if isinstance(space, Box):
-        return Box(space.low.flatten(), space.high.flatten(), dtype=space.dtype)
-    if isinstance(space, Discrete):
-        return Box(low=0, high=1, shape=(space.n,), dtype=space.dtype)
-    if isinstance(space, Tuple):
-        space = [flatten_space(s) for s in space.spaces]
-        return Box(
-            low=np.concatenate([s.low for s in space]),
-            high=np.concatenate([s.high for s in space]),
-            dtype=np.result_type(*[s.dtype for s in space]),
-        )
-    if isinstance(space, Dict):
-        space = [flatten_space(s) for s in space.spaces.values()]
-        return Box(
-            low=np.concatenate([s.low for s in space]),
-            high=np.concatenate([s.high for s in space]),
-            dtype=np.result_type(*[s.dtype for s in space]),
-        )
-    if isinstance(space, MultiBinary):
-        return Box(low=0, high=1, shape=(space.n,), dtype=space.dtype)
-    if isinstance(space, MultiDiscrete):
-        return Box(low=np.zeros_like(space.nvec), high=space.nvec, dtype=space.dtype)
-    raise NotImplementedError
+    raise NotImplementedError(f"Unknown space: `{space}`")
+
+
+@flatten_space.register(Box)
+def flatten_space_box(space):
+    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):
+    return Box(low=0, high=1, shape=(flatdim(space),), dtype=space.dtype)
+
+
+@flatten_space.register(Tuple)
+def flatten_space_tuple(space):
+    space = [flatten_space(s) for s in space.spaces]
+    return Box(
+        low=np.concatenate([s.low for s in space]),
+        high=np.concatenate([s.high for s in space]),
+        dtype=np.result_type(*[s.dtype for s in space]),
+    )
+
+
+@flatten_space.register(Dict)
+def flatten_space_dict(space):
+    space = [flatten_space(s) for s in space.spaces.values()]
+    return Box(
+        low=np.concatenate([s.low for s in space]),
+        high=np.concatenate([s.high for s in space]),
+        dtype=np.result_type(*[s.dtype for s in space]),
+    )


Test patch:
diff --git a/gym/spaces/tests/test_utils.py b/gym/spaces/tests/test_utils.py
--- a/gym/spaces/tests/test_utils.py
+++ b/gym/spaces/tests/test_utils.py
@@ -6,67 +6,40 @@
 from gym.spaces import Box, Dict, Discrete, MultiBinary, MultiDiscrete, Tuple, utils
 
 
-@pytest.mark.parametrize(
-    ["space", "flatdim"],
-    [
-        (Discrete(3), 3),
-        (Box(low=0.0, high=np.inf, shape=(2, 2)), 4),
-        (Tuple([Discrete(5), Discrete(10)]), 15),
-        (
-            Tuple(
-                [
-                    Discrete(5),
-                    Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-                ]
-            ),
-            7,
-        ),
-        (Tuple((Discrete(5), Discrete(2), Discrete(2))), 9),
-        (MultiDiscrete([2, 2, 100]), 3),
-        (MultiBinary(10), 10),
-        (
-            Dict(
-                {
-                    "position": Discrete(5),
-                    "velocity": Box(
-                        low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
-                    ),
-                }
+spaces = [
+    Discrete(3),
+    Box(low=0.0, high=np.inf, shape=(2, 2)),
+    Box(low=0.0, high=np.inf, shape=(2, 2), dtype=np.float16),
+    Tuple([Discrete(5), Discrete(10)]),
+    Tuple(
+        [
+            Discrete(5),
+            Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
+        ]
+    ),
+    Tuple((Discrete(5), Discrete(2), Discrete(2))),
+    MultiDiscrete([2, 2, 10]),
+    MultiBinary(10),
+    Dict(
+        {
+            "position": Discrete(5),
+            "velocity": Box(
+                low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
             ),
-            7,
-        ),
-    ],
-)
+        }
+    ),
+]
+
+flatdims = [3, 4, 4, 15, 7, 9, 14, 10, 7]
+
+
+@pytest.mark.parametrize(["space", "flatdim"], zip(spaces, flatdims))
 def test_flatdim(space, flatdim):
     dim = utils.flatdim(space)
     assert dim == flatdim, "Expected {} to equal {}".format(dim, flatdim)
 
 
-@pytest.mark.parametrize(
-    "space",
-    [
-        Discrete(3),
-        Box(low=0.0, high=np.inf, shape=(2, 2)),
-        Tuple([Discrete(5), Discrete(10)]),
-        Tuple(
-            [
-                Discrete(5),
-                Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-            ]
-        ),
-        Tuple((Discrete(5), Discrete(2), Discrete(2))),
-        MultiDiscrete([2, 2, 100]),
-        MultiBinary(10),
-        Dict(
-            {
-                "position": Discrete(5),
-                "velocity": Box(
-                    low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
-                ),
-            }
-        ),
-    ],
-)
+@pytest.mark.parametrize("space", spaces)
 def test_flatten_space_boxes(space):
     flat_space = utils.flatten_space(space)
     assert isinstance(flat_space, Box), "Expected {} to equal {}".format(
@@ -77,31 +50,7 @@ def test_flatten_space_boxes(space):
     assert single_dim == flatdim, "Expected {} to equal {}".format(single_dim, flatdim)
 
 
-@pytest.mark.parametrize(
-    "space",
-    [
-        Discrete(3),
-        Box(low=0.0, high=np.inf, shape=(2, 2)),
-        Tuple([Discrete(5), Discrete(10)]),
-        Tuple(
-            [
-                Discrete(5),
-                Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-            ]
-        ),
-        Tuple((Discrete(5), Discrete(2), Discrete(2))),
-        MultiDiscrete([2, 2, 100]),
-        MultiBinary(10),
-        Dict(
-            {
-                "position": Discrete(5),
-                "velocity": Box(
-                    low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
-                ),
-            }
-        ),
-    ],
-)
+@pytest.mark.parametrize("space", spaces)
 def test_flat_space_contains_flat_points(space):
     some_samples = [space.sample() for _ in range(10)]
     flattened_samples = [utils.flatten(space, sample) for sample in some_samples]
@@ -112,31 +61,7 @@ def test_flat_space_contains_flat_points(space):
         )
 
 
-@pytest.mark.parametrize(
-    "space",
-    [
-        Discrete(3),
-        Box(low=0.0, high=np.inf, shape=(2, 2)),
-        Tuple([Discrete(5), Discrete(10)]),
-        Tuple(
-            [
-                Discrete(5),
-                Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-            ]
-        ),
-        Tuple((Discrete(5), Discrete(2), Discrete(2))),
-        MultiDiscrete([2, 2, 100]),
-        MultiBinary(10),
-        Dict(
-            {
-                "position": Discrete(5),
-                "velocity": Box(
-                    low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
-                ),
-            }
-        ),
-    ],
-)
+@pytest.mark.parametrize("space", spaces)
 def test_flatten_dim(space):
     sample = utils.flatten(space, space.sample())
     (single_dim,) = sample.shape
@@ -144,31 +69,7 @@ def test_flatten_dim(space):
     assert single_dim == flatdim, "Expected {} to equal {}".format(single_dim, flatdim)
 
 
-@pytest.mark.parametrize(
-    "space",
-    [
-        Discrete(3),
-        Box(low=0.0, high=np.inf, shape=(2, 2)),
-        Tuple([Discrete(5), Discrete(10)]),
-        Tuple(
-            [
-                Discrete(5),
-                Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-            ]
-        ),
-        Tuple((Discrete(5), Discrete(2), Discrete(2))),
-        MultiDiscrete([2, 2, 100]),
-        MultiBinary(10),
-        Dict(
-            {
-                "position": Discrete(5),
-                "velocity": Box(
-                    low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32
-                ),
-            }
-        ),
-    ],
-)
+@pytest.mark.parametrize("space", spaces)
 def test_flatten_roundtripping(space):
     some_samples = [space.sample() for _ in range(10)]
     flattened_samples = [utils.flatten(space, sample) for sample in some_samples]
@@ -214,38 +115,22 @@ def compare_nested(left, right):
    internally by numpy when np.concatenate is called.
 """
 
+expected_flattened_dtypes = [
+    np.int64,
+    np.float32,
+    np.float16,
+    np.int64,
+    np.float64,
+    np.int64,
+    np.int64,
+    np.int8,
+    np.float64,
+]
+
 
 @pytest.mark.parametrize(
     ["original_space", "expected_flattened_dtype"],
-    [
-        (Discrete(3), np.int64),
-        (Box(low=0.0, high=np.inf, shape=(2, 2)), np.float32),
-        (Box(low=0.0, high=np.inf, shape=(2, 2), dtype=np.float16), np.float16),
-        (Tuple([Discrete(5), Discrete(10)]), np.int64),
-        (
-            Tuple(
-                [
-                    Discrete(5),
-                    Box(low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float32),
-                ]
-            ),
-            np.float64,
-        ),
-        (Tuple((Discrete(5), Discrete(2), Discrete(2))), np.int64),
-        (MultiDiscrete([2, 2, 100]), np.int64),
-        (MultiBinary(10), np.int8),
-        (
-            Dict(
-                {
-                    "position": Discrete(5),
-                    "velocity": Box(
-                        low=np.array([0, 0]), high=np.array([1, 5]), dtype=np.float16
-                    ),
-                }
-            ),
-            np.float64,
-        ),
-    ],
+    zip(spaces, expected_flattened_dtypes),
 )
 def test_dtypes(original_space, expected_flattened_dtype):
     flattened_space = utils.flatten_space(original_space)
@@ -292,3 +177,82 @@ def compare_sample_types(original_space, original_sample, unflattened_sample):
             "original_sample's dtype. unflattened_sample: "
             "{} original_sample: {}".format(unflattened_sample, original_sample)
         )
+
+
+samples = [
+    2,
+    np.array([[1.0, 3.0], [5.0, 8.0]], dtype=np.float32),
+    np.array([[1.0, 3.0], [5.0, 8.0]], dtype=np.float16),
+    (3, 7),
+    (2, np.array([0.5, 3.5], dtype=np.float32)),
+    (3, 0, 1),
+    np.array([0, 1, 7], dtype=np.int64),
+    np.array([0, 1, 1, 0, 0, 0, 1, 1, 1, 1], dtype=np.int8),
+    OrderedDict(
+        [("position", 3), ("velocity", np.array([0.5, 3.5], dtype=np.float32))]
+    ),
+]
+
+
+expected_flattened_samples = [
+    np.array([0, 0, 1], dtype=np.int64),
+    np.array([1.0, 3.0, 5.0, 8.0], dtype=np.float32),
+    np.array([1.0, 3.0, 5.0, 8.0], dtype=np.float16),
+    np.array([0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], dtype=np.int64),
+    np.array([0, 0, 1, 0, 0, 0.5, 3.5], dtype=np.float64),
+    np.array([0, 0, 0, 1, 0, 1, 0, 0, 1], dtype=np.int64),
+    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),
+]
+
+
+@pytest.mark.parametrize(
+    ["space", "sample", "expected_flattened_sample"],
+    zip(spaces, samples, expected_flattened_samples),
+)
+def test_flatten(space, sample, expected_flattened_sample):
+    assert sample in space
+
+    flattened_sample = utils.flatten(space, sample)
+    assert flattened_sample.shape == expected_flattened_sample.shape
+    assert flattened_sample.dtype == expected_flattened_sample.dtype
+    assert np.all(flattened_sample == expected_flattened_sample)
+
+
+@pytest.mark.parametrize(
+    ["space", "flattened_sample", "expected_sample"],
+    zip(spaces, expected_flattened_samples, samples),
+)
+def test_unflatten(space, flattened_sample, expected_sample):
+    sample = utils.unflatten(space, flattened_sample)
+    assert compare_nested(sample, expected_sample)
+
+
+expected_flattened_spaces = [
+    Box(low=0, high=1, shape=(3,), dtype=np.int64),
+    Box(low=0.0, high=np.inf, shape=(4,), dtype=np.float32),
+    Box(low=0.0, high=np.inf, shape=(4,), dtype=np.float16),
+    Box(low=0, high=1, shape=(15,), dtype=np.int64),
+    Box(
+        low=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        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=(9,), dtype=np.int64),
+    Box(low=0, high=1, shape=(14,), dtype=np.int64),
+    Box(low=0, high=1, shape=(10,), dtype=np.int8),
+    Box(
+        low=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64),
+        high=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 5.0], dtype=np.float64),
+        dtype=np.float64,
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    ["space", "expected_flattened_space"], zip(spaces, expected_flattened_spaces)
+)
+def test_flatten_space(space, expected_flattened_space):
+    flattened_space = utils.flatten_space(space)
+    assert flattened_space == expected_flattened_space