Repository: pytorch/rl

Issue:
[BUG] `TransformedEnv` should call `transform.eval()` when a transform is added
Currently, `TransformedEnv` adds transforms "as-is", and only casts them to device. 
We should ensure that transforms that are added once the module is instantiated are being set to eval mode (if needed).

## To reproduce
```python
base_env = GymEnv("Pendulum-v1")
r = RewardScaling(0, 1)
t = TransformedEnv(base_env, r)
assert not t.transform.training  # passes
t.append(RewardScaling)
assert not t.transform[1].training  # breaks
```

## Solution
Since `append_transform` and `insert_transform` untimately call `Compose.append` and `Compose.insert`, modifying those methods in `Compose` should be sufficient.
We'd need to make sure that (1) When `Compose` is instantiated, it is automatically set to eval mode (and add that to the docstring!), (2) that `append` and `insert` both turn the transform to eval mode _if needed_

## Testing
New tests should be written in [test/test_transform.py](test/test_transform.py) to make sure that this behaviour is appropriate. We should test that all transforms are in eval if `train()` is not called, that they're all in training mode if `train` is called, that this is the case if the transformed env is created like this:
```python
base_env = GymEnv("Pendulum-v1")
t = TransformedEnv(base_env)
```
or this
```python
base_env = GymEnv("Pendulum-v1")
r = RewardScaling(0, 1)
t = TransformedEnv(base_env, r)
```



Discussion:

Resolving patch:
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
--- a/torchrl/envs/transforms/transforms.py
+++ b/torchrl/envs/transforms/transforms.py
@@ -266,8 +266,8 @@ class TransformedEnv(_EnvClass):
     Args:
         env (_EnvClass): original environment to be transformed_in.
         transform (Transform, optional): transform to apply to the tensordict resulting
-            from `env.step(td)`. If none is provided, an empty Compose placeholder is
-            used.
+            from `env.step(td)`. If none is provided, an empty Compose
+            placeholder in an eval mode is used.
         cache_specs (bool, optional): if True, the specs will be cached once
             and for all after the first call (i.e. the specs will be
             transformed_in only once). If the transform changes during
@@ -299,6 +299,7 @@ def __init__(
             transform.set_parent(self)
         else:
             transform = transform.to(device)
+        transform.eval()
         self.transform = transform
 
         self._last_obs = None
@@ -435,7 +436,7 @@ def append_transform(self, transform: Transform) -> None:
     def insert_transform(self, index: int, transform: Transform) -> None:
         if not isinstance(transform, Transform):
             raise ValueError(
-                "TransformedEnv.append_transform expected a transform but received an object of "
+                "TransformedEnv.insert_transform expected a transform but received an object of "
                 f"type {type(transform)} instead."
             )
         transform = transform.to(self.device)
@@ -599,6 +600,7 @@ def append(self, transform):
                 "Compose.append expected a transform but received an object of "
                 f"type {type(transform)} instead."
             )
+        transform.eval()
         self.transforms.append(transform)
         transform.set_parent(self)
 
@@ -617,6 +619,7 @@ def insert(self, index: int, transform: Transform) -> None:
         self.empty_cache()
         if index < 0:
             index = index + len(self.transforms)
+        transform.eval()
         self.transforms.insert(index, transform)
         transform.set_parent(self)
 
@@ -1666,6 +1669,9 @@ class VecNorm(Transform):
     the normalization layer is queried it will update the values for all
     processes that share the same reference.
 
+    To use VecNorm at inference time and avoid updating the values with the new
+    observations, one should substitute this layer by `vecnorm.to_observation_norm()`.
+
     Args:
         keys_in (iterable of str, optional): keys to be updated.
             default: ["next_observation", "reward"]
@@ -1768,32 +1774,52 @@ def _update(self, key, value, N) -> torch.Tensor:
         _ssq = self._td.get(key + "_ssq")
         _count = self._td.get(key + "_count")
 
