torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/envs/libs/habitat.py

@@ -0,0 +1,138 @@
+# 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 functools
+import importlib.util
+
+import torch
+from torchrl._utils import _make_ordinal_device
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.envs.common import EnvBase
+from torchrl.envs.libs.gym import GymEnv, set_gym_backend
+from torchrl.envs.utils import _classproperty
+
+_has_habitat = importlib.util.find_spec("habitat") is not None
+
+
+def _wrap_import_error(fun):
+    @functools.wraps(fun)
+    def new_fun(*args, **kwargs):
+        if not _has_habitat:
+            raise ImportError(
+                "Habitat could not be loaded. Consider installing "
+                "it or solving the import bugs (see attached error message). "
+                "Refer to TorchRL's knowledge base in the documentation to "
+                "debug habitat installation."
+            )
+        return fun(*args, **kwargs)
+
+    return new_fun
+
+
+@_wrap_import_error
+def _get_available_envs():
+    for env in GymEnv.available_envs:
+        if env.startswith("Habitat"):
+            yield env
+
+
+class HabitatEnv(GymEnv):
+    """A wrapper for habitat envs.
+
+    This class currently serves as placeholder and compatibility security.
+    It behaves exactly like the GymEnv wrapper.
+
+    Doc: https://aihabitat.org/docs/
+
+    GitHub: https://github.com/facebookresearch/habitat-lab
+
+    URL: https://aihabitat.org/habitat3/
+
+    Paper: https://ai.meta.com/static-resource/habitat3
+
+    Args:
+        env_name (str): The environment to execute.
+        categorical_action_encoding (bool, optional): if ``True``, categorical
+            specs will be converted to the TorchRL equivalent (:class:`torchrl.data.Categorical`),
+            otherwise a one-hot encoding will be used (:class:`torchrl.data.OneHot`).
+            Defaults to ``False``.
+
+    Keyword Args:
+        from_pixels (bool, optional): if ``True``, an attempt to return the pixel
+            observations from the env will be performed. By default, these observations
+            will be written under the ``"pixels"`` entry.
+            The method being used varies
+            depending on the gym version and may involve a ``wrappers.pixel_observation.PixelObservationWrapper``.
+            Defaults to ``False``.
+        pixels_only (bool, optional): if ``True``, only the pixel observations will
+            be returned (by default under the ``"pixels"`` entry in the output tensordict).
+            If ``False``, observations (eg, states) and pixels will be returned
+            whenever ``from_pixels=True``. Defaults to ``True``.
+        frame_skip (int, optional): if provided, indicates for how many steps the
+            same action is to be repeated. The observation returned will be the
+            last observation of the sequence, whereas the reward will be the sum
+            of rewards across steps.
+        device (torch.device, optional): if provided, the device on which the simulation
+            will occur. Defaults to ``torch.device("cuda:0")``.
+        batch_size (torch.Size, optional): the batch size of the environment.
+            Should match the leading dimensions of all observations, done states,
+            rewards, actions and infos.
+            Defaults to ``torch.Size([])``.
+        allow_done_after_reset (bool, optional): if ``True``, it is tolerated
+            for envs to be ``done`` just after :meth:`reset` is called.
+            Defaults to ``False``.
+
+    Attributes:
+        available_envs (List[str]): a list of environments to build.
+
+    Examples:
+        >>> from torchrl.envs import HabitatEnv
+        >>> env = HabitatEnv("HabitatRenderPick-v0", from_pixels=True)
+        >>> env.rollout(3)
+
+    """
+
+    @_wrap_import_error
+    @set_gym_backend("gym")
+    def __init__(self, env_name, **kwargs):
+        import habitat  # noqa
+        import habitat.gym  # noqa
+
+        device_num = torch.device(kwargs.pop("device", 0)).index
+        kwargs["override_options"] = [
+            f"habitat.simulator.habitat_sim_v0.gpu_device_id={device_num}",
+            "habitat.simulator.concur_render=False",
+        ]
+        super().__init__(env_name=env_name, **kwargs)
+
+    @_classproperty
+    def available_envs(cls):
+        if not _has_habitat:
+            return []
+        return list(_get_available_envs())
+
+    def _build_gym_env(self, env, pixels_only):
+        if self.from_pixels:
+            env.reset()
+        return super()._build_gym_env(env, pixels_only)
+
+    def to(self, device: DEVICE_TYPING) -> EnvBase:
+        device = _make_ordinal_device(torch.device(device))
+        if device.type != "cuda":
+            raise ValueError("The device must be of type cuda for Habitat.")
+        device_num = device.index
+        kwargs = {"override_options": []}
+        for arg in self._constructor_kwargs.get("override_options", []):
+            if arg.startswith("habitat.simulator.habitat_sim_v0.gpu_device_id"):
+                arg = f"habitat.simulator.habitat_sim_v0.gpu_device_id={device_num}"
+                kwargs["override_options"].append(arg)
+            else:
+                kwargs["override_options"].append(arg)
+
+        self._env.close()
+        del self._env
+        self.rebuild_with_kwargs(**kwargs)
+        return super().to(device)

