Fix type hints for DQN (#1354)

* Fix type hints for DQN

* [ci skip] Remove commented line

* Refine types

* Fix vectorized obs detection

* Fix for pytype

* Fix check at load time to create replay buffer

* One config file to rule them all

* Delete unused config

---------

Co-authored-by: Quentin Gallouédec <45557362+qgallouedec@users.noreply.github.com>

.coveragerc

@@ -1,15 +0,0 @@
-[run]
-branch = False
-omit =
-    tests/*
-    setup.py
-    # Require graphical interface
-    stable_baselines3/common/results_plotter.py
-    # Require ffmpeg
-    stable_baselines3/common/vec_env/vec_video_recorder.py
-
-[report]
-exclude_lines =
-    pragma: no cover
-    raise NotImplementedError()
-    if typing.TYPE_CHECKING:

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.8.0a12 (WIP)
+Release 1.8.0a13 (WIP)
 --------------------------
 
 .. warning::
@@ -61,6 +61,7 @@ Others:
 - Moved from ``setup.cg`` to ``pyproject.toml`` configuration file
 - Switched from ``flake8`` to ``ruff``
 - Upgraded AutoROM to latest version
+- Fixed ``stable_baselines3/dqn/*.py`` type hints
 - Added ``extra_no_roms`` option for package installation without Atari Roms
 
 Documentation:

pyproject.toml

@@ -61,8 +61,6 @@ exclude = """(?x)(
     | stable_baselines3/common/vec_env/vec_normalize.py$
     | stable_baselines3/common/vec_env/vec_transpose.py$
     | stable_baselines3/common/vec_env/vec_video_recorder.py$
-    | stable_baselines3/dqn/dqn.py$
-    | stable_baselines3/dqn/policies.py$
     | stable_baselines3/her/her_replay_buffer.py$
     | stable_baselines3/ppo/ppo.py$
     | stable_baselines3/sac/policies.py$
@@ -92,3 +90,18 @@ filterwarnings = [
 markers = [
     "expensive: marks tests as expensive (deselect with '-m \"not expensive\"')"
 ]
+
+[tool.coverage.run]
+disable_warnings = ["couldnt-parse"]
+branch = false
+omit = [
+  "tests/*",
+  "setup.py",
+  # Require graphical interface
+  "stable_baselines3/common/results_plotter.py",
+  # Require ffmpeg
+  "stable_baselines3/common/vec_env/vec_video_recorder.py",
+]
+
+[tool.coverage.report]
+exclude_lines = [ "pragma: no cover", "raise NotImplementedError()", "if typing.TYPE_CHECKING:"]

stable_baselines3/common/base_class.py

@@ -87,6 +87,7 @@ class BaseAlgorithm(ABC):
 
     # Policy aliases (see _get_policy_from_name())
     policy_aliases: Dict[str, Type[BasePolicy]] = {}
+    policy: BasePolicy
 
     def __init__(
         self,
@@ -118,9 +119,9 @@ class BaseAlgorithm(ABC):
         self._vec_normalize_env = unwrap_vec_normalize(env)
         self.verbose = verbose
         self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
-        self.observation_space = None  # type: Optional[spaces.Space]
-        self.action_space = None  # type: Optional[spaces.Space]
-        self.n_envs = None
+        self.observation_space: spaces.Space
+        self.action_space: spaces.Space
+        self.n_envs: int
         self.num_timesteps = 0
         # Used for updating schedules
         self._total_timesteps = 0
@@ -129,7 +130,6 @@ class BaseAlgorithm(ABC):
         self.seed = seed
         self.action_noise: Optional[ActionNoise] = None
         self.start_time = None
-        self.policy = None
         self.learning_rate = learning_rate
         self.tensorboard_log = tensorboard_log
         self.lr_schedule = None  # type: Optional[Schedule]

stable_baselines3/common/off_policy_algorithm.py

@@ -125,13 +125,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         self.gradient_steps = gradient_steps
         self.action_noise = action_noise
         self.optimize_memory_usage = optimize_memory_usage
-        if replay_buffer_class is None:
-            if isinstance(self.observation_space, spaces.Dict):
-                self.replay_buffer_class = DictReplayBuffer
-            else:
-                self.replay_buffer_class = ReplayBuffer
-        else:
-            self.replay_buffer_class = replay_buffer_class
+        self.replay_buffer_class = replay_buffer_class
         self.replay_buffer_kwargs = replay_buffer_kwargs or {}
         self._episode_storage = None
 
@@ -139,7 +133,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         self.train_freq = train_freq
 
         self.actor = None  # type: Optional[th.nn.Module]
-        self.replay_buffer = None  # type: Optional[ReplayBuffer]
+        self.replay_buffer: Optional[ReplayBuffer] = None
         # Update policy keyword arguments
         if sde_support:
             self.policy_kwargs["use_sde"] = self.use_sde
@@ -174,6 +168,12 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         self._setup_lr_schedule()
         self.set_random_seed(self.seed)
 
+        if self.replay_buffer_class is None:
+            if isinstance(self.observation_space, spaces.Dict):
+                self.replay_buffer_class = DictReplayBuffer
+            else:
+                self.replay_buffer_class = ReplayBuffer
+
         if self.replay_buffer is None:
             # Make a local copy as we should not pickle
             # the environment when using HerReplayBuffer

stable_baselines3/common/policies.py

@@ -64,7 +64,7 @@ class BaseModel(nn.Module):
         action_space: spaces.Space,
         features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
         features_extractor_kwargs: Optional[Dict[str, Any]] = None,
-        features_extractor: Optional[nn.Module] = None,
+        features_extractor: Optional[BaseFeaturesExtractor] = None,
         normalize_images: bool = True,
         optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Dict[str, Any]] = None,
@@ -84,7 +84,7 @@ class BaseModel(nn.Module):
 
         self.optimizer_class = optimizer_class
         self.optimizer_kwargs = optimizer_kwargs
-        self.optimizer = None  # type: Optional[th.optim.Optimizer]
+        self.optimizer: th.optim.Optimizer
 
         self.features_extractor_class = features_extractor_class
         self.features_extractor_kwargs = features_extractor_kwargs
@@ -207,6 +207,26 @@ class BaseModel(nn.Module):
         """
         self.train(mode)
 
+    def is_vectorized_observation(self, observation: Union[np.ndarray, Dict[str, np.ndarray]]) -> bool:
+        """
+        Check whether or not the observation is vectorized,
+        apply transposition to image (so that they are channel-first) if needed.
+        This is used in DQN when sampling random action (epsilon-greedy policy)
+
+        :param observation: the input observation to check
+        :return: whether the given observation is vectorized or not
+        """
+        vectorized_env = False
+        if isinstance(observation, dict):
+            for key, obs in observation.items():
+                obs_space = self.observation_space.spaces[key]
+                vectorized_env = vectorized_env or is_vectorized_observation(maybe_transpose(obs, obs_space), obs_space)
+        else:
+            vectorized_env = is_vectorized_observation(
+                maybe_transpose(observation, self.observation_space), self.observation_space
+            )
+        return vectorized_env
+
     def obs_to_tensor(self, observation: Union[np.ndarray, Dict[str, np.ndarray]]) -> Tuple[th.Tensor, bool]:
         """
         Convert an input observation to a PyTorch tensor that can be fed to a model.

stable_baselines3/dqn/dqn.py

@@ -9,9 +9,8 @@ from torch.nn import functional as F
 from stable_baselines3.common.buffers import ReplayBuffer
 from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
 from stable_baselines3.common.policies import BasePolicy
-from stable_baselines3.common.preprocessing import maybe_transpose
 from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
-from stable_baselines3.common.utils import get_linear_fn, get_parameters_by_name, is_vectorized_observation, polyak_update
+from stable_baselines3.common.utils import get_linear_fn, get_parameters_by_name, polyak_update
 from stable_baselines3.dqn.policies import CnnPolicy, DQNPolicy, MlpPolicy, MultiInputPolicy
 
 SelfDQN = TypeVar("SelfDQN", bound="DQN")
@@ -93,7 +92,7 @@ class DQN(OffPolicyAlgorithm):
         seed: Optional[int] = None,
         device: Union[th.device, str] = "auto",
         _init_setup_model: bool = True,
-    ):
+    ) -> None:
         super().__init__(
             policy,
             env,
@@ -129,8 +128,9 @@ class DQN(OffPolicyAlgorithm):
         # "epsilon" for the epsilon-greedy exploration
         self.exploration_rate = 0.0
         # Linear schedule will be defined in `_setup_model()`
-        self.exploration_schedule = None
-        self.q_net, self.q_net_target = None, None
+        self.exploration_schedule: Schedule
+        self.q_net: th.nn.Module
+        self.q_net_target: th.nn.Module
 
         if _init_setup_model:
             self._setup_model()
@@ -160,6 +160,8 @@ class DQN(OffPolicyAlgorithm):
             self.target_update_interval = max(self.target_update_interval // self.n_envs, 1)
 
     def _create_aliases(self) -> None:
+        # For type checker:
+        assert isinstance(self.policy, DQNPolicy)
         self.q_net = self.policy.q_net
         self.q_net_target = self.policy.q_net_target
 
@@ -186,7 +188,7 @@ class DQN(OffPolicyAlgorithm):
         losses = []
         for _ in range(gradient_steps):
             # Sample replay buffer
-            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
+            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)  # type: ignore[union-attr]
 
             with th.no_grad():
                 # Compute the next Q-values using the target network
@@ -239,7 +241,7 @@ class DQN(OffPolicyAlgorithm):
             (used in recurrent policies)
         """
         if not deterministic and np.random.rand() < self.exploration_rate:
-            if is_vectorized_observation(maybe_transpose(observation, self.observation_space), self.observation_space):
+            if self.policy.is_vectorized_observation(observation):
                 if isinstance(observation, dict):
                     n_batch = observation[list(observation.keys())[0]].shape[0]
                 else:

stable_baselines3/dqn/policies.py

@@ -31,12 +31,12 @@ class QNetwork(BasePolicy):
         self,
         observation_space: spaces.Space,
         action_space: spaces.Space,
-        features_extractor: nn.Module,
+        features_extractor: BaseFeaturesExtractor,
         features_dim: int,
         net_arch: Optional[List[int]] = None,
         activation_fn: Type[nn.Module] = nn.ReLU,
         normalize_images: bool = True,
-    ):
+    ) -> None:
         super().__init__(
             observation_space,
             action_space,
@@ -49,7 +49,6 @@ class QNetwork(BasePolicy):
 
         self.net_arch = net_arch
         self.activation_fn = activation_fn
-        self.features_extractor = features_extractor
         self.features_dim = features_dim
         action_dim = self.action_space.n  # number of actions
         q_net = create_mlp(self.features_dim, action_dim, self.net_arch, self.activation_fn)
@@ -62,6 +61,8 @@ class QNetwork(BasePolicy):
         :param obs: Observation
         :return: The estimated Q-Value for each action.
         """
+        # For type checker:
+        assert isinstance(self.features_extractor, BaseFeaturesExtractor)
         return self.q_net(self.extract_features(obs, self.features_extractor))
 
     def _predict(self, observation: th.Tensor, deterministic: bool = True) -> th.Tensor:
@@ -116,7 +117,7 @@ class DQNPolicy(BasePolicy):
         normalize_images: bool = True,
         optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> None:
         super().__init__(
             observation_space,
             action_space,
@@ -144,7 +145,8 @@ class DQNPolicy(BasePolicy):
             "normalize_images": normalize_images,
         }
 
-        self.q_net, self.q_net_target = None, None
+        self.q_net: QNetwork
+        self.q_net_target: QNetwork
         self._build(lr_schedule)
 
     def _build(self, lr_schedule: Schedule) -> None:
@@ -163,7 +165,11 @@ class DQNPolicy(BasePolicy):
         self.q_net_target.set_training_mode(False)
 
         # Setup optimizer with initial learning rate
-        self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs)
+        self.optimizer = self.optimizer_class(  # type: ignore[call-arg]
+            self.parameters(),
+            lr=lr_schedule(1),
+            **self.optimizer_kwargs,
+        )
 
     def make_q_net(self) -> QNetwork:
         # Make sure we always have separate networks for features extractors etc
@@ -237,7 +243,7 @@ class CnnPolicy(DQNPolicy):
         normalize_images: bool = True,
         optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> None:
         super().__init__(
             observation_space,
             action_space,
@@ -282,7 +288,7 @@ class MultiInputPolicy(DQNPolicy):
         normalize_images: bool = True,
         optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
         optimizer_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> None:
         super().__init__(
             observation_space,
             action_space,

stable_baselines3/version.txt

@@ -1 +1 @@
-1.8.0a12
+1.8.0a13

tests/test_utils.py

@@ -186,9 +186,10 @@ def test_evaluate_policy(direct_policy: bool):
 
     assert model.policy is not None
     policy = model.policy if direct_policy else model
-    policy.n_callback_calls = 0
+
+    policy.n_callback_calls = 0  # type: ignore[assignment, attr-defined]
     _, episode_lengths = evaluate_policy(
-        policy,
+        policy,  # type: ignore[arg-type]
         model.get_env(),
         n_eval_episodes,
         deterministic=True,
@@ -198,21 +199,26 @@ def test_evaluate_policy(direct_policy: bool):
         return_episode_rewards=True,
     )
 
-    n_steps = sum(episode_lengths)
+    n_steps = sum(episode_lengths)  # type: ignore[arg-type]
     assert n_steps == n_steps_per_episode * n_eval_episodes
-    assert n_steps == policy.n_callback_calls
+    assert n_steps == policy.n_callback_calls  # type: ignore[attr-defined]
 
     # Reaching a mean reward of zero is impossible with the Pendulum env
     with pytest.raises(AssertionError):
-        evaluate_policy(policy, model.get_env(), n_eval_episodes, reward_threshold=0.0)
+        evaluate_policy(policy, model.get_env(), n_eval_episodes, reward_threshold=0.0)  # type: ignore[arg-type]
 
-    episode_rewards, _ = evaluate_policy(policy, model.get_env(), n_eval_episodes, return_episode_rewards=True)
-    assert len(episode_rewards) == n_eval_episodes
+    episode_rewards, _ = evaluate_policy(
+        policy,  # type: ignore[arg-type]
+        model.get_env(),
+        n_eval_episodes,
+        return_episode_rewards=True,
+    )
+    assert len(episode_rewards) == n_eval_episodes  # type: ignore[arg-type]
 
     # Test that warning is given about no monitor
     eval_env = gym.make("Pendulum-v1")
     with pytest.warns(UserWarning):
-        _ = evaluate_policy(policy, eval_env, n_eval_episodes)
+        _ = evaluate_policy(policy, eval_env, n_eval_episodes)  # type: ignore[arg-type]
 
 
 class ZeroRewardWrapper(gym.RewardWrapper):
@@ -512,6 +518,55 @@ def test_is_vectorized_observation():
         is_vectorized_observation(dict_obs, dict_space)
 
 
+def test_policy_is_vectorized_obs():
+    """
+    Additional tests to check `policy.is_vectorized()`
+    which handle transposing image to channel-first if needed.
+
+    We check for basic cases, the rest is handled
+    by is_vectorized_observation() helper.
+    """
+    policy = sb3.DQN("MlpPolicy", "CartPole-v1").policy
+
+    box_space = spaces.Box(-1, 1, shape=(2,))
+    box_obs = np.ones((1, *box_space.shape))
+    policy.observation_space = box_space
+    assert policy.is_vectorized_observation(box_obs)
+    assert not policy.is_vectorized_observation(np.ones(box_space.shape))
+
+    discrete_space = spaces.Discrete(2)
+    discrete_obs = np.ones((3,), dtype=np.int8)
+    policy.observation_space = discrete_space
+    assert not policy.is_vectorized_observation(np.ones((), dtype=np.int8))
+
+    dict_space = spaces.Dict({"box": box_space, "discrete": discrete_space})
+    dict_obs = {"box": box_obs, "discrete": discrete_obs}
+    policy.observation_space = dict_space
+    assert policy.is_vectorized_observation(dict_obs)
+    dict_obs = {"box": np.ones(box_space.shape), "discrete": np.ones((), dtype=np.int8)}
+    assert not policy.is_vectorized_observation(dict_obs)
+
+    # Image space are channel-first (done automatically in SB3 using VecTranspose)
+    # but observation passed is channel last
+    image_space = spaces.Box(low=0, high=255, shape=(3, 32, 32), dtype=np.uint8)
+
+    image_channel_first = image_space.sample()
+    image_channel_last = np.transpose(image_channel_first, (1, 2, 0))
+    policy.observation_space = image_space
+    assert not policy.is_vectorized_observation(image_channel_first)
+    assert not policy.is_vectorized_observation(image_channel_last)
+    assert policy.is_vectorized_observation(image_channel_first[np.newaxis])
+    assert policy.is_vectorized_observation(image_channel_last[np.newaxis])
+
+    # Same with dict obs
+    dict_space = spaces.Dict({"image": image_space})
+    policy.observation_space = dict_space
+    assert not policy.is_vectorized_observation({"image": image_channel_first})
+    assert not policy.is_vectorized_observation({"image": image_channel_last})
+    assert policy.is_vectorized_observation({"image": image_channel_first[np.newaxis]})
+    assert policy.is_vectorized_observation({"image": image_channel_last[np.newaxis]})
+
+
 def test_check_shape_equal():
     space1 = spaces.Box(low=0, high=1, shape=(2, 2))
     space2 = spaces.Box(low=-1, high=1, shape=(2, 2))