Add enum for environment commands

Author: Jonathan Harper

ml-agents/mlagents/trainers/subprocess_env_manager.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Dict, NamedTuple, List, Any, Optional, Callable, Set
 import cloudpickle
+import enum
 
 from mlagents_envs.environment import UnityEnvironment
 from mlagents_envs.exception import (
@@ -33,13 +34,22 @@
 logger = logging.getLogger("mlagents.trainers")
 
 
-class EnvironmentCommand(NamedTuple):
-    name: str
+class EnvironmentCommand(enum.Enum):
+    STEP = 1
+    EXTERNAL_BRAINS = 2
+    GET_PROPERTIES = 3
+    RESET = 4
+    CLOSE = 5
+    ENV_EXITED = 6
+
+
+class EnvironmentRequest(NamedTuple):
+    cmd: EnvironmentCommand
     payload: Any = None
 
 
 class EnvironmentResponse(NamedTuple):
-    name: str
+    cmd: EnvironmentCommand
     worker_id: int
     payload: Any
 
@@ -58,17 +68,17 @@ def __init__(self, process: Process, worker_id: int, conn: Connection):
         self.previous_all_action_info: Dict[str, ActionInfo] = {}
         self.waiting = False
 
-    def send(self, name: str, payload: Any = None) -> None:
+    def send(self, cmd: EnvironmentCommand, payload: Any = None) -> None:
         try:
-            cmd = EnvironmentCommand(name, payload)
-            self.conn.send(cmd)
+            req = EnvironmentRequest(cmd, payload)
+            self.conn.send(req)
         except (BrokenPipeError, EOFError):
             raise UnityCommunicationException("UnityEnvironment worker: send failed.")
 
     def recv(self) -> EnvironmentResponse:
         try:
             response: EnvironmentResponse = self.conn.recv()
-            if response.name == "env_close":
+            if response.cmd == EnvironmentCommand.ENV_EXITED:
                 env_exception: Exception = response.payload
                 raise env_exception
             return response
@@ -77,7 +87,7 @@ def recv(self) -> EnvironmentResponse:
 
     def close(self):
         try:
-            self.conn.send(EnvironmentCommand("close"))
+            self.conn.send(EnvironmentRequest(EnvironmentCommand.CLOSE))
         except (BrokenPipeError, EOFError):
             logger.debug(
                 f"UnityEnvWorker {self.worker_id} got exception trying to close."
@@ -102,7 +112,7 @@ def worker(
     engine_configuration_channel.set_configuration(engine_configuration)
     env: BaseEnv = None
 
-    def _send_response(cmd_name, payload):
+    def _send_response(cmd_name: EnvironmentCommand, payload: Any) -> None:
         parent_conn.send(EnvironmentResponse(cmd_name, worker_id, payload))
 
     def _generate_all_results() -> AllStepResult:
@@ -124,9 +134,9 @@ def external_brains():
             worker_id, [shared_float_properties, engine_configuration_channel]
         )
         while True:
-            cmd: EnvironmentCommand = parent_conn.recv()
-            if cmd.name == "step":
-                all_action_info = cmd.payload
+            req: EnvironmentRequest = parent_conn.recv()
+            if req.cmd == EnvironmentCommand.STEP:
+                all_action_info = req.payload
                 for brain_name, action_info in all_action_info.items():
                     if len(action_info.action) != 0:
                         env.set_actions(brain_name, action_info.action)
@@ -138,20 +148,24 @@ def external_brains():
                 # the data transferred.
                 # TODO get gauges from the workers and merge them in the main process too.
                 step_response = StepResponse(all_step_result, get_timer_root())
-                step_queue.put(EnvironmentResponse("step", worker_id, step_response))
+                step_queue.put(
+                    EnvironmentResponse(
+                        EnvironmentCommand.STEP, worker_id, step_response
+                    )
+                )
                 reset_timers()
-            elif cmd.name == "external_brains":
-                _send_response("external_brains", external_brains())
-            elif cmd.name == "get_properties":
+            elif req.cmd == EnvironmentCommand.EXTERNAL_BRAINS:
+                _send_response(EnvironmentCommand.EXTERNAL_BRAINS, external_brains())
+            elif req.cmd == EnvironmentCommand.GET_PROPERTIES:
                 reset_params = shared_float_properties.get_property_dict_copy()
-                _send_response("get_properties", reset_params)
-            elif cmd.name == "reset":
-                for k, v in cmd.payload.items():
+                _send_response(EnvironmentCommand.GET_PROPERTIES, reset_params)
+            elif req.cmd == EnvironmentCommand.RESET:
+                for k, v in req.payload.items():
                     shared_float_properties.set_property(k, v)
                 env.reset()
                 all_step_result = _generate_all_results()
-                _send_response("reset", all_step_result)
-            elif cmd.name == "close":
+                _send_response(EnvironmentCommand.RESET, all_step_result)
+            elif req.cmd == EnvironmentCommand.CLOSE:
                 break
     except (
         KeyboardInterrupt,
@@ -160,8 +174,10 @@ def external_brains():
         UnityEnvironmentException,
     ) as ex:
         logger.info(f"UnityEnvironment worker {worker_id}: environment stopping.")
-        step_queue.put(EnvironmentResponse("env_close", worker_id, ex))
-        _send_response("env_close", ex)
+        step_queue.put(
+            EnvironmentResponse(EnvironmentCommand.ENV_EXITED, worker_id, ex)
+        )
+        _send_response(EnvironmentCommand.ENV_EXITED, ex)
     finally:
         # If this worker has put an item in the step queue that hasn't been processed by the EnvManager, the process
         # will hang until the item is processed. We avoid this behavior by using Queue.cancel_join_thread()
@@ -222,7 +238,7 @@ def _queue_steps(self) -> None:
             if not env_worker.waiting:
                 env_action_info = self._take_step(env_worker.previous_step)
                 env_worker.previous_all_action_info = env_action_info
-                env_worker.send("step", env_action_info)
+                env_worker.send(EnvironmentCommand.STEP, env_action_info)
                 env_worker.waiting = True
 
     def _step(self) -> List[EnvironmentStep]:
@@ -236,8 +252,8 @@ def _step(self) -> List[EnvironmentStep]:
         while len(worker_steps) < 1:
             try:
                 while True:
-                    step = self.step_queue.get_nowait()
-                    if step.name == "env_close":
+                    step: EnvironmentResponse = self.step_queue.get_nowait()
+                    if step.cmd == EnvironmentCommand.ENV_EXITED:
                         env_exception: Exception = step.payload
                         raise env_exception
                     self.env_workers[step.worker_id].waiting = False
@@ -257,20 +273,20 @@ def _reset_env(self, config: Optional[Dict] = None) -> List[EnvironmentStep]:
                 self.env_workers[step.worker_id].waiting = False
         # First enqueue reset commands for all workers so that they reset in parallel
         for ew in self.env_workers:
-            ew.send("reset", config)
+            ew.send(EnvironmentCommand.RESET, config)
         # Next (synchronously) collect the reset observations from each worker in sequence
         for ew in self.env_workers:
             ew.previous_step = EnvironmentStep(ew.recv().payload, ew.worker_id, {})
         return list(map(lambda ew: ew.previous_step, self.env_workers))
 
     @property
     def external_brains(self) -> Dict[AgentGroup, BrainParameters]:
-        self.env_workers[0].send("external_brains")
+        self.env_workers[0].send(EnvironmentCommand.EXTERNAL_BRAINS)
         return self.env_workers[0].recv().payload
 
     @property
     def get_properties(self) -> Dict[AgentGroup, float]:
-        self.env_workers[0].send("get_properties")
+        self.env_workers[0].send(EnvironmentCommand.GET_PROPERTIES)
         return self.env_workers[0].recv().payload
 
     def close(self) -> None:

ml-agents/mlagents/trainers/tests/test_subprocess_env_manager.py

@@ -8,6 +8,7 @@
     SubprocessEnvManager,
     EnvironmentResponse,
     StepResponse,
+    EnvironmentCommand,
 )
 from mlagents.trainers.env_manager import EnvironmentStep
 from mlagents_envs.base_env import BaseEnv
@@ -38,7 +39,9 @@ def __init__(self, worker_id, resp=None):
 
 
 def create_worker_mock(worker_id, step_queue, env_factor, engine_c):
-    return MockEnvWorker(worker_id, EnvironmentResponse("reset", worker_id, worker_id))
+    return MockEnvWorker(
+        worker_id, EnvironmentResponse(EnvironmentCommand.RESET, worker_id, worker_id)
+    )
 
 
 class SubprocessEnvManagerTest(unittest.TestCase):
@@ -71,7 +74,9 @@ def test_reset_passes_reset_params(self, mock_create_worker):
         )
         params = {"test": "params"}
         manager._reset_env(params)