torchrl/envs/libs/isaac_lab.py

@@ -0,0 +1,87 @@
+# 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 torchrl.envs.libs.gym import GymWrapper
+
+
+class IsaacLabWrapper(GymWrapper):
+    """A wrapper for IsaacLab environments.
+
+    Args:
+        env (scripts_isaaclab.envs.ManagerBasedRLEnv or equivalent): the environment instance to wrap.
+        categorical_action_encoding (bool, optional): if ``True``, categorical
+            specs will be converted to the TorchRL equivalent (:class:`torchrl.data.Categorical`),
+            otherwise a one-hot encoding will be used (:class:`torchrl.data.OneHot`).
+            Defaults to ``False``.
+        allow_done_after_reset (bool, optional): if ``True``, it is tolerated
+            for envs to be ``done`` just after :meth:`reset` is called.
+            Defaults to ``False``.
+
+    For other arguments, see the :class:`torchrl.envs.GymWrapper` documentation.
+
+    Refer to `the Isaac Lab doc for installation instructions <https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/pip_installation.html>`_.
+
+    Example:
+        >>> # This code block ensures that the Isaac app is started in headless mode
+        >>> from scripts_isaaclab.app import AppLauncher
+        >>> import argparse
+
+        >>> parser = argparse.ArgumentParser(description="Train an RL agent with TorchRL.")
+        >>> AppLauncher.add_app_launcher_args(parser)
+        >>> args_cli, hydra_args = parser.parse_known_args(["--headless"])
+        >>> app_launcher = AppLauncher(args_cli)
+
+        >>> # Imports and env
+        >>> import gymnasium as gym
+        >>> import isaaclab_tasks  # noqa: F401
+        >>> from isaaclab_tasks.manager_based.classic.ant.ant_env_cfg import AntEnvCfg
+        >>> from torchrl.envs.libs.isaac_lab import IsaacLabWrapper
+
+        >>> env = gym.make("Isaac-Ant-v0", cfg=AntEnvCfg())
+        >>> env = IsaacLabWrapper(env)
+
+    """
+
+    def __init__(
+        self,
+        env: isaaclab.envs.ManagerBasedRLEnv,  # noqa: F821
+        *,
+        categorical_action_encoding: bool = False,
+        allow_done_after_reset: bool = True,
+        convert_actions_to_numpy: bool = False,
+        device: torch.device | None = None,
+        **kwargs,
+    ):
+        if device is None:
+            device = torch.device("cuda:0")
+        super().__init__(
+            env,
+            device=device,
+            categorical_action_encoding=categorical_action_encoding,
+            allow_done_after_reset=allow_done_after_reset,
+            convert_actions_to_numpy=convert_actions_to_numpy,
+            **kwargs,
+        )
+
+    def seed(self, seed: int | None):
+        self._set_seed(seed)
+
+    def _output_transform(self, step_outputs_tuple):  # noqa: F811
+        # IsaacLab will modify the `terminated` and `truncated` tensors
+        #  in-place. We clone them here to make sure data doesn't inadvertently get modified.
+        # The variable naming follows torchrl's convention here.
+        observations, reward, terminated, truncated, info = step_outputs_tuple
+        done = terminated | truncated
+        reward = reward.unsqueeze(-1)  # to get to (num_envs, 1)
+        return (
+            observations,
+            reward,
+            terminated.clone(),
+            truncated.clone(),
+            done.clone(),
+            info,
+        )

torchrl/envs/libs/isaacgym.py

