[Refactor] Non-daemonic processes in PEnv

ghstack-source-id: c625cd3
Pull-Request: #3225

Author: vmoens

torchrl/envs/batched_envs.py

@@ -196,6 +196,9 @@ class BatchedEnvBase(EnvBase):
             one of the environment has dynamic specs.
 
               .. note:: Learn more about dynamic specs and environments :ref:`here <dynamic_envs>`.
+        daemon (bool, optional): whether the processes should be daemonized.
+            This is only applicable to parallel environments such as :class:`~torchrl.envs.ParallelEnv`.
+            Defaults to ``False``.
 
     .. note::
         One can pass keyword arguments to each sub-environments using the following
@@ -311,6 +314,7 @@ def __init__(
         mp_start_method: str | None = None,
         use_buffers: bool | None = None,
         consolidate: bool = True,
+        daemon: bool = False,
     ):
         super().__init__(device=device)
         self.serial_for_single = serial_for_single
@@ -320,6 +324,7 @@ def __init__(
         self._cache_in_keys = None
         self._use_buffers = use_buffers
         self.consolidate = consolidate
+        self.daemon = daemon
 
         self._single_task = callable(create_env_fn) or (len(set(create_env_fn)) == 1)
         if callable(create_env_fn):
@@ -1503,7 +1508,7 @@ def look_for_cuda(tensor, has_cuda=has_cuda):
                         }
                     )
                 process = proc_fun(target=func, kwargs=kwargs[idx])
-                process.daemon = True
+                process.daemon = self.daemon
                 process.start()
                 child_pipe.close()
                 self.parent_channels.append(parent_pipe)