[Refactor] Make env creator optional for Ray

ghstack-source-id: 6444d2c
Pull-Request: #3227

Author: vmoens

torchrl/collectors/distributed/ray.py

@@ -268,6 +268,10 @@ class RayCollector(DataCollectorBase):
             :class:`~torchrl.weight_update.weight_sync_schemes.WeightSyncScheme` instances.
             This is the recommended way to configure weight synchronization. If not provided,
             defaults to ``{"policy": RayWeightSyncScheme()}``.
+        use_env_creator (bool, optional): if ``True``, the environment constructor functions will be wrapped
+            in :class:`~torchrl.envs.EnvCreator`. This is useful for multiprocessed settings where shared memory
+            needs to be managed, but Ray has its own object storage mechanism, so this is typically not needed.
+            Defaults to ``False``.
 
     Examples:
         >>> from torch import nn
@@ -332,6 +336,7 @@ def __init__(
         | Callable[[], WeightUpdaterBase]
         | None = None,
         weight_sync_schemes: dict[str, WeightSyncScheme] | None = None,
+        use_env_creator: bool = False,
     ):
         self.frames_per_batch = frames_per_batch
         if remote_configs is None:
@@ -406,9 +411,10 @@ def check_list_length_consistency(*lists):
         create_env_fn, collector_kwargs, remote_configs = out_lists
         num_collectors = len(create_env_fn)
 
-        for i in range(len(create_env_fn)):
-            if not isinstance(create_env_fn[i], (EnvBase, EnvCreator)):
-                create_env_fn[i] = EnvCreator(create_env_fn[i])
+        if use_env_creator:
+            for i in range(len(create_env_fn)):
+                if not isinstance(create_env_fn[i], (EnvBase, EnvCreator)):
+                    create_env_fn[i] = EnvCreator(create_env_fn[i])
 
         # If ray available, try to connect to an existing Ray cluster or start one and connect to it.
         if not _has_ray: