Add timeout for communicator exchange

When we initially connect to the environment using RPCCommunicator,
the connection is polled so we don't hang forever on `.recv()` when
the environment wasn't launched or failed.  However we don't currently
have any similar check for the exchanges mid-training-run.

This change applies the same timeout from initialization to each exchange,
and extends the default `timeout_wait` to 60 seconds to generally improve
the chances we won't have a mismatch between environment launch time and
the trainer timeout.

Tested on: single-env and multi-env cases.  Killed 1 environment process
manually and saw that the model was saved appropriately and all processes
closed.

Author: Jonathan Harper

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

@@ -51,7 +51,7 @@ def __init__(
         seed: int = 0,
         docker_training: bool = False,
         no_graphics: bool = False,
-        timeout_wait: int = 30,
+        timeout_wait: int = 60,
         args: Optional[List[str]] = None,
     ):
         """

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

@@ -82,14 +82,22 @@ def check_port(self, port):
         finally:
             s.close()
 
-    def initialize(self, inputs: UnityInputProto) -> UnityOutputProto:
+    def poll_for_timeout(self):
+        """
+        Polls the GRPC parent connection for data, to be used before calling recv.  This prevents
+        us from hanging indefinitely in the case where the environment process has died or was not
+        launched.
+        """
         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"
                 "\t The Agents are linked to the appropriate Brains\n"
                 "\t The environment and the Python interface have compatible versions."
             )
+
+    def initialize(self, inputs: UnityInputProto) -> UnityOutputProto:
+        self.poll_for_timeout()
         aca_param = self.unity_to_external.parent_conn.recv().unity_output
         message = UnityMessageProto()
         message.header.status = 200
@@ -103,6 +111,7 @@ def exchange(self, inputs: UnityInputProto) -> Optional[UnityOutputProto]:
         message.header.status = 200
         message.unity_input.CopyFrom(inputs)
         self.unity_to_external.parent_conn.send(message)
+        self.poll_for_timeout()
         output = self.unity_to_external.parent_conn.recv()
         if output.header.status != 200:
             return None

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

@@ -3,7 +3,7 @@
 import cloudpickle
 
 from mlagents.envs.environment import UnityEnvironment
-from mlagents.envs.exception import UnityCommunicationException
+from mlagents.envs.exception import UnityCommunicationException, UnityTimeOutException
 from multiprocessing import Process, Pipe, Queue
 from multiprocessing.connection import Connection
 from queue import Empty as EmptyQueueException
@@ -116,7 +116,7 @@ def _send_response(cmd_name, payload):
                 _send_response("reset", all_brain_info)
             elif cmd.name == "close":
                 break
-    except (KeyboardInterrupt, UnityCommunicationException):
+    except (KeyboardInterrupt, UnityCommunicationException, UnityTimeOutException):
         logger.info(f"UnityEnvironment worker {worker_id}: environment stopping.")
         step_queue.put(EnvironmentResponse("env_close", worker_id, None))
     finally: