Enable mypy in precommit checks (#2177)

* WIP precommit on top level

* update CI

* circleci fixes

* intentionally fail black

* use --show-diff-on-failure in CI

* fix command order

* rebreak a file

* apply black

* WIP enable mypy

* run mypy on each package

* fix trainer_metrics mypy errors

* more mypy errors

* more mypy

* Fix some partially typed functions

* types for take_action_outputs

* fix formatting

* cleanup

* generate stubs for proto objects

* fix ml-agents-env mypy errors

* disallow-incomplete-defs for gym-unity

* Add CI notes to CONTRIBUTING.md

.gitignore

 
 # Python virtual environment
 venv/
+.mypy_cache/

.pre-commit-config.yaml

-exclude: '^$'
-fail_fast: false
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.701'
+    hooks:
+    -   id: mypy
+        name: mypy-ml-agents
+        files: "ml-agents/.*"
+        args: [--ignore-missing-imports, --disallow-incomplete-defs]
+    -   id: mypy
+        name: mypy-ml-agents-envs
+        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]
+    -   id: mypy
+        name: mypy-gym-unity
+        files: "gym-unity/.*"
+        args: [--ignore-missing-imports, --disallow-incomplete-defs]

CONTRIBUTING.md

 machine learning algorithms. Feel free to submit these environments with a
 PR explaining the nature of the environment and task.
 
