Improve mypy coverage by adding --namespace-packages (#3049)

.pre-commit-config.yaml

             )$
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.740
+    rev: v0.750
     hooks:
     -   id: mypy
         name: mypy-ml-agents
         files: "ml-agents-envs/.*"
         # Exclude protobuf files and don't follow them when imported
         exclude: ".*_pb2.py"
-        args: [--ignore-missing-imports, --disallow-incomplete-defs]
+        args: [--ignore-missing-imports, --disallow-incomplete-defs,  --namespace-packages]
-        args: [--ignore-missing-imports, --disallow-incomplete-defs]
+        args: [--ignore-missing-imports, --disallow-incomplete-defs,  --namespace-packages]
 
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v2.4.0

.pylintrc

     # E0401: Unable to import...
     # E0611: No name '...' in module '...'
     # need to look into these, probably namespace packages
-    E0401, E0611
+    E0401, E0611,
+
+    # This was causing false positives
+    # Appears to be https://github.com/PyCQA/pylint/issues/2981
+    W0201,
+

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

 """
 
 from abc import ABC, abstractmethod
-from typing import List, NamedTuple, Tuple, Optional, Union, Dict, NewType
+from typing import List, NamedTuple, Tuple, Optional, Union, Dict
-AgentId = NewType("AgentId", int)
-AgentGroup = NewType("AgentGroup", str)
+AgentId = int
+AgentGroup = str
 
 
 class StepResult(NamedTuple):

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

     AgentId,
 )
 from mlagents.envs.timers import timed, hierarchical_timer
-from .exception import (
+from mlagents.envs.exception import (
     UnityEnvironmentException,
     UnityCommunicationException,
     UnityActionException,
+from mlagents.envs.communicator_objects.command_pb2 import STEP, RESET
 from mlagents.envs.rpc_utils import (
     agent_group_spec_from_proto,
     batched_step_result_from_proto,
             action = action.astype(expected_type)
 
         if agent_group not in self._env_actions:
-            self._env_actions[agent_group] = self._empty_action(
-                spec, self._env_state[agent_group].n_agents()
+            self._env_actions[agent_group] = spec.create_empty_action(
+                self._env_state[agent_group].n_agents()
             )
         try:
             index = np.where(self._env_state[agent_group].agent_id == agent_id)[0][0]
 
     @staticmethod
     def _parse_side_channel_message(
-        side_channels: Dict[int, SideChannel], data: bytearray
+        side_channels: Dict[int, SideChannel], data: bytes
     ) -> None:
         offset = 0
         while offset < len(data):
             for i in range(n_agents):
                 action = AgentActionProto(vector_actions=vector_action[b][i])
                 rl_in.agent_actions[b].value.extend([action])
-                rl_in.command = 0
+                rl_in.command = STEP
-        rl_in.command = 1
+        rl_in.command = RESET
         rl_in.side_channel = bytes(self._generate_side_channel_data(self.side_channels))
         return self.wrap_unity_input(rl_in)
 

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

     NONE as COMPRESSION_TYPE_NONE,
     PNG as COMPRESSION_TYPE_PNG,
 )
+from mlagents.envs.communicator_objects.space_type_pb2 import discrete, continuous
 
 
 class MockCommunicator(Communicator):
         bp = BrainParametersProto(
             vector_action_size=[2],
             vector_action_descriptions=["", ""],
-            vector_action_space_type=int(not self.is_discrete),
+            vector_action_space_type=discrete if self.is_discrete else continuous,
             brain_name=self.brain_name,
             is_training=True,
         )

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

 import logging
 import numpy as np
 import io
-from typing import List, Tuple
+from typing import cast, List, Tuple, Union, Collection
 from PIL import Image
 
 logger = logging.getLogger("mlagents.envs")
         if brain_param_proto.vector_action_space_type == 0
         else ActionType.CONTINUOUS
     )
-    action_shape = None
-        action_shape = brain_param_proto.vector_action_size[0]
+        action_shape: Union[
+            int, Tuple[int, ...]
+        ] = brain_param_proto.vector_action_size[0]
     else:
         action_shape = tuple(brain_param_proto.vector_action_size)
     return AgentGroupSpec(observation_shape, action_type, action_shape)
 
 @timed
 def _process_visual_observation(
-    obs_index: int, shape: Tuple[int, int, int], agent_info_list: List[AgentInfoProto]
+    obs_index: int,
+    shape: Tuple[int, int, int],
+    agent_info_list: Collection[
+        AgentInfoProto
+    ],  # pylint: disable=unsubscriptable-object
 ) -> np.ndarray:
     if len(agent_info_list) == 0:
         return np.zeros((0, shape[0], shape[1], shape[2]), dtype=np.float32)
 
 @timed
 def _process_vector_observation(
-    obs_index: int, shape: Tuple[int, ...], agent_info_list: List[AgentInfoProto]
+    obs_index: int,
+    shape: Tuple[int, ...],
+    agent_info_list: Collection[
+        AgentInfoProto
+    ],  # pylint: disable=unsubscriptable-object
 ) -> np.ndarray:
     if len(agent_info_list) == 0:
         return np.zeros((0, shape[0]), dtype=np.float32)
 
 @timed
 def batched_step_result_from_proto(
-    agent_info_list: List[AgentInfoProto], group_spec: AgentGroupSpec
+    agent_info_list: Collection[
+        AgentInfoProto
+    ],  # pylint: disable=unsubscriptable-object
+    group_spec: AgentGroupSpec,
+            obs_shape = cast(Tuple[int, int, int], obs_shape)
             obs_list += [
                 _process_visual_observation(obs_index, obs_shape, agent_info_list)
             ]

ml-agents-envs/mlagents/envs/side_channel/engine_configuration_channel.py

     def channel_type(self) -> int:
         return SideChannelType.EngineSettings
 
-    def on_message_received(self, data: bytearray) -> None:
+    def on_message_received(self, data: bytes) -> None:
         """
         Is called by the environment to the side channel. Can be called
         multiple times per step if multiple messages are meant for that

