Repository: pytorch/rl

Issue:
[Feature Request] Create a tensorspec for the R3M output
## Motivation

R3M (torchrl.envs.transforms.r3m) is a generic transform layer that will embed an image in a tensor, easier to work with using a policy or value function.
In TorchRL, each env has am `observation_spec` attribute that contains `tensor_spec` that describe what the observations are (they behave like dictionaries, and all of their keys must start with `"next_"`).

Here is how the observation_spec should behave:
1. We should add one TensorSpec for the output of R3M (the vector). It should have the size that is expected given the resnet used (i.e. an unbounded continuous tensorspec of size 2048 IIRC for resnet50).
2. We should delete the tensorspec of the pixels if it disappears (which will happen unless `tensor_pixels_keys` is set). e.g. when doing `R3MTransform("resnet50", keys_in=["next_pixels"])`, we should not see the "next_pixels" key in the `observation_spec` of the TransformedEnv once R3M has been applied since this key-value pair has disappeared. Instead, we should see the spec of "next_r3m_vec".
3. If the value of "next_pixels" is kept somewhere (using `tensor_pixels_keys`) we should duplicate the value of the `"next_pixels"`'s TensorSpec onto the key(s) indicated by `tensor_pixels_keys`.






Discussion:

Resolving patch:
diff --git a/torchrl/envs/transforms/r3m.py b/torchrl/envs/transforms/r3m.py
--- a/torchrl/envs/transforms/r3m.py
+++ b/torchrl/envs/transforms/r3m.py
@@ -5,6 +5,11 @@
 from torch.nn import Identity
 
 from torchrl.data import TensorDict, DEVICE_TYPING
+from torchrl.data.tensor_specs import (
+    TensorSpec,
+    CompositeSpec,
+    NdUnboundedContinuousTensorSpec,
+)
 from torchrl.envs.transforms import (
     ToTensorImage,
     Compose,
@@ -75,6 +80,24 @@ def _apply_transform(self, obs: torch.Tensor) -> None:
             out = out.view(*shape, *out.shape[1:])
         return out
 
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        if not isinstance(observation_spec, CompositeSpec):
+            raise ValueError("_R3MNet can only infer CompositeSpec")
+
+        keys = [key for key in observation_spec._specs.keys() if key in self.keys_in]
+        device = observation_spec[keys[0]].device
+
+        if self.del_keys:
+            for key_in in keys:
+                del observation_spec[key_in]
+
+        for key_out in self.keys_out:
+            observation_spec[key_out] = NdUnboundedContinuousTensorSpec(
+                shape=torch.Size([self.outdim]), device=device
+            )
+
+        return observation_spec
+
     @staticmethod
     def _load_weights(model_name, r3m_instance, dir_prefix):
         if model_name not in ("r3m_50", "r3m_34", "r3m_18"):


Test patch:
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -16,6 +16,7 @@
     NdBoundedTensorSpec,
     CompositeSpec,
     UnboundedContinuousTensorSpec,
+    NdUnboundedContinuousTensorSpec,
 )
 from torchrl.data import TensorDict
 from torchrl.envs import EnvCreator, SerialEnv
@@ -37,6 +38,7 @@
 )
 from torchrl.envs.libs.gym import _has_gym, GymEnv
 from torchrl.envs.transforms import VecNorm, TransformedEnv
+from torchrl.envs.transforms.r3m import _R3MNet
 from torchrl.envs.transforms.transforms import (
     _has_tv,
     NoopResetEnv,
@@ -1293,6 +1295,75 @@ def test_r3m_parallel(self, model, device):
         transformed_env.close()
         del transformed_env
 
+    @pytest.mark.parametrize("del_keys", [True, False])
+    @pytest.mark.parametrize(
+        "in_keys",
+        [["next_pixels"], ["next_pixels_1", "next_pixels_2", "next_pixels_3"]],
+    )
+    @pytest.mark.parametrize(
+        "out_keys",
+        [["next_r3m_vec"], ["next_r3m_vec_1", "next_r3m_vec_2", "next_r3m_vec_3"]],
+    )
+    def test_r3mnet_transform_observation_spec(
+        self, in_keys, out_keys, del_keys, device, model
+    ):
+        r3m_net = _R3MNet(in_keys, out_keys, model, del_keys)
+
+        observation_spec = CompositeSpec(
+            **{key: NdBoundedTensorSpec(-1, 1, (3, 16, 16), device) for key in in_keys}
+        )
+        if del_keys:
+            exp_ts = CompositeSpec(
+                **{
+                    key: NdUnboundedContinuousTensorSpec(r3m_net.outdim, device)
+                    for key in out_keys
+                }
+            )
+
+            observation_spec_out = r3m_net.transform_observation_spec(observation_spec)
+
+            for key in in_keys:
+                assert key not in observation_spec_out
+            for key in out_keys:
+                assert observation_spec_out[key].shape == exp_ts[key].shape
+                assert observation_spec_out[key].device == exp_ts[key].device
+                assert observation_spec_out[key].dtype == exp_ts[key].dtype
+        else:
+            ts_dict = {}
+            for key in in_keys:
+                ts_dict[key] = observation_spec[key]
+            for key in out_keys:
+                ts_dict[key] = NdUnboundedContinuousTensorSpec(r3m_net.outdim, device)
+            exp_ts = CompositeSpec(**ts_dict)
+
+            observation_spec_out = r3m_net.transform_observation_spec(observation_spec)
+
+            for key in in_keys + out_keys:
+                assert observation_spec_out[key].shape == exp_ts[key].shape
+                assert observation_spec_out[key].dtype == exp_ts[key].dtype
+                assert observation_spec_out[key].device == exp_ts[key].device
+
+    @pytest.mark.parametrize("tensor_pixels_key", [None, ["funny_key"]])
+    def test_r3m_spec_against_real(self, model, tensor_pixels_key, device):
+        keys_in = ["next_pixels"]
+        keys_out = ["next_vec"]
+        r3m = R3MTransform(
+            model,
+            keys_in=keys_in,
+            keys_out=keys_out,
+            tensor_pixels_keys=tensor_pixels_key,
+        )
+        base_env = DiscreteActionConvMockEnvNumpy().to(device)
+        transformed_env = TransformedEnv(base_env, r3m)
+        expected_keys = (
+            list(transformed_env.input_spec.keys())
+            + list(transformed_env.observation_spec.keys())
+            + [key.strip("next_") for key in transformed_env.observation_spec.keys()]
+            + ["reward"]
+            + ["done"]
+        )
+        assert set(expected_keys) == set(transformed_env.rollout(3).keys())
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()