-## Style Guide
+## Continuous Integration (CI)
-When performing changes to the codebase, please ensure that all python code is reformatted using the [black](https://github.com/ambv/black) formatter. For C#, we will soon be requirements for style and formatting.
+We run CircleCI on all PRs; all tests must be passing before the PR is merged.
+
+Several static checks are run on the codebase using the [pre-commit framework](https://pre-commit.com/) during CI. To execute the same checks locally, install `pre-commit` and run `pre-commit run --all-files`. Some hooks (for example, `black`) will output the corrected version of the code; others (like `mypy`) may require more effort to fix.
+
+### Code style
+All python code should be formatted with [`black`](https://github.com/ambv/black). Style and formatting for C# may be enforced later.
+
+### Python type annotations
+We use [`mypy`](http://mypy-lang.org/) to perform static type checking on python code. Currently not all code is annotated but we will increase coverage over time. If you are adding or refactoring code, please
+1. Add type annotations to the new or refactored code.
+2. Make sure that code calling or called by the modified code also has type annotations.
+
+The [type hint cheat sheet](https://mypy.readthedocs.io/en/stable/cheat_sheet_py3.html) provides a good introduction to adding type hints.

gym-unity/gym_unity/envs/unity_env.py

     def __init__(
         self,
         environment_filename: str,
-        worker_id=0,
-        use_visual=False,
-        uint8_visual=False,
-        multiagent=False,
-        flatten_branched=False,
-        no_graphics=False,
-        allow_multiple_visual_obs=False,
+        worker_id: int = 0,
+        use_visual: bool = False,
+        uint8_visual: bool = False,
+        multiagent: bool = False,
+        flatten_branched: bool = False,
+        no_graphics: bool = False,
+        allow_multiple_visual_obs: bool = False,
     ):
         """
         Environment initialization

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

         return brain_info
 
 
-def safe_concat_lists(l1: Optional[List], l2: Optional[List]):
-    if l1 is None and l2 is None:
-        return None
-    if l1 is None and l2 is not None:
-        return l2.copy()
-    if l1 is not None and l2 is None:
-        return l1.copy()
+def safe_concat_lists(l1: Optional[List], l2: Optional[List]) -> Optional[List]:
+    if l1 is None:
+        if l2 is None:
+            return None
+        else:
+            return l2.copy()
-        copy = l1.copy()
-        copy.extend(l2)
-        return copy
+        if l2 is None:
+            return l1.copy()
+        else:
+            copy = l1.copy()
+            copy.extend(l2)
+            return copy
 
 
 def safe_concat_np_ndarray(a1: Optional[np.ndarray], a2: Optional[np.ndarray]):

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

 import logging
+from typing import Optional
 
 from .communicator_objects import UnityOutput, UnityInput
 
         :return: UnityOutput: The initialization output sent by Unity
         """
 
-    def exchange(self, inputs: UnityInput) -> UnityOutput:
+    def exchange(self, inputs: UnityInput) -> Optional[UnityOutput]:
         """
         Used to send an input and receive an output from the Environment
         :param inputs: The UnityInput that needs to be sent the Environment

ml-agents-envs/mlagents/envs/communicator_objects/unity_to_external_pb2.py

+# -*- coding: utf-8 -*-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
 # source: mlagents/envs/communicator_objects/unity_to_external.proto
 

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

                 "{1}.\nPlease go to https://github.com/Unity-Technologies/ml-agents to download the latest version "
                 "of ML-Agents.".format(self._version_, self._unity_version)
             )
-        self._n_agents = {}
-        self._global_done = None
+        self._n_agents: Dict[str, int] = {}
+        self._global_done: Optional[bool] = None
-        self._brains = {}
-        self._brain_names = []
-        self._external_brain_names = []
+        self._brains: Dict[str, BrainParameters] = {}
+        self._brain_names: List[str] = []
+        self._external_brain_names: List[str] = []
         for brain_param in aca_params.brain_parameters:
             self._brain_names += [brain_param.brain_name]
             self._brains[brain_param.brain_name] = BrainParameters.from_proto(
         custom_action = {} if custom_action is None else custom_action
 
         # Check that environment is loaded, and episode is currently running.
-        if self._loaded and not self._global_done and self._global_done is not None:
+        if not self._loaded:
+            raise UnityEnvironmentException("No Unity environment is loaded.")
+        elif self._global_done:
+            raise UnityActionException(
+                "The episode is completed. Reset the environment with 'reset()'"
+            )
+        elif self.global_done is None:
+            raise UnityActionException(
+                "You cannot conduct step without first calling reset. "
+                "Reset the environment with 'reset()'"
+            )
+        else:
             if isinstance(vector_action, self.SINGLE_BRAIN_ACTION_TYPES):
                 if self._num_external_brains == 1:
                     vector_action = {self._external_brain_names[0]: vector_action}
             for _b in self._external_brain_names:
                 self._n_agents[_b] = len(state[0][_b].agents)
             return state[0]
-        elif not self._loaded:
-            raise UnityEnvironmentException("No Unity environment is loaded.")
-        elif self._global_done:
-            raise UnityActionException(
-                "The episode is completed. Reset the environment with 'reset()'"
-            )
-        elif self.global_done is None:
-            raise UnityActionException(
-                "You cannot conduct step without first calling reset. "
-                "Reset the environment with 'reset()'"
-            )
 
     def close(self):
         """
         arr = [float(x) for x in arr]
         return arr
 
-    def _get_state(self, output: UnityRLOutput) -> (AllBrainInfo, bool):
+    def _get_state(self, output: UnityRLOutput) -> Tuple[AllBrainInfo, bool]:
         """
         Collects experience information from all external brains in environment at current step.
         :return: a dictionary of BrainInfo objects.
 
     def _generate_step_input(
         self, vector_action, memory, text_action, value, custom_action
-    ) -> UnityRLInput:
+    ) -> UnityInput:
         rl_in = UnityRLInput()
         for b in vector_action:
             n_agents = self._n_agents[b]
 
     def _generate_reset_input(
         self, training, config, custom_reset_parameters
-    ) -> UnityRLInput:
+    ) -> UnityInput:
         rl_in = UnityRLInput()
         rl_in.is_training = training
         rl_in.environment_parameters.CopyFrom(EnvironmentParametersProto())
         return self.communicator.initialize(inputs).rl_initialization_output
 
     @staticmethod
-    def wrap_unity_input(rl_input: UnityRLInput) -> UnityOutput:
+    def wrap_unity_input(rl_input: UnityRLInput) -> UnityInput:
         result = UnityInput()
         result.rl_input.CopyFrom(rl_input)
         return result

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

 import logging
 import grpc
+from typing import Optional
 
 import socket
 from multiprocessing import Pipe
         self.unity_to_external.parent_conn.recv()
         return aca_param
 
-    def exchange(self, inputs: UnityInput) -> UnityOutput:
+    def exchange(self, inputs: UnityInput) -> Optional[UnityOutput]:
         message = UnityMessage()
         message.header.status = 200
         message.unity_input.CopyFrom(inputs)

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

 import logging
 import socket
 import struct
+from typing import Optional
 
 from .communicator import Communicator
 from .communicator_objects import UnityMessage, UnityOutput, UnityInput
     def _communicator_send(self, message):
         self._conn.send(struct.pack("I", len(message)) + message)
 
-    def exchange(self, inputs: UnityInput) -> UnityOutput:
+    def exchange(self, inputs: UnityInput) -> Optional[UnityOutput]:
         message = UnityMessage()
         message.header.status = 200
         message.unity_input.CopyFrom(inputs)

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

     def __init__(
         self, env_factory: Callable[[int], BaseUnityEnvironment], n_env: int = 1
     ):
-        self.envs = []
-        self.env_agent_counts = {}
+        self.envs: List[UnityEnvWorker] = []
+        self.env_agent_counts: Dict[str, List[int]] = {}
         self.waiting = False
         for worker_id in range(n_env):
             self.envs.append(self.create_worker(worker_id, env_factory))

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

 
 from mlagents.envs.subprocess_environment import *
 from mlagents.envs import UnityEnvironmentException, BrainInfo
+from mlagents.envs.base_unity_environment import BaseUnityEnvironment
 
 
 def mock_env_factory(worker_id: int):

ml-agents/mlagents/trainers/action_info.py

 from typing import NamedTuple, Any, Dict, Optional
 
+ActionInfoOutputs = Optional[Dict[str, Any]]
+
 
 class ActionInfo(NamedTuple):
     action: Any
-    outputs: Optional[Dict[str, Any]]
+    outputs: ActionInfoOutputs

ml-agents/mlagents/trainers/bc/online_trainer.py

 import numpy as np
 
 from mlagents.envs import AllBrainInfo
+from mlagents.trainers import ActionInfoOutputs
 from mlagents.trainers.bc.trainer import BCTrainer
 
 logger = logging.getLogger("mlagents.trainers")
         )
 
     def add_experiences(
-        self, curr_info: AllBrainInfo, next_info: AllBrainInfo, take_action_outputs
-    ):
+        self,
+        curr_info: AllBrainInfo,
+        next_info: AllBrainInfo,
+        take_action_outputs: ActionInfoOutputs,
+    ) -> None:
         """
         Adds experiences to each agent's experience history.
         :param curr_info: Current AllBrainInfo (Dictionary of all current brains and corresponding BrainInfo).
             curr_info, next_info, take_action_outputs
         )
 
-    def process_experiences(self, current_info: AllBrainInfo, next_info: AllBrainInfo):
+    def process_experiences(
+        self, current_info: AllBrainInfo, next_info: AllBrainInfo
+    ) -> None:
         """
         Checks agent histories for processing condition, and processes them as necessary.
         Processing involves calculating value and advantage targets for model updating step.

ml-agents/mlagents/trainers/bc/trainer.py

 import tensorflow as tf
 
 from mlagents.envs import AllBrainInfo
+from mlagents.trainers import ActionInfoOutputs
 from mlagents.trainers.bc.policy import BCPolicy
 from mlagents.trainers.buffer import Buffer
 from mlagents.trainers.trainer import Trainer
         return
 
     def add_experiences(
-        self, curr_info: AllBrainInfo, next_info: AllBrainInfo, take_action_outputs
-    ):
+        self,
+        curr_info: AllBrainInfo,
+        next_info: AllBrainInfo,
+        take_action_outputs: ActionInfoOutputs,
+    ) -> None:
         """
         Adds experiences to each agent's experience history.
         :param curr_info: Current AllBrainInfo (Dictionary of all current brains and corresponding BrainInfo).
                         self.episode_steps[agent_id] = 0
                     self.episode_steps[agent_id] += 1
 
-    def process_experiences(self, current_info: AllBrainInfo, next_info: AllBrainInfo):
+    def process_experiences(
+        self, current_info: AllBrainInfo, next_info: AllBrainInfo
+    ) -> None:
         """
         Checks agent histories for processing condition, and processes them as necessary.
         Processing involves calculating value and advantage targets for model updating step.

ml-agents/mlagents/trainers/demo_loader.py

 from mlagents.trainers.buffer import Buffer
 from mlagents.envs.brain import BrainParameters, BrainInfo
 from mlagents.envs.communicator_objects import *
-from google.protobuf.internal.decoder import _DecodeVarint32
+from google.protobuf.internal.decoder import _DecodeVarint32  # type: ignore
+
 
 logger = logging.getLogger("mlagents.trainers")
 

ml-agents/mlagents/trainers/learn.py

 import numpy as np
 import yaml
 from docopt import docopt
-from typing import Optional, Callable
+from typing import Any, Callable, Dict, Optional
 
 
 from mlagents.trainers.trainer_controller import TrainerController
 from mlagents.envs.subprocess_environment import SubprocessUnityEnvironment
 
 
-def run_training(sub_id: int, run_seed: int, run_options, process_queue):
+def run_training(
+    sub_id: int, run_seed: int, run_options: Dict[str, Any], process_queue: Queue
+) -> None:
     """
     Launches training session.
     :param process_queue: Queue used to send signal back to main.
     return env_path
 
 
-def load_config(trainer_config_path):
+def load_config(trainer_config_path: str) -> Dict[str, Any]:
     try:
         with open(trainer_config_path) as data_file:
             trainer_config = yaml.safe_load(data_file)

ml-agents/mlagents/trainers/policy.py

 import logging
+from typing import Any, Dict
+
 import numpy as np
 import tensorflow as tf
 
 from mlagents.envs import BrainInfo
+
 
 logger = logging.getLogger("mlagents.trainers")
 
                 )
             self.saver.restore(self.sess, ckpt.model_checkpoint_path)
 
-    def evaluate(self, brain_info: BrainInfo):
+    def evaluate(self, brain_info: BrainInfo) -> Dict[str, Any]:
         """
         Evaluates policy for the agent experiences provided.
         :param brain_info: BrainInfo input to network.

ml-agents/mlagents/trainers/ppo/trainer.py

 
 import logging
 from collections import deque
+from typing import Any, List
 
 import numpy as np
 import tensorflow as tf
 from mlagents.trainers.ppo.policy import PPOPolicy
 from mlagents.trainers.trainer import Trainer
-
+from mlagents.trainers.action_info import ActionInfoOutputs
 
 logger = logging.getLogger("mlagents.trainers")
 
         :BrainInfo next_info: A t+1 BrainInfo.
         :return: curr_info: Reconstructed BrainInfo to match agents of next_info.
         """
-        visual_observations = [[]]
+        visual_observations: List[List[Any]] = [
+            []
+        ]  # TODO add types to brain.py methods
         vector_observations = []
         text_observations = []
         memories = []
         self,
         curr_all_info: AllBrainInfo,
         next_all_info: AllBrainInfo,
-        take_action_outputs,
-    ):
+        take_action_outputs: ActionInfoOutputs,
+    ) -> None:
         """
         Adds experiences to each agent's experience history.
         :param curr_all_info: Dictionary of all current brains and corresponding BrainInfo.
                     self.episode_steps[agent_id] += 1
         self.trainer_metrics.end_experience_collection_timer()
 
-    def process_experiences(self, current_info: AllBrainInfo, new_info: AllBrainInfo):
+    def process_experiences(
+        self, current_info: AllBrainInfo, new_info: AllBrainInfo
+    ) -> None:
         """
         Checks agent histories for processing condition, and processes them as necessary.
         Processing involves calculating value and advantage targets for model updating step.

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

     tc = basic_trainer_controller(brain_info_mock)
     bad_config = dummy_bad_config()
 
-    try:
+    with pytest.raises(UnityEnvironmentException):
-        assert (
-            1 == 0,
-            "Initialize trainers with bad config did not raise an exception.",
-        )
-    except UnityEnvironmentException:
-        pass
 
 
 def trainer_controller_with_start_learning_mocks():

ml-agents/mlagents/trainers/trainer.py

 import numpy as np
 
 from mlagents.envs import UnityException, AllBrainInfo, BrainInfo
-from mlagents.trainers import ActionInfo
+from mlagents.trainers import ActionInfo, ActionInfoOutputs
 from mlagents.trainers import TrainerMetrics
 
 LOGGER = logging.getLogger("mlagents.trainers")
         return action
 
     def add_experiences(
-        self, curr_info: AllBrainInfo, next_info: AllBrainInfo, take_action_outputs
-    ):
+        self,
+        curr_info: AllBrainInfo,
+        next_info: AllBrainInfo,
+        take_action_outputs: ActionInfoOutputs,
+    ) -> None:
         """
         Adds experiences to each agent's experience history.
         :param curr_info: Current AllBrainInfo.
         raise UnityTrainerException("The add_experiences method was not implemented.")
 
-    def process_experiences(self, current_info: AllBrainInfo, next_info: AllBrainInfo):
+    def process_experiences(
+        self, current_info: AllBrainInfo, next_info: AllBrainInfo
+    ) -> None:
         """
         Checks agent histories for processing condition, and processes them as necessary.
         Processing involves calculating value and advantage targets for model updating step.