ml-agents-envs/mlagents/envs/side_channel/float_properties_channel.py

 from mlagents.envs.side_channel.side_channel import SideChannel, SideChannelType
 import struct
-from typing import Tuple, Optional, List
+from typing import Dict, Tuple, Optional, List
 
 
 class FloatPropertiesChannel(SideChannel):
     set_property, get_property and list_properties.
     """
 
-    def __init__(self):
-        self._float_properties = {}
+    def __init__(self) -> None:
+        self._float_properties: Dict[str, float] = {}
         super().__init__()
 
     @property
-    def on_message_received(self, data: bytearray) -> None:
+    def on_message_received(self, data: bytes) -> None:
         """
         Is called by the environment to the side channel. Can be called
         multiple times per step if multiple messages are meant for that
         Returns a list of all the string identifiers of the properties
         currently present in the Unity Environment.
         """
-        return self._float_properties.keys()
+        return list(self._float_properties.keys())
+
+    def get_property_dict_copy(self) -> Dict[str, float]:
+        """
+        Returns a copy of the float properties.
+        :return:
+        """
+        return dict(self._float_properties)
 
     @staticmethod
     def serialize_float_prop(key: str, value: float) -> bytearray:
         return result
 
     @staticmethod
-    def deserialize_float_prop(data: bytearray) -> Tuple[str, float]:
+    def deserialize_float_prop(data: bytes) -> Tuple[str, float]:
         offset = 0
         encoded_key_len = struct.unpack_from("<i", data, offset)[0]
         offset = offset + 4

ml-agents-envs/mlagents/envs/side_channel/raw_bytes_channel.py

     def channel_type(self) -> int:
         return SideChannelType.RawBytesChannelStart + self._channel_id
 
-    def on_message_received(self, data: bytearray) -> None:
+    def on_message_received(self, data: bytes) -> None:
         """
         Is called by the environment to the side channel. Can be called
         multiple times per step if multiple messages are meant for that

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

         self.message_queue.append(data)
 
     @abstractmethod
