Hotfix for train signature

Author: araffin

sbx/sac/sac.py

@@ -179,20 +179,21 @@ def learn(
             progress_bar=progress_bar,
         )
 
-    def train(self, batch_size: int, gradient_steps: int):
+    def train(self, gradient_steps: int, batch_size: int) -> None:
+        assert self.replay_buffer is not None
         # Sample all at once for efficiency (so we can jit the for loop)
         data = self.replay_buffer.sample(batch_size * gradient_steps, env=self._vec_normalize_env)
 
         if isinstance(data.observations, dict):
-            keys = list(self.observation_space.keys())
+            keys = list(self.observation_space.keys())  # type: ignore[attr-defined]
             obs = np.concatenate([data.observations[key].numpy() for key in keys], axis=1)
             next_obs = np.concatenate([data.next_observations[key].numpy() for key in keys], axis=1)
         else:
             obs = data.observations.numpy()
             next_obs = data.next_observations.numpy()
 
         # Convert to numpy
-        data = ReplayBufferSamplesNp(
+        data = ReplayBufferSamplesNp(  # type: ignore[assignment]
             obs,
             data.actions.numpy(),
             next_obs,

sbx/td3/td3.py

@@ -120,20 +120,21 @@ def learn(
             progress_bar=progress_bar,
         )
 
-    def train(self, batch_size, gradient_steps):
+    def train(self, gradient_steps: int, batch_size: int) -> None:
+        assert self.replay_buffer is not None
         # Sample all at once for efficiency (so we can jit the for loop)
         data = self.replay_buffer.sample(batch_size * gradient_steps, env=self._vec_normalize_env)
 
         if isinstance(data.observations, dict):
-            keys = list(self.observation_space.keys())
+            keys = list(self.observation_space.keys())  # type: ignore[attr-defined]
             obs = np.concatenate([data.observations[key].numpy() for key in keys], axis=1)
             next_obs = np.concatenate([data.next_observations[key].numpy() for key in keys], axis=1)
         else:
             obs = data.observations.numpy()
             next_obs = data.next_observations.numpy()
 
         # Convert to numpy
-        data = ReplayBufferSamplesNp(
+        data = ReplayBufferSamplesNp(  # type: ignore[assignment]
             obs,
             data.actions.numpy(),
             next_obs,

sbx/tqc/tqc.py

@@ -180,20 +180,21 @@ def learn(
             progress_bar=progress_bar,
         )
 
-    def train(self, batch_size, gradient_steps):
+    def train(self, gradient_steps: int, batch_size: int) -> None:
+        assert self.replay_buffer is not None
         # Sample all at once for efficiency (so we can jit the for loop)
         data = self.replay_buffer.sample(batch_size * gradient_steps, env=self._vec_normalize_env)
 
         if isinstance(data.observations, dict):
-            keys = list(self.observation_space.keys())
+            keys = list(self.observation_space.keys())  # type: ignore[attr-defined]
             obs = np.concatenate([data.observations[key].numpy() for key in keys], axis=1)
             next_obs = np.concatenate([data.next_observations[key].numpy() for key in keys], axis=1)
         else:
             obs = data.observations.numpy()
             next_obs = data.next_observations.numpy()
 
         # Convert to numpy
-        data = ReplayBufferSamplesNp(
+        data = ReplayBufferSamplesNp(  # type: ignore[assignment]
             obs,
             data.actions.numpy(),
             next_obs,

sbx/version.txt

@@ -1 +1 @@
-0.9.1
+0.10.0