From 26180853534795297f3a07c66090026bed20eb15 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 6 Nov 2024 09:09:59 +0000
Subject: [PATCH 01/10] init

---
 .../linux_libs/scripts_habitat/run_test.sh    |  3 +++
 torchrl/envs/libs/habitat.py                  | 19 ++++++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index a60fffd8f45..bc80e03c11e 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -42,6 +42,9 @@ conda deactivate && conda activate ./env
 
 
 # this workflow only tests the libs
+mkdir data
+python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
+
 python -c "import habitat;import habitat.gym"
 python -c """from torchrl.envs.libs.habitat import HabitatEnv
 env = HabitatEnv('HabitatRenderPick-v0')
diff --git a/torchrl/envs/libs/habitat.py b/torchrl/envs/libs/habitat.py
index 4180c42b2dc..45d50f3f87b 100644
--- a/torchrl/envs/libs/habitat.py
+++ b/torchrl/envs/libs/habitat.py
@@ -9,7 +9,7 @@
 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.libs.gym import GymEnv, set_gym_backend, GymWrapper
 from torchrl.envs.utils import _classproperty
 
 _has_habitat = importlib.util.find_spec("habitat") is not None
@@ -106,6 +106,23 @@ def __init__(self, env_name, **kwargs):
         ]
         super().__init__(env_name=env_name, **kwargs)
 
+    @classmethod
+    def from_config(cls, cfg):
+        """Creates a HabitatEnv from the config.
+
+        Examples:
+            >>> config = habitat.get_config(
+            ...     "benchmark/nav/objectnav/objectnav_hssd-hab.yaml"
+            ... )
+            >>> env = HabitatEnv(config)
+
+        """
+        import habitat.gym
+
+        wrapper = cls.__new__(cls)
+        wrapper.__dict__.update(GymWrapper(habitat.gym.make_gym_from_config(cfg)).__dict__)
+        return wrapper
+
     @_classproperty
     def available_envs(cls):
         if not _has_habitat:

From 1818e95f1c62df4b9a03ecdf38f5f5e90b12d104 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 6 Nov 2024 09:36:55 +0000
Subject: [PATCH 02/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/run_test.sh | 3 ++-
 test/test_libs.py                                       | 5 +++++
 torchrl/envs/libs/habitat.py                            | 4 +++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index bc80e03c11e..bfc093dea6f 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -43,7 +43,8 @@ conda deactivate && conda activate ./env
 
 # this workflow only tests the libs
 mkdir data
-python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
+#python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
+python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data
 
 python -c "import habitat;import habitat.gym"
 python -c """from torchrl.envs.libs.habitat import HabitatEnv
diff --git a/test/test_libs.py b/test/test_libs.py
index b3ba8d54c3d..9e941deb477 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -1566,6 +1566,11 @@ def test_habitat(self, envname):
         _ = env.rollout(3)
         check_env_specs(env)
 
+    def test_from_config(self):
+        env = HabitatEnv.from_config(cfg)
+        check_env_specs(env)
+        assert isinstance(env, HabitatEnv)
+
     @pytest.mark.parametrize("from_pixels", [True, False])
     def test_habitat_render(self, envname, from_pixels):
         env = HabitatEnv(envname, from_pixels=from_pixels)
diff --git a/torchrl/envs/libs/habitat.py b/torchrl/envs/libs/habitat.py
index 45d50f3f87b..d0403d21cd4 100644
--- a/torchrl/envs/libs/habitat.py
+++ b/torchrl/envs/libs/habitat.py
@@ -120,7 +120,9 @@ def from_config(cls, cfg):
         import habitat.gym
 
         wrapper = cls.__new__(cls)
-        wrapper.__dict__.update(GymWrapper(habitat.gym.make_gym_from_config(cfg)).__dict__)
+        env = habitat.gym.make_gym_from_config(cfg)
+        env.reset()  # Prevents crash from TorchRL GymWrapper
+        wrapper.__dict__.update(GymWrapper(env).__dict__)
         return wrapper
 
     @_classproperty