@@ -0,0 +1,203 @@
+# 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.util
+import itertools
+import warnings
+from typing import Any
+
+import numpy as np
+import torch
+from tensordict import TensorDictBase
+from torchrl.data.tensor_specs import Composite
+from torchrl.envs.libs.gym import GymWrapper
+from torchrl.envs.utils import _classproperty, make_composite_from_td
+
+_has_isaac = importlib.util.find_spec("isaacgym") is not None
+
+
+class IsaacGymWrapper(GymWrapper):
+    """Wrapper for IsaacGymEnvs environments.
+
+    The original library can be found `here <https://github.com/NVIDIA-Omniverse/IsaacGymEnvs>`_
+    and is based on IsaacGym which can be downloaded `through NVIDIA's webpage <https://developer.nvidia.com/isaac-gym>_`.
+
+    .. note:: IsaacGym environments cannot be executed consecutively, ie. instantiating one
+        environment after another (even if it has been cleared) will cause
+        CUDA memory issues. We recommend creating one environment per process only.
+        If you need more than one environment, the best way to achieve that is
+        to spawn them across processes.
+
+    .. note:: IsaacGym works on CUDA devices by essence. Make sure your machine
+        has GPUs available and the required setup for IsaacGym (eg, Ubuntu 20.04).
+
+    """
+
+    @property
+    def lib(self):
+        import isaacgym
+
+        return isaacgym
+
+    def __init__(
+        self, env: isaacgymenvs.tasks.base.vec_task.Env, **kwargs  # noqa: F821
+    ):
+        warnings.warn(
+            "IsaacGym environment support is an experimental feature that may change in the future."
+        )
+        super().__init__(
+            env, torch.device(env.device), batch_size=torch.Size([]), **kwargs
+        )
+        if not hasattr(self, "task"):
+            # by convention in IsaacGymEnvs
+            self.task = env.__name__
+
+    def _make_specs(self, env: gym.Env) -> None:  # noqa: F821
+        super()._make_specs(env, batch_size=self.batch_size)
+        self.full_done_spec = Composite(
+            {
+                key: spec.squeeze(-1)
+                for key, spec in self.full_done_spec.items(True, True)
+            },
+            shape=self.batch_size,
+        )
+
+        self.observation_spec["obs"] = self.observation_spec["observation"]
+        del self.observation_spec["observation"]
+
+        data = self.rollout(3).get("next")[..., 0]
+        del data[self.reward_key]
+        for done_key in self.done_keys:
+            try:
+                del data[done_key]
+            except KeyError:
+                continue
+        specs = make_composite_from_td(data)
+
+        obs_spec = self.observation_spec
+        obs_spec.unlock_(recurse=True)
+        obs_spec.update(specs)
+        obs_spec.lock_(recurse=True)
+
+    def _output_transform(self, output):
+        obs, reward, done, info = output
+        if self.from_pixels:
+            obs["pixels"] = self._env.render(mode="rgb_array")
+        return obs, reward, done ^ done, done, done, info
+
+    def _reset_output_transform(self, reset_data):
+        reset_data.pop("reward", None)
+        if self.from_pixels:
+            reset_data["pixels"] = self._env.render(mode="rgb_array")
+        return reset_data, {}
+
+    @classmethod
+    def _make_envs(cls, *, task, num_envs, device, seed=None, headless=False, **kwargs):
+        import isaacgym  # noqa
+        import isaacgymenvs  # noqa
+
+        _ = kwargs.pop("from_pixels", None)
+        envs = isaacgymenvs.make(
+            seed=seed,
+            task=task,
+            num_envs=num_envs,
+            sim_device=str(device),
+            rl_device=str(device),
+            headless=headless,
+            **kwargs,
+        )
+        return envs
+
+    def _set_seed(self, seed: int | None) -> None:
+        # as of #665c32170d84b4be66722eea405a1e08b6e7f761 the seed points nowhere in gym.make for IsaacGymEnvs
+        ...
+
+    def read_action(self, action):
+        """Reads the action obtained from the input TensorDict and transforms it in the format expected by the contained environment.
+
+        Args:
+            action (Tensor or TensorDict): an action to be taken in the environment
+
+        Returns: an action in a format compatible with the contained environment.
+
+        """
+        return action
+
+    def read_done(
+        self,
+        terminated: bool | None = None,
+        truncated: bool | None = None,
+        done: bool | None = None,
+    ) -> tuple[bool, bool, bool]:
+        if terminated is not None:
+            terminated = terminated.bool()
+        if truncated is not None:
+            truncated = truncated.bool()
+        if done is not None:
+            done = done.bool()
+        return terminated, truncated, done, done.any()
+
+    def read_reward(self, total_reward):
+        return total_reward
+
+    def read_obs(
+        self, observations: dict[str, Any] | torch.Tensor | np.ndarray
+    ) -> dict[str, Any]:
+        """Reads an observation from the environment and returns an observation compatible with the output TensorDict.
+
+        Args:
+            observations (observation under a format dictated by the inner env): observation to be read.
+
+        """
+        if isinstance(observations, dict):
+            if "state" in observations and "observation" not in observations:
+                # we rename "state" in "observation" as "observation" is the conventional name
+                # for single observation in torchrl.
+                # naming it 'state' will result in envs that have a different name for the state vector
+                # when queried with and without pixels
+                observations["observation"] = observations.pop("state")
+        if not isinstance(observations, (TensorDictBase, dict)):
+            (key,) = itertools.islice(self.observation_spec.keys(True, True), 1)
+            observations = {key: observations}
+        return observations
+
+
+class IsaacGymEnv(IsaacGymWrapper):
+    """A TorchRL Env interface for IsaacGym environments.
+
+    See :class:`~.IsaacGymWrapper` for more information.
+
+    Examples:
+        >>> env = IsaacGymEnv(task="Ant", num_envs=2000, device="cuda:0")
+        >>> rollout = env.rollout(3)
+        >>> assert env.batch_size == (2000,)
+
+    """
+
+    @_classproperty
+    def available_envs(cls):
+        if not _has_isaac:
+            return []
+
+        import isaacgymenvs  # noqa
+
+        return list(isaacgymenvs.tasks.isaacgym_task_map.keys())
+
+    def __init__(self, task=None, *, env=None, num_envs, device, **kwargs):
+        if env is not None and task is not None:
+            raise RuntimeError("Cannot provide both `task` and `env` arguments.")
+        elif env is not None:
+            task = env
+        from_pixels = kwargs.pop("from_pixels", False)
+        envs = self._make_envs(
+            task=task,
+            num_envs=num_envs,
+            device=device,
+            virtual_screen_capture=False,
+            **kwargs,
+        )
+        self.task = task
+        super().__init__(envs, from_pixels=from_pixels, **kwargs)

