[BugFix] Fix mp_start_method for ParallelEnv with single_for_serial (p…

…ytorch#2007)
ai4co · Mar 11, 2024 · 2b8450c · 2b8450c
1 parent 358475a
commit 2b8450c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/test/test_env.py b/test/test_env.py
@@ -447,13 +447,22 @@ def test_parallel_devices(
                 env.shared_tensordict_parent.device.type == torch.device(edevice).type
             )
 
-    def test_serial_for_single(self, maybe_fork_ParallelEnv):
-        env = ParallelEnv(1, ContinuousActionVecMockEnv, serial_for_single=True)
+    @pytest.mark.parametrize("start_method", [None, "fork"])
+    def test_serial_for_single(self, maybe_fork_ParallelEnv, start_method):
+        env = ParallelEnv(
+            1,
+            ContinuousActionVecMockEnv,
+            serial_for_single=True,
+            mp_start_method=start_method,
+        )
         assert isinstance(env, SerialEnv)
-        env = maybe_fork_ParallelEnv(1, ContinuousActionVecMockEnv)
+        env = ParallelEnv(1, ContinuousActionVecMockEnv, mp_start_method=start_method)
         assert isinstance(env, ParallelEnv)
-        env = maybe_fork_ParallelEnv(
-            2, ContinuousActionVecMockEnv, serial_for_single=True
+        env = ParallelEnv(
+            2,
+            ContinuousActionVecMockEnv,
+            serial_for_single=True,
+            mp_start_method=start_method,
         )
         assert isinstance(env, ParallelEnv)
 

diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -115,6 +115,8 @@ def __call__(cls, *args, **kwargs):
         serial_for_single = kwargs.pop("serial_for_single", False)
         if serial_for_single:
             num_workers = kwargs.get("num_workers", None)
+            # Remove start method from kwargs
+            kwargs.pop("mp_start_method", None)
             if num_workers is None:
                 num_workers = args[0]
             if num_workers == 1: