We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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`.





--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchrl/envs/transforms/r3m.py`
Content:
```
1 from typing import List, Optional, Union
2 
3 import torch
4 from torch.hub import load_state_dict_from_url
5 from torch.nn import Identity
6 
7 from torchrl.data import TensorDict, DEVICE_TYPING
8 from torchrl.envs.transforms import (
9     ToTensorImage,
10     Compose,
11     ObservationNorm,
12     Resize,
13     Transform,
14     CatTensors,
15     FlattenObservation,
16     UnsqueezeTransform,
17 )
18 
19 try:
20     from torchvision import models
21 
22     _has_tv = True
23 except ImportError:
24     _has_tv = False
25 
26 __all__ = ["R3MTransform"]
27 
28 
29 class _R3MNet(Transform):
30 
31     inplace = False
32 
33     def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):
34         if not _has_tv:
35             raise ImportError(
36                 "Tried to instantiate R3M without torchvision. Make sure you have "
37                 "torchvision installed in your environment."
38             )
39         if model_name == "resnet18":
40             self.model_name = "r3m_18"
41             self.outdim = 512
42             convnet = models.resnet18(pretrained=False)
43         elif model_name == "resnet34":
44             self.model_name = "r3m_34"
45             self.outdim = 512
46             convnet = models.resnet34(pretrained=False)
47         elif model_name == "resnet50":
48             self.model_name = "r3m_50"
49             self.outdim = 2048
50             convnet = models.resnet50(pretrained=False)
51         else:
52             raise NotImplementedError(
53                 f"model {model_name} is currently not supported by R3M"
54             )
55         convnet.fc = Identity()
56         super().__init__(keys_in=in_keys, keys_out=out_keys)
57         self.convnet = convnet
58         self.del_keys = del_keys
59 
60     def _call(self, tensordict):
61         tensordict_view = tensordict.view(-1)
62         super()._call(tensordict_view)
63         if self.del_keys:
64             tensordict.exclude(*self.keys_in, inplace=True)
65         return tensordict
66 
67     @torch.no_grad()
68     def _apply_transform(self, obs: torch.Tensor) -> None:
69         shape = None
70         if obs.ndimension() > 4:
71             shape = obs.shape[:-3]
72             obs = obs.flatten(0, -4)
73         out = self.convnet(obs)
74         if shape is not None:
75             out = out.view(*shape, *out.shape[1:])
76         return out
77 
78     @staticmethod
79     def _load_weights(model_name, r3m_instance, dir_prefix):
80         if model_name not in ("r3m_50", "r3m_34", "r3m_18"):
81             raise ValueError(
82                 "model_name should be one of 'r3m_50', 'r3m_34' or 'r3m_18'"
83             )
84         # url = "https://download.pytorch.org/models/rl/r3m/" + model_name
85         url = "https://pytorch.s3.amazonaws.com/models/rl/r3m/" + model_name + ".pt"
86         d = load_state_dict_from_url(
87             url,
88             progress=True,
89             map_location=next(r3m_instance.parameters()).device,
90             model_dir=dir_prefix,
91         )
92         td = TensorDict(d["r3m"], []).unflatten_keys(".")
93         td_flatten = td["module"]["convnet"].flatten_keys(".")
94         state_dict = td_flatten.to_dict()
95         r3m_instance.convnet.load_state_dict(state_dict)
96 
97     def load_weights(self, dir_prefix=None):
98         self._load_weights(self.model_name, self, dir_prefix)
99 
100 
101 def _init_first(fun):
102     def new_fun(self, *args, **kwargs):
103         if not self.initialized:
104             self._init()
105         return fun(self, *args, **kwargs)
106 
107     return new_fun
108 
109 
110 class R3MTransform(Compose):
111     """R3M Transform class.
112 
113     R3M provides pre-trained ResNet weights aimed at facilitating visual
114     embedding for robotic tasks. The models are trained using Ego4d.
115     See the paper:
116         R3M: A Universal Visual Representation for Robot Manipulation (Suraj Nair,
117             Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, Abhinav Gupta)
118             https://arxiv.org/abs/2203.12601
119 
120     Args:
121         model_name (str): one of resnet50, resnet34 or resnet18
122         keys_in (list of str, optional): list of input keys. If left empty, the
123             "next_pixels" key is assumed.
124         keys_out (list of str, optional): list of output keys. If left empty,
125              "next_r3m_vec" is assumed.
126         size (int, optional): Size of the image to feed to resnet.
127             Defaults to 244.
128         download (bool, optional): if True, the weights will be downloaded using
129             the torch.hub download API (i.e. weights will be cached for future use).
130             Defaults to False.
131         download_path (str, optional): path where to download the models.
132             Default is None (cache path determined by torch.hub utils).
133         tensor_pixels_keys (list of str, optional): Optionally, one can keep the
134             original images (as collected from the env) in the output tensordict.
135             If no value is provided, this won't be collected.
136     """
137 
138     @classmethod
139     def __new__(cls, *args, **kwargs):
140         cls._is_3d = None
141         cls.initialized = False
142         cls._device = None
143         cls._dtype = None
144         return super().__new__(cls)
145 
146     def __init__(
147         self,
148         model_name: str,
149         keys_in: List[str] = None,
150         keys_out: List[str] = None,
151         size: int = 244,
152         stack_images: bool = True,
153         download: bool = False,
154         download_path: Optional[str] = None,
155         tensor_pixels_keys: List[str] = None,
156     ):
157         super().__init__()
158         self.keys_in = keys_in
159         self.download = download
160         self.download_path = download_path
161         self.model_name = model_name
162         self.keys_out = keys_out
163         self.size = size
164         self.stack_images = stack_images
165         self.tensor_pixels_keys = tensor_pixels_keys
166 
167     def _init(self):
168         keys_in = self.keys_in
169         model_name = self.model_name
170         keys_out = self.keys_out
171         size = self.size
172         stack_images = self.stack_images
173         tensor_pixels_keys = self.tensor_pixels_keys
174 
175         # ToTensor
176         transforms = []
177         if tensor_pixels_keys:
178             for i in range(len(keys_in)):
179                 transforms.append(
180                     CatTensors(
181                         keys_in=[keys_in[i]],
182                         out_key=tensor_pixels_keys[i],
183                         del_keys=False,
184                     )
185                 )
186 
187         totensor = ToTensorImage(
188             unsqueeze=False,
189             keys_in=keys_in,
190         )
191         transforms.append(totensor)
192 
193         # Normalize
194         mean = [0.485, 0.456, 0.406]
195         std = [0.229, 0.224, 0.225]
196         normalize = ObservationNorm(
197             keys_in=keys_in,
198             loc=torch.tensor(mean).view(3, 1, 1),
199             scale=torch.tensor(std).view(3, 1, 1),
200             standard_normal=True,
201         )
202         transforms.append(normalize)
203 
204         # Resize: note that resize is a no-op if the tensor has the desired size already
205         resize = Resize(size, size, keys_in=keys_in)
206         transforms.append(resize)
207 
208         # R3M
209         if keys_out is None:
210             if stack_images:
211                 keys_out = ["next_r3m_vec"]
212             else:
213                 keys_out = [f"next_r3m_vec_{i}" for i in range(len(keys_in))]
214         elif stack_images and len(keys_out) != 1:
215             raise ValueError(
216                 f"key_out must be of length 1 if stack_images is True. Got keys_out={keys_out}"
217             )
218         elif not stack_images and len(keys_out) != len(keys_in):
219             raise ValueError(
220                 "key_out must be of length equal to keys_in if stack_images is False."
221             )
222 
223         if stack_images and len(keys_in) > 1:
224             if self.is_3d:
225                 unsqueeze = UnsqueezeTransform(
226                     keys_in=keys_in,
227                     keys_out=keys_in,
228                     unsqueeze_dim=-4,
229                 )
230                 transforms.append(unsqueeze)
231 
232             cattensors = CatTensors(
233                 keys_in,
234                 keys_out[0],
235                 dim=-4,
236             )
237             network = _R3MNet(
238                 in_keys=keys_out,
239                 out_keys=keys_out,
240                 model_name=model_name,
241                 del_keys=False,
242             )
243             flatten = FlattenObservation(-2, -1, keys_out)
244             transforms = [*transforms, cattensors, network, flatten]
245         else:
246             network = _R3MNet(
247                 in_keys=keys_in,
248                 out_keys=keys_out,
249                 model_name=model_name,
250                 del_keys=True,
251             )
252             transforms = [*transforms, normalize, network]
253 
254         for transform in transforms:
255             self.append(transform)
256         if self.download:
257             self[-1].load_weights(dir_prefix=self.download_path)
258         self.initialized = True
259 
260         if self._device is not None:
261             self.to(self._device)
262         if self._dtype is not None:
263             self.to(self._dtype)
264 
265     @property
266     def is_3d(self):
267         if self._is_3d is None:
268             parent = self.parent
269             for key in parent.observation_spec.keys():
270                 self._is_3d = len(parent.observation_spec[key].shape) == 3
271                 break
272         return self._is_3d
273 
274     def to(self, dest: Union[DEVICE_TYPING, torch.dtype]):
275         if isinstance(dest, torch.dtype):
276             self._dtype = dest
277         else:
278             self._device = dest
279         return super().to(dest)
280 
281     @property
282     def device(self):
283         return self._device
284 
285     @property
286     def dtype(self):
287         return self._dtype
288 
289     forward = _init_first(Compose.forward)
290     transform_action_spec = _init_first(Compose.transform_action_spec)
291     transform_observation_spec = _init_first(Compose.transform_observation_spec)
292     transform_input_spec = _init_first(Compose.transform_input_spec)
293     transform_reward_spec = _init_first(Compose.transform_reward_spec)
294     reset = _init_first(Compose.reset)
295     init = _init_first(Compose.init)
296
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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

{"golden_diff": "diff --git a/torchrl/envs/transforms/r3m.py b/torchrl/envs/transforms/r3m.py\n--- a/torchrl/envs/transforms/r3m.py\n+++ b/torchrl/envs/transforms/r3m.py\n@@ -5,6 +5,11 @@\n from torch.nn import Identity\n \n from torchrl.data import TensorDict, DEVICE_TYPING\n+from torchrl.data.tensor_specs import (\n+    TensorSpec,\n+    CompositeSpec,\n+    NdUnboundedContinuousTensorSpec,\n+)\n from torchrl.envs.transforms import (\n     ToTensorImage,\n     Compose,\n@@ -75,6 +80,24 @@\n             out = out.view(*shape, *out.shape[1:])\n         return out\n \n+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:\n+        if not isinstance(observation_spec, CompositeSpec):\n+            raise ValueError(\"_R3MNet can only infer CompositeSpec\")\n+\n+        keys = [key for key in observation_spec._specs.keys() if key in self.keys_in]\n+        device = observation_spec[keys[0]].device\n+\n+        if self.del_keys:\n+            for key_in in keys:\n+                del observation_spec[key_in]\n+\n+        for key_out in self.keys_out:\n+            observation_spec[key_out] = NdUnboundedContinuousTensorSpec(\n+                shape=torch.Size([self.outdim]), device=device\n+            )\n+\n+        return observation_spec\n+\n     @staticmethod\n     def _load_weights(model_name, r3m_instance, dir_prefix):\n         if model_name not in (\"r3m_50\", \"r3m_34\", \"r3m_18\"):\n", "issue": "[Feature Request] Create a tensorspec for the R3M output\n## Motivation\r\n\r\nR3M (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.\r\nIn 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_\"`).\r\n\r\nHere is how the observation_spec should behave:\r\n1. 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).\r\n2. 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\".\r\n3. 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`.\r\n\r\n\r\n\r\n\n", "before_files": [{"content": "from typing import List, Optional, Union\n\nimport torch\nfrom torch.hub import load_state_dict_from_url\nfrom torch.nn import Identity\n\nfrom torchrl.data import TensorDict, DEVICE_TYPING\nfrom torchrl.envs.transforms import (\n    ToTensorImage,\n    Compose,\n    ObservationNorm,\n    Resize,\n    Transform,\n    CatTensors,\n    FlattenObservation,\n    UnsqueezeTransform,\n)\n\ntry:\n    from torchvision import models\n\n    _has_tv = True\nexcept ImportError:\n    _has_tv = False\n\n__all__ = [\"R3MTransform\"]\n\n\nclass _R3MNet(Transform):\n\n    inplace = False\n\n    def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):\n        if not _has_tv:\n            raise ImportError(\n                \"Tried to instantiate R3M without torchvision. Make sure you have \"\n                \"torchvision installed in your environment.\"\n            )\n        if model_name == \"resnet18\":\n            self.model_name = \"r3m_18\"\n            self.outdim = 512\n            convnet = models.resnet18(pretrained=False)\n        elif model_name == \"resnet34\":\n            self.model_name = \"r3m_34\"\n            self.outdim = 512\n            convnet = models.resnet34(pretrained=False)\n        elif model_name == \"resnet50\":\n            self.model_name = \"r3m_50\"\n            self.outdim = 2048\n            convnet = models.resnet50(pretrained=False)\n        else:\n            raise NotImplementedError(\n                f\"model {model_name} is currently not supported by R3M\"\n            )\n        convnet.fc = Identity()\n        super().__init__(keys_in=in_keys, keys_out=out_keys)\n        self.convnet = convnet\n        self.del_keys = del_keys\n\n    def _call(self, tensordict):\n        tensordict_view = tensordict.view(-1)\n        super()._call(tensordict_view)\n        if self.del_keys:\n            tensordict.exclude(*self.keys_in, inplace=True)\n        return tensordict\n\n    @torch.no_grad()\n    def _apply_transform(self, obs: torch.Tensor) -> None:\n        shape = None\n        if obs.ndimension() > 4:\n            shape = obs.shape[:-3]\n            obs = obs.flatten(0, -4)\n        out = self.convnet(obs)\n        if shape is not None:\n            out = out.view(*shape, *out.shape[1:])\n        return out\n\n    @staticmethod\n    def _load_weights(model_name, r3m_instance, dir_prefix):\n        if model_name not in (\"r3m_50\", \"r3m_34\", \"r3m_18\"):\n            raise ValueError(\n                \"model_name should be one of 'r3m_50', 'r3m_34' or 'r3m_18'\"\n            )\n        # url = \"https://download.pytorch.org/models/rl/r3m/\" + model_name\n        url = \"https://pytorch.s3.amazonaws.com/models/rl/r3m/\" + model_name + \".pt\"\n        d = load_state_dict_from_url(\n            url,\n            progress=True,\n            map_location=next(r3m_instance.parameters()).device,\n            model_dir=dir_prefix,\n        )\n        td = TensorDict(d[\"r3m\"], []).unflatten_keys(\".\")\n        td_flatten = td[\"module\"][\"convnet\"].flatten_keys(\".\")\n        state_dict = td_flatten.to_dict()\n        r3m_instance.convnet.load_state_dict(state_dict)\n\n    def load_weights(self, dir_prefix=None):\n        self._load_weights(self.model_name, self, dir_prefix)\n\n\ndef _init_first(fun):\n    def new_fun(self, *args, **kwargs):\n        if not self.initialized:\n            self._init()\n        return fun(self, *args, **kwargs)\n\n    return new_fun\n\n\nclass R3MTransform(Compose):\n    \"\"\"R3M Transform class.\n\n    R3M provides pre-trained ResNet weights aimed at facilitating visual\n    embedding for robotic tasks. The models are trained using Ego4d.\n    See the paper:\n        R3M: A Universal Visual Representation for Robot Manipulation (Suraj Nair,\n            Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, Abhinav Gupta)\n            https://arxiv.org/abs/2203.12601\n\n    Args:\n        model_name (str): one of resnet50, resnet34 or resnet18\n        keys_in (list of str, optional): list of input keys. If left empty, the\n            \"next_pixels\" key is assumed.\n        keys_out (list of str, optional): list of output keys. If left empty,\n             \"next_r3m_vec\" is assumed.\n        size (int, optional): Size of the image to feed to resnet.\n            Defaults to 244.\n        download (bool, optional): if True, the weights will be downloaded using\n            the torch.hub download API (i.e. weights will be cached for future use).\n            Defaults to False.\n        download_path (str, optional): path where to download the models.\n            Default is None (cache path determined by torch.hub utils).\n        tensor_pixels_keys (list of str, optional): Optionally, one can keep the\n            original images (as collected from the env) in the output tensordict.\n            If no value is provided, this won't be collected.\n    \"\"\"\n\n    @classmethod\n    def __new__(cls, *args, **kwargs):\n        cls._is_3d = None\n        cls.initialized = False\n        cls._device = None\n        cls._dtype = None\n        return super().__new__(cls)\n\n    def __init__(\n        self,\n        model_name: str,\n        keys_in: List[str] = None,\n        keys_out: List[str] = None,\n        size: int = 244,\n        stack_images: bool = True,\n        download: bool = False,\n        download_path: Optional[str] = None,\n        tensor_pixels_keys: List[str] = None,\n    ):\n        super().__init__()\n        self.keys_in = keys_in\n        self.download = download\n        self.download_path = download_path\n        self.model_name = model_name\n        self.keys_out = keys_out\n        self.size = size\n        self.stack_images = stack_images\n        self.tensor_pixels_keys = tensor_pixels_keys\n\n    def _init(self):\n        keys_in = self.keys_in\n        model_name = self.model_name\n        keys_out = self.keys_out\n        size = self.size\n        stack_images = self.stack_images\n        tensor_pixels_keys = self.tensor_pixels_keys\n\n        # ToTensor\n        transforms = []\n        if tensor_pixels_keys:\n            for i in range(len(keys_in)):\n                transforms.append(\n                    CatTensors(\n                        keys_in=[keys_in[i]],\n                        out_key=tensor_pixels_keys[i],\n                        del_keys=False,\n                    )\n                )\n\n        totensor = ToTensorImage(\n            unsqueeze=False,\n            keys_in=keys_in,\n        )\n        transforms.append(totensor)\n\n        # Normalize\n        mean = [0.485, 0.456, 0.406]\n        std = [0.229, 0.224, 0.225]\n        normalize = ObservationNorm(\n            keys_in=keys_in,\n            loc=torch.tensor(mean).view(3, 1, 1),\n            scale=torch.tensor(std).view(3, 1, 1),\n            standard_normal=True,\n        )\n        transforms.append(normalize)\n\n        # Resize: note that resize is a no-op if the tensor has the desired size already\n        resize = Resize(size, size, keys_in=keys_in)\n        transforms.append(resize)\n\n        # R3M\n        if keys_out is None:\n            if stack_images:\n                keys_out = [\"next_r3m_vec\"]\n            else:\n                keys_out = [f\"next_r3m_vec_{i}\" for i in range(len(keys_in))]\n        elif stack_images and len(keys_out) != 1:\n            raise ValueError(\n                f\"key_out must be of length 1 if stack_images is True. Got keys_out={keys_out}\"\n            )\n        elif not stack_images and len(keys_out) != len(keys_in):\n            raise ValueError(\n                \"key_out must be of length equal to keys_in if stack_images is False.\"\n            )\n\n        if stack_images and len(keys_in) > 1:\n            if self.is_3d:\n                unsqueeze = UnsqueezeTransform(\n                    keys_in=keys_in,\n                    keys_out=keys_in,\n                    unsqueeze_dim=-4,\n                )\n                transforms.append(unsqueeze)\n\n            cattensors = CatTensors(\n                keys_in,\n                keys_out[0],\n                dim=-4,\n            )\n            network = _R3MNet(\n                in_keys=keys_out,\n                out_keys=keys_out,\n                model_name=model_name,\n                del_keys=False,\n            )\n            flatten = FlattenObservation(-2, -1, keys_out)\n            transforms = [*transforms, cattensors, network, flatten]\n        else:\n            network = _R3MNet(\n                in_keys=keys_in,\n                out_keys=keys_out,\n                model_name=model_name,\n                del_keys=True,\n            )\n            transforms = [*transforms, normalize, network]\n\n        for transform in transforms:\n            self.append(transform)\n        if self.download:\n            self[-1].load_weights(dir_prefix=self.download_path)\n        self.initialized = True\n\n        if self._device is not None:\n            self.to(self._device)\n        if self._dtype is not None:\n            self.to(self._dtype)\n\n    @property\n    def is_3d(self):\n        if self._is_3d is None:\n            parent = self.parent\n            for key in parent.observation_spec.keys():\n                self._is_3d = len(parent.observation_spec[key].shape) == 3\n                break\n        return self._is_3d\n\n    def to(self, dest: Union[DEVICE_TYPING, torch.dtype]):\n        if isinstance(dest, torch.dtype):\n            self._dtype = dest\n        else:\n            self._device = dest\n        return super().to(dest)\n\n    @property\n    def device(self):\n        return self._device\n\n    @property\n    def dtype(self):\n        return self._dtype\n\n    forward = _init_first(Compose.forward)\n    transform_action_spec = _init_first(Compose.transform_action_spec)\n    transform_observation_spec = _init_first(Compose.transform_observation_spec)\n    transform_input_spec = _init_first(Compose.transform_input_spec)\n    transform_reward_spec = _init_first(Compose.transform_reward_spec)\n    reset = _init_first(Compose.reset)\n    init = _init_first(Compose.init)\n", "path": "torchrl/envs/transforms/r3m.py"}], "after_files": [{"content": "from typing import List, Optional, Union\n\nimport torch\nfrom torch.hub import load_state_dict_from_url\nfrom torch.nn import Identity\n\nfrom torchrl.data import TensorDict, DEVICE_TYPING\nfrom torchrl.data.tensor_specs import (\n    TensorSpec,\n    CompositeSpec,\n    NdUnboundedContinuousTensorSpec,\n)\nfrom torchrl.envs.transforms import (\n    ToTensorImage,\n    Compose,\n    ObservationNorm,\n    Resize,\n    Transform,\n    CatTensors,\n    FlattenObservation,\n    UnsqueezeTransform,\n)\n\ntry:\n    from torchvision import models\n\n    _has_tv = True\nexcept ImportError:\n    _has_tv = False\n\n__all__ = [\"R3MTransform\"]\n\n\nclass _R3MNet(Transform):\n\n    inplace = False\n\n    def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):\n        if not _has_tv:\n            raise ImportError(\n                \"Tried to instantiate R3M without torchvision. Make sure you have \"\n                \"torchvision installed in your environment.\"\n            )\n        if model_name == \"resnet18\":\n            self.model_name = \"r3m_18\"\n            self.outdim = 512\n            convnet = models.resnet18(pretrained=False)\n        elif model_name == \"resnet34\":\n            self.model_name = \"r3m_34\"\n            self.outdim = 512\n            convnet = models.resnet34(pretrained=False)\n        elif model_name == \"resnet50\":\n            self.model_name = \"r3m_50\"\n            self.outdim = 2048\n            convnet = models.resnet50(pretrained=False)\n        else:\n            raise NotImplementedError(\n                f\"model {model_name} is currently not supported by R3M\"\n            )\n        convnet.fc = Identity()\n        super().__init__(keys_in=in_keys, keys_out=out_keys)\n        self.convnet = convnet\n        self.del_keys = del_keys\n\n    def _call(self, tensordict):\n        tensordict_view = tensordict.view(-1)\n        super()._call(tensordict_view)\n        if self.del_keys:\n            tensordict.exclude(*self.keys_in, inplace=True)\n        return tensordict\n\n    @torch.no_grad()\n    def _apply_transform(self, obs: torch.Tensor) -> None:\n        shape = None\n        if obs.ndimension() > 4:\n            shape = obs.shape[:-3]\n            obs = obs.flatten(0, -4)\n        out = self.convnet(obs)\n        if shape is not None:\n            out = out.view(*shape, *out.shape[1:])\n        return out\n\n    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:\n        if not isinstance(observation_spec, CompositeSpec):\n            raise ValueError(\"_R3MNet can only infer CompositeSpec\")\n\n        keys = [key for key in observation_spec._specs.keys() if key in self.keys_in]\n        device = observation_spec[keys[0]].device\n\n        if self.del_keys:\n            for key_in in keys:\n                del observation_spec[key_in]\n\n        for key_out in self.keys_out:\n            observation_spec[key_out] = NdUnboundedContinuousTensorSpec(\n                shape=torch.Size([self.outdim]), device=device\n            )\n\n        return observation_spec\n\n    @staticmethod\n    def _load_weights(model_name, r3m_instance, dir_prefix):\n        if model_name not in (\"r3m_50\", \"r3m_34\", \"r3m_18\"):\n            raise ValueError(\n                \"model_name should be one of 'r3m_50', 'r3m_34' or 'r3m_18'\"\n            )\n        # url = \"https://download.pytorch.org/models/rl/r3m/\" + model_name\n        url = \"https://pytorch.s3.amazonaws.com/models/rl/r3m/\" + model_name + \".pt\"\n        d = load_state_dict_from_url(\n            url,\n            progress=True,\n            map_location=next(r3m_instance.parameters()).device,\n            model_dir=dir_prefix,\n        )\n        td = TensorDict(d[\"r3m\"], []).unflatten_keys(\".\")\n        td_flatten = td[\"module\"][\"convnet\"].flatten_keys(\".\")\n        state_dict = td_flatten.to_dict()\n        r3m_instance.convnet.load_state_dict(state_dict)\n\n    def load_weights(self, dir_prefix=None):\n        self._load_weights(self.model_name, self, dir_prefix)\n\n\ndef _init_first(fun):\n    def new_fun(self, *args, **kwargs):\n        if not self.initialized:\n            self._init()\n        return fun(self, *args, **kwargs)\n\n    return new_fun\n\n\nclass R3MTransform(Compose):\n    \"\"\"R3M Transform class.\n\n    R3M provides pre-trained ResNet weights aimed at facilitating visual\n    embedding for robotic tasks. The models are trained using Ego4d.\n    See the paper:\n        R3M: A Universal Visual Representation for Robot Manipulation (Suraj Nair,\n            Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, Abhinav Gupta)\n            https://arxiv.org/abs/2203.12601\n\n    Args:\n        model_name (str): one of resnet50, resnet34 or resnet18\n        keys_in (list of str, optional): list of input keys. If left empty, the\n            \"next_pixels\" key is assumed.\n        keys_out (list of str, optional): list of output keys. If left empty,\n             \"next_r3m_vec\" is assumed.\n        size (int, optional): Size of the image to feed to resnet.\n            Defaults to 244.\n        download (bool, optional): if True, the weights will be downloaded using\n            the torch.hub download API (i.e. weights will be cached for future use).\n            Defaults to False.\n        download_path (str, optional): path where to download the models.\n            Default is None (cache path determined by torch.hub utils).\n        tensor_pixels_keys (list of str, optional): Optionally, one can keep the\n            original images (as collected from the env) in the output tensordict.\n            If no value is provided, this won't be collected.\n    \"\"\"\n\n    @classmethod\n    def __new__(cls, *args, **kwargs):\n        cls._is_3d = None\n        cls.initialized = False\n        cls._device = None\n        cls._dtype = None\n        return super().__new__(cls)\n\n    def __init__(\n        self,\n        model_name: str,\n        keys_in: List[str] = None,\n        keys_out: List[str] = None,\n        size: int = 244,\n        stack_images: bool = True,\n        download: bool = False,\n        download_path: Optional[str] = None,\n        tensor_pixels_keys: List[str] = None,\n    ):\n        super().__init__()\n        self.keys_in = keys_in\n        self.download = download\n        self.download_path = download_path\n        self.model_name = model_name\n        self.keys_out = keys_out\n        self.size = size\n        self.stack_images = stack_images\n        self.tensor_pixels_keys = tensor_pixels_keys\n\n    def _init(self):\n        keys_in = self.keys_in\n        model_name = self.model_name\n        keys_out = self.keys_out\n        size = self.size\n        stack_images = self.stack_images\n        tensor_pixels_keys = self.tensor_pixels_keys\n\n        # ToTensor\n        transforms = []\n        if tensor_pixels_keys:\n            for i in range(len(keys_in)):\n                transforms.append(\n                    CatTensors(\n                        keys_in=[keys_in[i]],\n                        out_key=tensor_pixels_keys[i],\n                        del_keys=False,\n                    )\n                )\n\n        totensor = ToTensorImage(\n            unsqueeze=False,\n            keys_in=keys_in,\n        )\n        transforms.append(totensor)\n\n        # Normalize\n        mean = [0.485, 0.456, 0.406]\n        std = [0.229, 0.224, 0.225]\n        normalize = ObservationNorm(\n            keys_in=keys_in,\n            loc=torch.tensor(mean).view(3, 1, 1),\n            scale=torch.tensor(std).view(3, 1, 1),\n            standard_normal=True,\n        )\n        transforms.append(normalize)\n\n        # Resize: note that resize is a no-op if the tensor has the desired size already\n        resize = Resize(size, size, keys_in=keys_in)\n        transforms.append(resize)\n\n        # R3M\n        if keys_out is None:\n            if stack_images:\n                keys_out = [\"next_r3m_vec\"]\n            else:\n                keys_out = [f\"next_r3m_vec_{i}\" for i in range(len(keys_in))]\n        elif stack_images and len(keys_out) != 1:\n            raise ValueError(\n                f\"key_out must be of length 1 if stack_images is True. Got keys_out={keys_out}\"\n            )\n        elif not stack_images and len(keys_out) != len(keys_in):\n            raise ValueError(\n                \"key_out must be of length equal to keys_in if stack_images is False.\"\n            )\n\n        if stack_images and len(keys_in) > 1:\n            if self.is_3d:\n                unsqueeze = UnsqueezeTransform(\n                    keys_in=keys_in,\n                    keys_out=keys_in,\n                    unsqueeze_dim=-4,\n                )\n                transforms.append(unsqueeze)\n\n            cattensors = CatTensors(\n                keys_in,\n                keys_out[0],\n                dim=-4,\n            )\n            network = _R3MNet(\n                in_keys=keys_out,\n                out_keys=keys_out,\n                model_name=model_name,\n                del_keys=False,\n            )\n            flatten = FlattenObservation(-2, -1, keys_out)\n            transforms = [*transforms, cattensors, network, flatten]\n        else:\n            network = _R3MNet(\n                in_keys=keys_in,\n                out_keys=keys_out,\n                model_name=model_name,\n                del_keys=True,\n            )\n            transforms = [*transforms, normalize, network]\n\n        for transform in transforms:\n            self.append(transform)\n        if self.download:\n            self[-1].load_weights(dir_prefix=self.download_path)\n        self.initialized = True\n\n        if self._device is not None:\n            self.to(self._device)\n        if self._dtype is not None:\n            self.to(self._dtype)\n\n    @property\n    def is_3d(self):\n        if self._is_3d is None:\n            parent = self.parent\n            for key in parent.observation_spec.keys():\n                self._is_3d = len(parent.observation_spec[key].shape) == 3\n                break\n        return self._is_3d\n\n    def to(self, dest: Union[DEVICE_TYPING, torch.dtype]):\n        if isinstance(dest, torch.dtype):\n            self._dtype = dest\n        else:\n            self._device = dest\n        return super().to(dest)\n\n    @property\n    def device(self):\n        return self._device\n\n    @property\n    def dtype(self):\n        return self._dtype\n\n    forward = _init_first(Compose.forward)\n    transform_action_spec = _init_first(Compose.transform_action_spec)\n    transform_observation_spec = _init_first(Compose.transform_observation_spec)\n    transform_input_spec = _init_first(Compose.transform_input_spec)\n    transform_reward_spec = _init_first(Compose.transform_reward_spec)\n    reset = _init_first(Compose.reset)\n    init = _init_first(Compose.init)\n", "path": "torchrl/envs/transforms/r3m.py"}]}