[BugFix] Prevent transform parent from being reassigned (#641)

* Prevent transform parent from being reassigned

* Added missed resets, fixed flake8 error

* Prevent transform.parent from modifying original transforms

* Added test to ensure Transform.parent doesn't alter original _parents

Co-authored-by: Jason Kutarnia <jasonkut@fb.com>
Co-authored-by: vmoens <vincentmoens@gmail.com>

Author: 3 people

test/test_transforms.py

@@ -327,6 +327,27 @@ def test_nested_transformed_env():
     assert children[1] == t2
 
 
+def test_transform_parent():
+    base_env = ContinuousActionVecMockEnv()
+    t1 = RewardScaling(0, 1)
+    t2 = RewardScaling(0, 2)
+    env = TransformedEnv(TransformedEnv(base_env, t1), t2)
+    t3 = RewardClipping(0.1, 0.5)
+    env.append_transform(t3)
+
+    t1_parent_gt = t1._parent
+    t2_parent_gt = t2._parent
+    t3_parent_gt = t3._parent
+
+    _ = t1.parent
+    _ = t2.parent
+    _ = t3.parent
+
+    assert t1_parent_gt == t1._parent
+    assert t2_parent_gt == t2._parent
+    assert t3_parent_gt == t3._parent
+
+
 class TestTransforms:
     @pytest.mark.skipif(not _has_tv, reason="no torchvision")
     @pytest.mark.parametrize("interpolation", ["bilinear", "bicubic"])

torchrl/envs/transforms/transforms.py

@@ -204,8 +204,13 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}(keys={self.keys_in})"
 
     def set_parent(self, parent: Union[Transform, EnvBase]) -> None:
+        if self.__dict__["_parent"] is not None:
+            raise AttributeError("parent of transform already set")
         self.__dict__["_parent"] = parent
 
+    def reset_parent(self) -> None:
+        self.__dict__["_parent"] = None
+
     @property
     def parent(self) -> EnvBase:
         if not hasattr(self, "_parent"):
@@ -226,15 +231,20 @@ def parent(self) -> EnvBase:
                 raise ValueError(
                     f"Compose parent was of type {type(compose_parent)} but expected TransformedEnv."
                 )
+            if compose_parent.transform is not compose:
+                comp_parent_trans = copy(compose_parent.transform)
+                comp_parent_trans.reset_parent()
+            else:
+                comp_parent_trans = None
             out = TransformedEnv(
                 compose_parent.base_env,
-                transform=compose_parent.transform
-                if compose_parent.transform is not compose
-                else None,
+                transform=comp_parent_trans,
             )
-            for transform in compose.transforms:
-                if transform is self:
+            for orig_trans in compose.transforms:
+                if orig_trans is self:
                     break
+                transform = copy(orig_trans)
+                transform.reset_parent()
                 out.append_transform(transform)
         elif isinstance(parent, TransformedEnv):
             out = TransformedEnv(parent.base_env)
@@ -287,9 +297,16 @@ def __init__(
                 # we don't use isinstance as some transforms may be subclassed from
                 # Compose but with other features that we don't want to loose.
                 transform = [transform]
+            else:
+                for t in transform:
+                    t.reset_parent()
             env_transform = env.transform
             if type(env_transform) is not Compose:
+                env_transform.reset_parent()
                 env_transform = [env_transform]
+            else:
+                for t in env_transform:
+                    t.reset_parent()
             transform = Compose(*env_transform, *transform).to(device)
         else:
             self._set_env(env, device)
@@ -474,9 +491,10 @@ def append_transform(self, transform: Transform) -> None:
         transform = transform.to(self.device)
         if not isinstance(self.transform, Compose):
             prev_transform = self.transform
+            prev_transform.reset_parent()
             self.transform = Compose()
             self.transform.append(prev_transform)
-            self.transform.set_parent(self)
+
         self.transform.append(transform)
 
     def insert_transform(self, index: int, transform: Transform) -> None:
@@ -538,8 +556,6 @@ def to(self, device: DEVICE_TYPING) -> TransformedEnv:
     def __setattr__(self, key, value):
         propobj = getattr(self.__class__, key, None)
 
-        if isinstance(value, Transform):
-            value.set_parent(self)
         if isinstance(propobj, property):
             ancestors = list(__class__.__mro__)[::-1]
             while isinstance(propobj, property):