fix type errors

araffin · Feb 26, 2025 · 815e5be · 815e5be
1 parent 621b64d
commit 815e5be
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/sbx/common/off_policy_algorithm.py b/sbx/common/off_policy_algorithm.py
@@ -89,7 +89,7 @@ def _maybe_reset_params(self) -> None:
         ):
             # Note: we are not resetting the entropy coeff
             assert isinstance(self.qf_learning_rate, float)
-            self.key = self.policy.build(self.key, self.lr_schedule, self.qf_learning_rate)
+            self.key = self.policy.build(self.key, self.lr_schedule, self.qf_learning_rate)  # type: ignore[operator]
             self.reset_idx += 1
 
     def _get_torch_save_params(self):

diff --git a/sbx/common/type_aliases.py b/sbx/common/type_aliases.py
@@ -14,8 +14,8 @@ class BatchNormTrainState(TrainState):  # type: ignore[misc]
 
 
 class ReplayBufferSamplesNp(NamedTuple):
-    observations: Union[np.ndarray, dict[str, np.ndarray]]
+    observations: np.ndarray
     actions: np.ndarray
-    next_observations: Union[np.ndarray, dict[str, np.ndarray]]
+    next_observations: np.ndarray
     dones: np.ndarray
     rewards: np.ndarray