ml-agents/mlagents/trainers/trainer_controller.py

 from mlagents.envs import AllBrainInfo, BrainParameters
 from mlagents.envs.base_unity_environment import BaseUnityEnvironment
 from mlagents.envs.exception import UnityEnvironmentException
-from mlagents.trainers import Trainer
+from mlagents.trainers import Trainer, TrainerMetrics
 from mlagents.trainers.ppo.trainer import PPOTrainer
 from mlagents.trainers.bc.offline_trainer import OfflineBCTrainer
 from mlagents.trainers.bc.online_trainer import OnlineBCTrainer
         for brain_name in self.trainers.keys():
             self.trainers[brain_name].export_model()
 
-    def initialize_trainers(self, trainer_config: Dict[str, Dict[str, str]]):
+    def initialize_trainers(self, trainer_config: Dict[str, Any]) -> None:
         """
         Initialization of the trainers
         :param trainer_config: The configurations of the trainers
             )
             trainer_parameters["keep_checkpoints"] = self.keep_checkpoints
             if brain_name in trainer_config:
-                _brain_key = brain_name
+                _brain_key: Any = brain_name
-                for k in trainer_config[_brain_key]:
-                    trainer_parameters[k] = trainer_config[_brain_key][k]
+                trainer_parameters.update(trainer_config[_brain_key])
             trainer_parameters_dict[brain_name] = trainer_parameters.copy()
         for brain_name in self.external_brains:
             if trainer_parameters_dict[brain_name]["trainer"] == "offline_bc":
                 "permissions are set correctly.".format(model_path)
             )
 
-    def _reset_env(self, env: BaseUnityEnvironment):
+    def _reset_env(self, env: BaseUnityEnvironment) -> AllBrainInfo:
         """Resets the environment.
 
         Returns:
         else:
             return env.reset(train_mode=self.fast_simulation)
 
