torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/transforms/utils.py

@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def check_finite(tensor: torch.Tensor):
+    """Raise an error if a tensor has non-finite elements."""
+    if not tensor.isfinite().all():
+        raise ValueError("Encountered a non-finite tensor.")
+
+
+def _init_first(fun):
+    def new_fun(self, *args, **kwargs):
+        if not self.initialized:
+            self._init()
+        return fun(self, *args, **kwargs)
+
+    return new_fun
+
+
+class _set_missing_tolerance:
+    """Context manager to change the transform tolerance to missing values.
+
+    If a transform has a missing_tolerance of True, it will not raise an error if a key is missing during reset.
+
+    This is implemented via :meth:`~torchrl.envs.transforms.Transform.set_missing_tolerance`.
+
+    The way this is handled is that, if `_reset` calls the default `_call` method, it will not raise an error if an input key is missing.
+
+    For custom `_reset` methods, you should implement this yourself:
+
+    Exmples:
+        >>> def _reset(self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase) -> TensorDictBase:
+        ...     with _set_missing_tolerance(self, True):
+        ...         tensordict_reset = self.foo(tensordict, tensordict_reset)
+        ...     return tensordict_reset
+        >>> def foo(self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase) -> TensorDictBase:
+        ...     if self.input_keys[0] not in tensordict_reset and self.missing_tolerance:
+        ...         return tensordict_reset
+        ...     else:
+        ...         # your code here
+
+    Because `missing_tolerance` will be turned off during calls to `_step`, you can be sure that an appropriate KeyError will be raised
+    if the input key is missing at that time.
+
+    """
+
+    def __init__(self, transform, mode):
+        self.transform = transform
+        self.mode = mode
+
+    def __enter__(self):
+        self.exit_mode = self.transform.missing_tolerance
+        if self.mode != self.exit_mode:
+            self.transform.set_missing_tolerance(self.mode)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.mode != self.exit_mode:
+            self.transform.set_missing_tolerance(self.exit_mode)
+
+
+def _get_reset(reset_key, tensordict):
+    _reset = tensordict.get(reset_key, None)
+    # reset key must be unraveled already
+    parent_td = (
+        tensordict.get(reset_key[:-1], None)
+        if isinstance(reset_key, tuple)
+        else tensordict
+    )
+    if parent_td is None:
+        # we do this just in case the nested td wasn't found
+        parent_td = tensordict
+    if _reset is None:
+        _reset = torch.ones(
+            (),
+            dtype=torch.bool,
+            device=parent_td.device,
+        ).expand(parent_td.batch_size)
+    if _reset.ndim > parent_td.ndim:
+        _reset = _reset.flatten(parent_td.ndim, -1).any(-1)
+    return _reset
+
+
+def _stateless_param(param):
+    is_param = isinstance(param, nn.Parameter)
+    param = param.data.to("meta")
+    if is_param:
+        return nn.Parameter(param, requires_grad=False)
+    return param

torchrl/envs/transforms/vc1.py

