The environment can now close right after being instantiated (#4071)

Author: vincentpierre

com.unity.ml-agents/Runtime/Academy.cs

@@ -366,6 +366,8 @@ void InitializeEnvironment()
                         port = port
                     }
                 );
+                Communicator.QuitCommandReceived += OnQuitCommandReceived;
+                Communicator.ResetCommandReceived += OnResetCommand;
             }
 
             if (Communicator != null)
@@ -397,12 +399,6 @@ void InitializeEnvironment()
                     );
                     Communicator = null;
                 }
-
-                if (Communicator != null)
-                {
-                    Communicator.QuitCommandReceived += OnQuitCommandReceived;
-                    Communicator.ResetCommandReceived += OnResetCommand;
-                }
             }
 
             // If a communicator is enabled/provided, then we assume we are in

com.unity.ml-agents/Runtime/Communicator/RpcCommunicator.cs

@@ -214,10 +214,16 @@ UnityInputProto Initialize(UnityOutputProto unityOutput,
 
             m_Client = new UnityToExternalProto.UnityToExternalProtoClient(channel);
             var result = m_Client.Exchange(WrapMessage(unityOutput, 200));
-            unityInput = m_Client.Exchange(WrapMessage(null, 200)).UnityInput;
+            var inputMessage = m_Client.Exchange(WrapMessage(null, 200));
+            unityInput = inputMessage.UnityInput;
 #if UNITY_EDITOR
             EditorApplication.playModeStateChanged += HandleOnPlayModeChanged;
 #endif
+            if (result.Header.Status != 200 || inputMessage.Header.Status != 200)
+            {
+                m_IsOpen = false;
+                QuitCommandReceived?.Invoke();
+            }
             return result.UnityInput;
 #else
             throw new UnityAgentsException(