update _worker test

Author: huangshiyu13

openrl/envs/vec_env/async_venv.py

@@ -734,7 +734,7 @@ def _worker(
     index: int,
     env_fn: callable,
     pipe: Connection,
-    parent_pipe: Connection,
+    parent_pipe: Optional[Connection],
     shared_memory: bool,
     error_queue: Queue,
     auto_reset: bool = True,
@@ -757,11 +757,14 @@ def prepare_obs(observation):
             )
             observation = None
         return observation
-
-    parent_pipe.close()
+    if parent_pipe is not None:
+        parent_pipe.close()
     try:
         while True:
             command, data = pipe.recv()
+            print(command)
+
+
 
             if command == "reset":
                 result = env.reset(**data)

tests/test_env/test_vec_env/test_async_env.py

@@ -24,7 +24,8 @@
 
 from openrl.envs.toy_envs import make_toy_envs
 from openrl.envs.vec_env.async_venv import AsyncVectorEnv
-
+from openrl.envs.vec_env.async_venv import _worker
+import multiprocessing as mp
 
 class CustomEnvCompatibility(EnvCompatibility):
     def reset(self, **kwargs):
@@ -48,18 +49,68 @@ def assert_env_name(env, env_name):
         assert env.metadata["name"].__name__ == env_name
 
 
+# @pytest.mark.unittest
+# def test_async_env():
+#     env_name = "IdentityEnv"
+#     env = AsyncVectorEnv(init_envs(), shared_memory=True)
+#     assert (
+#         env._env_name == env_name
+#     ), "AsyncVectorEnv should have the same metadata as the wrapped env"
+#     env.exec_func(assert_env_name, indices=None, env_name=env_name)
+#     env.call("render")
+#     env_name_new = "IdentityEnvNew"
+#     env.set_attr("metadata", {"name": env_name_new})
+#     env.exec_func(assert_env_name, indices=None, env_name=env_name_new)
+
+def main_control(parent_pipe,child_pipe):
+    child_pipe.close()
+
+    parent_pipe.send(("reset", {"seed":0}))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+    parent_pipe.send(("step", [0]))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+    parent_pipe.send(("_call", ("render",[],{})))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+    parent_pipe.send(("_setattr", ("metadata", {"name": "IdentityEnvNew"})))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+    parent_pipe.send(("_func_exec",(assert_env_name,None,[],{"env_name":"IdentityEnvNew"})))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+    parent_pipe.send(("close",None))
+    result, success = parent_pipe.recv()
+    assert success, result
+
+
 @pytest.mark.unittest
-def test_async_env():
-    env_name = "IdentityEnv"
-    env = AsyncVectorEnv(init_envs(), shared_memory=True)
-    assert (
-        env._env_name == env_name
-    ), "AsyncVectorEnv should have the same metadata as the wrapped env"
-    env.exec_func(assert_env_name, indices=None, env_name=env_name)
-    env.call("render")
-    env_name_new = "IdentityEnvNew"
-    env.set_attr("metadata", {"name": env_name_new})
-    env.exec_func(assert_env_name, indices=None, env_name=env_name_new)
+def test_worker():
+    for auto_reset in [True,False]:
+        ctx = mp.get_context(None)
+        parent_pipe, child_pipe = ctx.Pipe()
+
+        error_queue = ctx.Queue()
+
+        process = ctx.Process(
+            target=main_control,
+            name="test",
+            args=(
+                parent_pipe,
+                child_pipe
+            ),
+        )
+        process.daemon = True
+        process.start()
+        _worker(0, init_envs()[0], child_pipe, None, False, error_queue, auto_reset)
+
+
 
 
 if __name__ == "__main__":