[Feature] IsaacLab wrapper

ghstack-source-id: d688505
Pull-Request-resolved: #2937

Author: vmoens

.github/unittest/linux_libs/scripts_gym/setup_env.sh

@@ -10,7 +10,6 @@ set -e
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 # Avoid error: "fatal: unsafe repository"
 apt-get update && apt-get install -y git wget gcc g++
-
 apt-get install -y libglfw3 libgl1-mesa-glx libosmesa6 libglew-dev libsdl2-dev libsdl2-2.0-0
 apt-get install -y libglvnd0 libgl1 libglx0 libegl1 libgles2 xvfb libegl-dev libx11-dev freeglut3-dev
 

.github/workflows/test-linux-libs.yml

@@ -230,6 +230,93 @@ jobs:
         ./.github/unittest/linux_libs/scripts_gym/batch_scripts.sh
         ./.github/unittest/linux_libs/scripts_gym/post_process.sh
 
+  unittests-isaaclab:
+    strategy:
+      matrix:
+        python_version: ["3.10"]
+        cuda_arch_version: ["12.8"]
+    if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    with:
+      repository: pytorch/rl
+      runner: "linux.g5.4xlarge.nvidia.gpu"
+      gpu-arch-type: cuda
+      gpu-arch-version: "12.8"
+      docker-image: "nvidia/cuda:12.4.0-devel-ubuntu22.04"
+      timeout: 120
+      script: |
+        if [[ "${{ github.ref }}" =~ release/* ]]; then
+          export RELEASE=1
+          export TORCH_VERSION=stable
+        else
+          export RELEASE=0
+          export TORCH_VERSION=nightly
+        fi
+
+        set -euo pipefail
+        export PYTHON_VERSION="3.10"
+        export CU_VERSION="12.8"
+        export TAR_OPTIONS="--no-same-owner"
+        export UPLOAD_CHANNEL="nightly"
+        export TF_CPP_MIN_LOG_LEVEL=0
+        export BATCHED_PIPE_TIMEOUT=60
+        export TD_GET_DEFAULTS_TO_NONE=1
+
+        nvidia-smi
+        
+        # Setup
+        apt-get update && apt-get install -y git wget gcc g++
+        apt-get install -y libglfw3 libgl1-mesa-glx libosmesa6 libglew-dev libsdl2-dev libsdl2-2.0-0
+        apt-get install -y libglvnd0 libgl1 libglx0 libegl1 libgles2 xvfb libegl-dev libx11-dev freeglut3-dev
+
+        git config --global --add safe.directory '*'
+        root_dir="$(git rev-parse --show-toplevel)"
+        conda_dir="${root_dir}/conda"
+        env_dir="${root_dir}/env"
+        lib_dir="${env_dir}/lib"
+
+        cd "${root_dir}"
+
+        # install conda
+        printf "* Installing conda\n"
+        wget -O miniconda.sh "http://repo.continuum.io/miniconda/Miniconda3-latest-${os}-x86_64.sh"
+        bash ./miniconda.sh -b -f -p "${conda_dir}"
+        eval "$(${conda_dir}/bin/conda shell.bash hook)"
+
+
+        conda create -n env_isaaclab python=3.10 -y
+        conda activate env_isaaclab
+        pip install --upgrade pip
+        pip install 'isaacsim[all,extscache]==4.5.0' --extra-index-url https://pypi.nvidia.com
+        git clone [email protected]:isaac-sim/IsaacLab.git
+        conda install "conda-forge::cmake>3.22" -y
+        cd IsaacLab
+        ./isaaclab.sh --install
+        cd ../
+        
+        # install tensordict
+        if [[ "$RELEASE" == 0 ]]; then
+          conda install "anaconda::cmake>=3.22" -y
+          pip3 install "pybind11[global]"
+          pip3 install git+https://github.com/pytorch/tensordict.git
+        else
+          pip3 install tensordict
+        fi
+
+        # smoke test
+        python -c "import tensordict"
+        
+        printf "* Installing torchrl\n"
+        python setup.py develop
+        python -c "import torchrl"
+        
+        # Install pytest
+        pip install pytest pytest-cov pytest-mock pytest-instafail pytest-rerunfailures pytest-error-for-skips pytest-asyncio
+
+        # Run tests
+        pytest test/test_libs.py -k isaac
+
+
   unittests-jumanji:
     strategy:
       matrix:

docs/source/reference/envs.rst

@@ -1417,6 +1417,7 @@ the following function will return ``1`` when queried:
     HabitatEnv
     IsaacGymEnv
     IsaacGymWrapper
+    IsaacLabWrapper
     JumanjiEnv
     JumanjiWrapper
     MeltingpotEnv

test/test_libs.py

@@ -32,32 +32,6 @@
 import pytest
 import torch
 
-if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import (
-        _make_multithreaded_env,
-        CARTPOLE_VERSIONED,
-        get_available_devices,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-        retry,
-        rollout_consistency_assertion,
-    )
-else:
-    from _utils_internal import (
-        _make_multithreaded_env,
-        CARTPOLE_VERSIONED,
-        get_available_devices,
-        get_default_devices,
-        HALFCHEETAH_VERSIONED,
-        PENDULUM_VERSIONED,
-        PONG_VERSIONED,
-        rand_reset,
-        retry,
-        rollout_consistency_assertion,
-    )
 from packaging import version
 from tensordict import (
     assert_allclose_td,
@@ -155,6 +129,33 @@
     ValueOperator,
 )
 
+if os.getenv("PYTORCH_TEST_FBCODE"):
+    from pytorch.rl.test._utils_internal import (
+        _make_multithreaded_env,
+        CARTPOLE_VERSIONED,
+        get_available_devices,
+        get_default_devices,
+        HALFCHEETAH_VERSIONED,
+        PENDULUM_VERSIONED,
+        PONG_VERSIONED,
+        rand_reset,
+        retry,
+        rollout_consistency_assertion,
+    )
+else:
+    from _utils_internal import (
+        _make_multithreaded_env,
+        CARTPOLE_VERSIONED,
+        get_available_devices,
+        get_default_devices,
+        HALFCHEETAH_VERSIONED,
+        PENDULUM_VERSIONED,
+        PONG_VERSIONED,
+        rand_reset,
+        retry,
+        rollout_consistency_assertion,
+    )
+
 _has_d4rl = importlib.util.find_spec("d4rl") is not None
 
 _has_mo = importlib.util.find_spec("mo_gymnasium") is not None
@@ -166,6 +167,9 @@
 _has_minari = importlib.util.find_spec("minari") is not None
 
 _has_gymnasium = importlib.util.find_spec("gymnasium") is not None
+
+_has_isaaclab = importlib.util.find_spec("isaaclab") is not None
+
 _has_gym_regular = importlib.util.find_spec("gym") is not None
 if _has_gymnasium:
     set_gym_backend("gymnasium").set()
@@ -4541,6 +4545,38 @@ def test_render(self, rollout_steps):
         assert not torch.equal(rollout_penultimate_image, image_from_env)
 
 
+@pytest.mark.skipif(not _has_isaaclab, reason="Isaaclab not found")
+class TestIsaacLab:
+    @pytest.fixture(scope="class")
+    def env(self):
+        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)
+        yield env
+
+    def test_isaaclab(self, env):
+        assert env.batch_size == (4096,)
+        assert env._is_batched
+        env.check_env_specs(break_when_any_done="both")
+
+    def test_isaac_collector(self, env):
+        col = SyncDataCollector(
+            env, env.rand_action, frames_per_batch=1000, total_frames=100_000_000
+        )
+        for _ in col:
+            break
+
+    def test_isaaclab_reset(self):
+        # Make a rollout that will stop as soon as a trajectory reaches a done state
+        r = env.rollout(1_000_000)
+        # Check that done obs are None
+        assert (r["next", "policy"][r["next", "done"].squeeze(-1)] == np.nan).all()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/__init__.py

@@ -20,6 +20,7 @@
     HabitatEnv,
     IsaacGymEnv,
     IsaacGymWrapper,
+    IsaacLabWrapper,
     JumanjiEnv,
     JumanjiWrapper,
     MeltingpotEnv,
@@ -131,6 +132,7 @@
     "ActionDiscretizer",
     "ActionMask",
     "VecNormV2",
+    "IsaacLabWrapper",
     "AutoResetEnv",
     "AutoResetTransform",
     "AsyncEnvPool",

torchrl/envs/gym_like.py

@@ -515,6 +515,12 @@ def validated(self, value):
     def _reset(
         self, tensordict: TensorDictBase | None = None, **kwargs
     ) -> TensorDictBase:
+        if (
+            tensordict is not None
+            and "_reset" in tensordict
+            and not tensordict["_reset"].all()
+        ):
+            raise RuntimeError("Partial resets are not handled at this level.")
         obs, info = self._reset_output_transform(self._env.reset(**kwargs))
 
         source = self.read_obs(obs)

torchrl/envs/libs/__init__.py

@@ -16,6 +16,7 @@
     set_gym_backend,
 )
 from .habitat import HabitatEnv
+from .isaac_lab import IsaacLabWrapper
 from .isaacgym import IsaacGymEnv, IsaacGymWrapper
 from .jumanji import JumanjiEnv, JumanjiWrapper
 from .meltingpot import MeltingpotEnv, MeltingpotWrapper
@@ -32,22 +33,20 @@
     "BraxWrapper",
     "DMControlEnv",
     "DMControlWrapper",
-    "MultiThreadedEnv",
-    "MultiThreadedEnvWrapper",
-    "gym_backend",
     "GymEnv",
     "GymWrapper",
-    "MOGymEnv",
-    "MOGymWrapper",
-    "register_gym_spec_conversion",
-    "set_gym_backend",
     "HabitatEnv",
     "IsaacGymEnv",
     "IsaacGymWrapper",
+    "IsaacLabWrapper",
     "JumanjiEnv",
     "JumanjiWrapper",
+    "MOGymEnv",
+    "MOGymWrapper",
     "MeltingpotEnv",
     "MeltingpotWrapper",
+    "MultiThreadedEnv",
+    "MultiThreadedEnvWrapper",
     "OpenMLEnv",
     "OpenSpielEnv",
     "OpenSpielWrapper",
@@ -60,4 +59,7 @@
     "UnityMLAgentsWrapper",
     "VmasEnv",
     "VmasWrapper",
+    "gym_backend",
+    "register_gym_spec_conversion",
+    "set_gym_backend",
 ]

torchrl/envs/libs/gym.py

@@ -53,6 +53,7 @@
 
 _has_mo = importlib.util.find_spec("mo_gymnasium") is not None
 _has_sb3 = importlib.util.find_spec("stable_baselines3") is not None
+_has_isaaclab = importlib.util.find_spec("isaaclab") is not None
 _has_minigrid = importlib.util.find_spec("minigrid") is not None
 
 
@@ -793,6 +794,7 @@ class PixelObservationWrapper:
 
 class _GymAsyncMeta(_EnvPostInit):
     def __call__(cls, *args, **kwargs):
+        missing_obs_value = kwargs.pop("missing_obs_value", None)
         instance: GymWrapper = super().__call__(*args, **kwargs)
 
         # before gym 0.22, there was no final_observation
@@ -803,6 +805,15 @@ def __call__(cls, *args, **kwargs):
                 VecGymEnvTransform,
             )
 
+            if _has_isaaclab:
+                from isaaclab.envs import ManagerBasedRLEnv
+
+                kwargs = {}
+                if missing_obs_value is not None:
+                    kwargs["missing_obs_value"] = missing_obs_value
+                if isinstance(instance._env.unwrapped, ManagerBasedRLEnv):
+                    return TransformedEnv(instance, VecGymEnvTransform(**kwargs))
+
             if _has_sb3:
                 from stable_baselines3.common.vec_env.base_vec_env import VecEnv
 
@@ -845,7 +856,10 @@ def __call__(cls, *args, **kwargs):
                         instance.observation_spec, backend=backend
                     )
                 )
-            return TransformedEnv(instance, VecGymEnvTransform())
+            kwargs = {}
+            if missing_obs_value is not None:
+                kwargs["missing_obs_value"] = missing_obs_value
+            return TransformedEnv(instance, VecGymEnvTransform(**kwargs))
         return instance
 
 
@@ -892,6 +906,10 @@ class GymWrapper(GymLikeEnv, metaclass=_GymAsyncMeta):
             env step function. Set this to ``False`` if the environment is evaluated
             on GPU, such as IsaacLab.
             Defaults to ``True``.
+        missing_obs_value (Any, optional): default value to use as placeholder for missing observations, when
+            the environment is auto-resetting and missing observations cannot be found in the info dictionary
+            (e.g., with IsaacLab). This argument is passed to :class:`~torchrl.envs.VecGymEnvTransform` by
+            the metaclass.
 
     Attributes:
         available_envs (List[str]): a list of environments to build.
@@ -1069,14 +1087,17 @@ def _post_init(self):
 
     @property
     def _is_batched(self):
+        tuple_of_classes = ()
         if _has_sb3:
             from stable_baselines3.common.vec_env.base_vec_env import VecEnv
 
-            tuple_of_classes = (VecEnv,)
-        else:
-            tuple_of_classes = ()
+            tuple_of_classes = tuple_of_classes + (VecEnv,)
+        if _has_isaaclab:
+            from isaaclab.envs import ManagerBasedRLEnv
+
+            tuple_of_classes = tuple_of_classes + (ManagerBasedRLEnv,)
         return isinstance(
-            self._env, tuple_of_classes + (gym_backend("vector").VectorEnv,)
+            self._env.unwrapped, tuple_of_classes + (gym_backend("vector").VectorEnv,)
         )
 
     @implement_for("gym")
@@ -1562,7 +1583,10 @@ def _replace_reset(self, reset, kwargs):  # noqa
     def _replace_reset(self, reset, kwargs):  # noqa
         import gymnasium as gym
 
-        if self._env.autoreset_mode == gym.vector.AutoresetMode.DISABLED:
+        if (
+            getattr(self._env, "autoreset_mode", None)
+            == gym.vector.AutoresetMode.DISABLED
+        ):
             options = {"reset_mask": reset.view(self.batch_size).numpy()}
             kwargs.setdefault("options", {}).update(options)
         return kwargs