diff --git a/.coveragerc b/.coveragerc deleted file mode 100644 index 8ee0fb9..0000000 --- a/.coveragerc +++ /dev/null @@ -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: diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 5c4dc42..4a0c2e3 100644 --- a/docs/misc/changelog.rst +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 7941679..2f98aa4 100644 --- a/pyproject.toml +++ b/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:"] diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 0ab03d8..15d93e2 100644 --- a/stable_baselines3/common/base_class.py +++ b/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] diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 943f2ab..63afa90 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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 diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index e07aefa..45d7930 100644 --- a/stable_baselines3/common/policies.py +++ b/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. diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 61560e2..015da5d 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/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: diff --git a/stable_baselines3/dqn/policies.py b/stable_baselines3/dqn/policies.py index 22e6d0a..6991b8a 100644 --- a/stable_baselines3/dqn/policies.py +++ b/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, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 15159be..3c66ede 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.8.0a12 +1.8.0a13 \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 0bab41e..c4399a8 100644 --- a/tests/test_utils.py +++ b/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))