Add wrappers.vector.TransformObs/Action single obs/action space argument

[howardh]

Description

Fixes #1287

Type of change

Please delete options that are not relevant.

• Documentation only change (no code changed)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[pseudo-rnd-thoughts]

Thanks for the PR,
I think the __init__ can be significantly simpler,

if action_space is None:
    if single_action_space is not None:
        self.single_action_space = single_action_space
        self.action_space = batch_space(single_action_space, self.num_envs)
else:
    self.action_space = action_space
    if single_action_space is not None:
        self.single_action_space = single_action_space
        if action_space != batch_space(single_action_space, self.num_envs):
            warn("The action space and the batched single action space don't match as expected.")

self.func = func

The vector wrapper will automatically do the variable forwarding so we don't need the new function either

Could you update and change for the transform obs wrapper

[howardh]

The extra stuff I'm doing with self._single_action_space_error is to preserve the current behaviour. If someone has code that is using the wrapper and they only assign the action_space and never use the single_action_space, then this allows them to continue as they were with no breaking changes. What you propose does simplify the check, but it will give the user a warning regardless of whether they need the single_action_space or not, which I believe is not desired behaviour.

Also, I think you want is this:

if action_space is None:
    if single_action_space is not None:
        self.single_action_space = single_action_space
        self.action_space = batch_space(single_action_space, self.num_envs)
else:
    self.action_space = action_space
    if single_action_space is not None:
        self.single_action_space = single_action_space
if self.action_space != batch_space(self.single_action_space, self.num_envs):
    warn("The action space and the batched single action space don't match as expected.")

self.func = func

[howardh]

If you're okay with the behaviour change, then I'll make the changes and resubmit.

[pseudo-rnd-thoughts]

If you're okay with the behaviour change, then I'll make the changes and resubmit.

Yes, I think it should be fine

[pseudo-rnd-thoughts]

Thanks for the update, looks good.
If you could add these changes then good to merge

[pseudo-rnd-thoughts]

Thanks for the changes, looks good