Add observation_space override and float32 cast to vector NormalizeOb… (#1528)

Author: JonahFSD

gymnasium/wrappers/vector/stateful_observation.py

@@ -12,6 +12,8 @@
 import gymnasium as gym
 from gymnasium.core import ObsType
 from gymnasium.logger import warn
+from gymnasium.spaces import Box
+from gymnasium.vector.utils import batch_space
 from gymnasium.vector.vector_env import (
     AutoresetMode,
     VectorEnv,
@@ -78,6 +80,15 @@ def __init__(self, env: VectorEnv, epsilon: float = 1e-8):
         else:
             assert self.env.metadata["autoreset_mode"] in {AutoresetMode.NEXT_STEP}
 
+        new_single_space = Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=self.single_observation_space.shape,
+            dtype=np.float32,
+        )
+        self.single_observation_space = new_single_space
+        self.observation_space = batch_space(new_single_space, self.num_envs)
+
         self.obs_rms = RunningMeanStd(
             shape=self.single_observation_space.shape,
             dtype=self.single_observation_space.dtype,
@@ -120,6 +131,7 @@ def observations(self, observations: ObsType) -> ObsType:
         """
         if self._update_running_mean:
             self.obs_rms.update(observations)
-        return (observations - self.obs_rms.mean) / np.sqrt(
-            self.obs_rms.var + self.epsilon
-        )
+        return (
+            (observations - self.obs_rms.mean)
+            / np.sqrt(self.obs_rms.var + self.epsilon)
+        ).astype(np.float32)

tests/wrappers/vector/test_normalize_observation.py

@@ -100,3 +100,18 @@ def test_update_running_mean():
 
     assert np.any(copied_rms_mean != env.obs_rms.mean)
     assert np.any(copied_rms_var != env.obs_rms.var)
+
+
+def test_observation_space_and_dtype():
+    vec_env = SyncVectorEnv([create_env for _ in range(2)])
+    vec_env = wrappers.vector.NormalizeObservation(vec_env)
+
+    assert vec_env.single_observation_space.dtype == np.float32
+    assert np.all(vec_env.single_observation_space.low == -np.inf)
+    assert np.all(vec_env.single_observation_space.high == np.inf)
+
+    obs, _ = vec_env.reset(seed=123)
+    assert obs.dtype == np.float32
+
+    obs, *_ = vec_env.step(vec_env.action_space.sample())
+    assert obs.dtype == np.float32