[BugFix] Compose cloning fix (#899)

Author: vmoens

test/test_transforms.py

@@ -6059,6 +6059,42 @@ def test_smoke_compose_transform(transform):
     Compose(transform())
 
 
+@pytest.mark.parametrize("transform", transforms)
+def test_clone_parent(transform):
+    base_env1 = ContinuousActionVecMockEnv()
+    base_env2 = ContinuousActionVecMockEnv()
+    env = TransformedEnv(base_env1, transform())
+    env_clone = TransformedEnv(base_env2, env.transform.clone())
+
+    assert env_clone.transform.parent.base_env is not base_env1
+    assert env_clone.transform.parent.base_env is base_env2
+    assert env.transform.parent.base_env is not base_env2
+    assert env.transform.parent.base_env is base_env1
+
+
+@pytest.mark.parametrize("transform", transforms)
+def test_clone_parent_compose(transform):
+    base_env1 = ContinuousActionVecMockEnv()
+    base_env2 = ContinuousActionVecMockEnv()
+    env = TransformedEnv(base_env1, Compose(ToTensorImage(), transform()))
+    t = env.transform.clone()
+
+    assert t.parent is None
+    assert t[0].parent is None
+    assert t[1].parent is None
+
+    env_clone = TransformedEnv(base_env2, Compose(ToTensorImage(), *t))
+
+    assert env_clone.transform[0].parent.base_env is not base_env1
+    assert env_clone.transform[0].parent.base_env is base_env2
+    assert env.transform[0].parent.base_env is not base_env2
+    assert env.transform[0].parent.base_env is base_env1
+    assert env_clone.transform[1].parent.base_env is not base_env1
+    assert env_clone.transform[1].parent.base_env is base_env2
+    assert env.transform[1].parent.base_env is not base_env2
+    assert env.transform[1].parent.base_env is base_env1
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/transforms/transforms.py

@@ -265,7 +265,10 @@ def reset_parent(self) -> None:
 
     def clone(self):
         self_copy = copy(self)
-        self_copy.reset_parent()
+        state = copy(self.__dict__)
+        state["_container"] = None
+        state["_parent"] = None
+        self_copy.__dict__.update(state)
         return self_copy
 
     @property
@@ -778,7 +781,11 @@ def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> Compose:
         return super().to(dest)
 
     def __iter__(self):
-        return iter(self.transforms)
+        # We clone the transforms to be able to do
+        # env2 = TransformedEnv(base_env, *env1.transform.clone())
+        # which will otherwise result in an error because all the transforms
+        # from the Compose already have a container.
+        yield from (t.clone() for t in self.transforms)
 
     def __len__(self):
         return len(self.transforms)
@@ -794,6 +801,17 @@ def empty_cache(self):
             t.empty_cache()
         super().empty_cache()
 
+    def reset_parent(self):
+        for t in self.transforms:
+            t.reset_parent()
+        super().reset_parent()
+
+    def clone(self):
+        transforms = []
+        for t in self.transforms:
+            transforms.append(t.clone())
+        return Compose(*transforms)
+
 
 class ToTensorImage(ObservationTransform):
     """Transforms a numpy-like image (3 x W x H) to a pytorch image (3 x W x H).