update to match upstream folder changes

ray-project · Jan 22, 2021 · 6b6c548 · 6b6c548
1 parent ae8a2e2
commit 6b6c548
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 6 deletions.
diff --git a/rllib/tests/test_unity3d_env.py → rllib/env/wrappers/tests/test_unity3d_env.py b/rllib/tests/test_unity3d_env.py → rllib/env/wrappers/tests/test_unity3d_env.py
@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 
-from ray.rllib.env.unity3d_env import Unity3DEnv
+from ray.rllib.env.wrappers.unity3d_env import Unity3DEnv
 
 
 @patch("mlagents_envs.environment.UnityEnvironment")

diff --git a/rllib/env/wrappers/unity3d_env.py b/rllib/env/wrappers/unity3d_env.py
@@ -27,7 +27,12 @@ class Unity3DEnv(MultiAgentEnv):
     inside an RLlib PolicyClient for cloud/distributed training of Unity games.
     """
 
-    _BASE_PORT = 5004
+    # Default base port when connecting directly to the Editor
+    _BASE_PORT_EDITOR = 5004
+    # Default base port when connecting to a compiled environment
+    _BASE_PORT_ENVIRONMENT = 5005
+    # The worker_id for each environment instance
+    _WORKER_ID = 0
 
     def __init__(self,
                  file_name: str = None,
@@ -73,18 +78,24 @@ def __init__(self,
             # environments (num_workers >> 1). Otherwise, would lead to port
             # conflicts sometimes.
             time.sleep(random.randint(1, 10))
-            port_ = port or self._BASE_PORT
-            self._BASE_PORT += 1
+            port_ = port or (self._BASE_PORT_ENVIRONMENT
+                             if file_name else self._BASE_PORT_EDITOR)
+            # cache the worker_id and
+            # increase it for the next environment
+            worker_id_ = Unity3DEnv._WORKER_ID if file_name else 0
+            Unity3DEnv._WORKER_ID += 1
             try:
                 self.unity_env = UnityEnvironment(
                     file_name=file_name,
-                    worker_id=0,
+                    worker_id=worker_id_,
                     base_port=port_,
                     seed=seed,
                     no_graphics=no_graphics,
                     timeout_wait=timeout_wait,
                 )
-                print("Created UnityEnvironment for port {}".format(port_))
+                print(
+                    "Created UnityEnvironment for port {}".format(port_ +
+                                                                  worker_id_))
             except mlagents_envs.exception.UnityWorkerInUseException:
                 pass
             else: