Shorten timeout duration for environment close

com.unity.ml-agents/CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
  - Fixed an issue where specifying `vis_encode_type` was required only for SAC. (#3677)
  - The way that UnityEnvironment decides the port was changed. If no port is specified, the behavior will depend on the `file_name` parameter. If it is `None`, 5004 (the editor port) will be used; otherwise 5005 (the base environment port) will be used.
  - Fixed an issue where switching models using `SetModel()` during training would use an excessive amount of memory. (#3664)
+ - Environment subprocesses now close immediately on timeout or wrong API version. (#3679)
 
 ## [0.15.0-preview] - 2020-03-18
 ### Major Changes

ml-agents-envs/mlagents_envs/environment.py

@@ -153,12 +153,12 @@ def __init__(
             aca_output = self.send_academy_parameters(rl_init_parameters_in)
             aca_params = aca_output.rl_initialization_output
         except UnityTimeOutException:
-            self._close()
+            self._close(0)
             raise
 
         unity_communicator_version = aca_params.communication_version
         if unity_communicator_version != UnityEnvironment.API_VERSION:
-            self._close()
+            self._close(0)
             raise UnityEnvironmentException(
                 f"The communication API version is not compatible between Unity and python. "
                 f"Python API: {UnityEnvironment.API_VERSION}, Unity API: {unity_communicator_version}.\n "
@@ -239,7 +239,7 @@ def validate_environment_path(env_path: str) -> Optional[str]:
     def executable_launcher(self, file_name, docker_training, no_graphics, args):
         launch_string = self.validate_environment_path(file_name)
         if launch_string is None:
-            self._close()
+            self._close(0)
             raise UnityEnvironmentException(
                 f"Couldn't launch the {file_name} environment. Provided filename does not match any environments."
             )
@@ -444,13 +444,21 @@ def close(self):
         else:
             raise UnityEnvironmentException("No Unity environment is loaded.")
 
-    def _close(self):
+    def _close(self, timeout: Optional[int] = None) -> None:
+        """
+        Close the communicator and environment subprocess (if necessary).
+
+        :int timeout: [Optional] Number of seconds to wait for the environment to shut down before
+            force-killing it.  Defaults to `self.timeout_wait`.
+        """
+        if timeout is None:
+            timeout = self.timeout_wait
         self._loaded = False
         self.communicator.close()
         if self.proc1 is not None:
             # Wait a bit for the process to shutdown, but kill it if it takes too long
             try:
-                self.proc1.wait(timeout=self.timeout_wait)
+                self.proc1.wait(timeout=timeout)
                 signal_name = self.returncode_to_signal_name(self.proc1.returncode)
                 signal_name = f" ({signal_name})" if signal_name else ""
                 return_info = f"Environment shut down with return code {self.proc1.returncode}{signal_name}."