v1: replace all isinstance checks from BatchProtocol to Batch

Seriously improves performance of Batch constructor

Author: MischaPanch

test/base/test_collector.py

@@ -72,7 +72,7 @@ def forward(
         if self.dict_state:
             if self.action_shape:
                 action_shape = self.action_shape
-            elif isinstance(batch.obs, BatchProtocol):
+            elif isinstance(batch.obs, Batch):
                 action_shape = len(batch.obs["index"])
             else:
                 action_shape = len(batch.obs)

tianshou/data/batch.py

@@ -983,7 +983,7 @@ def __cat(self, batches: Sequence[dict | Self], lens: list[int]) -> None:
                     self.__dict__[key][sum_lens[i] : sum_lens[i + 1]] = value
 
     def cat_(self, batches: BatchProtocol | Sequence[dict | BatchProtocol]) -> None:
-        if isinstance(batches, BatchProtocol | dict):
+        if isinstance(batches, Batch | dict):
             batches = [batches]
         # check input format
         batch_list = []
@@ -1069,7 +1069,7 @@ def stack_(self, batches: Sequence[dict | BatchProtocol], axis: int = 0) -> None
             {
                 batch_key
                 for batch_key, obj in batch.items()
-                if not (isinstance(obj, BatchProtocol) and len(obj.get_keys()) == 0)
+                if not (isinstance(obj, Batch) and len(obj.get_keys()) == 0)
             }
             for batch in batches
         ]
@@ -1080,7 +1080,7 @@ def stack_(self, batches: Sequence[dict | BatchProtocol], axis: int = 0) -> None
             if all(isinstance(element, torch.Tensor) for element in value):
                 self.__dict__[shared_key] = torch.stack(value, axis)
             # third often
-            elif all(isinstance(element, BatchProtocol | dict) for element in value):
+            elif all(isinstance(element, Batch | dict) for element in value):
                 self.__dict__[shared_key] = Batch.stack(value, axis)
             else:  # most often case is np.ndarray
                 try:
@@ -1114,7 +1114,7 @@ def stack_(self, batches: Sequence[dict | BatchProtocol], axis: int = 0) -> None
                 value = batch.get(key)
                 # TODO: fix code/annotations s.t. the ignores can be removed
                 if (
-                    isinstance(value, BatchProtocol)  # type: ignore
+                    isinstance(value, Batch)  # type: ignore
                     and len(value.get_keys()) == 0  # type: ignore
                 ):
                     continue  # type: ignore
@@ -1288,7 +1288,7 @@ def set_array_at_key(
                     ) from exception
             else:
                 existing_entry = self[key]
-                if isinstance(existing_entry, BatchProtocol):
+                if isinstance(existing_entry, Batch):
                     raise ValueError(
                         f"Cannot set sequence at key {key} because it is a nested batch, "
                         f"can only set a subsequence of an array.",
@@ -1312,7 +1312,7 @@ def hasnull(self) -> bool:
 
         def is_any_true(boolean_batch: BatchProtocol) -> bool:
             for val in boolean_batch.values():
-                if isinstance(val, BatchProtocol):
+                if isinstance(val, Batch):
                     if is_any_true(val):
                         return True
                 else:
@@ -1375,7 +1375,7 @@ def _apply_batch_values_func_recursively(
     """
     result = batch if inplace else deepcopy(batch)
     for key, val in batch.__dict__.items():
-        if isinstance(val, BatchProtocol):
+        if isinstance(val, Batch):
             result[key] = _apply_batch_values_func_recursively(val, values_transform, inplace=False)
         else:
             result[key] = values_transform(val)

tianshou/data/buffer/her.py

@@ -150,12 +150,12 @@ def rewrite_transitions(self, indices: np.ndarray) -> None:
         ep_obs = self[unique_ep_indices].obs
         # to satisfy mypy
         # TODO: add protocol covering these batches
-        assert isinstance(ep_obs, BatchProtocol)
+        assert isinstance(ep_obs, Batch)
         ep_rew = self[unique_ep_indices].rew
         if self._save_obs_next:
             ep_obs_next = self[unique_ep_indices].obs_next
             # to satisfy mypy
-            assert isinstance(ep_obs_next, BatchProtocol)
+            assert isinstance(ep_obs_next, Batch)
             future_obs = self[future_t[unique_ep_close_indices]].obs_next
         else:
             future_obs = self[self.next(future_t[unique_ep_close_indices])].obs
@@ -172,7 +172,7 @@ def rewrite_transitions(self, indices: np.ndarray) -> None:
             ep_rew[:, her_ep_indices] = self._compute_reward(ep_obs_next)[:, her_ep_indices]
         else:
             tmp_ep_obs_next = self[self.next(unique_ep_indices)].obs
-            assert isinstance(tmp_ep_obs_next, BatchProtocol)
+            assert isinstance(tmp_ep_obs_next, Batch)
             ep_rew[:, her_ep_indices] = self._compute_reward(tmp_ep_obs_next)[:, her_ep_indices]
 
         # Sanity check
@@ -181,7 +181,7 @@ def rewrite_transitions(self, indices: np.ndarray) -> None:
         assert ep_rew.shape == unique_ep_indices.shape
 
         # Re-write meta
-        assert isinstance(self._meta.obs, BatchProtocol)
+        assert isinstance(self._meta.obs, Batch)
         self._meta.obs[unique_ep_indices] = ep_obs
         if self._save_obs_next:
             self._meta.obs_next[unique_ep_indices] = ep_obs_next  # type: ignore

tianshou/policy/modelbased/psrl.py

@@ -236,7 +236,8 @@ def learn(self, batch: RolloutBatchProtocol, *args: Any, **kwargs: Any) -> TPSRL
         for minibatch in batch.split(size=1):
             obs, act, obs_next = minibatch.obs, minibatch.act, minibatch.obs_next
             obs_next = cast(np.ndarray, obs_next)
-            assert not isinstance(obs, BatchProtocol), "Observations cannot be Batches here"
+            assert not isinstance(obs, Batch), "Observations cannot be Batches here"
+            obs = cast(np.ndarray, obs)
             trans_count[obs, act, obs_next] += 1
             rew_sum[obs, act] += minibatch.rew
             rew_square_sum[obs, act] += minibatch.rew**2