Make various 'space' functions easier to extend for custom Spaces

[lebrice]

Adds better support for "custom" Space subclasses, by making it possible to "extend" the various gym functions that take a Space as an argument by registering new handlers that type.

This is achieved by converting the given functions into singledispatch functions:

• gym.spaces.utils.flatten
• gym.spaces.utils.flatdim
• gym.spaces.utils.unflatten
• gym.spaces.utils.flatten_space
• gym.vector.utils.numpy_utils.concatenate
• gym.vector.utils.numpy_utils.create_empty_array
• gym.vector.utils..shared_memory.create_shared_memory
• gym.vector.utils..shared_memory.read_from_shared_memory
• gym.vector.utils..shared_memory.write_to_shared_memory

The ordering of the positional arguments for the following functions was changed, since the type of the first argument (in this case space) is the one used by singledispatch to choose the appropriate callable to use. There is however code in each of these to re-order the arguments in case they weren't passed in the right order, making this change fully backward-compatible.

• gym.vector.utils.numpy_utils.concatenate
• gym.vector.utils..shared_memory.read_from_shared_memory
• gym.vector.utils..shared_memory.write_to_shared_memory

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

[jkterry1]

So after discussion with @benblack769 this actually seems like a generally good idea. However, this detailed tests need to be added before I can merge this and Pr currently does not have them.

[jkterry1]

Reviewer: @benblack769

[jkterry1]

@lebrice tests are failing and there are merge conflicts. Are you still interested in working on this PR?

[lebrice]

Hey @jkterry1, Yeah I will update this when I find a moment, should be relatively soon.

[lebrice]

Alright so turns out that @tristandeleu actually already added part of what this PR was meant for in #2328 (at least for the flatten/unflatten/flatdim utils).

This PR also transforms the other utility functions into singledispatch callables, and also fixes a potential bug:

• Calling these generic utility functions without a Space positional argument would raise an error.

My opinion is that there shouldn't be any other tests needed specifically for this PR, since it's simply a refactoring of the utility functions, each of which already has tests. This PR also adds a test for the positional argument issue I just described.

Let me know what you think.

[tristandeleu]

Alright so turns out that @tristandeleu actually already added part of what this PR was meant for in #2328 (at least for the flatten/unflatten/flatdim utils).

Apologies for the changes, I should have noted them here.

I would also be in favor of eventually changing the order of the arguments in all the functions in gym/vector/utils to get the space as a first argument, and make all the changes internally in VectorEnv. This would indeed break backward-compatibility though, as noted by @lebrice, but these are largely internal functions.

[lebrice]

This would indeed break backward-compatibility though, as noted by @lebrice, but these are largely internal functions.

@tristandeleu This is actually not the case! I've added a "wrapper" function around all the introduced generic functions, which makes this totally backward-compatible! As an example, I didn't need to change anything VectorEnvs for this to work! :)

[tristandeleu]

Yes sorry I saw the wrapper function. What I was suggesting was to change the order of the arguments everywhere (and remove the wrapper function). The PR in its current form is absolutely fine, but in the future it would be cleaner (and possibly less confusing) to have all these functions have space as their first argument.

[jkterry1]

@lebrice is this PR currently in a state where you'd like it to be reviewed for being merged?

[jkterry1]

@lebrice could you please also fix tests and change argument orders per Tristan's comment?

[lebrice]

PR is in a good state and ready for review, however I guess I could first change the order of arguments as you mentioned. I'll also keep it backward-compatible as well and raise a warning or something in case the old ordering is used with positional args.

[jkterry1]

@lebrice did you change the argument order and can you please fix the merge conflicts?

[lebrice]

Not yet, will do. You can close this if you want, or mark it as draft, and I'll mark it as "ready for review" when it's good to go

[vwxyzjn]

Hey, this looks really good in refactoring the code. One thing that pops to me was the use of _space_as_first_positional_argument. It seems pretty hacky: shouldn’t we just tell the users to put space as the first positional argument?

[jkterry1]

Closed in favor of #2525