-    def start_learning(self, env: BaseUnityEnvironment, trainer_config):
+    def start_learning(
+        self, env: BaseUnityEnvironment, trainer_config: Dict[str, Any]
+    ) -> None:
         # TODO: Should be able to start learning at different lesson numbers
         # for each curriculum.
         if self.meta_curriculum is not None:
             self._write_training_metrics()
             self._export_graph()
 
-    def take_step(self, env: BaseUnityEnvironment, curr_info: AllBrainInfo):
+    def take_step(
+        self, env: BaseUnityEnvironment, curr_info: AllBrainInfo
+    ) -> AllBrainInfo:
         if self.meta_curriculum:
             # Get the sizes of the reward buffers.
             reward_buff_sizes = {

ml-agents/mlagents/trainers/trainer_metrics.py

 import logging
 import csv
 from time import time
+from typing import List, Optional
 
 LOGGER = logging.getLogger("mlagents.trainers")
 FIELD_NAMES = [
         """
         self.path = path
         self.brain_name = brain_name
-        self.rows = []
-        self.time_start_experience_collection = None
+        self.rows: List[List[Optional[str]]] = []
+        self.time_start_experience_collection: Optional[float] = None
-        self.last_buffer_length = None
-        self.last_mean_return = None
-        self.time_policy_update_start = None
-        self.delta_last_experience_collection = None
-        self.delta_policy_update = None
+        self.last_buffer_length: Optional[int] = None
+        self.last_mean_return: Optional[float] = None
+        self.time_policy_update_start: Optional[float] = None
+        self.delta_last_experience_collection: Optional[float] = None
+        self.delta_policy_update: Optional[float] = None
-    def start_experience_collection_timer(self):
+    def start_experience_collection_timer(self) -> None:
         """
         Inform Metrics class that experience collection is starting. Intended to be idempotent
         """
-    def end_experience_collection_timer(self):
+    def end_experience_collection_timer(self) -> None:
         """
         Inform Metrics class that experience collection is done.
         """
                 self.delta_last_experience_collection += curr_delta
         self.time_start_experience_collection = None
 
-    def add_delta_step(self, delta: float):
+    def add_delta_step(self, delta: float) -> None:
         """
         Inform Metrics class about time to step in environment.
         """
             self.delta_last_experience_collection = delta
 
-    def start_policy_update_timer(self, number_experiences: int, mean_return: float):
+    def start_policy_update_timer(
+        self, number_experiences: int, mean_return: float
+    ) -> None:
         """
         Inform Metrics class that policy update has started.
         :int number_experiences: Number of experiences in Buffer at this point.
         self.last_mean_return = mean_return
         self.time_policy_update_start = time()
 
-    def _add_row(self, delta_train_start):
-        row = [self.brain_name]
+    def _add_row(self, delta_train_start: float) -> None:
+        row: List[Optional[str]] = [self.brain_name]
         row.extend(
             format(c, ".3f") if isinstance(c, float) else c
             for c in [
         self.delta_last_experience_collection = None
         self.rows.append(row)
 
-    def end_policy_update(self):
+    def end_policy_update(self) -> None:
         """
         Inform Metrics class that policy update has started.
         """
         )
         self._add_row(delta_train_start)
 
-    def write_training_metrics(self):
+    def write_training_metrics(self) -> None:
         """
         Write Training Metrics to CSV
         """

protobuf-definitions/README.md

 
 `pip install grpcio-tools`
 
+`pip install mypy-protobuf`
+
 `nuget install Grpc.Tools` into known directory.
 
 ### Installing Protobuf Compiler

protobuf-definitions/make.bat

 # generate proto objects in python and C#
 
 protoc --proto_path=proto --csharp_out=$DST_DIR_C $SRC_DIR/*.proto
-protoc --proto_path=proto --python_out=$DST_DIR_P $SRC_DIR/*.proto 
+protoc --proto_path=proto --python_out=$DST_DIR_P --mypy_out=$DST_DIR_P $SRC_DIR/*.proto
 
 # grpc 
 

ml-agents-envs/mlagents/envs/communicator_objects/agent_action_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedScalarFieldContainer as google___protobuf___internal___containers___RepeatedScalarFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.custom_action_pb2 import (
+    CustomAction as mlagents___envs___communicator_objects___custom_action_pb2___CustomAction,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class AgentActionProto(google___protobuf___message___Message):
+    vector_actions = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[float]
+    text_actions = ... # type: typing___Text
+    memories = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[float]
+    value = ... # type: float
+
+    @property
+    def custom_action(self) -> mlagents___envs___communicator_objects___custom_action_pb2___CustomAction: ...
+
+    def __init__(self,
+        vector_actions : typing___Optional[typing___Iterable[float]] = None,
+        text_actions : typing___Optional[typing___Text] = None,
+        memories : typing___Optional[typing___Iterable[float]] = None,
+        value : typing___Optional[float] = None,
+        custom_action : typing___Optional[mlagents___envs___communicator_objects___custom_action_pb2___CustomAction] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> AgentActionProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_action"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"custom_action",u"memories",u"text_actions",u"value",u"vector_actions"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_action",b"custom_action"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"custom_action",b"memories",b"text_actions",b"value",b"vector_actions"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/agent_info_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedScalarFieldContainer as google___protobuf___internal___containers___RepeatedScalarFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.custom_observation_pb2 import (
+    CustomObservation as mlagents___envs___communicator_objects___custom_observation_pb2___CustomObservation,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class AgentInfoProto(google___protobuf___message___Message):
+    stacked_vector_observation = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[float]
+    visual_observations = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[bytes]
+    text_observation = ... # type: typing___Text
+    stored_vector_actions = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[float]
+    stored_text_actions = ... # type: typing___Text
+    memories = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[float]
+    reward = ... # type: float
+    done = ... # type: bool
+    max_step_reached = ... # type: bool
+    id = ... # type: int
+    action_mask = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[bool]
+
+    @property
+    def custom_observation(self) -> mlagents___envs___communicator_objects___custom_observation_pb2___CustomObservation: ...
+
+    def __init__(self,
+        stacked_vector_observation : typing___Optional[typing___Iterable[float]] = None,
+        visual_observations : typing___Optional[typing___Iterable[bytes]] = None,
+        text_observation : typing___Optional[typing___Text] = None,
+        stored_vector_actions : typing___Optional[typing___Iterable[float]] = None,
+        stored_text_actions : typing___Optional[typing___Text] = None,
+        memories : typing___Optional[typing___Iterable[float]] = None,
+        reward : typing___Optional[float] = None,
+        done : typing___Optional[bool] = None,
+        max_step_reached : typing___Optional[bool] = None,
+        id : typing___Optional[int] = None,
+        action_mask : typing___Optional[typing___Iterable[bool]] = None,
+        custom_observation : typing___Optional[mlagents___envs___communicator_objects___custom_observation_pb2___CustomObservation] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> AgentInfoProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_observation"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"action_mask",u"custom_observation",u"done",u"id",u"max_step_reached",u"memories",u"reward",u"stacked_vector_observation",u"stored_text_actions",u"stored_vector_actions",u"text_observation",u"visual_observations"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_observation",b"custom_observation"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"action_mask",b"custom_observation",b"done",b"id",b"max_step_reached",b"memories",b"reward",b"stacked_vector_observation",b"stored_text_actions",b"stored_vector_actions",b"text_observation",b"visual_observations"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/brain_parameters_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedCompositeFieldContainer as google___protobuf___internal___containers___RepeatedCompositeFieldContainer,
+    RepeatedScalarFieldContainer as google___protobuf___internal___containers___RepeatedScalarFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.resolution_proto_pb2 import (
+    ResolutionProto as mlagents___envs___communicator_objects___resolution_proto_pb2___ResolutionProto,
+)
+
+from mlagents.envs.communicator_objects.space_type_proto_pb2 import (
+    SpaceTypeProto as mlagents___envs___communicator_objects___space_type_proto_pb2___SpaceTypeProto,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class BrainParametersProto(google___protobuf___message___Message):
+    vector_observation_size = ... # type: int
+    num_stacked_vector_observations = ... # type: int
+    vector_action_size = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[int]
+    vector_action_descriptions = ... # type: google___protobuf___internal___containers___RepeatedScalarFieldContainer[typing___Text]
+    vector_action_space_type = ... # type: mlagents___envs___communicator_objects___space_type_proto_pb2___SpaceTypeProto
+    brain_name = ... # type: typing___Text
+    is_training = ... # type: bool
+
+    @property
+    def camera_resolutions(self) -> google___protobuf___internal___containers___RepeatedCompositeFieldContainer[mlagents___envs___communicator_objects___resolution_proto_pb2___ResolutionProto]: ...
+
+    def __init__(self,
+        vector_observation_size : typing___Optional[int] = None,
+        num_stacked_vector_observations : typing___Optional[int] = None,
+        vector_action_size : typing___Optional[typing___Iterable[int]] = None,
+        camera_resolutions : typing___Optional[typing___Iterable[mlagents___envs___communicator_objects___resolution_proto_pb2___ResolutionProto]] = None,
+        vector_action_descriptions : typing___Optional[typing___Iterable[typing___Text]] = None,
+        vector_action_space_type : typing___Optional[mlagents___envs___communicator_objects___space_type_proto_pb2___SpaceTypeProto] = None,
+        brain_name : typing___Optional[typing___Text] = None,
+        is_training : typing___Optional[bool] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> BrainParametersProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"brain_name",u"camera_resolutions",u"is_training",u"num_stacked_vector_observations",u"vector_action_descriptions",u"vector_action_size",u"vector_action_space_type",u"vector_observation_size"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"brain_name",b"camera_resolutions",b"is_training",b"num_stacked_vector_observations",b"vector_action_descriptions",b"vector_action_size",b"vector_action_space_type",b"vector_observation_size"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/command_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.descriptor import (
+    EnumDescriptor as google___protobuf___descriptor___EnumDescriptor,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    List as typing___List,
+    Tuple as typing___Tuple,
+    cast as typing___cast,
+)
+
+
+class CommandProto(int):
+    DESCRIPTOR: google___protobuf___descriptor___EnumDescriptor = ...
+    @classmethod
+    def Name(cls, number: int) -> str: ...
+    @classmethod
+    def Value(cls, name: str) -> CommandProto: ...
+    @classmethod
+    def keys(cls) -> typing___List[str]: ...
+    @classmethod
+    def values(cls) -> typing___List[CommandProto]: ...
+    @classmethod
+    def items(cls) -> typing___List[typing___Tuple[str, CommandProto]]: ...
+STEP = typing___cast(CommandProto, 0)
+RESET = typing___cast(CommandProto, 1)
+QUIT = typing___cast(CommandProto, 2)

ml-agents-envs/mlagents/envs/communicator_objects/custom_action_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+
+class CustomAction(google___protobuf___message___Message):
+
+    def __init__(self,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> CustomAction: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/custom_observation_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+
+class CustomObservation(google___protobuf___message___Message):
+
+    def __init__(self,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> CustomObservation: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/custom_reset_parameters_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+
+class CustomResetParameters(google___protobuf___message___Message):
+
+    def __init__(self,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> CustomResetParameters: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/demonstration_meta_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class DemonstrationMetaProto(google___protobuf___message___Message):
+    api_version = ... # type: int
+    demonstration_name = ... # type: typing___Text
+    number_steps = ... # type: int
+    number_episodes = ... # type: int
+    mean_reward = ... # type: float
+
+    def __init__(self,
+        api_version : typing___Optional[int] = None,
+        demonstration_name : typing___Optional[typing___Text] = None,
+        number_steps : typing___Optional[int] = None,
+        number_episodes : typing___Optional[int] = None,
+        mean_reward : typing___Optional[float] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> DemonstrationMetaProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"api_version",u"demonstration_name",u"mean_reward",u"number_episodes",u"number_steps"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"api_version",b"demonstration_name",b"mean_reward",b"number_episodes",b"number_steps"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/engine_configuration_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class EngineConfigurationProto(google___protobuf___message___Message):
+    width = ... # type: int
+    height = ... # type: int
+    quality_level = ... # type: int
+    time_scale = ... # type: float
+    target_frame_rate = ... # type: int
+    show_monitor = ... # type: bool
+
+    def __init__(self,
+        width : typing___Optional[int] = None,
+        height : typing___Optional[int] = None,
+        quality_level : typing___Optional[int] = None,
+        time_scale : typing___Optional[float] = None,
+        target_frame_rate : typing___Optional[int] = None,
+        show_monitor : typing___Optional[bool] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> EngineConfigurationProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"height",u"quality_level",u"show_monitor",u"target_frame_rate",u"time_scale",u"width"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"height",b"quality_level",b"show_monitor",b"target_frame_rate",b"time_scale",b"width"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/environment_parameters_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.custom_reset_parameters_pb2 import (
+    CustomResetParameters as mlagents___envs___communicator_objects___custom_reset_parameters_pb2___CustomResetParameters,
+)
+
+from typing import (
+    Mapping as typing___Mapping,
+    MutableMapping as typing___MutableMapping,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class EnvironmentParametersProto(google___protobuf___message___Message):
+    class FloatParametersEntry(google___protobuf___message___Message):
+        key = ... # type: typing___Text
+        value = ... # type: float
+
+        def __init__(self,
+            key : typing___Optional[typing___Text] = None,
+            value : typing___Optional[float] = None,
+            ) -> None: ...
+        @classmethod
+        def FromString(cls, s: bytes) -> EnvironmentParametersProto.FloatParametersEntry: ...
+        def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        if sys.version_info >= (3,):
+            def ClearField(self, field_name: typing_extensions___Literal[u"key",u"value"]) -> None: ...
+        else:
+            def ClearField(self, field_name: typing_extensions___Literal[b"key",b"value"]) -> None: ...
+
+
+    @property
+    def float_parameters(self) -> typing___MutableMapping[typing___Text, float]: ...
+
+    @property
+    def custom_reset_parameters(self) -> mlagents___envs___communicator_objects___custom_reset_parameters_pb2___CustomResetParameters: ...
+
+    def __init__(self,
+        float_parameters : typing___Optional[typing___Mapping[typing___Text, float]] = None,
+        custom_reset_parameters : typing___Optional[mlagents___envs___communicator_objects___custom_reset_parameters_pb2___CustomResetParameters] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> EnvironmentParametersProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_reset_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"custom_reset_parameters",u"float_parameters"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"custom_reset_parameters",b"custom_reset_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"custom_reset_parameters",b"float_parameters"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/header_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class Header(google___protobuf___message___Message):
+    status = ... # type: int
+    message = ... # type: typing___Text
+
+    def __init__(self,
+        status : typing___Optional[int] = None,
+        message : typing___Optional[typing___Text] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> Header: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"message",u"status"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"message",b"status"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/resolution_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class ResolutionProto(google___protobuf___message___Message):
+    width = ... # type: int
+    height = ... # type: int
+    gray_scale = ... # type: bool
+
+    def __init__(self,
+        width : typing___Optional[int] = None,
+        height : typing___Optional[int] = None,
+        gray_scale : typing___Optional[bool] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> ResolutionProto: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"gray_scale",u"height",u"width"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"gray_scale",b"height",b"width"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/space_type_proto_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.descriptor import (
+    EnumDescriptor as google___protobuf___descriptor___EnumDescriptor,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    List as typing___List,
+    Tuple as typing___Tuple,
+    cast as typing___cast,
+)
+
+
+class SpaceTypeProto(int):
+    DESCRIPTOR: google___protobuf___descriptor___EnumDescriptor = ...
+    @classmethod
+    def Name(cls, number: int) -> str: ...
+    @classmethod
+    def Value(cls, name: str) -> SpaceTypeProto: ...
+    @classmethod
+    def keys(cls) -> typing___List[str]: ...
+    @classmethod
+    def values(cls) -> typing___List[SpaceTypeProto]: ...
+    @classmethod
+    def items(cls) -> typing___List[typing___Tuple[str, SpaceTypeProto]]: ...
+discrete = typing___cast(SpaceTypeProto, 0)
+continuous = typing___cast(SpaceTypeProto, 1)

ml-agents-envs/mlagents/envs/communicator_objects/unity_input_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.unity_rl_initialization_input_pb2 import (
+    UnityRLInitializationInput as mlagents___envs___communicator_objects___unity_rl_initialization_input_pb2___UnityRLInitializationInput,
+)
+
+from mlagents.envs.communicator_objects.unity_rl_input_pb2 import (
+    UnityRLInput as mlagents___envs___communicator_objects___unity_rl_input_pb2___UnityRLInput,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityInput(google___protobuf___message___Message):
+
+    @property
+    def rl_input(self) -> mlagents___envs___communicator_objects___unity_rl_input_pb2___UnityRLInput: ...
+
+    @property
+    def rl_initialization_input(self) -> mlagents___envs___communicator_objects___unity_rl_initialization_input_pb2___UnityRLInitializationInput: ...
+
+    def __init__(self,
+        rl_input : typing___Optional[mlagents___envs___communicator_objects___unity_rl_input_pb2___UnityRLInput] = None,
+        rl_initialization_input : typing___Optional[mlagents___envs___communicator_objects___unity_rl_initialization_input_pb2___UnityRLInitializationInput] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityInput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"rl_initialization_input",u"rl_input"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"rl_initialization_input",u"rl_input"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"rl_initialization_input",b"rl_initialization_input",u"rl_input",b"rl_input"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"rl_initialization_input",b"rl_input"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_message_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.header_pb2 import (
+    Header as mlagents___envs___communicator_objects___header_pb2___Header,
+)
+
+from mlagents.envs.communicator_objects.unity_input_pb2 import (
+    UnityInput as mlagents___envs___communicator_objects___unity_input_pb2___UnityInput,
+)
+
+from mlagents.envs.communicator_objects.unity_output_pb2 import (
+    UnityOutput as mlagents___envs___communicator_objects___unity_output_pb2___UnityOutput,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityMessage(google___protobuf___message___Message):
+
+    @property
+    def header(self) -> mlagents___envs___communicator_objects___header_pb2___Header: ...
+
+    @property
+    def unity_output(self) -> mlagents___envs___communicator_objects___unity_output_pb2___UnityOutput: ...
+
+    @property
+    def unity_input(self) -> mlagents___envs___communicator_objects___unity_input_pb2___UnityInput: ...
+
+    def __init__(self,
+        header : typing___Optional[mlagents___envs___communicator_objects___header_pb2___Header] = None,
+        unity_output : typing___Optional[mlagents___envs___communicator_objects___unity_output_pb2___UnityOutput] = None,
+        unity_input : typing___Optional[mlagents___envs___communicator_objects___unity_input_pb2___UnityInput] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityMessage: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"header",u"unity_input",u"unity_output"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"header",u"unity_input",u"unity_output"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"header",b"header",u"unity_input",b"unity_input",u"unity_output",b"unity_output"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"header",b"unity_input",b"unity_output"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_output_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.unity_rl_initialization_output_pb2 import (
+    UnityRLInitializationOutput as mlagents___envs___communicator_objects___unity_rl_initialization_output_pb2___UnityRLInitializationOutput,
+)
+
+from mlagents.envs.communicator_objects.unity_rl_output_pb2 import (
+    UnityRLOutput as mlagents___envs___communicator_objects___unity_rl_output_pb2___UnityRLOutput,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityOutput(google___protobuf___message___Message):
+
+    @property
+    def rl_output(self) -> mlagents___envs___communicator_objects___unity_rl_output_pb2___UnityRLOutput: ...
+
+    @property
+    def rl_initialization_output(self) -> mlagents___envs___communicator_objects___unity_rl_initialization_output_pb2___UnityRLInitializationOutput: ...
+
+    def __init__(self,
+        rl_output : typing___Optional[mlagents___envs___communicator_objects___unity_rl_output_pb2___UnityRLOutput] = None,
+        rl_initialization_output : typing___Optional[mlagents___envs___communicator_objects___unity_rl_initialization_output_pb2___UnityRLInitializationOutput] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityOutput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"rl_initialization_output",u"rl_output"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"rl_initialization_output",u"rl_output"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"rl_initialization_output",b"rl_initialization_output",u"rl_output",b"rl_output"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"rl_initialization_output",b"rl_output"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_rl_initialization_input_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from typing import (
+    Optional as typing___Optional,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityRLInitializationInput(google___protobuf___message___Message):
+    seed = ... # type: int
+
+    def __init__(self,
+        seed : typing___Optional[int] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityRLInitializationInput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"seed"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"seed"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_rl_initialization_output_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedCompositeFieldContainer as google___protobuf___internal___containers___RepeatedCompositeFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.brain_parameters_proto_pb2 import (
+    BrainParametersProto as mlagents___envs___communicator_objects___brain_parameters_proto_pb2___BrainParametersProto,
+)
+
+from mlagents.envs.communicator_objects.environment_parameters_proto_pb2 import (
+    EnvironmentParametersProto as mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityRLInitializationOutput(google___protobuf___message___Message):
+    name = ... # type: typing___Text
+    version = ... # type: typing___Text
+    log_path = ... # type: typing___Text
+
+    @property
+    def brain_parameters(self) -> google___protobuf___internal___containers___RepeatedCompositeFieldContainer[mlagents___envs___communicator_objects___brain_parameters_proto_pb2___BrainParametersProto]: ...
+
+    @property
+    def environment_parameters(self) -> mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto: ...
+
+    def __init__(self,
+        name : typing___Optional[typing___Text] = None,
+        version : typing___Optional[typing___Text] = None,
+        log_path : typing___Optional[typing___Text] = None,
+        brain_parameters : typing___Optional[typing___Iterable[mlagents___envs___communicator_objects___brain_parameters_proto_pb2___BrainParametersProto]] = None,
+        environment_parameters : typing___Optional[mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityRLInitializationOutput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"environment_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"brain_parameters",u"environment_parameters",u"log_path",u"name",u"version"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"environment_parameters",b"environment_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"brain_parameters",b"environment_parameters",b"log_path",b"name",b"version"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_rl_input_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedCompositeFieldContainer as google___protobuf___internal___containers___RepeatedCompositeFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.agent_action_proto_pb2 import (
+    AgentActionProto as mlagents___envs___communicator_objects___agent_action_proto_pb2___AgentActionProto,
+)
+
+from mlagents.envs.communicator_objects.command_proto_pb2 import (
+    CommandProto as mlagents___envs___communicator_objects___command_proto_pb2___CommandProto,
+)
+
+from mlagents.envs.communicator_objects.environment_parameters_proto_pb2 import (
+    EnvironmentParametersProto as mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Mapping as typing___Mapping,
+    MutableMapping as typing___MutableMapping,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityRLInput(google___protobuf___message___Message):
+    class ListAgentActionProto(google___protobuf___message___Message):
+
+        @property
+        def value(self) -> google___protobuf___internal___containers___RepeatedCompositeFieldContainer[mlagents___envs___communicator_objects___agent_action_proto_pb2___AgentActionProto]: ...
+
+        def __init__(self,
+            value : typing___Optional[typing___Iterable[mlagents___envs___communicator_objects___agent_action_proto_pb2___AgentActionProto]] = None,
+            ) -> None: ...
+        @classmethod
+        def FromString(cls, s: bytes) -> UnityRLInput.ListAgentActionProto: ...
+        def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        if sys.version_info >= (3,):
+            def ClearField(self, field_name: typing_extensions___Literal[u"value"]) -> None: ...
+        else:
+            def ClearField(self, field_name: typing_extensions___Literal[b"value"]) -> None: ...
+
+    class AgentActionsEntry(google___protobuf___message___Message):
+        key = ... # type: typing___Text
+
+        @property
+        def value(self) -> UnityRLInput.ListAgentActionProto: ...
+
+        def __init__(self,
+            key : typing___Optional[typing___Text] = None,
+            value : typing___Optional[UnityRLInput.ListAgentActionProto] = None,
+            ) -> None: ...
+        @classmethod
+        def FromString(cls, s: bytes) -> UnityRLInput.AgentActionsEntry: ...
+        def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        if sys.version_info >= (3,):
+            def HasField(self, field_name: typing_extensions___Literal[u"value"]) -> bool: ...
+            def ClearField(self, field_name: typing_extensions___Literal[u"key",u"value"]) -> None: ...
+        else:
+            def HasField(self, field_name: typing_extensions___Literal[u"value",b"value"]) -> bool: ...
+            def ClearField(self, field_name: typing_extensions___Literal[b"key",b"value"]) -> None: ...
+
+    is_training = ... # type: bool
+    command = ... # type: mlagents___envs___communicator_objects___command_proto_pb2___CommandProto
+
+    @property
+    def agent_actions(self) -> typing___MutableMapping[typing___Text, UnityRLInput.ListAgentActionProto]: ...
+
+    @property
+    def environment_parameters(self) -> mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto: ...
+
+    def __init__(self,
+        agent_actions : typing___Optional[typing___Mapping[typing___Text, UnityRLInput.ListAgentActionProto]] = None,
+        environment_parameters : typing___Optional[mlagents___envs___communicator_objects___environment_parameters_proto_pb2___EnvironmentParametersProto] = None,
+        is_training : typing___Optional[bool] = None,
+        command : typing___Optional[mlagents___envs___communicator_objects___command_proto_pb2___CommandProto] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityRLInput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def HasField(self, field_name: typing_extensions___Literal[u"environment_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[u"agent_actions",u"command",u"environment_parameters",u"is_training"]) -> None: ...
+    else:
+        def HasField(self, field_name: typing_extensions___Literal[u"environment_parameters",b"environment_parameters"]) -> bool: ...
+        def ClearField(self, field_name: typing_extensions___Literal[b"agent_actions",b"command",b"environment_parameters",b"is_training"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_rl_output_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.internal.containers import (
+    RepeatedCompositeFieldContainer as google___protobuf___internal___containers___RepeatedCompositeFieldContainer,
+)
+
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+
+from mlagents.envs.communicator_objects.agent_info_proto_pb2 import (
+    AgentInfoProto as mlagents___envs___communicator_objects___agent_info_proto_pb2___AgentInfoProto,
+)
+
+from typing import (
+    Iterable as typing___Iterable,
+    Mapping as typing___Mapping,
+    MutableMapping as typing___MutableMapping,
+    Optional as typing___Optional,
+    Text as typing___Text,
+)
+
+from typing_extensions import (
+    Literal as typing_extensions___Literal,
+)
+
+
+class UnityRLOutput(google___protobuf___message___Message):
+    class ListAgentInfoProto(google___protobuf___message___Message):
+
+        @property
+        def value(self) -> google___protobuf___internal___containers___RepeatedCompositeFieldContainer[mlagents___envs___communicator_objects___agent_info_proto_pb2___AgentInfoProto]: ...
+
+        def __init__(self,
+            value : typing___Optional[typing___Iterable[mlagents___envs___communicator_objects___agent_info_proto_pb2___AgentInfoProto]] = None,
+            ) -> None: ...
+        @classmethod
+        def FromString(cls, s: bytes) -> UnityRLOutput.ListAgentInfoProto: ...
+        def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        if sys.version_info >= (3,):
+            def ClearField(self, field_name: typing_extensions___Literal[u"value"]) -> None: ...
+        else:
+            def ClearField(self, field_name: typing_extensions___Literal[b"value"]) -> None: ...
+
+    class AgentInfosEntry(google___protobuf___message___Message):
+        key = ... # type: typing___Text
+
+        @property
+        def value(self) -> UnityRLOutput.ListAgentInfoProto: ...
+
+        def __init__(self,
+            key : typing___Optional[typing___Text] = None,
+            value : typing___Optional[UnityRLOutput.ListAgentInfoProto] = None,
+            ) -> None: ...
+        @classmethod
+        def FromString(cls, s: bytes) -> UnityRLOutput.AgentInfosEntry: ...
+        def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+        if sys.version_info >= (3,):
+            def HasField(self, field_name: typing_extensions___Literal[u"value"]) -> bool: ...
+            def ClearField(self, field_name: typing_extensions___Literal[u"key",u"value"]) -> None: ...
+        else:
+            def HasField(self, field_name: typing_extensions___Literal[u"value",b"value"]) -> bool: ...
+            def ClearField(self, field_name: typing_extensions___Literal[b"key",b"value"]) -> None: ...
+
+    global_done = ... # type: bool
+
+    @property
+    def agentInfos(self) -> typing___MutableMapping[typing___Text, UnityRLOutput.ListAgentInfoProto]: ...
+
+    def __init__(self,
+        global_done : typing___Optional[bool] = None,
+        agentInfos : typing___Optional[typing___Mapping[typing___Text, UnityRLOutput.ListAgentInfoProto]] = None,
+        ) -> None: ...
+    @classmethod
+    def FromString(cls, s: bytes) -> UnityRLOutput: ...
+    def MergeFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    def CopyFrom(self, other_msg: google___protobuf___message___Message) -> None: ...
+    if sys.version_info >= (3,):
+        def ClearField(self, field_name: typing_extensions___Literal[u"agentInfos",u"global_done"]) -> None: ...
+    else:
+        def ClearField(self, field_name: typing_extensions___Literal[b"agentInfos",b"global_done"]) -> None: ...

ml-agents-envs/mlagents/envs/communicator_objects/unity_to_external_pb2.pyi

+# @generated by generate_proto_mypy_stubs.py.  Do not edit!
+import sys
+from google.protobuf.message import (
+    Message as google___protobuf___message___Message,
+)
+