[BugFix] Fix slow and flaky non-tensor parallel env test

ghstack-source-id: fcb5caa56e05176958b3468a7d6f69e363cfe558
Pull-Request-resolved: #2926

Author: vmoens

test/test_env.py

@@ -3716,14 +3716,15 @@ def test_serial(self, bwad, use_buffers):
 
     @pytest.mark.parametrize("bwad", [True, False])
     @pytest.mark.parametrize("use_buffers", [False, True])
-    def test_parallel(self, bwad, use_buffers):
+    def test_parallel(self, bwad, use_buffers, maybe_fork_ParallelEnv):
         N = 50
-        env = ParallelEnv(2, EnvWithMetadata, use_buffers=use_buffers)
+        env = maybe_fork_ParallelEnv(2, EnvWithMetadata, use_buffers=use_buffers)
         try:
             r = env.rollout(N, break_when_any_done=bwad)
             assert r.get("non_tensor").tolist() == [list(range(N))] * 2
         finally:
             env.close(raise_if_closed=False)
+            del env
 
     class AddString(Transform):
         def __init__(self):

torchrl/envs/batched_envs.py

@@ -2232,8 +2232,14 @@ def _shutdown_workers(self) -> None:
 
             for channel in self.parent_channels:
                 channel.close()
+            start_time = time.time()
+            while (
+                any(proc.is_alive() for proc in self._workers)
+                and (time.time() - start_time) < self._timeout
+            ):
+                time.sleep(0.01)
             for proc in self._workers:
-                proc.join(timeout=self._timeout)
+                proc.join()
         finally:
             for proc in self._workers:
                 if proc.is_alive():
@@ -2731,16 +2737,13 @@ def _run_worker_pipe_direct(
             if not initialized:
                 raise RuntimeError("call 'init' before closing")
             env.close()
-            del (
-                env,
-                data,
-            )
             mp_event.set()
             child_pipe.close()
             if verbose:
                 torchrl_logger.info(f"{pid} closed")
+            del (env, data, child_pipe, mp_event)
             gc.collect()
-            break
+            return
 
         elif cmd == "load_state_dict":
             env.load_state_dict(data)