-        manager.env_workers[0].send.assert_called_with("reset", (params))
+        manager.env_workers[0].send.assert_called_with(
+            EnvironmentCommand.RESET, (params)
+        )
 
     @mock.patch(
         "mlagents.trainers.subprocess_env_manager.SubprocessEnvManager.create_worker"
@@ -85,7 +90,7 @@ def test_reset_collects_results_from_all_envs(self, mock_create_worker):
         params = {"test": "params"}
         res = manager._reset_env(params)
         for i, env in enumerate(manager.env_workers):
-            env.send.assert_called_with("reset", (params))
+            env.send.assert_called_with(EnvironmentCommand.RESET, (params))
             env.recv.assert_called()
             # Check that the "last steps" are set to the value returned for each step
             self.assertEqual(
@@ -103,8 +108,8 @@ def test_step_takes_steps_for_all_non_waiting_envs(self, mock_create_worker):
         )
         manager.step_queue = Mock()
         manager.step_queue.get_nowait.side_effect = [
-            EnvironmentResponse("step", 0, StepResponse(0, None)),
-            EnvironmentResponse("step", 1, StepResponse(1, None)),
+            EnvironmentResponse(EnvironmentCommand.STEP, 0, StepResponse(0, None)),
+            EnvironmentResponse(EnvironmentCommand.STEP, 1, StepResponse(1, None)),
             EmptyQueue(),
         ]
         step_mock = Mock()
@@ -117,7 +122,7 @@ def test_step_takes_steps_for_all_non_waiting_envs(self, mock_create_worker):
         res = manager._step()
         for i, env in enumerate(manager.env_workers):
             if i < 2:
-                env.send.assert_called_with("step", step_mock)
+                env.send.assert_called_with(EnvironmentCommand.STEP, step_mock)
                 manager.step_queue.get_nowait.assert_called()
                 # Check that the "last steps" are set to the value returned for each step
                 self.assertEqual(