[Feature] Too many deepcopy in transforms.py (#625)

* Remove optional deepcopy in tranform.py

* lint

* Refactoring/Renaming deepcopy cleanup related tests

* lint

Co-authored-by: vmoens <vincentmoens@gmail.com>

Author: romainjln

test/test_transforms.py

@@ -5,6 +5,7 @@
 import argparse
 from copy import copy, deepcopy
 
+import numpy as np
 import pytest
 import torch
 from _utils_internal import get_available_devices, retry
@@ -22,6 +23,7 @@
     NdUnboundedContinuousTensorSpec,
     TensorDict,
     UnboundedContinuousTensorSpec,
+    BoundedTensorSpec,
 )
 from torchrl.envs import (
     BinarizeReward,
@@ -313,6 +315,50 @@ def test_added_transforms_are_in_eval_mode():
     assert t.transform[1].training
 
 
+class TestTransformedEnv:
+    def test_independent_obs_specs_from_shared_env(self):
+        obs_spec = CompositeSpec(
+            next_observation=BoundedTensorSpec(minimum=0, maximum=10)
+        )
+        base_env = ContinuousActionVecMockEnv(observation_spec=obs_spec)
+        t1 = TransformedEnv(base_env, transform=ObservationNorm(loc=3, scale=2))
+        t2 = TransformedEnv(base_env, transform=ObservationNorm(loc=1, scale=6))
+
+        t1_obs_spec = t1.observation_spec
+        t2_obs_spec = t2.observation_spec
+
+        assert t1_obs_spec["next_observation"].space.minimum == 3
+        assert t1_obs_spec["next_observation"].space.maximum == 23
+
+        assert t2_obs_spec["next_observation"].space.minimum == 1
+        assert t2_obs_spec["next_observation"].space.maximum == 61
+
+        assert base_env.observation_spec["next_observation"].space.minimum == 0
+        assert base_env.observation_spec["next_observation"].space.maximum == 10
+
+    def test_independent_reward_specs_from_shared_env(self):
+        reward_spec = UnboundedContinuousTensorSpec()
+        base_env = ContinuousActionVecMockEnv(reward_spec=reward_spec)
+        t1 = TransformedEnv(
+            base_env, transform=RewardClipping(clamp_min=0, clamp_max=4)
+        )
+        t2 = TransformedEnv(
+            base_env, transform=RewardClipping(clamp_min=-2, clamp_max=2)
+        )
+
+        t1_reward_spec = t1.reward_spec
+        t2_reward_spec = t2.reward_spec
+
+        assert t1_reward_spec.space.minimum == 0
+        assert t1_reward_spec.space.maximum == 4
+
+        assert t2_reward_spec.space.minimum == -2
+        assert t2_reward_spec.space.maximum == 2
+
+        assert base_env.reward_spec.space.minimum == -np.inf
+        assert base_env.reward_spec.space.maximum == np.inf
+
+
 def test_nested_transformed_env():
     base_env = ContinuousActionVecMockEnv()
     t1 = RewardScaling(0, 1)

torchrl/envs/transforms/transforms.py

@@ -757,7 +757,7 @@ def _apply_transform(self, observation: torch.FloatTensor) -> torch.Tensor:
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
-        observation_spec = self._pixel_observation(deepcopy(observation_spec))
+        observation_spec = self._pixel_observation(observation_spec)
         observation_spec.shape = torch.Size(
             [
                 *observation_spec.shape[:-3],
@@ -913,7 +913,6 @@ 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)
@@ -970,20 +969,17 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
                     for key, _obs_spec in observation_spec._specs.items()
                 }
             )
-        else:
-            _observation_spec = deepcopy(observation_spec)
 
-        space = _observation_spec.space
+        space = observation_spec.space
         if isinstance(space, ContinuousBox):
             space.minimum = self._apply_transform(space.minimum)
             space.maximum = self._apply_transform(space.maximum)
-            _observation_spec.shape = space.minimum.shape
+            observation_spec.shape = space.minimum.shape
         else:
-            _observation_spec.shape = self._apply_transform(
-                torch.zeros(_observation_spec.shape)
+            observation_spec.shape = self._apply_transform(
+                torch.zeros(observation_spec.shape)
             ).shape
 
-        observation_spec = _observation_spec
         return observation_spec
 
     def __repr__(self) -> str:
@@ -1036,7 +1032,6 @@ def set_parent(self, parent: Union[Transform, EnvBase]) -> None:
 
     @_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):
@@ -1139,17 +1134,17 @@ def _transform_spec(self, spec: TensorSpec) -> None:
 
     def transform_input_spec(self, input_spec: TensorSpec) -> TensorSpec:
         for key in self.keys_inv_in:
-            input_spec = self._transform_spec(deepcopy(input_spec[key]))
+            input_spec = self._transform_spec(input_spec[key])
         return input_spec
 
     def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
         if "reward" in self.keys_in:
-            reward_spec = self._transform_spec(deepcopy(reward_spec))
+            reward_spec = self._transform_spec(reward_spec)
         return reward_spec
 
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
-        observation_spec = self._transform_spec(deepcopy(observation_spec))
+        observation_spec = self._transform_spec(observation_spec)
         return observation_spec
 
     def __repr__(self) -> str:
@@ -1213,7 +1208,6 @@ 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)
@@ -1303,7 +1297,6 @@ 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)
@@ -1756,9 +1749,8 @@ def _inv_apply_transform(self, action: torch.Tensor) -> torch.Tensor:
         return action
 
     def tranform_input_spec(self, input_spec: CompositeSpec):
-        input_spec_out = deepcopy(input_spec)
-        input_spec_out["action"] = self.transform_action_spec(input_spec_out["action"])
-        return input_spec_out
+        input_spec["action"] = self.transform_action_spec(input_spec["action"])
+        return input_spec
 
     def __repr__(self) -> str:
         return (