Update

[ghstack-poisoned]

Author: Vincent Moens

test/test_transforms.py

@@ -14590,7 +14590,6 @@ def test_trans_serial_env_check(self):
         base_env = SerialEnv(
             3,
             [partial(CountingEnv, 6), partial(CountingEnv, 7), partial(CountingEnv, 8)],
-            batch_locked=False,
         )
         condition = lambda td: ((td.get("step_count") % 2) == 0).squeeze(-1)
         policy_odd = self._create_policy_odd(base_env)
@@ -14604,7 +14603,6 @@ def test_trans_parallel_env_check(self):
         base_env = ParallelEnv(
             3,
             [partial(CountingEnv, 6), partial(CountingEnv, 7), partial(CountingEnv, 8)],
-            batch_locked=False,
             mp_start_method=mp_ctx,
         )
         condition = lambda td: ((td.get("step_count") % 2) == 0).squeeze(-1)

torchrl/envs/__init__.py

@@ -58,6 +58,7 @@
     CenterCrop,
     ClipTransform,
     Compose,
+    ConditionalPolicySwitch,
     ConditionalSkip,
     Crop,
     DeviceCastTransform,
@@ -137,6 +138,7 @@
     "AutoResetTransform",
     "AsyncEnvPool",
     "ProcessorAsyncEnvPool",
+    "ConditionalPolicySwitch",
     "ThreadingAsyncEnvPool",
     "BatchSizeTransform",
     "BinarizeReward",

torchrl/envs/transforms/__init__.py

@@ -20,6 +20,7 @@
     CenterCrop,
     ClipTransform,
     Compose,
+    ConditionalPolicySwitch,
     ConditionalSkip,
     Crop,
     DeviceCastTransform,
@@ -83,6 +84,7 @@
     "CatFrames",
     "CatTensors",
     "CenterCrop",
+    "ConditionalPolicySwitch",
     "ClipTransform",
     "Compose",
     "ConditionalSkip",

torchrl/envs/transforms/transforms.py

@@ -11398,12 +11398,6 @@ def _reset(
             parent: TransformedEnv = self.parent
             reset_td_save = None
             if not cond.all():
-                if parent.base_env.batch_locked:
-                    raise RuntimeError(
-                        "Cannot run partial steps in a batched locked environment. "
-                        "Hint: Parallel and Serial envs can be unlocked through a keyword argument in "
-                        "the constructor."
-                    )
                 reset_td_save = tensordict_reset.copy()
                 tensordict_reset = tensordict_reset[cond]
                 tensordict = tensordict[cond]