Allow mypy to reject incomplete defs for mlagents-envs (#2585)

This wasn't working before because of several remaining partially defined
function definitions.

.pre-commit-config.yaml

@@ -15,8 +15,7 @@ repos:
         files: "ml-agents-envs/.*"
         # Exclude protobuf files and don't follow them when imported
         exclude: ".*_pb2.py"
-        # TODO get disallow-incomplete-defs working
-        args: [--ignore-missing-imports, --follow-imports=silent]
+        args: [--ignore-missing-imports, --disallow-incomplete-defs]
     -   id: mypy
         name: mypy-gym-unity
         files: "gym-unity/.*"

ml-agents-envs/mlagents/envs/base_unity_environment.py

@@ -1,19 +1,26 @@
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Dict, Optional, Any
 
 from mlagents.envs.brain import AllBrainInfo, BrainParameters
 
 
 class BaseUnityEnvironment(ABC):
     @abstractmethod
     def step(
-        self, vector_action=None, memory=None, text_action=None, value=None
+        self,
+        vector_action: Optional[Dict] = None,
+        memory: Optional[Dict] = None,
+        text_action: Optional[Dict] = None,
+        value: Optional[Dict] = None,
     ) -> AllBrainInfo:
         pass
 
     @abstractmethod
     def reset(
-        self, config=None, train_mode=True, custom_reset_parameters=None
+        self,
+        config: Optional[Dict] = None,
+        train_mode: bool = True,
+        custom_reset_parameters: Any = None,
     ) -> AllBrainInfo:
         pass
 

ml-agents-envs/mlagents/envs/brain.py

@@ -2,12 +2,81 @@
 import numpy as np
 import io
 
+from mlagents.envs.communicator_objects.agent_info_proto_pb2 import AgentInfoProto
+from mlagents.envs.communicator_objects.brain_parameters_proto_pb2 import (
+    BrainParametersProto,
+)
 from typing import Dict, List, Optional
 from PIL import Image
 
 logger = logging.getLogger("mlagents.envs")
 
 
+class BrainParameters:
+    def __init__(
+        self,
+        brain_name: str,
+        vector_observation_space_size: int,
+        num_stacked_vector_observations: int,
+        camera_resolutions: List[Dict],
+        vector_action_space_size: List[int],
+        vector_action_descriptions: List[str],
+        vector_action_space_type: int,
+    ):
+        """
+        Contains all brain-specific parameters.
+        """
+        self.brain_name = brain_name
+        self.vector_observation_space_size = vector_observation_space_size
+        self.num_stacked_vector_observations = num_stacked_vector_observations
+        self.number_visual_observations = len(camera_resolutions)
+        self.camera_resolutions = camera_resolutions
+        self.vector_action_space_size = vector_action_space_size
+        self.vector_action_descriptions = vector_action_descriptions
+        self.vector_action_space_type = ["discrete", "continuous"][
+            vector_action_space_type
+        ]
+
+    def __str__(self):
+        return """Unity brain name: {}
+        Number of Visual Observations (per agent): {}
+        Vector Observation space size (per agent): {}
+        Number of stacked Vector Observation: {}
+        Vector Action space type: {}
+        Vector Action space size (per agent): {}
+        Vector Action descriptions: {}""".format(
+            self.brain_name,
+            str(self.number_visual_observations),
+            str(self.vector_observation_space_size),
+            str(self.num_stacked_vector_observations),
+            self.vector_action_space_type,
+            str(self.vector_action_space_size),
+            ", ".join(self.vector_action_descriptions),
+        )
+
+    @staticmethod
+    def from_proto(brain_param_proto: BrainParametersProto) -> "BrainParameters":
+        """
+        Converts brain parameter proto to BrainParameter object.
+        :param brain_param_proto: protobuf object.
+        :return: BrainParameter object.
+        """
+        resolution = [
+            {"height": x.height, "width": x.width, "blackAndWhite": x.gray_scale}
+            for x in brain_param_proto.camera_resolutions
+        ]
+        brain_params = BrainParameters(
+            brain_param_proto.brain_name,
+            brain_param_proto.vector_observation_size,
+            brain_param_proto.num_stacked_vector_observations,
+            resolution,
+            list(brain_param_proto.vector_action_size),
+            list(brain_param_proto.vector_action_descriptions),
+            brain_param_proto.vector_action_space_type,
+        )
+        return brain_params
+
+
 class BrainInfo:
     def __init__(
         self,
@@ -101,7 +170,11 @@ def process_pixels(image_bytes: bytes, gray_scale: bool) -> np.ndarray:
         return s
 
     @staticmethod
-    def from_agent_proto(worker_id: int, agent_info_list, brain_params):
+    def from_agent_proto(
+        worker_id: int,
+        agent_info_list: List[AgentInfoProto],
+        brain_params: BrainParameters,
+    ) -> "BrainInfo":
         """
         Converts list of agent infos to BrainInfo.
         """
@@ -190,7 +263,9 @@ def safe_concat_lists(l1: Optional[List], l2: Optional[List]) -> Optional[List]:
             return copy
 
 
-def safe_concat_np_ndarray(a1: Optional[np.ndarray], a2: Optional[np.ndarray]):
+def safe_concat_np_ndarray(
+    a1: Optional[np.ndarray], a2: Optional[np.ndarray]
+) -> Optional[np.ndarray]:
     if a1 is not None and a1.size != 0:
         if a2 is not None and a2.size != 0:
             return np.append(a1, a2, axis=0)
@@ -203,68 +278,3 @@ def safe_concat_np_ndarray(a1: Optional[np.ndarray], a2: Optional[np.ndarray]):
 
 # Renaming of dictionary of brain name to BrainInfo for clarity
 AllBrainInfo = Dict[str, BrainInfo]
-
-
-class BrainParameters:
-    def __init__(
-        self,
-        brain_name: str,
-        vector_observation_space_size: int,
-        num_stacked_vector_observations: int,
-        camera_resolutions: List[Dict],
-        vector_action_space_size: List[int],
-        vector_action_descriptions: List[str],
-        vector_action_space_type: int,
-    ):
-        """
-        Contains all brain-specific parameters.
-        """
-        self.brain_name = brain_name
-        self.vector_observation_space_size = vector_observation_space_size
-        self.num_stacked_vector_observations = num_stacked_vector_observations
-        self.number_visual_observations = len(camera_resolutions)
-        self.camera_resolutions = camera_resolutions
-        self.vector_action_space_size = vector_action_space_size
-        self.vector_action_descriptions = vector_action_descriptions
-        self.vector_action_space_type = ["discrete", "continuous"][
-            vector_action_space_type
-        ]
-
-    def __str__(self):
-        return """Unity brain name: {}
-        Number of Visual Observations (per agent): {}
-        Vector Observation space size (per agent): {}
-        Number of stacked Vector Observation: {}
-        Vector Action space type: {}
-        Vector Action space size (per agent): {}
-        Vector Action descriptions: {}""".format(
-            self.brain_name,
-            str(self.number_visual_observations),
-            str(self.vector_observation_space_size),
-            str(self.num_stacked_vector_observations),
-            self.vector_action_space_type,
-            str(self.vector_action_space_size),
-            ", ".join(self.vector_action_descriptions),
-        )
-
-    @staticmethod
-    def from_proto(brain_param_proto):
-        """
-        Converts brain parameter proto to BrainParameter object.
-        :param brain_param_proto: protobuf object.
-        :return: BrainParameter object.
-        """
-        resolution = [
-            {"height": x.height, "width": x.width, "blackAndWhite": x.gray_scale}
-            for x in brain_param_proto.camera_resolutions
-        ]
-        brain_params = BrainParameters(
-            brain_param_proto.brain_name,
-            brain_param_proto.vector_observation_size,
-            brain_param_proto.num_stacked_vector_observations,
-            resolution,
-            list(brain_param_proto.vector_action_size),
-            list(brain_param_proto.vector_action_descriptions),
-            brain_param_proto.vector_action_space_type,
-        )
-        return brain_params

ml-agents-envs/mlagents/envs/env_manager.py

@@ -23,7 +23,9 @@ def step(self) -> List[EnvironmentStep]:
         pass
 
     @abstractmethod
-    def reset(self, config=None, train_mode=True) -> List[EnvironmentStep]:
+    def reset(
+        self, config: Dict = None, train_mode: bool = True
+    ) -> List[EnvironmentStep]:
         pass
 
     @property

ml-agents-envs/mlagents/envs/environment.py

@@ -4,7 +4,7 @@
 import numpy as np
 import os
 import subprocess
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 
 from mlagents.envs.base_unity_environment import BaseUnityEnvironment
 from mlagents.envs.timers import timed, hierarchical_timer
@@ -41,7 +41,7 @@
 class UnityEnvironment(BaseUnityEnvironment):
     SCALAR_ACTION_TYPES = (int, np.int32, np.int64, float, np.float32, np.float64)
     SINGLE_BRAIN_ACTION_TYPES = SCALAR_ACTION_TYPES + (list, np.ndarray)
-    SINGLE_BRAIN_TEXT_TYPES = (str, list, np.ndarray)
+    SINGLE_BRAIN_TEXT_TYPES = list
 
     def __init__(
         self,
@@ -309,7 +309,10 @@ def __str__(self):
         )
 
     def reset(
-        self, config=None, train_mode=True, custom_reset_parameters=None
+        self,
+        config: Dict = None,
+        train_mode: bool = True,
+        custom_reset_parameters: Any = None,
     ) -> AllBrainInfo:
         """
         Sends a signal to reset the unity environment.
@@ -355,11 +358,11 @@ def reset(
     @timed
     def step(
         self,
-        vector_action=None,
-        memory=None,
-        text_action=None,
-        value=None,
-        custom_action=None,
+        vector_action: Dict[str, np.ndarray] = None,
+        memory: Optional[Dict[str, np.ndarray]] = None,
+        text_action: Optional[Dict[str, List[str]]] = None,
+        value: Optional[Dict[str, np.ndarray]] = None,
+        custom_action: Dict[str, Any] = None,
     ) -> AllBrainInfo:
         """
         Provides the environment with an action, moves the environment dynamics forward accordingly,
@@ -495,8 +498,6 @@ def step(
                 else:
                     if text_action[brain_name] is None:
                         text_action[brain_name] = [""] * n_agent
-                    if isinstance(text_action[brain_name], str):
-                        text_action[brain_name] = [text_action[brain_name]] * n_agent
                 if brain_name not in custom_action:
                     custom_action[brain_name] = [None] * n_agent
                 else:
@@ -585,7 +586,7 @@ def _close(self):
             self.proc1.kill()
 
     @classmethod
-    def _flatten(cls, arr) -> List[float]:
+    def _flatten(cls, arr: Any) -> List[float]:
         """
         Converts arrays to list.
         :param arr: numpy vector.
@@ -619,7 +620,12 @@ def _get_state(self, output: UnityRLOutput) -> AllBrainInfo:
 
     @timed
     def _generate_step_input(
-        self, vector_action, memory, text_action, value, custom_action
+        self,
+        vector_action: Dict[str, np.ndarray],
+        memory: Dict[str, np.ndarray],
+        text_action: Dict[str, list],
+        value: Dict[str, np.ndarray],
+        custom_action: Dict[str, list],
     ) -> UnityInput:
         rl_in = UnityRLInput()
         for b in vector_action:
@@ -643,7 +649,7 @@ def _generate_step_input(
         return self.wrap_unity_input(rl_in)
 
     def _generate_reset_input(
-        self, training, config, custom_reset_parameters
+        self, training: bool, config: Dict, custom_reset_parameters: Any
     ) -> UnityInput:
         rl_in = UnityRLInput()
         rl_in.is_training = training

ml-agents-envs/mlagents/envs/sampler_class.py

@@ -21,8 +21,7 @@ def __init__(
         min_value: Union[int, float],
         max_value: Union[int, float],
         seed: Optional[int] = None,
-        **kwargs
-    ) -> None:
+    ):
         """
         :param min_value: minimum value of the range to be sampled uniformly from
         :param max_value: maximum value of the range to be sampled uniformly from
@@ -49,11 +48,8 @@ class MultiRangeUniformSampler(Sampler):
     """
 
     def __init__(
-        self,
-        intervals: List[List[Union[int, float]]],
-        seed: Optional[int] = None,
-        **kwargs
-    ) -> None:
+        self, intervals: List[List[Union[int, float]]], seed: Optional[int] = None
+    ):
         """
         :param intervals: List of intervals to draw uniform samples from
         :param seed: Random seed used for making uniform draws from the specified intervals
@@ -88,8 +84,7 @@ def __init__(
         mean: Union[float, int],
         st_dev: Union[float, int],
         seed: Optional[int] = None,
-        **kwargs
-    ) -> None:
+    ):
         """
         :param mean: Specifies the mean of the gaussian distribution to draw from
         :param st_dev: Specifies the standard devation of the gaussian distribution to draw from

ml-agents-envs/mlagents/envs/subprocess_env_manager.py

@@ -44,7 +44,7 @@ def __init__(self, process: Process, worker_id: int, conn: Connection):
         self.previous_all_action_info: Dict[str, ActionInfo] = {}
         self.waiting = False
 
-    def send(self, name: str, payload=None):
+    def send(self, name: str, payload: Any = None) -> None:
         try:
             cmd = EnvironmentCommand(name, payload)
             self.conn.send(cmd)
@@ -68,7 +68,7 @@ def close(self):
 
 def worker(
     parent_conn: Connection, step_queue: Queue, pickled_env_factory: str, worker_id: int
-):
+) -> None:
     env_factory: Callable[[int], UnityEnvironment] = cloudpickle.loads(
         pickled_env_factory
     )
@@ -183,7 +183,10 @@ def step(self) -> List[EnvironmentStep]:
         return step_infos
 
     def reset(
-        self, config=None, train_mode=True, custom_reset_parameters=None
+        self,
+        config: Optional[Dict] = None,
+        train_mode: bool = True,
+        custom_reset_parameters: Any = None,
     ) -> List[EnvironmentStep]:
         while any([ew.waiting for ew in self.env_workers]):
             if not self.step_queue.empty():

ml-agents-envs/mlagents/envs/tests/test_subprocess_env_manager.py

@@ -11,7 +11,7 @@
 from mlagents.envs.base_unity_environment import BaseUnityEnvironment
 
 
-def mock_env_factory(worker_id: int):
+def mock_env_factory(worker_id):
     return mock.create_autospec(spec=BaseUnityEnvironment)