@@ -0,0 +1,307 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import importlib
+import os
+import subprocess
+from functools import partial
+
+import torch
+from tensordict import TensorDictBase
+from torch import nn
+from torchrl._utils import logger as torchrl_logger
+
+from torchrl.data.tensor_specs import Composite, DEVICE_TYPING, TensorSpec, Unbounded
+from torchrl.envs.transforms.transforms import (
+    CenterCrop,
+    Compose,
+    ObservationNorm,
+    Resize,
+    ToTensorImage,
+    Transform,
+)
+from torchrl.envs.transforms.utils import _set_missing_tolerance
+
+_has_vc = importlib.util.find_spec("vc_models") is not None
+
+
+class VC1Transform(Transform):
+    """VC1 Transform class.
+
+    VC1 provides pre-trained ResNet weights aimed at facilitating visual
+    embedding for robotic tasks. The models are trained using Ego4d.
+
+    See the paper:
+        VC1: A Universal Visual Representation for Robot Manipulation (Suraj Nair,
+            Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, Abhinav Gupta)
+            https://arxiv.org/abs/2203.12601
+
+    The VC1Transform is created in a lazy manner: the object will be initialized
+    only when an attribute (a spec or the forward method) will be queried.
+    The reason for this is that the :obj:`_init()` method requires some attributes of
+    the parent environment (if any) to be accessed: by making the class lazy we
+    can ensure that the following code snippet works as expected:
+
+    Examples:
+        >>> transform = VC1Transform("default", in_keys=["pixels"])
+        >>> env.append_transform(transform)
+        >>> # the forward method will first call _init which will look at env.observation_spec
+        >>> env.reset()
+
+    Args:
+        in_keys (list of NestedKeys): list of input keys. If left empty, the
+            "pixels" key is assumed.
+        out_keys (list of NestedKeys, optional): list of output keys. If left empty,
+             "VC1_vec" is assumed.
+        model_name (str): One of ``"large"``, ``"base"`` or any other compatible
+            model name (see the `github repo <https://github.com/facebookresearch/eai-vc>`_ for more info). Defaults to ``"default"``
+            which provides a small, untrained model for testing.
+        del_keys (bool, optional): If ``True`` (default), the input key will be
+            discarded from the returned tensordict.
+    """
+
+    inplace = False
+    IMPORT_ERROR = (
+        "Could not load vc_models. You can install it via "
+        "VC1Transform.install_vc_models()."
+    )
+
+    def __init__(self, in_keys, out_keys, model_name, del_keys: bool = True):
+        if model_name == "default":
+            self.make_noload_model()
+            model_name = "vc1_vitb_noload"
+        self.model_name = model_name
+        self.del_keys = del_keys
+
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+        self._init()
+
+    def _init(self):
+        try:
+            from vc_models.models.vit import model_utils
+        except ModuleNotFoundError as err:
+            raise ModuleNotFoundError(self.IMPORT_ERROR) from err
+
+        if self.model_name == "base":
+            model_name = model_utils.VC1_BASE_NAME
+        elif self.model_name == "large":
+            model_name = model_utils.VC1_LARGE_NAME
+        else:
+            model_name = self.model_name
+
+        model, embd_size, model_transforms, model_info = model_utils.load_model(
+            model_name
+        )
+        self.model = model
+        self.embd_size = embd_size
+        self.model_transforms = self._map_tv_to_torchrl(model_transforms)
+
+    def _map_tv_to_torchrl(
+        self,
+        model_transforms,
+        in_keys=None,
+    ):
+        if in_keys is None:
+            in_keys = self.in_keys
+        from torchvision import transforms
+
+        if isinstance(model_transforms, transforms.Resize):
+            size = model_transforms.size
+            if isinstance(size, int):
+                size = (size, size)
+            return Resize(
+                *size,
+                in_keys=in_keys,
+            )
+        elif isinstance(model_transforms, transforms.CenterCrop):
+            size = model_transforms.size
+            if isinstance(size, int):
+                size = (size,)
+            return CenterCrop(
+                *size,
+                in_keys=in_keys,
+            )
+        elif isinstance(model_transforms, transforms.Normalize):
+            return ObservationNorm(
+                in_keys=in_keys,
+                loc=torch.as_tensor(model_transforms.mean).reshape(3, 1, 1),
+                scale=torch.as_tensor(model_transforms.std).reshape(3, 1, 1),
+                standard_normal=True,
+            )
+        elif isinstance(model_transforms, transforms.ToTensor):
+            return ToTensorImage(
+                in_keys=in_keys,
+            )
+        elif isinstance(model_transforms, transforms.Compose):
+            transform_list = []
+            for t in model_transforms.transforms:
+
+                if isinstance(t, transforms.ToTensor):
+                    transform_list.insert(0, t)
+                else:
+                    transform_list.append(t)
+            if len(transform_list) == 0:
+                raise RuntimeError("Did not find any transform.")
+            for i, t in enumerate(transform_list):
+                if i == 0:
+                    transform_list[i] = self._map_tv_to_torchrl(t)
+                else:
+                    transform_list[i] = self._map_tv_to_torchrl(t)
+            return Compose(*transform_list)
+        else:
+            raise NotImplementedError(type(model_transforms))
+
+    def _call(self, next_tensordict):
+        if not self.del_keys:
+            in_keys = [
+                in_key
+                for in_key, out_key in zip(self.in_keys, self.out_keys)
+                if in_key != out_key
+            ]
+            saved_td = next_tensordict.select(*in_keys)
+        with next_tensordict.view(-1) as tensordict_view:
+            super()._call(self.model_transforms(tensordict_view))
+        if self.del_keys:
+            next_tensordict.exclude(*self.in_keys, inplace=True)
+        else:
+            # reset in_keys
+            next_tensordict.update(saved_td)
+        return next_tensordict
+
+    forward = _call
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        # TODO: Check this makes sense
+        with _set_missing_tolerance(self, True):
+            tensordict_reset = self._call(tensordict_reset)
+        return tensordict_reset
+
+    @torch.no_grad()
+    def _apply_transform(self, obs: torch.Tensor) -> None:
+        shape = None
+        if obs.ndimension() > 4:
+            shape = obs.shape[:-3]
+            obs = obs.flatten(0, -4)
+        out = self.model(obs)
+        if shape is not None:
+            out = out.view(*shape, *out.shape[1:])
+        return out
+
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        if not isinstance(observation_spec, Composite):
+            raise ValueError("VC1Transform can only infer Composite")
+
+        keys = [key for key in observation_spec.keys(True, True) if key in self.in_keys]
+        device = observation_spec[keys[0]].device
+        dim = observation_spec[keys[0]].shape[:-3]
+
+        observation_spec = observation_spec.clone()
+        if self.del_keys:
+            for in_key in keys:
+                del observation_spec[in_key]
+
+        for out_key in self.out_keys:
+            observation_spec[out_key] = Unbounded(
+                shape=torch.Size([*dim, self.embd_size]), device=device
+            )
+
+        return observation_spec
+
+    def to(self, dest: DEVICE_TYPING | torch.dtype):
+        if isinstance(dest, torch.dtype):
+            self._dtype = dest
+        else:
+            self._device = dest
+        return super().to(dest)
+
+    @property
+    def device(self):
+        return self._device
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @classmethod
+    def install_vc_models(cls, auto_exit=False):
+        try:
+            from vc_models import models  # noqa: F401
+
+            torchrl_logger.info("vc_models found, no need to install.")
+        except ModuleNotFoundError:
+            HOME = os.environ.get("HOME")
+            vcdir = HOME + "/.cache/torchrl/eai-vc"
+            parentdir = os.path.dirname(os.path.abspath(vcdir))
+            os.makedirs(parentdir, exist_ok=True)
+            try:
+                from git import Repo
+            except ModuleNotFoundError as err:
+                raise ModuleNotFoundError(
+                    "Could not load git. Make sure that `git` has been installed "
+                    "in your virtual environment."
+                ) from err
+            Repo.clone_from("https://github.com/facebookresearch/eai-vc.git", vcdir)
+            os.chdir(vcdir + "/vc_models")
+            subprocess.call(["python", "setup.py", "develop"])
+            if not auto_exit:
+                input(
+                    "VC1 has been successfully installed. Exit this python run and "
+                    "relaunch it again. Press Enter to exit..."
+                )
+                exit()
+
+    @classmethod
+    def make_noload_model(cls):
+        """Creates an naive model at a custom destination."""
+        import vc_models
+
+        models_filepath = os.path.dirname(os.path.abspath(vc_models.__file__))
+        cfg_path = os.path.join(
+            models_filepath, "conf", "model", "vc1_vitb_noload.yaml"
+        )
+        if os.path.exists(cfg_path):
+            return
+        config = """_target_: vc_models.models.load_model
+model:
+  _target_: vc_models.models.vit.vit.load_mae_encoder
+  checkpoint_path:
+  model:
+    _target_: torchrl.envs.transforms.vc1._vit_base_patch16
+    img_size: 224
+    use_cls: True
+    drop_path_rate: 0.0
+transform:
+  _target_: vc_models.transforms.vit_transforms
+metadata:
+  algo: mae
+  model: vit_base_patch16
+  data:
+    - ego
+    - imagenet
+    - inav
+  comment: 182_epochs
+"""
+        with open(cfg_path, "w") as file:
+            file.write(config)
+
+
+def _vit_base_patch16(**kwargs):
+    from vc_models.models.vit.vit import VisionTransformer
+
+    model = VisionTransformer(
+        patch_size=16,
+        embed_dim=16,
+        depth=4,
+        num_heads=4,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    return model