-        if self.training:
-            _sum = self._td.get(key + "_sum")
-            value_sum = _sum_left(value, _sum)
-            _sum *= self.decay
-            _sum += value_sum
-            self._td.set_(key + "_sum", _sum, no_check=True)
+        _sum = self._td.get(key + "_sum")
+        value_sum = _sum_left(value, _sum)
+        _sum *= self.decay
+        _sum += value_sum
+        self._td.set_(key + "_sum", _sum, no_check=True)
 
-            _ssq = self._td.get(key + "_ssq")
-            value_ssq = _sum_left(value.pow(2), _ssq)
-            _ssq *= self.decay
-            _ssq += value_ssq
-            self._td.set_(key + "_ssq", _ssq, no_check=True)
+        _ssq = self._td.get(key + "_ssq")
+        value_ssq = _sum_left(value.pow(2), _ssq)
+        _ssq *= self.decay
+        _ssq += value_ssq
+        self._td.set_(key + "_ssq", _ssq, no_check=True)
 
-            _count = self._td.get(key + "_count")
-            _count *= self.decay
-            _count += N
-            self._td.set_(key + "_count", _count, no_check=True)
-        else:
-            _sum = self._td.get(key + "_sum")
-            _ssq = self._td.get(key + "_ssq")
-            _count = self._td.get(key + "_count")
+        _count = self._td.get(key + "_count")
+        _count *= self.decay
+        _count += N
+        self._td.set_(key + "_count", _count, no_check=True)
 
         mean = _sum / _count
         std = (_ssq / _count - mean.pow(2)).clamp_min(self.eps).sqrt()
         return (value - mean) / std.clamp_min(self.eps)
 
+    def to_observation_norm(self) -> Union[Compose, ObservationNorm]:
+        """Converts VecNorm into an ObservationNorm class that can be used
+        at inference time.
+
+        """
+        out = []
+        for key in self.keys_in:
+            _sum = self._td.get(key + "_sum")
+            _ssq = self._td.get(key + "_ssq")
+            _count = self._td.get(key + "_count")
+            mean = _sum / _count
+            std = (_ssq / _count - mean.pow(2)).clamp_min(self.eps).sqrt()
+
+            _out = ObservationNorm(
+                loc=mean,
+                scale=std,
+                standard_normal=True,
+                keys_in=self.keys_in,
+            )
+            if len(self.keys_in) == 1:
+                return _out
+            else:
+                out += ObservationNorm
+        return Compose(*out)
+
     @staticmethod
     def build_td_for_shared_vecnorm(
         env: _EnvClass,
diff --git a/torchrl/version.py b/torchrl/version.py
new file mode 100644
--- /dev/null
+++ b/torchrl/version.py
@@ -0,0 +1,2 @@
+__version__ = "0.0.1c0+00ddc39"
+git_version = "00ddc39f2b2c5fcbcea6f1fc310a3a8e98069213"


Test patch:
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -356,6 +356,29 @@ def test_vecnorm(parallel, thr=0.2, N=200):  # 10000):
         env_t.close()
 
 
+def test_added_transforms_are_in_eval_mode_trivial():
+    base_env = ContinuousActionVecMockEnv()
+    t = TransformedEnv(base_env)
+    assert not t.transform.training
+
+    t.train()
+    assert t.transform.training
+
+
+def test_added_transforms_are_in_eval_mode():
+    base_env = ContinuousActionVecMockEnv()
+    r = RewardScaling(0, 1)
+    t = TransformedEnv(base_env, r)
+    assert not t.transform.training
+    t.append_transform(RewardScaling(0, 1))
+    assert not t.transform[1].training
+
+    t.train()
+    assert t.transform.training
+    assert t.transform[0].training
+    assert t.transform[1].training
+
+
 class TestTransforms:
     @pytest.mark.skipif(not _has_tv, reason="no torchvision")
     @pytest.mark.parametrize("interpolation", ["bilinear", "bicubic"])