Add TypedDictSpace (spaces.Dict subclass)

[lebrice]

TypedDictSpace

(This is a draft PR, just looking for feedback on the idea at the moment)

---

TODOS

• Depends on the change that would be introduced by Make various 'space' functions easier to extend for custom Spaces #2093 (making batch_space and concatenate single-dispatch functions), since the dtype argument would be lost when using the current version of these functions with a TypedDictSpace.
• Gather some feedback
• Add tests

---

Adds TypedDictSpace, a subclass of gym.spaces.Dict, which adds the following:

• allows custom dtypes other than dict;
• uses type annotations on the class to set default and required items (à-la
  typing.TypedDict)
• allows indexing using attributes access (space.x := space["x"])

Examples:

• Using it just like a regular spaces.Dict:

>>> from gym.spaces import Box
>>> s = TypedDictSpace(x=Box(0, 1, (4,), dtype=np.float64))
>>> s
TypedDictSpace(x:Box(0.0, 1.0, (4,), float64))
>>> s.seed(123)
>>> s.sample()
{'x': array([0.70787616, 0.3698764 , 0.29010696, 0.10647454])}

• Using it like a TypedDict: (This equivalent to the above)

>>> class VisionSpace(TypedDictSpace):
...     x: Box = Box(0, 1, (4,), dtype=np.float64)  
>>> s = VisionSpace()
>>> s
VisionSpace(x:Box(0.0, 1.0, (4,), float64))
>>> s.seed(123)
>>> s.sample()
{'x': array([0.70787616, 0.3698764 , 0.29010696, 0.10647454])}

• You can also overwrite the values from the type annotations by passing them to the
  constructor:

>>> s = VisionSpace(x=spaces.Box(0, 2, (3,), dtype=np.int64))
>>> s
VisionSpace(x:Box(0, 2, (3,), int64))
>>> s.seed(123)
>>> s.sample()
{'x': array([2, 1, 0])}

• Using custom dtypes

Can use any type here, as long as it can receive the samples from each space as
keyword arguments.
One good example of this is to use a dataclass as the custom dtype.
You are strongly encouraged to use a dtype that inherits from the Mapping class
from collections.abc, so that samples form your space can be handled similarly to
regular dictionaries.

>>> from collections import OrderedDict
>>> s = TypedDictSpace(x=spaces.Box(0, 1, (4,), dtype=int), dtype=OrderedDict)
>>> s
TypedDictSpace(x:Box(0, 1, (4,), int64), dtype=<class 'collections.OrderedDict'>)
>>> s.seed(123)
>>> s.sample()
OrderedDict([('x', array([1, 0, 0, 0]))])

• Required items:

If an annotation on the class doesn't have a default value, then it is treated as a
required argument:

>>> class FooSpace(TypedDictSpace):
...     a: spaces.Box = spaces.Box(0, 1, (4,), int)
...     b: spaces.Discrete
>>> s = FooSpace()  # doesn't work!
Traceback (most recent call last):
  ...
TypeError: Space of type <class 'sequoia.common.spaces.typed_dict.FooSpace'> requires a 'b' item!
>>> s = FooSpace(b=spaces.Discrete(5))
>>> s
FooSpace(a:Box(0, 1, (4,), int64), b:Discrete(5))

NOTE: spaces can also inherit from each other!

>>> class ImageSegmentationSpace(VisionSpace):
...     bounding_box: Box
... 
>>> s = ImageSegmentationSpace(
...     x=spaces.Box(0, 1, (32, 32, 3), dtype=float),
...     bounding_box=spaces.Box(0, 32, (4, 2), dtype=int),
... )
>>> s
ImageSegmentationSpace(x:Box(0.0, 1.0, (32, 32, 3), float64), bounding_box:Box(0, 32, (4, 2), int64))

Let me know what you think! :)

Signed-off-by: Fabrice Normandin fabrice.normandin@gmail.com

[cclauss]

Looks cool! Please rebase now that we have fixed the automated testing.

[jkterry1]

reviewer @Bam4d

@lebrice could you please on master rebase so CI tests will run?

[Bam4d]

@jkterry1 This is a pretty neat addition and seems well thought out. Thanks also @lebrice for keeping this up to date.

[jkterry1]

@lebrice could you please fix flake so I can merge?

[lebrice]

Oh! Please don't merge this, I thought I marked the PR as draft, it isnt ready yet (still need to add tests)
I'll mark this as "ready for review" once the tests are added.

Thanks!

[jkterry1]

Hey, after some review of this project I'm going to close this for the purposes of keeping track of whats going on. If you get this into a state where you'd like us to be seriously considered, please let me know.