Add argument features_extractor to ActorCriticPolicy.extract_features

docs/misc/changelog.rst

@@ -3,7 +3,7 @@
 Changelog
 ==========
 
-Release 2.2.0a6 (WIP)
+Release 2.2.0a7 (WIP)
 --------------------------
 
 Breaking Changes:
@@ -52,6 +52,7 @@ Others:
 - Fixed ``stable_baselines3/common/buffers.py`` type hints
 - Fixed ``stable_baselines3/her/her_replay_buffer.py`` type hints
 - Buffers do no call an additional ``.copy()`` when storing new transitions
+- Fixed ``ActorCriticPolicy.extract_features()`` signature by adding an optional ``features_extractor`` argument
 
 Documentation:
 ^^^^^^^^^^^^^^

stable_baselines3/common/policies.py

@@ -123,9 +123,9 @@ def extract_features(self, obs: th.Tensor, features_extractor: BaseFeaturesExtra
         """
         Preprocess the observation if needed and extract features.
 
-         :param obs: The observation
-         :param features_extractor: The features extractor to use.
-         :return: The extracted features
+        :param obs: Observation
+        :param features_extractor: The features extractor to use.
+        :return: The extracted features
         """
         preprocessed_obs = preprocess_obs(obs, self.observation_space, normalize_images=self.normalize_images)
         return features_extractor(preprocessed_obs)
@@ -642,16 +642,26 @@ def forward(self, obs: th.Tensor, deterministic: bool = False) -> Tuple[th.Tenso
         actions = actions.reshape((-1, *self.action_space.shape))
         return actions, values, log_prob
 
-    def extract_features(self, obs: th.Tensor) -> Union[th.Tensor, Tuple[th.Tensor, th.Tensor]]:
+    def extract_features(
+        self, obs: th.Tensor, features_extractor: Optional[BaseFeaturesExtractor] = None
+    ) -> Union[th.Tensor, Tuple[th.Tensor, th.Tensor]]:
         """
         Preprocess the observation if needed and extract features.
 
         :param obs: Observation
-        :return: the output of the features extractor(s)
+        :param features_extractor: The features extractor to use. If None, then ``self.features_extractor`` is used.
+        :return: The extracted features. If features extractor is not shared, returns a tuple with the
+            features for the actor and the features for the critic.
         """
         if self.share_features_extractor:
-            return super().extract_features(obs, self.features_extractor)
+            return super().extract_features(obs, self.features_extractor if features_extractor is None else features_extractor)
         else:
+            if features_extractor is not None:
+                warnings.warn(
+                    "Provided features_extractor will be ignored because the features extractor is not shared.",
+                    UserWarning,
+                )
+
             pi_features = super().extract_features(obs, self.pi_features_extractor)
             vf_features = super().extract_features(obs, self.vf_features_extractor)
             return pi_features, vf_features

stable_baselines3/version.txt

@@ -1 +1 @@
-2.2.0a6
+2.2.0a7