Add timeout wait param

# Conflicts:
#	ml-agents-envs/setup.py

ml-agents/mlagents/envs/environment.py

@@ -26,7 +26,8 @@ class UnityEnvironment(object):
 
     def __init__(self, file_name=None, worker_id=0,
                  base_port=5005, seed=0,
-                 docker_training=False, no_graphics=False):
+                 docker_training=False, no_graphics=False,
+                 timeout_wait=30):
         """
         Starts a new unity environment and establishes a connection with the environment.
         Notice: Currently communication between Unity and Python takes place over an open socket without authentication.
@@ -35,8 +36,9 @@ def __init__(self, file_name=None, worker_id=0,
         :string file_name: Name of Unity environment binary.
         :int base_port: Baseline port number to connect to Unity environment over. worker_id increments over this.
         :int worker_id: Number to add to communication port (5005) [0]. Used for asynchronous agent scenarios.
-        :param docker_training: Informs this class whether the process is being run within a container.
-        :param no_graphics: Whether to run the Unity simulator in no-graphics mode
+        :bool docker_training: Informs this class whether the process is being run within a container.
+        :bool no_graphics: Whether to run the Unity simulator in no-graphics mode
+        :int timeout_wait: Time (in seconds) to wait for connection from environment.
         """
 
         atexit.register(self._close)
@@ -45,7 +47,7 @@ def __init__(self, file_name=None, worker_id=0,
         self._version_ = "API-6"
         self._loaded = False  # If true, this means the environment was successfully loaded
         self.proc1 = None  # The process that is started. If None, no process was started
-        self.communicator = self.get_communicator(worker_id, base_port)
+        self.communicator = self.get_communicator(worker_id, base_port, timeout_wait)
 
         # If the environment name is None, a new environment will not be launched
         # and the communicator will directly try to connect to an existing unity environment.
@@ -207,8 +209,8 @@ def executable_launcher(self, file_name, docker_training, no_graphics):
                                               stderr=subprocess.PIPE,
                                               shell=True)
 
-    def get_communicator(self, worker_id, base_port):
-        return RpcCommunicator(worker_id, base_port)
+    def get_communicator(self, worker_id, base_port, timeout_wait):
+        return RpcCommunicator(worker_id, base_port, timeout_wait)
         # return SocketCommunicator(worker_id, base_port)
 
     def __str__(self):

ml-agents/mlagents/envs/rpc_communicator.py

@@ -27,7 +27,7 @@ def Exchange(self, request, context):
 
 
 class RpcCommunicator(Communicator):
-    def __init__(self, worker_id=0, base_port=5005):
+    def __init__(self, worker_id=0, base_port=5005, timeout_wait=30):
         """
         Python side of the grpc communication. Python is the server and Unity the client
 
@@ -37,6 +37,7 @@ def __init__(self, worker_id=0, base_port=5005):
         """
         self.port = base_port + worker_id
         self.worker_id = worker_id
+        self.timeout_wait = timeout_wait
         self.server = None
         self.unity_to_external = None
         self.is_open = False
@@ -74,7 +75,7 @@ def check_port(self, port):
             s.close()
 
     def initialize(self, inputs: UnityInput) -> UnityOutput:
-        if not self.unity_to_external.parent_conn.poll(30):
+        if not self.unity_to_external.parent_conn.poll(self.timeout_wait):
             raise UnityTimeOutException(
                 "The Unity environment took too long to respond. Make sure that :\n"
                 "\t The environment does not need user interaction to launch\n"