v2: add_exploration_noise - raise error on wrong type instead of doing nothing

Author: MischaPanch

tianshou/algorithm/modelfree/bdqn.py

@@ -1,4 +1,4 @@
-from typing import Any, cast
+from typing import cast
 
 import gymnasium as gym
 import numpy as np
@@ -65,7 +65,6 @@ def forward(
         batch: ObsBatchProtocol,
         state: dict | BatchProtocol | np.ndarray | None = None,
         model: torch.nn.Module | None = None,
-        **kwargs: Any,
     ) -> ModelOutputBatchProtocol:
         if model is None:
             model = self.model
@@ -84,8 +83,9 @@ def add_exploration_noise(
         batch: ObsBatchProtocol,
     ) -> TArrOrActBatch:
         eps = self.eps_training if self.is_within_training_step else self.eps_inference
-        # TODO: This looks problematic; the non-array case is silently ignored
-        if isinstance(act, np.ndarray) and not np.isclose(eps, 0.0):
+        if not np.isclose(eps, 0.0):
+            return act
+        if isinstance(act, np.ndarray):
             bsz = len(act)
             rand_mask = np.random.rand(bsz) < eps
             rand_act = np.random.randint(
@@ -96,7 +96,11 @@ def add_exploration_noise(
             if hasattr(batch.obs, "mask"):
                 rand_act += batch.obs.mask
             act[rand_mask] = rand_act[rand_mask]
-        return act
+            return act
+        else:
+            raise NotImplementedError(
+                f"Currently only numpy arrays are supported, got {type(act)=}."
+            )
 
 
 class BDQN(QLearningOffPolicyAlgorithm[BDQNPolicy]):

tianshou/algorithm/modelfree/dqn.py

@@ -156,8 +156,10 @@ def add_exploration_noise(
         batch: ObsBatchProtocol,
     ) -> TArrOrActBatch:
         eps = self.eps_training if self.is_within_training_step else self.eps_inference
-        # TODO: This looks problematic; the non-array case is silently ignored
-        if isinstance(act, np.ndarray) and not np.isclose(eps, 0.0):
+        eps = self.eps_training if self.is_within_training_step else self.eps_inference
+        if not np.isclose(eps, 0.0):
+            return act
+        if isinstance(act, np.ndarray):
             batch_size = len(act)
             rand_mask = np.random.rand(batch_size) < eps
             self.action_space = cast(Discrete, self.action_space)  # for mypy
@@ -167,7 +169,10 @@ def add_exploration_noise(
                 q += batch.obs.mask
             rand_act = q.argmax(axis=1)
             act[rand_mask] = rand_act[rand_mask]
-        return act
+            return act
+        raise NotImplementedError(
+            f"Currently only numpy array is supported for action, but got {type(act)}"
+        )
 
 
 TDQNPolicy = TypeVar("TDQNPolicy", bound=DiscreteQLearningPolicy)