[CI] Fix libs

ghstack-source-id: 3e9f8f9
Pull-Request: #3203

Author: vmoens

.github/unittest/linux_libs/scripts_isaaclab/isaac.sh

@@ -46,36 +46,45 @@ eval "$(${conda_dir}/bin/conda shell.bash hook)"
 conda create --prefix ${env_dir} python=3.10 -y
 conda activate ${env_dir}
 
+# Set LD_LIBRARY_PATH to prioritize conda environment libraries early
+export LD_LIBRARY_PATH=${lib_dir}:${LD_LIBRARY_PATH:-}
+
+# Ensure libexpat is at the correct version BEFORE installing other packages
+conda install -c conda-forge expat -y
+
+# Reinstall Python to ensure it links against the correct expat
+conda install --force-reinstall python=3.10 -y
+
 # Pin pytorch to 2.5.1 for IsaacLab
 conda install pytorch==2.5.1 torchvision==0.20.1 pytorch-cuda=12.4 -c pytorch -c nvidia -y
 
-conda run -p ${env_dir} pip install --upgrade pip
-conda run -p ${env_dir} pip install 'isaacsim[all,extscache]==4.5.0' --extra-index-url https://pypi.nvidia.com
+python -m pip install --upgrade pip
+python -m pip install 'isaacsim[all,extscache]==4.5.0' --extra-index-url https://pypi.nvidia.com
 conda install conda-forge::"cmake>3.22" -y
 
 git clone https://github.com/isaac-sim/IsaacLab.git
 cd IsaacLab
-conda run -p ${env_dir} ./isaaclab.sh --install sb3
+./isaaclab.sh --install sb3
 cd ../
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then
   conda install "anaconda::cmake>=3.22" -y
-  conda run -p ${env_dir} python -m pip install "pybind11[global]"
-  conda run -p ${env_dir} python -m pip install git+https://github.com/pytorch/tensordict.git
+  python -m pip install "pybind11[global]"
+  python -m pip install git+https://github.com/pytorch/tensordict.git
 else
-  conda run -p ${env_dir} python -m pip install tensordict
+  python -m pip install tensordict
 fi
 
 # smoke test
-conda run -p ${env_dir} python -c "import tensordict"
+python -c "import tensordict"
 
 printf "* Installing torchrl\n"
-conda run -p ${env_dir} python -m pip install -e . --no-build-isolation
-conda run -p ${env_dir} python -c "import torchrl"
+python -m pip install -e . --no-build-isolation
+python -c "import torchrl"
 
 # Install pytest
-conda run -p ${env_dir} python -m pip install pytest pytest-cov pytest-mock pytest-instafail pytest-rerunfailures pytest-error-for-skips pytest-asyncio
+python -m pip install pytest pytest-cov pytest-mock pytest-instafail pytest-rerunfailures pytest-error-for-skips pytest-asyncio
 
 # Run tests
-conda run -p ${env_dir} python -m pytest test/test_libs.py -k isaac -s
+python -m pytest test/test_libs.py -k isaac -s

.github/unittest/linux_libs/scripts_minari/environment.yml

@@ -25,6 +25,6 @@ dependencies:
     - gymnasium-robotics
     - minari[create]
     - jax>=0.7.0
-    - mujoco<3.3.6      
+    - mujoco>=2.3.0,<3.0.0
     - mujoco-py<2.2,>=2.1
     - minigrid

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

@@ -156,10 +156,10 @@ jobs:
         export BATCHED_PIPE_TIMEOUT=60
         export TD_GET_DEFAULTS_TO_NONE=1
 
-        bash .github/unittest/linux_libs/scripts_d4rl/setup_env.sh
-        bash .github/unittest/linux_libs/scripts_d4rl/install.sh
-        bash .github/unittest/linux_libs/scripts_d4rl/run_test.sh
-        bash .github/unittest/linux_libs/scripts_d4rl/post_process.sh
+        bash .github/unittest/linux_libs/scripts_envpool/setup_env.sh
+        bash .github/unittest/linux_libs/scripts_envpool/install.sh
+        bash .github/unittest/linux_libs/scripts_envpool/run_test.sh
+        bash .github/unittest/linux_libs/scripts_envpool/post_process.sh
 
   unittests-gendgrl:
     strategy:

test/test_libs.py

@@ -73,6 +73,8 @@
 from torchrl.data.datasets.roboset import RobosetExperienceReplay
 from torchrl.data.datasets.vd4rl import VD4RLExperienceReplay
 from torchrl.data.replay_buffers import SamplerWithoutReplacement
+from torchrl.data.replay_buffers.samplers import SliceSampler
+from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import (
     CatTensors,
@@ -82,6 +84,7 @@
     EnvCreator,
     RemoveEmptySpecs,
     RenameTransform,
+    StepCounter,
 )
 from torchrl.envs.batched_envs import SerialEnv
 from torchrl.envs.libs.brax import _has_brax, BraxEnv, BraxWrapper
@@ -2790,7 +2793,6 @@ class TestVmas:
     @pytest.mark.parametrize("scenario_name", VmasWrapper.available_envs)
     @pytest.mark.parametrize("continuous_actions", [True, False])
     def test_all_vmas_scenarios(self, scenario_name, continuous_actions):
-        
         env = VmasEnv(
             scenario=scenario_name,
             continuous_actions=continuous_actions,
@@ -3455,6 +3457,8 @@ def test_d4rl_dummy(self, task):
     @pytest.mark.parametrize("split_trajs", [True, False])
     @pytest.mark.parametrize("from_env", [True, False])
     def test_dataset_build(self, task, split_trajs, from_env):
+        import d4rl  # noqa: F401
+
         t0 = time.time()
         data = D4RLExperienceReplay(
             task, split_trajs=split_trajs, from_env=from_env, batch_size=2
@@ -5144,6 +5148,17 @@ def test_isaaclab(self, env):
         env.check_env_specs(break_when_any_done="both")
         torchrl_logger.info("Check succeeded!")
 
+    def test_isaaclab_rb(self, env):
+        env = env.append_transform(StepCounter())
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(50, ndim=2), sampler=SliceSampler(num_slices=5)
+        )
+        rb.extend(env.rollout(20))
+        # check that rb["step_count"].flatten() is made of sequences of 4 consecutive numbers
+        flat_ranges = rb["step_count"].flatten() % 4
+        arange = torch.arange(flat_ranges.numel(), device=flat_ranges.device) % 4
+        assert (flat_ranges == arange).all()
+
     def test_isaac_collector(self, env):
         col = SyncDataCollector(
             env, env.rand_action, frames_per_batch=1000, total_frames=100_000_000

torchrl/data/datasets/d4rl.py

@@ -346,6 +346,7 @@ def _get_dataset_from_env(self, name, env_kwargs):
         """
         if env_kwargs:
             raise RuntimeError("env_kwargs cannot be passed with using from_env=True")
+        import d4rl  # noqa: F401
         import gym
 
         # we do a local import to avoid circular import issues