[BugFix] Add transform_observation_spec _R3MNet (#443)

* Add transform_observation_spec function and related test cases for _R3MNet

* Improve unit tests: 1) check observation_spec from rollouts; 2) integrate transform_observation_spec tests inside TestR3M test suite; 3) fix some errors

Co-authored-by: Zhe Sun <zhesun@fb.com>

Author: ymwdalex

.gitignore

@@ -169,3 +169,6 @@ version.py
 torchrl/version.py
 wandb
 outputs
+
+# PyCharm
+.idea

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()

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"):