-    def on_message_received(self, data: bytearray) -> None:
+    def on_message_received(self, data: bytes) -> None:
         """
         Is called by the environment to the side channel. Can be called
         multiple times per step if multiple messages are meant for that

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

 from typing import List, Tuple
 from mlagents.envs.communicator_objects.agent_info_pb2 import AgentInfoProto
-from mlagents.envs.communicator_objects.observation_pb2 import ObservationProto
+from mlagents.envs.communicator_objects.observation_pb2 import (
+    ObservationProto,
+    NONE,
+    PNG,
+)
 from mlagents.envs.communicator_objects.brain_parameters_pb2 import BrainParametersProto
 import numpy as np
 from mlagents.envs.base_env import AgentGroupSpec, ActionType
         for obs_index in range(len(shape)):
             obs_proto = ObservationProto()
             obs_proto.shape.extend(list(shape[obs_index]))
-            obs_proto.compression_type = 0
+            obs_proto.compression_type = NONE
             obs_proto.float_data.data.extend([0.1] * np.prod(shape[obs_index]))
             obs_proto_list.append(obs_proto)
         ap.observations.extend(obs_proto_list)
 def generate_compressed_proto_obs(in_array: np.ndarray) -> ObservationProto:
     obs_proto = ObservationProto()
     obs_proto.compressed_data = generate_compressed_data(in_array)
-    obs_proto.compression_type = 1
+    obs_proto.compression_type = PNG
     obs_proto.shape.extend(in_array.shape)
     return obs_proto
 

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

     val = sender.get_property("prop1")
     assert val == 1.0
 
+    assert receiver.get_property_dict_copy() == {"prop1": 1.0, "prop2": 2.0}
+    assert receiver.get_property_dict_copy() == sender.get_property_dict_copy()
+
 
 def test_raw_bytes():
     sender = RawBytesChannel()

ml-agents/mlagents/trainers/agent_processor.py

-from typing import List
+from typing import List, Union
 
 from mlagents.trainers.buffer import AgentBuffer, BufferException
 
     def append_to_update_buffer(
         self,
         update_buffer: AgentBuffer,
-        agent_id: str,
+        agent_id: Union[int, str],
         key_list: List[str] = None,
         batch_size: int = None,
         training_length: int = None,

ml-agents/mlagents/trainers/brain.py

 from mlagents.envs.communicator_objects.brain_parameters_pb2 import BrainParametersProto
 from mlagents.envs.communicator_objects.observation_pb2 import ObservationProto
 from mlagents.envs.timers import hierarchical_timer, timed
-from typing import Dict, List, NamedTuple
+from typing import Dict, List, NamedTuple, Collection
 from PIL import Image
 
 logger = logging.getLogger("mlagents.envs")
     @timed
     def from_agent_proto(
         worker_id: int,
-        agent_info_list: List[AgentInfoProto],
+        agent_info_list: Collection[
+            AgentInfoProto
+        ],  # pylint: disable=unsubscriptable-object
         brain_params: BrainParameters,
     ) -> "BrainInfo":
         """
 
     @staticmethod
     def _process_visual_observations(
-        brain_params: BrainParameters, agent_info_list: List[AgentInfoProto]
+        brain_params: BrainParameters,
+        agent_info_list: Collection[
+            AgentInfoProto
+        ],  # pylint: disable=unsubscriptable-object
     ) -> List[np.ndarray]:
 
         visual_observation_protos: List[List[ObservationProto]] = []
 
     @staticmethod
     def _process_vector_observations(
-        brain_params: BrainParameters, agent_info_list: List[AgentInfoProto]
+        brain_params: BrainParameters,
+        agent_info_list: Collection[
+            AgentInfoProto
+        ],  # pylint: disable=unsubscriptable-object
     ) -> np.ndarray:
         if len(agent_info_list) == 0:
             vector_obs = np.zeros(

ml-agents/mlagents/trainers/models.py

         self.running_variance: Optional[tf.Variable] = None
         self.update_normalization: Optional[tf.Operation] = None
         self.value: Optional[tf.Tensor] = None
+        self.all_log_probs: Optional[tf.Tensor] = None
+        self.output: Optional[tf.Tensor] = None
+        self.selected_actions: Optional[tf.Tensor] = None
+        self.action_holder: Optional[tf.Tensor] = None
 
     @staticmethod
     def create_global_steps():

ml-agents/mlagents/trainers/rl_trainer.py

         raise UnityTrainerException(
             "The add_rewards_outputs method was not implemented."
         )
+
+    def advance(self):
+        """
+        Eventually logic from TrainerController.advance() will live here.
+        """
+        self.clear_update_buffer()

ml-agents/mlagents/trainers/sac/models.py

         self.q2_memory_in: Optional[tf.Tensor] = None
         self.q1_memory_out: Optional[tf.Tensor] = None
         self.q2_memory_out: Optional[tf.Tensor] = None
-        self.action_holder: Optional[tf.Tensor] = None
         self.prev_action: Optional[tf.Tensor] = None
         self.action_masks: Optional[tf.Tensor] = None
         self.external_action_in: Optional[tf.Tensor] = None
-        self.all_log_probs: Optional[tf.Tensor] = None
-        self.selected_actions: Optional[tf.Tensor] = None
-        self.output: Optional[tf.Tensor] = None
         self.output_oh: Optional[tf.Tensor] = None
         self.output_pre: Optional[tf.Tensor] = None
 

ml-agents/mlagents/trainers/simple_env_manager.py

 
     @property
     def get_properties(self) -> Dict[str, float]:
-        reset_params = {}
-        for k in self.shared_float_properties.list_properties():
-            reset_params[k] = self.shared_float_properties.get_property(k)
-        return reset_params
+        return self.shared_float_properties.get_property_dict_copy()
 
     def close(self):
         self.env.close()

ml-agents/mlagents/trainers/subprocess_env_manager.py

             elif cmd.name == "external_brains":
                 _send_response("external_brains", external_brains())
             elif cmd.name == "get_properties":
-                reset_params = {}
-                for k in shared_float_properties.list_properties():
-                    reset_params[k] = shared_float_properties.get_property(k)
-
+                reset_params = shared_float_properties.get_property_dict_copy()
                 _send_response("get_properties", reset_params)
             elif cmd.name == "reset":
                 for k, v in cmd.payload.items():

ml-agents/mlagents/trainers/tests/test_trainer_controller.py

         new_step_info.previous_all_brain_info[brain_name],
         new_step_info.current_all_brain_info[brain_name],
     )
