Fixes for Numpy v2

Author: araffin

docs/misc/changelog.rst

@@ -13,7 +13,8 @@ Breaking Changes:
 
 New Features:
 ^^^^^^^^^^^^^
-- Added support for NumPy v2.0 (via Torch)
+- Added support for NumPy v2.0: ``VecNormalize`` now cast normalized rewards to float32,
+updated bit flipping env to avoid overflow issues too
 - Added official support for Python 3.12
 
 Bug Fixes:

stable_baselines3/common/envs/bit_flipping_env.py

@@ -75,14 +75,17 @@ def convert_if_needed(self, state: np.ndarray) -> Union[int, np.ndarray]:
         :param state:
         :return:
         """
+
         if self.discrete_obs_space:
+            # Convert from int8 to int32 for NumPy 2.0
+            state = state.astype(np.int32)
             # The internal state is the binary representation of the
             # observed one
             return int(sum(state[i] * 2**i for i in range(len(state))))
 
         if self.image_obs_space:
             size = np.prod(self.image_shape)
-            image = np.concatenate((state * 255, np.zeros(size - len(state), dtype=np.uint8)))
+            image = np.concatenate((state.astype(np.uint8) * 255, np.zeros(size - len(state), dtype=np.uint8)))
             return image.reshape(self.image_shape).astype(np.uint8)
         return state
 

stable_baselines3/common/vec_env/vec_normalize.py

@@ -254,7 +254,9 @@ def normalize_reward(self, reward: np.ndarray) -> np.ndarray:
         """
         if self.norm_reward:
             reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
-        return reward
+        # Note: we cast to float32 as it correspond to Python default float type
+        # This cast is needed because `RunningMeanStd` keeps stats in float64
+        return reward.astype(np.float32)
 
     def unnormalize_obs(self, obs: Union[np.ndarray, dict[str, np.ndarray]]) -> Union[np.ndarray, dict[str, np.ndarray]]:
         # Avoid modifying by reference the original object