[BugFix] deepcopy specs before transforming (#461)

Author: vmoens

torchrl/envs/transforms/r3m.py

@@ -87,6 +87,7 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
         keys = [key for key in observation_spec._specs.keys() if key in self.keys_in]
         device = observation_spec[keys[0]].device
 
+        observation_spec = CompositeSpec(**observation_spec)
         if self.del_keys:
             for key_in in keys:
                 del observation_spec[key_in]
@@ -272,7 +273,7 @@ def _init(self):
                 model_name=model_name,
                 del_keys=True,
             )
-            transforms = [*transforms, normalize, network]
+            transforms = [*transforms, network]
 
         for transform in transforms:
             self.append(transform)

torchrl/envs/transforms/transforms.py

@@ -72,7 +72,7 @@
 def _apply_to_composite(function):
     def new_fun(self, observation_spec):
         if isinstance(observation_spec, CompositeSpec):
-            d = copy(observation_spec._specs)
+            d = observation_spec._specs
             for key_in, key_out in zip(self.keys_in, self.keys_out):
                 if key_in in observation_spec.keys():
                     d[key_out] = function(self, observation_spec[key_in])
@@ -506,7 +506,9 @@ def __getattr__(self, attr: str) -> Any:
         )
 
     def __repr__(self) -> str:
-        return f"TransformedEnv(env={self.base_env}, transform={self.transform})"
+        env_str = indent(f"env={self.base_env}", 4 * " ")
+        t_str = indent(f"transform={self.transform}", 4 * " ")
+        return f"TransformedEnv(\n{env_str},\n{t_str})"
 
     def _erase_metadata(self):
         if self.cache_specs:
@@ -621,7 +623,9 @@ def __getitem__(self, item: Union[int, slice, List]) -> Union:
         transform = self.transforms
         transform = transform[item]
         if not isinstance(transform, Transform):
-            return Compose(*self.transforms[item])
+            out = Compose(*self.transforms[item])
+            out.set_parent(self.parent)
+            return out
         return transform
 
     def dump(self, **kwargs) -> None:
@@ -737,7 +741,7 @@ def _apply_transform(self, observation: torch.FloatTensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
-        self._pixel_observation(observation_spec)
+        observation_spec = self._pixel_observation(deepcopy(observation_spec))
         observation_spec.shape = torch.Size(
             [
                 *observation_spec.shape[:-3],
@@ -747,13 +751,13 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
             ]
         )
         observation_spec.dtype = self.dtype
-        observation_spec = observation_spec
         return observation_spec
 
     def _pixel_observation(self, spec: TensorSpec) -> None:
-        if isinstance(spec, BoundedTensorSpec):
+        if isinstance(spec.space, ContinuousBox):
             spec.space.maximum = self._apply_transform(spec.space.maximum)
             spec.space.minimum = self._apply_transform(spec.space.minimum)
+        return spec
 
 
 class RewardClipping(Transform):
@@ -899,6 +903,7 @@ def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec = deepcopy(observation_spec)
         space = observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)
@@ -962,7 +967,8 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
                 }
             )
         else:
-            _observation_spec = observation_spec
+            _observation_spec = deepcopy(observation_spec)
+
         space = _observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)
@@ -1019,6 +1025,7 @@ def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec = deepcopy(observation_spec)
         space = observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)
@@ -1122,25 +1129,26 @@ def _transform_spec(self, spec: TensorSpec) -> None:
                 spec.shape = space.minimum.shape
             else:
                 spec.shape = self._apply_transform(torch.zeros(spec.shape)).shape
+        return spec
 
     def transform_action_spec(self, action_spec: TensorSpec) -> TensorSpec:
         if "action" in self.keys_inv_in:
-            self._transform_spec(action_spec)
+            action_spec = self._transform_spec(deepcopy(action_spec))
         return action_spec
 
     def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
         for key in self.keys_inv_in:
-            self._transform_spec(input_spec[key])
+            input_spec = self._transform_spec(deepcopy(input_spec[key]))
         return input_spec
 
     def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
         if "reward" in self.keys_in:
-            self._transform_spec(reward_spec)
+            reward_spec = self._transform_spec(deepcopy(reward_spec))
         return reward_spec
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
-        self._transform_spec(observation_spec)
+        observation_spec = self._transform_spec(deepcopy(observation_spec))
         return observation_spec
 
     def __repr__(self) -> str:
@@ -1207,6 +1215,7 @@ def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec = deepcopy(observation_spec)
         space = observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)
@@ -1295,6 +1304,7 @@ def _apply_transform(self, obs: torch.Tensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        observation_spec = deepcopy(observation_spec)
         space = observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)