Stop VecNormalize rewriting the observation space of the env it wraps

[DenisDrobyshev]

Summary

VecNormalize replaces an image sub-space with the normalized bounds, so that a policy built on top of it does not take the normalized floats for an image (#1214). For a Dict observation it writes that replacement into the space dict itself:

if isinstance(self.observation_space, spaces.Dict):
    self.obs_spaces = self.observation_space.spaces
    ...
    for key in self.obs_rms.keys():
        if is_image_space(self.obs_spaces[key]):
            self.observation_space.spaces[key] = spaces.Box(...)

VecEnvWrapper.__init__ does not copy the space, so self.observation_space is the wrapped VecEnv's space, which is in turn the underlying env's space. Building the wrapper therefore changes the environment it wraps:

from stable_baselines3.common.envs import SimpleMultiObsEnv
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize

venv = DummyVecEnv([lambda: SimpleMultiObsEnv()])
venv.observation_space
# Dict('img': Box(0, 255, (64, 64, 1), uint8), 'vec': Box(0.0, 1.0, (5,), float64))

VecNormalize(venv)

venv.observation_space
# Dict('img': Box(-10.0, 10.0, (64, 64, 1), float32), 'vec': Box(0.0, 1.0, (5,), float64))

The Box branch of the same constructor, a dozen lines below, rebinds self.observation_space instead of writing through it, so the non-dict case never had this.

A model built on the venv afterwards then picks a different feature extractor for the image — and it is built on the venv, not on the VecNormalize, which does not even have to be kept:

def extractors(venv):
    model = PPO("MultiInputPolicy", venv, n_steps=8, batch_size=8)
    return {k: type(v).__name__ for k, v in model.policy.features_extractor.extractors.items()}

venv = DummyVecEnv([lambda: SimpleMultiObsEnv()])
extractors(venv)     # {'img': 'NatureCNN', 'vec': 'Flatten'}
VecNormalize(venv)
extractors(venv)     # {'img': 'Flatten',   'vec': 'Flatten'}

The original bounds are not recoverable from the wrapper either: self.obs_spaces is assigned before the loop, but it aliases the same dict, so it reports the rewritten space as well.

Changes

• stable_baselines3/common/vec_env/vec_normalize.py — the Dict branch copies the space before editing it, which is what the Box branch already effectively does. One line.
• docs/misc/changelog.md — an entry under Bug Fixes.
• tests/test_vec_normalize.py — test_vec_normalize_keeps_the_wrapped_obs_space asserts the venv's space is unchanged, and that the wrapper's own image sub-space is still rewritten to float32.

Verification

tests/test_vec_normalize.py: 18 passed.

With the fix reverted and the test kept, it fails on the first assertion, reporting img as Box(-10.0, 10.0, (64, 64, 1), float32) where the environment declared Box(0, 255, (64, 64, 1), uint8). The test detects the defect rather than passing regardless.

ruff check, ruff format --check and mypy are clean on both files.

[DenisDrobyshev]

Adjacent suites on the same checkout, since the description said I would add the counts: tests/test_vec_normalize.py, tests/test_vec_envs.py, tests/test_dict_env.py, tests/test_vec_check_nan.py and tests/test_vec_extract_dict_obs.py together are 117 passed, 2 skipped.