Enable numpy typing check

d3rlpy/algos/qlearning/base.py

@@ -46,6 +46,7 @@
     sync_optimizer_state,
     train_api,
 )
+from ...types import NDArray
 from ..utility import (
     assert_action_space_with_dataset,
     assert_action_space_with_env,
@@ -243,7 +244,7 @@ def _func(x: torch.Tensor) -> torch.Tensor:
         # workaround until version 1.6
         self._impl.modules.unfreeze()
 
-    def predict(self, x: Observation) -> np.ndarray:
+    def predict(self, x: Observation) -> NDArray:
         """Returns greedy actions.
 
         .. code-block:: python
@@ -278,9 +279,9 @@ def predict(self, x: Observation) -> np.ndarray:
             if self._config.action_scaler:
                 action = self._config.action_scaler.reverse_transform(action)
 
-        return action.cpu().detach().numpy()
+        return action.cpu().detach().numpy()  # type: ignore
 
-    def predict_value(self, x: Observation, action: np.ndarray) -> np.ndarray:
+    def predict_value(self, x: Observation, action: NDArray) -> NDArray:
         """Returns predicted action-values.
 
         .. code-block:: python
@@ -332,9 +333,9 @@ def predict_value(self, x: Observation, action: np.ndarray) -> np.ndarray:
 
             value = self._impl.predict_value(torch_x, torch_action)
 
-        return value.cpu().detach().numpy()
+        return value.cpu().detach().numpy()  # type: ignore
 
-    def sample_action(self, x: Observation) -> np.ndarray:
+    def sample_action(self, x: Observation) -> NDArray:
         """Returns sampled actions.
 
         The sampled actions are identical to the output of `predict` method if
@@ -364,7 +365,7 @@ def sample_action(self, x: Observation) -> np.ndarray:
             if self._config.action_scaler:
                 action = self._config.action_scaler.reverse_transform(action)
 
-        return action.cpu().detach().numpy()
+        return action.cpu().detach().numpy()  # type: ignore
 
     def fit(
         self,
@@ -795,7 +796,7 @@ def collect(
             clip_episode = terminal or truncated
 
             # store observation
-            buffer.append(observation, action, reward)
+            buffer.append(observation, action, float(reward))
 
             # reset if terminated
             if clip_episode:

d3rlpy/algos/qlearning/explorers.py

@@ -1,10 +1,12 @@
 from abc import ABCMeta, abstractmethod
-from typing import Any, List, Optional, Union
+from typing import Union
 
 import numpy as np
-from typing_extensions import Protocol
 
-from ...preprocessing.action_scalers import ActionScaler, MinMaxActionScaler
+from ...dataset import Observation
+from ...interface import QLearningAlgoProtocol
+from ...preprocessing.action_scalers import MinMaxActionScaler
+from ...types import NDArray
 
 __all__ = [
     "Explorer",
@@ -14,24 +16,11 @@
 ]
 
 
-class _ActionProtocol(Protocol):
-    def predict(self, x: Union[np.ndarray, List[Any]]) -> np.ndarray:
-        ...
-
-    @property
-    def action_size(self) -> Optional[int]:
-        ...
-
-    @property
-    def action_scaler(self) -> Optional[ActionScaler]:
-        ...
-
-
 class Explorer(metaclass=ABCMeta):
     @abstractmethod
     def sample(
-        self, algo: _ActionProtocol, x: np.ndarray, step: int
-    ) -> np.ndarray:
+        self, algo: QLearningAlgoProtocol, x: Observation, step: int
+    ) -> NDArray:
         pass
 
 
@@ -48,11 +37,14 @@ def __init__(self, epsilon: float):
         self._epsilon = epsilon
 
     def sample(
-        self, algo: _ActionProtocol, x: np.ndarray, step: int
-    ) -> np.ndarray:
+        self, algo: QLearningAlgoProtocol, x: Observation, step: int
+    ) -> NDArray:
+        action_size = algo.action_size
+        assert action_size is not None
         greedy_actions = algo.predict(x)
-        random_actions = np.random.randint(algo.action_size, size=x.shape[0])
-        is_random = np.random.random(x.shape[0]) < self._epsilon
+        batch_size = greedy_actions.shape[0]
+        random_actions = np.random.randint(action_size, size=batch_size)
+        is_random = np.random.random(batch_size) < self._epsilon
         return np.where(is_random, random_actions, greedy_actions)
 
 
@@ -80,8 +72,8 @@ def __init__(
         self._duration = duration
 
     def sample(
-        self, algo: _ActionProtocol, x: np.ndarray, step: int
-    ) -> np.ndarray:
+        self, algo: QLearningAlgoProtocol, x: Observation, step: int
+    ) -> NDArray:
         """Returns :math:`\\epsilon`-greedy action.
 
         Args:
@@ -92,9 +84,12 @@ def sample(
         Returns:
             :math:`\\epsilon`-greedy action.
         """
+        action_size = algo.action_size
+        assert action_size is not None
         greedy_actions = algo.predict(x)
-        random_actions = np.random.randint(algo.action_size, size=x.shape[0])
-        is_random = np.random.random(x.shape[0]) < self.compute_epsilon(step)
+        batch_size = greedy_actions.shape[0]
+        random_actions = np.random.randint(action_size, size=batch_size)
+        is_random = np.random.random(batch_size) < self.compute_epsilon(step)
         return np.where(is_random, random_actions, greedy_actions)
 
     def compute_epsilon(self, step: int) -> float:
@@ -125,8 +120,8 @@ def __init__(self, mean: float = 0.0, std: float = 0.1):
         self._std = std
 
     def sample(
-        self, algo: _ActionProtocol, x: np.ndarray, step: int
-    ) -> np.ndarray:
+        self, algo: QLearningAlgoProtocol, x: Observation, step: int
+    ) -> NDArray:
         """Returns action with noise injection.
 
         Args:
@@ -139,12 +134,16 @@ def sample(
         action = algo.predict(x)
         noise = np.random.normal(self._mean, self._std, size=action.shape)
 
+        minimum: Union[float, NDArray]
+        maximum: Union[float, NDArray]
         if isinstance(algo.action_scaler, MinMaxActionScaler):
             # scale noise
+            assert algo.action_scaler.minimum is not None
+            assert algo.action_scaler.maximum is not None
             minimum = algo.action_scaler.minimum
             maximum = algo.action_scaler.maximum
         else:
             minimum = -1.0
             maximum = 1.0
 
-        return np.clip(action + noise, minimum, maximum)
+        return np.clip(action + noise, minimum, maximum)  # type: ignore

d3rlpy/algos/qlearning/random_policy.py

@@ -7,6 +7,7 @@
 from ...constants import ActionSpace
 from ...dataset import Observation, Shape
 from ...torch_utility import TorchMiniBatch
+from ...types import NDArray
 from .base import QLearningAlgoBase
 
 __all__ = [
@@ -54,10 +55,10 @@ def inner_create_impl(
     ) -> None:
         self._action_size = action_size
 
-    def predict(self, x: Observation) -> np.ndarray:
+    def predict(self, x: Observation) -> NDArray:
         return self.sample_action(x)
 
-    def sample_action(self, x: Observation) -> np.ndarray:
+    def sample_action(self, x: Observation) -> NDArray:
         x = np.asarray(x)
         action_shape = (x.shape[0], self._action_size)
 
@@ -79,7 +80,7 @@ def sample_action(self, x: Observation) -> np.ndarray:
 
         return action
 
-    def predict_value(self, x: Observation, action: np.ndarray) -> np.ndarray:
+    def predict_value(self, x: Observation, action: NDArray) -> NDArray:
         raise NotImplementedError
 
     def inner_update(self, batch: TorchMiniBatch) -> Dict[str, float]:
@@ -117,14 +118,14 @@ def inner_create_impl(
     ) -> None:
         self._action_size = action_size
 
-    def predict(self, x: Observation) -> np.ndarray:
+    def predict(self, x: Observation) -> NDArray:
         return self.sample_action(x)
 
-    def sample_action(self, x: Observation) -> np.ndarray:
+    def sample_action(self, x: Observation) -> NDArray:
         x = np.asarray(x)
         return np.random.randint(self._action_size, size=x.shape[0])
 
-    def predict_value(self, x: Observation, action: np.ndarray) -> np.ndarray:
+    def predict_value(self, x: Observation, action: NDArray) -> NDArray:
         raise NotImplementedError
 
     def inner_update(self, batch: TorchMiniBatch) -> Dict[str, float]:

d3rlpy/algos/transformer/action_samplers.py

@@ -3,6 +3,8 @@
 import numpy as np
 from typing_extensions import Protocol
 
+from ...types import NDArray
+
 __all__ = [
     "TransformerActionSampler",
     "IdentityTransformerActionSampler",
@@ -14,9 +16,7 @@
 class TransformerActionSampler(Protocol):
     r"""Interface of TransformerActionSampler."""
 
-    def __call__(
-        self, transformer_output: np.ndarray
-    ) -> Union[np.ndarray, int]:
+    def __call__(self, transformer_output: NDArray) -> Union[NDArray, int]:
         r"""Returns sampled action from Transformer output.
 
         Args:
@@ -35,9 +35,7 @@ class IdentityTransformerActionSampler(TransformerActionSampler):
     Sampled action is the exactly same as ``transformer_output``.
     """
 
-    def __call__(
-        self, transformer_output: np.ndarray
-    ) -> Union[np.ndarray, int]:
+    def __call__(self, transformer_output: NDArray) -> Union[NDArray, int]:
         return transformer_output
 
 
@@ -55,9 +53,7 @@ class SoftmaxTransformerActionSampler(TransformerActionSampler):
     def __init__(self, temperature: float = 1.0):
         self._temperature = temperature
 
-    def __call__(
-        self, transformer_output: np.ndarray
-    ) -> Union[np.ndarray, int]:
+    def __call__(self, transformer_output: NDArray) -> Union[NDArray, int]:
         assert transformer_output.ndim == 1
         logits = transformer_output / self._temperature
         x = np.exp(logits - np.max(logits))
@@ -73,8 +69,6 @@ class GreedyTransformerActionSampler(TransformerActionSampler):
     probability distribution.
     """
 
-    def __call__(
-        self, transformer_output: np.ndarray
-    ) -> Union[np.ndarray, int]:
+    def __call__(self, transformer_output: NDArray) -> Union[NDArray, int]:
         assert transformer_output.ndim == 1
         return int(np.argmax(transformer_output))

d3rlpy/algos/transformer/base.py

@@ -20,6 +20,7 @@
 )
 from ...metrics import evaluate_transformer_with_environment
 from ...torch_utility import TorchTrajectoryMiniBatch, train_api
+from ...types import NDArray
 from ..utility import (
     assert_action_space_with_dataset,
     build_scalers_with_trajectory_slicer,
@@ -102,7 +103,7 @@ class StatefulTransformerWrapper(Generic[TTransformerImpl, TTransformerConfig]):
     _action_sampler: TransformerActionSampler
     _return_rest: float
     _observations: Deque[Observation]
-    _actions: Deque[Union[np.ndarray, int]]
+    _actions: Deque[Union[NDArray, int]]
     _rewards: Deque[float]
     _returns_to_go: Deque[float]
     _timesteps: Deque[int]
@@ -128,7 +129,7 @@ def __init__(
         self._timesteps = deque([], maxlen=context_size)
         self._timestep = 1
 
-    def predict(self, x: Observation, reward: float) -> Union[np.ndarray, int]:
+    def predict(self, x: Observation, reward: float) -> Union[NDArray, int]:
         r"""Returns action.
 
         Args:
@@ -173,8 +174,9 @@ def algo(
     ) -> "TransformerAlgoBase[TTransformerImpl, TTransformerConfig]":
         return self._algo
 
-    def _get_pad_action(self) -> Union[int, np.ndarray]:
+    def _get_pad_action(self) -> Union[int, NDArray]:
         assert self._algo.impl
+        pad_action: Union[int, NDArray]
         if self._algo.get_action_type() == ActionSpace.CONTINUOUS:
             pad_action = np.zeros(self._algo.impl.action_size, dtype=np.float32)
         else:
@@ -186,7 +188,7 @@ class TransformerAlgoBase(
     Generic[TTransformerImpl, TTransformerConfig],
     LearnableBase[TTransformerImpl, TTransformerConfig],
 ):
-    def predict(self, inpt: TransformerInput) -> np.ndarray:
+    def predict(self, inpt: TransformerInput) -> NDArray:
         """Returns action.
 
         This is for internal use. For evaluation, use
@@ -213,7 +215,7 @@ def predict(self, inpt: TransformerInput) -> np.ndarray:
             if self._config.action_scaler:
                 action = self._config.action_scaler.reverse_transform(action)
 
-        return action.cpu().detach().numpy()
+        return action.cpu().detach().numpy()  # type: ignore
 
     def fit(
         self,