init

pytorch · vmoens · Jan 29, 2024 · Jan 29, 2024 · Jan 29, 2024 · Jan 29, 2024
commit 19c379fd61353bdb5fe00c1d0fe37484c823a8e3
diff --git a/examples/a2c/utils_atari.py b/examples/a2c/utils_atari.py
@@ -61,7 +61,7 @@ def make_base_env(
 
 def make_parallel_env(env_name, num_envs, device, is_test=False):
     env = ParallelEnv(
-        num_envs, EnvCreator(lambda: make_base_env(env_name, device=device))
+        num_envs, EnvCreator(lambda: make_base_env(env_name, device=device)), serial_for_single=True,
     )
     env = TransformedEnv(env)
     env.append_transform(ToTensorImage())

diff --git a/examples/cql/utils.py b/examples/cql/utils.py
@@ -79,7 +79,7 @@ def make_environment(cfg, train_num_envs=1, eval_num_envs=1):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         train_num_envs,
-        EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+        EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -88,7 +88,7 @@ def make_environment(cfg, train_num_envs=1, eval_num_envs=1):
     eval_env = TransformedEnv(
         ParallelEnv(
             eval_num_envs,
-            EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+            EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/examples/ddpg/utils.py b/examples/ddpg/utils.py
@@ -75,7 +75,7 @@ def make_environment(cfg):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         cfg.collector.env_per_collector,
-        EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+        EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -86,7 +86,7 @@ def make_environment(cfg):
     eval_env = TransformedEnv(
         ParallelEnv(
             cfg.collector.env_per_collector,
-            EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+            EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/examples/decision_transformer/utils.py b/examples/decision_transformer/utils.py
@@ -142,7 +142,7 @@ def make_env():
             return make_base_env(env_cfg)
 
     env = make_transformed_env(
-        ParallelEnv(num_envs, EnvCreator(make_env)),
+        ParallelEnv(num_envs, EnvCreator(make_env), serial_for_single=True),
         env_cfg,
         obs_loc,
         obs_std,

diff --git a/examples/discrete_sac/utils.py b/examples/discrete_sac/utils.py
@@ -76,7 +76,7 @@ def make_environment(cfg):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         cfg.collector.env_per_collector,
-        EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+        EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -87,7 +87,7 @@ def make_environment(cfg):
     eval_env = TransformedEnv(
         ParallelEnv(
             cfg.collector.env_per_collector,
-            EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+            EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/examples/distributed/collectors/single_machine/generic.py b/examples/distributed/collectors/single_machine/generic.py
@@ -100,7 +100,7 @@ def gym_make():
     if args.worker_parallelism == "collector" or num_workers == 1:
         action_spec = make_env().action_spec
     else:
-        make_env = ParallelEnv(num_workers, make_env)
+        make_env = ParallelEnv(num_workers, make_env, serial_for_single=True,)
         action_spec = make_env.action_spec
 
     if args.worker_parallelism == "collector" and num_workers > 1:

diff --git a/examples/distributed/collectors/single_machine/rpc.py b/examples/distributed/collectors/single_machine/rpc.py
@@ -96,7 +96,7 @@ def gym_make():
     if num_workers == 1:
         action_spec = make_env().action_spec
     else:
-        make_env = ParallelEnv(num_workers, make_env)
+        make_env = ParallelEnv(num_workers, make_env, serial_for_single=True,)
         action_spec = make_env.action_spec
 
     collector = RPCDataCollector(

diff --git a/examples/distributed/collectors/single_machine/sync.py b/examples/distributed/collectors/single_machine/sync.py
@@ -95,7 +95,7 @@ def gym_make():
     if args.worker_parallelism == "collector" or num_workers == 1:
         action_spec = make_env().action_spec
     else:
-        make_env = ParallelEnv(num_workers, make_env)
+        make_env = ParallelEnv(num_workers, make_env, serial_for_single=True,)
         action_spec = make_env.action_spec
 
     if args.worker_parallelism == "collector" and num_workers > 1:

diff --git a/examples/dreamer/dreamer_utils.py b/examples/dreamer/dreamer_utils.py
@@ -269,7 +269,7 @@ def parallel_env_constructor(
         create_env_fn=make_transformed_env,
         create_env_kwargs=None,
         pin_memory=cfg.pin_memory,
-        device=cfg.collector_device,
+        device=cfg.collector_device, serial_for_single=True,
     )
     if batch_transform:
         kwargs.update(

diff --git a/examples/iql/utils.py b/examples/iql/utils.py
@@ -83,7 +83,7 @@ def make_environment(cfg, train_num_envs=1, eval_num_envs=1):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         train_num_envs,
-        EnvCreator(lambda: env_maker(cfg)),
+        EnvCreator(lambda: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -92,7 +92,7 @@ def make_environment(cfg, train_num_envs=1, eval_num_envs=1):
     eval_env = TransformedEnv(
         ParallelEnv(
             eval_num_envs,
-            EnvCreator(lambda: env_maker(cfg)),
+            EnvCreator(lambda: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/examples/ppo/utils_atari.py b/examples/ppo/utils_atari.py
@@ -60,7 +60,7 @@ def make_base_env(
 
 def make_parallel_env(env_name, num_envs, device, is_test=False):
     env = ParallelEnv(
-        num_envs, EnvCreator(lambda: make_base_env(env_name, device=device))
+        num_envs, EnvCreator(lambda: make_base_env(env_name, device=device)), serial_for_single=True,
     )
     env = TransformedEnv(env)
     env.append_transform(ToTensorImage())

diff --git a/examples/redq/utils.py b/examples/redq/utils.py
@@ -746,7 +746,7 @@ def parallel_env_constructor(
     parallel_env = ParallelEnv(
         num_workers=cfg.collector.env_per_collector,
         create_env_fn=make_transformed_env,
-        create_env_kwargs=None,
+        create_env_kwargs=None, serial_for_single=True,
         pin_memory=False,
     )
     if batch_transform:

diff --git a/examples/sac/utils.py b/examples/sac/utils.py
@@ -67,7 +67,7 @@ def make_environment(cfg):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         cfg.collector.env_per_collector,
-        EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+        EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -76,7 +76,7 @@ def make_environment(cfg):
     eval_env = TransformedEnv(
         ParallelEnv(
             cfg.collector.env_per_collector,
-            EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+            EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/examples/td3/utils.py b/examples/td3/utils.py
@@ -77,7 +77,7 @@ def make_environment(cfg):
     """Make environments for training and evaluation."""
     parallel_env = ParallelEnv(
         cfg.collector.env_per_collector,
-        EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+        EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
     )
     parallel_env.set_seed(cfg.env.seed)
 
@@ -88,7 +88,7 @@ def make_environment(cfg):
     eval_env = TransformedEnv(
         ParallelEnv(
             cfg.collector.env_per_collector,
-            EnvCreator(lambda cfg=cfg: env_maker(cfg)),
+            EnvCreator(lambda cfg=cfg: env_maker(cfg)), serial_for_single=True,
         ),
         train_env.transform.clone(),
     )

diff --git a/test/test_env.py b/test/test_env.py
@@ -395,6 +395,17 @@ def test_parallel_devices(self, parallel, hetero, pdevice, edevice, bwad):
                 env.shared_tensordict_parent.device.type == torch.device(edevice).type
             )
 
+    def test_serial_for_single(self):
+        env = ParallelEnv(1, lambda: ContinuousActionVecMockEnv, serial_for_single=True)
+        assert isinstance(env, SerialEnv)
+        env.close()
+        env = ParallelEnv(1, lambda: ContinuousActionVecMockEnv)
+        assert isinstance(env, ParallelEnv)
+        env.close()
+        env = ParallelEnv(2, lambda: ContinuousActionVecMockEnv, serial_for_single=True)
+        assert isinstance(env, ParallelEnv)
+        env.close()
+
     @pytest.mark.parametrize("num_parallel_env", [1, 10])
     @pytest.mark.parametrize("env_batch_size", [[], (32,), (32, 1), (32, 0)])
     def test_env_with_batch_size(self, num_parallel_env, env_batch_size):

diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -25,7 +25,7 @@
 from torchrl._utils import _check_for_faulty_process, _ProcessNoWarn, VERBOSE
 from torchrl.data.tensor_specs import CompositeSpec
 from torchrl.data.utils import CloudpickleWrapper, contains_lazy_spec, DEVICE_TYPING
-from torchrl.envs.common import EnvBase
+from torchrl.envs.common import EnvBase, _EnvPostInit
 from torchrl.envs.env_creator import get_env_metadata
 
 # legacy
@@ -103,6 +103,18 @@ def new_fun(self, *args, **kwargs):
 
     return new_fun
 
+class _PEnvMeta(_EnvPostInit):
+    def __call__(cls, *args, **kwargs):
+        serial_for_single = kwargs.get("serial_for_single", False)
+        if serial_for_single:
+            num_workers = kwargs.get("num_workers", None)
+            if num_workers is None:
+                num_workers = args[0]
+            if num_workers == 1:
+                # We still use a serial to keep the shape unchanged
+                return SerialEnv(*args, **kwargs)
+        instance: EnvBase = super().__call__(*args, **kwargs)
+
 
 class _BatchedEnv(EnvBase):
     """Batched environments allow the user to query an arbitrary method / attribute of the environment running remotely.
@@ -120,6 +132,8 @@ class _BatchedEnv(EnvBase):
             If a single task is used, a callable should be used and not a list of identical callables:
             if a list of callable is provided, the environment will be executed as if multiple, diverse tasks were
             needed, which comes with a slight compute overhead;
+
+    Keyword Args:
         create_env_kwargs (dict or list of dicts, optional): kwargs to be used with the environments being created;
         share_individual_td (bool, optional): if ``True``, a different tensordict is created for every process/worker and a lazy
             stack is returned.
@@ -147,7 +161,9 @@ class _BatchedEnv(EnvBase):
             Defaults to 1 for safety: if none is indicated, launching multiple
             workers may charge the cpu load too much and harm performance.
             This parameter has no effect for the :class:`~SerialEnv` class.
-
+        serial_for_single (bool, optional): if ``True``, creating a parallel environment
+            with a single worker will return a :class:`~SerialEnv` instead.
+            This option has no effect with :class:`~SerialEnv`. Defaults to ``False``.
     """
 
     _verbose: bool = VERBOSE
@@ -162,6 +178,7 @@ def __init__(
         self,
         num_workers: int,
         create_env_fn: Union[Callable[[], EnvBase], Sequence[Callable[[], EnvBase]]],
+        *,
         create_env_kwargs: Union[dict, Sequence[dict]] = None,
         pin_memory: bool = False,
         share_individual_td: Optional[bool] = None,
@@ -172,8 +189,10 @@ def __init__(
         allow_step_when_done: bool = False,
         num_threads: int = None,
         num_sub_threads: int = 1,
+        serial_for_single: bool = False,
     ):
         super().__init__(device=device)
+        self.serial_for_single = serial_for_single
         self.is_closed = True
         if num_threads is None:
             num_threads = num_workers + 1  # 1 more thread for this proc
@@ -760,7 +779,7 @@ def to(self, device: DEVICE_TYPING):
         return self
 
 
-class ParallelEnv(_BatchedEnv):
+class ParallelEnv(_BatchedEnv, metaclass=_PEnvMeta):
     """Creates one environment per process.
 
     TensorDicts are passed via shared memory or memory map.