[Refactor] Removing inplace transform attribute (#871)

Author: vmoens

test/test_transforms.py

@@ -1259,31 +1259,6 @@ def test_observationnorm_uninitialized_stats_error(self):
         with pytest.raises(RuntimeError, match=err_msg):
             transform._apply_transform(torch.Tensor([1]))
 
-    @pytest.mark.parametrize("device", get_available_devices())
-    def test_observationnorm_infinite_stats_error(self, device):
-        base_env = ContinuousActionVecMockEnv(
-            observation_spec=CompositeSpec(
-                observation=BoundedTensorSpec(
-                    minimum=1, maximum=1, shape=torch.Size([1])
-                ),
-                observation_orig=BoundedTensorSpec(
-                    minimum=1, maximum=1, shape=torch.Size([1])
-                ),
-            ),
-            action_spec=BoundedTensorSpec(minimum=1, maximum=1, shape=torch.Size((1,))),
-            seed=0,
-        )
-        base_env.out_key = "observation"
-        t_env = TransformedEnv(
-            base_env,
-            transform=ObservationNorm(in_keys="observation"),
-        )
-        t_env.append_transform(ObservationNorm(in_keys="observation"))
-        err_msg = "Non-finite values found in"
-        with pytest.raises(RuntimeError, match=err_msg):
-            for transform in t_env.transform:
-                transform.init_stats(num_iter=100)
-
     def test_catframes_transform_observation_spec(self):
         N = 4
         key1 = "first key"
@@ -1649,7 +1624,6 @@ def test_binarized_reward(self, device, batch):
             device=device,
         )
         br(td)
-        assert td["reward"] is reward
         assert (td["reward"] != reward_copy).all()
         assert (td["misc"] == misc_copy).all()
         assert (torch.count_nonzero(td["reward"]) == torch.sum(reward_copy > 0)).all()

torchrl/envs/transforms/transforms.py

@@ -113,13 +113,6 @@ def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
         """Resets a tranform if it is stateful."""
         return tensordict
 
-    def _check_inplace(self) -> None:
-        if not hasattr(self, "inplace"):
-            raise AttributeError(
-                f"Transform of class {self.__class__.__name__} has no "
-                f"attribute inplace, consider implementing it."
-            )
-
     def init(self, tensordict) -> None:
         pass
 
@@ -134,11 +127,13 @@ def _apply_transform(self, obs: torch.Tensor) -> None:
 
     def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
         """Reads the input tensordict, and for the selected keys, applies the transform."""
-        self._check_inplace()
         for in_key, out_key in zip(self.in_keys, self.out_keys):
             if in_key in tensordict.keys(include_nested=True):
                 observation = self._apply_transform(tensordict.get(in_key))
-                tensordict.set(out_key, observation, inplace=self.inplace)
+                tensordict.set(
+                    out_key,
+                    observation,
+                )
         return tensordict
 
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
@@ -160,11 +155,13 @@ def _inv_apply_transform(self, obs: torch.Tensor) -> torch.Tensor:
             return obs
 
     def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
-        self._check_inplace()
         for in_key, out_key in zip(self.in_keys_inv, self.out_keys_inv):
             if in_key in tensordict.keys(include_nested=True):
                 observation = self._inv_apply_transform(tensordict.get(in_key))
-                tensordict.set(out_key, observation, inplace=self.inplace)
+                tensordict.set(
+                    out_key,
+                    observation,
+                )
         return tensordict
 
     def inv(self, tensordict: TensorDictBase) -> TensorDictBase:
@@ -607,8 +604,6 @@ def __del__(self):
 class ObservationTransform(Transform):
     """Abstract class for transformations of the observations."""
 
