Add timeout wait param (Develop) (#1700)

* Add timeout wait param
* Remove unnecessary function

ml-agents/mlagents/envs/environment.py

 
     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.
         :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)
         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 = RpcCommunicator(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.
                                               stdout=subprocess.PIPE,
                                               stderr=subprocess.PIPE,
                                               shell=True)
-
-    def get_communicator(self, worker_id, base_port):
-        return RpcCommunicator(worker_id, base_port)
-        # return SocketCommunicator(worker_id, base_port)
 
     def __str__(self):
         return '''Unity Academy name: {0}

ml-agents/mlagents/envs/rpc_communicator.py

 
 
 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
 
         """
         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
             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"