From ecaedad305ea625f8ba10c0ca33b382b94c704e7 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Thu, 7 Nov 2024 09:22:00 +0000
Subject: [PATCH 03/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/run_all.sh  | 8 +-------
 .github/unittest/linux_libs/scripts_habitat/run_test.sh | 2 ++
 .github/workflows/test-linux-habitat.yml                | 1 +
 3 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_all.sh b/.github/unittest/linux_libs/scripts_habitat/run_all.sh
index 69204455a69..3e109ee78bb 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_all.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_all.sh
@@ -5,14 +5,8 @@ set -v
 
 
 apt-get update && apt-get upgrade -y
-apt-get install -y vim git wget
+apt-get install -y g++ gcc vim git wget libglfw3 libgl1-mesa-glx libosmesa6 libglew-dev libglvnd0 libgl1 libglx0 libegl1 libgles2
 
-apt-get install -y libglfw3 libgl1-mesa-glx libosmesa6 libglew-dev
-apt-get install -y libglvnd0 libgl1 libglx0 libegl1 libgles2
-
-apt-get install -y g++ gcc
-#apt-get upgrade -y libstdc++6
-#apt-get install -y libgcc
 apt-get dist-upgrade -y
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index bfc093dea6f..66d7be4dd52 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -43,6 +43,8 @@ conda deactivate && conda activate ./env
 
 # this workflow only tests the libs
 mkdir data
+git lfs update
+
 #python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
 python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data
 
diff --git a/.github/workflows/test-linux-habitat.yml b/.github/workflows/test-linux-habitat.yml
index dc7a58aa116..b38ec8ac6ee 100644
--- a/.github/workflows/test-linux-habitat.yml
+++ b/.github/workflows/test-linux-habitat.yml
@@ -39,6 +39,7 @@ jobs:
           export TORCH_VERSION=nightly
         fi
 
+        export COMPOSE_IGNORE_ORPHANS=true
         # Set env vars from matrix
         export PYTHON_VERSION=${{ matrix.python_version }}
         # Commenting these out for now because the GPU test are not working inside docker

From afb573acf01b443fbc483fba9e32258a9e6db206 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Thu, 7 Nov 2024 09:26:00 +0000
Subject: [PATCH 04/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/run_test.sh | 2 +-
 .github/workflows/test-linux-habitat.yml                | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index 66d7be4dd52..bb75c7e761b 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -46,7 +46,7 @@ mkdir data
 git lfs update
 
 #python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
-python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data
+python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data --prune_lfs 0
 
 python -c "import habitat;import habitat.gym"
 python -c """from torchrl.envs.libs.habitat import HabitatEnv
diff --git a/.github/workflows/test-linux-habitat.yml b/.github/workflows/test-linux-habitat.yml
index b38ec8ac6ee..dc7a58aa116 100644
--- a/.github/workflows/test-linux-habitat.yml
+++ b/.github/workflows/test-linux-habitat.yml
@@ -39,7 +39,6 @@ jobs:
           export TORCH_VERSION=nightly
         fi
 
-        export COMPOSE_IGNORE_ORPHANS=true
         # Set env vars from matrix
         export PYTHON_VERSION=${{ matrix.python_version }}
         # Commenting these out for now because the GPU test are not working inside docker

From b5033cde494b9d4fd14e82db8445cfda23bc27a5 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Thu, 7 Nov 2024 10:00:53 +0000
Subject: [PATCH 05/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/run_test.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index bb75c7e761b..0cba133ae8e 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -46,7 +46,7 @@ mkdir data
 git lfs update
 
 #python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
-python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data --prune_lfs 0
+python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data --no-prune
 
 python -c "import habitat;import habitat.gym"
 python -c """from torchrl.envs.libs.habitat import HabitatEnv

From 97495f2dd7019d7dd2b4408da415f99fbc2184d8 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 11 Nov 2024 09:51:38 +0000
Subject: [PATCH 06/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/run_test.sh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
index 0cba133ae8e..f790e757d3d 100755
--- a/.github/unittest/linux_libs/scripts_habitat/run_test.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -45,8 +45,8 @@ conda deactivate && conda activate ./env
 mkdir data
 git lfs update
 
-#python -m habitat_sim.utils.datasets_download --uids rearrange_dataset_v1 rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data
-python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data --no-prune
+python -m habitat_sim.utils.datasets_download --uids rearrange_pick_dataset_v0 rearrange_task_assets --data-path ./data --no-prune
+#python -m habitat_sim.utils.datasets_download --uids rearrange_task_assets --data-path ./data --no-prune
 
 python -c "import habitat;import habitat.gym"
 python -c """from torchrl.envs.libs.habitat import HabitatEnv

From 53ae51f02fa8ad16e68a19f13b5fa2eedeff6d4f Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 4 Dec 2024 11:46:29 +0000
Subject: [PATCH 07/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/environment.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/unittest/linux_libs/scripts_habitat/environment.yml b/.github/unittest/linux_libs/scripts_habitat/environment.yml
index feab90cd052..256d8edbf64 100644
--- a/.github/unittest/linux_libs/scripts_habitat/environment.yml
+++ b/.github/unittest/linux_libs/scripts_habitat/environment.yml
@@ -18,3 +18,4 @@ dependencies:
     - scipy==1.9.1
     - hydra-core
     - ninja
+    - numpy<2.0

From 97722eac7fe656a9b0bad1399c2c0751af13f5be Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 4 Dec 2024 11:58:29 +0000
Subject: [PATCH 08/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/install.sh   | 4 ++--
 .github/unittest/linux_libs/scripts_habitat/setup_env.sh | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/install.sh b/.github/unittest/linux_libs/scripts_habitat/install.sh
index 6b948279a61..fac059bc4d5 100755
--- a/.github/unittest/linux_libs/scripts_habitat/install.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/install.sh
@@ -21,9 +21,9 @@ git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then
-  pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+  pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
-  pip3 install torch --index-url https://download.pytorch.org/whl/cu121
+  pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu121
 fi
 
 # install tensordict
diff --git a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
index 6ad970c3f47..e640f868322 100755
--- a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
@@ -72,4 +72,5 @@ git clone https://github.com/facebookresearch/habitat-lab.git
 cd habitat-lab
 pip3 install -e habitat-lab
 pip3 install -e habitat-baselines  # install habitat_baselines
+pip3 install "numpy<2.0"
 conda run python -m pip install "gym[atari,accept-rom-license]" pygame

From 53a52399e0078b9d56603fb08b9237ee80954a66 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 4 Dec 2024 11:59:56 +0000
Subject: [PATCH 09/10] amend

---
 test/test_libs.py            | 3 +++
 torchrl/envs/libs/habitat.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index 9e941deb477..4046dabfb8e 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -1567,6 +1567,9 @@ def test_habitat(self, envname):
         check_env_specs(env)
 
     def test_from_config(self):
+        import habitat
+
+        cfg = habitat.get_config("benchmark/nav/objectnav/objectnav_hssd-hab.yaml")
         env = HabitatEnv.from_config(cfg)
         check_env_specs(env)
         assert isinstance(env, HabitatEnv)
diff --git a/torchrl/envs/libs/habitat.py b/torchrl/envs/libs/habitat.py
index d0403d21cd4..3fc9672dabb 100644
--- a/torchrl/envs/libs/habitat.py
+++ b/torchrl/envs/libs/habitat.py
@@ -9,7 +9,7 @@
 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, GymWrapper
+from torchrl.envs.libs.gym import GymEnv, GymWrapper, set_gym_backend
 from torchrl.envs.utils import _classproperty
 
 _has_habitat = importlib.util.find_spec("habitat") is not None

From 19ae46a503ea481adb8617d5b84ea49cf6cb9687 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 4 Dec 2024 12:52:18 +0000
Subject: [PATCH 10/10] amend

---
 .github/unittest/linux_libs/scripts_habitat/setup_env.sh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
index e640f868322..075931be71c 100755
--- a/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
+++ b/.github/unittest/linux_libs/scripts_habitat/setup_env.sh
@@ -70,7 +70,7 @@ conda env update --file "${this_dir}/environment.yml" --prune
 conda install habitat-sim withbullet headless -c conda-forge -c aihabitat -y
 git clone https://github.com/facebookresearch/habitat-lab.git
 cd habitat-lab
-pip3 install -e habitat-lab
-pip3 install -e habitat-baselines  # install habitat_baselines
-pip3 install "numpy<2.0"
+echo "numpy<2.0" > constraints.txt
+pip3 install -e habitat-lab --constraint constraints.txt
+pip3 install -e habitat-baselines --constraint constraints.txt # install habitat_baselines
 conda run python -m pip install "gym[atari,accept-rom-license]" pygame