[Bug Report] single_observation_space not updated with vectorized TransformObservation

[howardh]

Describe the bug

gymnasium.wrappers.vector.TransformObservation does not set the observation space correctly. It has an argument for observation_space, which is correct, but it keeps the old single_observation_space. TransformAction appears to have the same issue.

I believe this can be a simple fix by just adding a single_observation_space argument to the constructor.

Code example

import gymnasium
from gymnasium.wrappers.vector import TransformObservation

env = gymnasium.make_vec("CartPole-v1", num_envs=8, vectorization_mode='sync')

old_obs_space = env.observation_space

env = TransformObservation(
    env,
    func = lambda obs: obs[:, 0:2],
    observation_space = gymnasium.spaces.Box(low=old_obs_space.low[:,0:2], high=old_obs_space.high[:,0:2], shape=(8, 2))
)

print(env.single_observation_space)
print(env.observation_space)

Expected output:

Box([-4.8               -inf], [4.8               inf], (2,), float32)
Box([[-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]], [[4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]], (8, 2), float32)

Actual output:

Box([-4.8               -inf -0.41887903        -inf], [4.8               inf 0.41887903        inf], (4,), float32)
Box([[-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]
 [-4.8 -inf]], [[4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]
 [4.8 inf]], (8, 2), float32)

System info

Reproduced on 1.0.0 and on the main branch as of today (commit fc74bb8)

Additional context

No response

Checklist

• I have checked that there is no similar issue in the repo

[pseudo-rnd-thoughts]

Thanks for raising this.
It looks like this is only an issue for TransformObservation and TransformAction.

Could you make a PR? My guess is that we would need the wrappers to specify either the single observation space (and compute the observation space with batch_space and num_envs) or the observation and single observation space