Fix type annotations

Author: araffin

stable_baselines3/common/env_checker.py

@@ -126,7 +126,7 @@ def _check_goal_env_compute_reward(
     env: gym.Env,
     reward: float,
     info: Dict[str, Any],
-):
+) -> None:
     """
     Check that reward is computed with `compute_reward`
     and that the implementation is vectorized.

stable_baselines3/common/env_util.py

@@ -25,7 +25,7 @@ def unwrap_wrapper(env: gym.Env, wrapper_class: Type[gym.Wrapper]) -> Optional[g
     return None
 
 
-def is_wrapped(env: Type[gym.Env], wrapper_class: Type[gym.Wrapper]) -> bool:
+def is_wrapped(env: gym.Env, wrapper_class: Type[gym.Wrapper]) -> bool:
     """
     Check if a given environment has been wrapped with a given wrapper.
 
@@ -73,13 +73,19 @@ def make_vec_env(
     :param wrapper_kwargs: Keyword arguments to pass to the ``Wrapper`` class constructor.
     :return: The wrapped environment
     """
-    env_kwargs = {} if env_kwargs is None else env_kwargs
-    vec_env_kwargs = {} if vec_env_kwargs is None else vec_env_kwargs
-    monitor_kwargs = {} if monitor_kwargs is None else monitor_kwargs
-    wrapper_kwargs = {} if wrapper_kwargs is None else wrapper_kwargs
+    env_kwargs = env_kwargs or {}
+    vec_env_kwargs = vec_env_kwargs or {}
+    monitor_kwargs = monitor_kwargs or {}
+    wrapper_kwargs = wrapper_kwargs or {}
+    assert vec_env_kwargs is not None  # for mypy
+
+    def make_env(rank: int) -> Callable[[], gym.Env]:
+        def _init() -> gym.Env:
+            # For type checker:
+            assert monitor_kwargs is not None
+            assert wrapper_kwargs is not None
+            assert env_kwargs is not None
 
-    def make_env(rank):
-        def _init():
             if isinstance(env_id, str):
                 env = gym.make(env_id, **env_kwargs)
             else:
@@ -91,7 +97,7 @@ def _init():
             # to have additional training information
             monitor_path = os.path.join(monitor_dir, str(rank)) if monitor_dir is not None else None
             # Create the monitor folder if needed
-            if monitor_path is not None:
+            if monitor_path is not None and monitor_dir is not None:
                 os.makedirs(monitor_dir, exist_ok=True)
             env = Monitor(env, filename=monitor_path, **monitor_kwargs)
             # Optionally, wrap the environment with the provided wrapper

stable_baselines3/common/torch_layers.py

@@ -57,7 +57,11 @@ class NatureCNN(BaseFeaturesExtractor):
         This corresponds to the number of unit for the last layer.
     """
 
-    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 512):
+    def __init__(self, observation_space: gym.Space, features_dim: int = 512):
+        assert isinstance(observation_space, gym.spaces.Box), (
+            "NatureCNN must be used with a gym.spaces.Box ",
+            f"observation space, not {observation_space}",
+        )
         super().__init__(observation_space, features_dim)
         # We assume CxHxW images (channels first)
         # Re-ordering will be done by pre-preprocessing or wrapper