torchrl/envs/libs/jax_utils.py

@@ -0,0 +1,166 @@
+# 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 dataclasses
+import importlib.util
+
+# import jax
+import numpy as np
+import torch
+
+# from jax import dlpack as jax_dlpack, numpy as jnp
+from tensordict import make_tensordict, TensorDictBase
+from torch.utils import dlpack as torch_dlpack
+from torchrl.data.tensor_specs import Composite, TensorSpec, Unbounded
+from torchrl.data.utils import numpy_to_torch_dtype_dict
+
+_has_jax = importlib.util.find_spec("jax") is not None
+
+
+def _tree_reshape(x, batch_size: torch.Size):
+    import jax
+
+    shape, n = batch_size, 1
+    return jax.tree_util.tree_map(lambda x: x.reshape(shape + x.shape[n:]), x)
+
+
+def _tree_flatten(x, batch_size: torch.Size):
+    import jax
+
+    shape, n = (batch_size.numel(),), len(batch_size)
+    return jax.tree_util.tree_map(lambda x: x.reshape(shape + x.shape[n:]), x)
+
+
+_dtype_conversion = {
+    np.dtype("uint16"): np.int16,
+    np.dtype("uint32"): np.int32,
+    np.dtype("uint64"): np.int64,
+}
+
+
+def _ndarray_to_tensor(value: jnp.ndarray | np.ndarray) -> torch.Tensor:  # noqa: F821
+    from jax import numpy as jnp
+
+    # JAX arrays generated by jax.vmap would have Numpy dtypes.
+    if value.dtype in _dtype_conversion:
+        value = value.view(_dtype_conversion[value.dtype])
+    if isinstance(value, jnp.ndarray):
+        dlpack_tensor = value.__dlpack__()
+    elif isinstance(value, np.ndarray):
+        dlpack_tensor = value.__dlpack__()
+    else:
+        raise NotImplementedError(f"unsupported data type {type(value)}")
+    out = torch_dlpack.from_dlpack(dlpack_tensor)
+    # dtype can be messed up by dlpack
+    return out.to(numpy_to_torch_dtype_dict[value.dtype])
+
+
+def _tensor_to_ndarray(value: torch.Tensor) -> jnp.ndarray:  # noqa: F821
+    from jax import dlpack as jax_dlpack
+
+    # Detach the tensor to remove gradients before converting to DLPack
+    value = value.contiguous().detach()
+    return jax_dlpack.from_dlpack(value)
+
+
+def _get_object_fields(obj) -> dict:
+    """Converts an object (named tuple or dataclass or dict) to a dict."""
+    if isinstance(obj, tuple) and hasattr(obj, "_fields"):  # named tuple
+        return dict(zip(obj._fields, obj))
+    elif dataclasses.is_dataclass(obj):
+        return {
+            field.name: getattr(obj, field.name) for field in dataclasses.fields(obj)
+        }
+    elif isinstance(obj, dict):
+        return obj
+    elif obj is None:
+        return {}
+    else:
+        raise NotImplementedError(f"unsupported data type {type(obj)}")
+
+
+def _object_to_tensordict(obj, device, batch_size) -> TensorDictBase:
+    """Converts a namedtuple or a dataclass to a TensorDict."""
+    from jax import numpy as jnp
+
+    t = {}
+    _fields = _get_object_fields(obj)
+    for name, value in _fields.items():
+        if isinstance(value, (np.number, int, float)):
+            t[name] = _ndarray_to_tensor(np.asarray([value])).to(device)
+        elif isinstance(value, (jnp.ndarray, np.ndarray)):
+            t[name] = _ndarray_to_tensor(value).to(device)
+        else:
+            nested = _object_to_tensordict(value, device, batch_size)
+            if nested is not None:
+                t[name] = nested
+    if len(t):
+        return make_tensordict(t, device=device, batch_size=batch_size)
+    # discard empty tensordicts
+    return None
+
+
+def _tensordict_to_object(tensordict: TensorDictBase, object_example, batch_size=None):
+    """Converts a TensorDict to a namedtuple or a dataclass."""
+    from jax import dlpack as jax_dlpack, numpy as jnp
+
+    if batch_size is None:
+        batch_size = []
+    t = {}
+    _fields = _get_object_fields(object_example)
+    for name, example in _fields.items():
+        value = tensordict.get(name, None)
+        if isinstance(value, TensorDictBase):
+            t[name] = _tensordict_to_object(value, example, batch_size=batch_size)
+        elif value is None:
+            if isinstance(example, dict):
+                t[name] = _tensordict_to_object({}, example, batch_size=batch_size)
+            else:
+                t[name] = None
+        else:
+            if value.dtype is torch.bool:
+                value = value.to(torch.uint8)
+            shape = value.shape
+            # We need to flatten to fix https://github.com/pytorch/rl/issues/2184
+            value = value.contiguous()
+            value = value.detach()
+            if value.ndim > 1:
+                value = value.flatten().clone()
+            else:
+                # Need this because otherwise an exception is raised
+                #  ValueError: INTERNAL: Address of buffer 1 must be a multiple of 10, but was 0x7efccec00824
+                value = value.clone()
+            value = jax_dlpack.from_dlpack(value)
+            if shape.numel() == 1 and not value.shape:
+                while value.shape != shape:
+                    value = jnp.expand_dims(value, 0)
+                if value.dtype != example.dtype:
+                    t[name] = value.view(example.dtype)
+                else:
+                    t[name] = value
+            else:
+                value = jnp.reshape(value, tuple(shape))
+                t[name] = value.view(example.dtype).reshape(
+                    (*batch_size, *example.shape)
+                )
+    return type(object_example)(**t)
+
+
+def _extract_spec(data: torch.Tensor | TensorDictBase, key=None) -> TensorSpec:
+    if isinstance(data, torch.Tensor):
+        shape = data.shape
+        if key in ("reward", "done"):
+            shape = (*shape, 1)
+        if data.dtype in (torch.float, torch.double, torch.half):
+            return Unbounded(shape=shape, dtype=data.dtype, device=data.device)
+        else:
+            return Unbounded(shape=shape, dtype=data.dtype, device=data.device)
+    elif isinstance(data, TensorDictBase):
+        return Composite(
+            {key: _extract_spec(value, key=key) for key, value in data.items()}
+        )
+    else:
+        raise TypeError(f"Unsupported data type {type(data)}")