-    inplace = False
-
     def __init__(
         self,
         in_keys: Optional[Sequence[str]] = None,
@@ -634,8 +629,6 @@ class Compose(Transform):
 
     """
 
-    inplace = False
-
     def __init__(self, *transforms: Transform):
         super().__init__(in_keys=[])
         self.transforms = nn.ModuleList(transforms)
@@ -773,8 +766,6 @@ class ToTensorImage(ObservationTransform):
         torch.Size([1, 1, 3, 10, 11]) torch.float32
     """
 
-    inplace = False
-
     def __init__(
         self,
         unsqueeze: bool = False,
@@ -827,8 +818,6 @@ class RewardClipping(Transform):
 
     """
 
-    inplace = True
-
     def __init__(
         self,
         clamp_min: float = None,
@@ -850,11 +839,11 @@ def __init__(
 
     def _apply_transform(self, reward: torch.Tensor) -> torch.Tensor:
         if self.clamp_max is not None and self.clamp_min is not None:
-            reward = reward.clamp_(self.clamp_min, self.clamp_max)
+            reward = reward.clamp(self.clamp_min, self.clamp_max)
         elif self.clamp_min is not None:
-            reward = reward.clamp_min_(self.clamp_min)
+            reward = reward.clamp_min(self.clamp_min)
         elif self.clamp_max is not None:
-            reward = reward.clamp_max_(self.clamp_max)
+            reward = reward.clamp_max(self.clamp_max)
         return reward
 
     def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
@@ -884,8 +873,6 @@ def __repr__(self) -> str:
 class BinarizeReward(Transform):
     """Maps the reward to a binary value (0 or 1) if the reward is null or non-null, respectively."""
 
-    inplace = True
-
     def __init__(
         self,
         in_keys: Optional[Sequence[str]] = None,
@@ -917,8 +904,6 @@ class Resize(ObservationTransform):
         interpolation (str): interpolation method
     """
 
-    inplace = False
-
     def __init__(
         self,
         w: int,
@@ -986,8 +971,6 @@ class CenterCrop(ObservationTransform):
         h (int, optional): resulting height. If None, then w is used (square crop).
     """
 
-    inplace = False
-
     def __init__(
         self,
         w: int,
@@ -1043,8 +1026,6 @@ class FlattenObservation(ObservationTransform):
         last_dim (int): last dimension of the dimensions to flatten.
     """
 
-    inplace = False
-
     def __init__(
         self,
         first_dim: int,
@@ -1115,7 +1096,6 @@ class UnsqueezeTransform(Transform):
     """
 
     invertible = True
-    inplace = False
 
     @classmethod
     def __new__(cls, *args, **kwargs):
@@ -1232,7 +1212,6 @@ class SqueezeTransform(UnsqueezeTransform):
     """
 
     invertible = True
-    inplace = False
 
     def __init__(
         self,
@@ -1269,8 +1248,6 @@ def inv(self, tensordict: TensorDictBase) -> TensorDictBase:
 class GrayScale(ObservationTransform):
     """Turns a pixel observation to grayscale."""
 
-    inplace = False
-
     def __init__(self, in_keys: Optional[Sequence[str]] = None):
         if in_keys is None:
             in_keys = IMAGE_KEYS
@@ -1342,8 +1319,6 @@ class ObservationNorm(ObservationTransform):
 
     """
 
-    inplace = True
-
     def __init__(
         self,
         loc: Optional[float, torch.Tensor] = None,
@@ -1471,7 +1446,6 @@ def raise_initialization_exception(module):
             raise RuntimeError("Non-finite values found in loc")
         if not torch.isfinite(scale).all():
             raise RuntimeError("Non-finite values found in scale")
-
         self.register_buffer("loc", loc)
         self.register_buffer("scale", scale.clamp_min(self.eps))
 
@@ -1662,8 +1636,6 @@ class RewardScaling(Transform):
             as it is done for standardization. Default is `False`.
     """
 
-    inplace = True
-
     def __init__(
         self,
         loc: Union[float, torch.Tensor],
@@ -1717,8 +1689,6 @@ def __repr__(self) -> str:
 class FiniteTensorDictCheck(Transform):
     """This transform will check that all the items of the tensordict are finite, and raise an exception if they are not."""
 
-    inplace = False
-
     def __init__(self):
         super().__init__(in_keys=[])
 
@@ -1741,7 +1711,6 @@ class DoubleToFloat(Transform):
     """
 
     invertible = True
-    inplace = False
 
     def __init__(
         self,
@@ -1835,7 +1804,6 @@ class CatTensors(Transform):
     """
 
     invertible = False
-    inplace = False
 
     def __init__(
         self,
@@ -1992,8 +1960,6 @@ class DiscreteActionProjection(Transform):
         tensor([1])
     """
 
-    inplace = False
-
     def __init__(self, max_n: int, m: int, action_key: str = "action"):
         super().__init__([action_key])
         self.max_n = max_n
@@ -2035,8 +2001,6 @@ class FrameSkipTransform(Transform):
 
     """
 
-    inplace = False
-
     def __init__(self, frame_skip: int = 1):
         super().__init__([])
         if frame_skip < 1:
@@ -2069,8 +2033,6 @@ class NoopResetEnv(Transform):
 
     """
 
-    inplace = True
-
     def __init__(self, noops: int = 30, random: bool = True):
         """Sample initial states by taking random number of no-ops on reset.
 
@@ -2169,8 +2131,6 @@ class TensorDictPrimer(Transform):
             is_shared=False)
     """
 
-    inplace = False
-
     def __init__(self, random=False, default_value=0.0, **kwargs):
         self.primers = kwargs
         self.random = random
@@ -2271,8 +2231,6 @@ class gSDENoise(Transform):
     See the :func:`~torchrl.modules.models.exploration.gSDEModule' for more info.
     """
 
-    inplace = False
-
     def __init__(
         self,
         state_dim=None,
@@ -2351,8 +2309,6 @@ class VecNorm(Transform):
 
     """
 
-    inplace = True
-
     def __init__(
         self,
         in_keys: Optional[Sequence[str]] = None,
@@ -2402,7 +2358,7 @@ def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
                 key, tensordict.get(key), N=max(1, tensordict.numel())
             )
 
-            tensordict.set_(key, new_val)
+            tensordict.set(key, new_val)
 
         if self.lock is not None:
             self.lock.release()
@@ -2582,8 +2538,6 @@ class RewardSum(Transform):
     this transform hos no effect.
     """
 
-    inplace = True
-
     def __init__(
         self,
         in_keys: Optional[Sequence[str]] = None,
@@ -2654,7 +2608,6 @@ def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
     def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
         """Updates the episode rewards with the step rewards."""
         # Sanity checks
-        self._check_inplace()
         for in_key in self.in_keys:
             if in_key not in tensordict.keys():
                 return tensordict
@@ -2727,7 +2680,6 @@ class StepCounter(Transform):
     """
 
     invertible = False
-    inplace = True
 
     def __init__(self, max_steps: Optional[int] = None):
         if max_steps is not None and max_steps < 1:
@@ -2799,8 +2751,6 @@ class ExcludeTransform(Transform):
 
     """
 
-    inplace = False
-
     def __init__(self, *excluded_keys):
         super().__init__(in_keys=[], in_keys_inv=[], out_keys=[], out_keys_inv=[])
         if not all(isinstance(item, str) for item in excluded_keys):
@@ -2840,8 +2790,6 @@ class SelectTransform(Transform):
 
     """
 
-    inplace = False
-
     def __init__(self, *selected_keys):
         super().__init__(in_keys=[], in_keys_inv=[], out_keys=[], out_keys_inv=[])
         if not all(isinstance(item, str) for item in selected_keys):
@@ -2887,7 +2835,6 @@ class TimeMaxPool(Transform):
         T (int, optional): Number of time steps over which to apply max pooling.
     """
 
-    inplace = False
     invertible = False
 
     def __init__(