-    trainer_mock.clear_update_buffer.assert_called_once()
+    trainer_mock.advance.assert_called_once()

ml-agents/mlagents/trainers/trainer.py

         )
         self.summary_writer = tf.summary.FileWriter(self.summary_path)
         self._reward_buffer: Deque[float] = deque(maxlen=reward_buff_cap)
-        self.policy: TFPolicy = None
+        self.policy: TFPolicy = None  # type: ignore  # this will always get set
         self.step: int = 0
 
     def check_param_keys(self):
         Uses demonstration_buffer to update model.
         """
         raise UnityTrainerException("The update_model method was not implemented.")
+
+    def advance(self) -> None:
+        pass

ml-agents/mlagents/trainers/trainer_controller.py

                     env.set_policy(brain_name, trainer.policy)
             else:
                 # Avoid memory leak during inference
-                trainer.clear_update_buffer()
+                # Eventually this whole block will take place in advance()
+                # But currently this only calls clear_update_buffer() in RLTrainer
+                # and nothing in the base class
+                trainer.advance()
         return len(new_step_infos)

ml-agents/mlagents/trainers/trainer_util.py

             _brain_key = trainer_config[_brain_key]
         trainer_parameters.update(trainer_config[_brain_key])
 
-    trainer = None
+    trainer: Trainer = None  # type: ignore  # will be set to one of these, or raise
     if trainer_parameters["trainer"] == "offline_bc":
         trainer = OfflineBCTrainer(
             brain_parameters, trainer_parameters, train_model, load_model, seed, run_id