amend

Author: vmoens

.github/unittest/linux_libs/scripts_roboset/environment.yml

@@ -19,3 +19,4 @@ dependencies:
     - hydra-core
     - huggingface_hub
     - tqdm
+    - h5py

setup.py

@@ -218,6 +218,7 @@ def _main(argv):
             "tqdm",
             "scikit-learn",
             "pandas",
+            "h5py",
         ],
         "marl": ["vmas>=1.2.10", "pettingzoo>=1.24.1"],
     }

torchrl/data/datasets/roboset.py

@@ -24,6 +24,8 @@
 from torchrl.data.replay_buffers.writers import Writer
 
 _has_tqdm = importlib.util.find_spec("tqdm", None) is not None
+_has_h5py = importlib.util.find_spec("h5py", None) is not None
+_has_hf_hub = importlib.util.find_spec("huggingface_hub", None) is not None
 
 _NAME_MATCH = KeyDependentDefaultDict(lambda key: key)
 _NAME_MATCH["observations"] = "observation"
@@ -158,6 +160,10 @@ def __init__(
         split_trajs: bool = False,
         **env_kwargs,
     ):
+        if not _has_h5py or not _has_hf_hub:
+            raise ImportError(
+                "h5py and huggingface_hub are required for Roboset datasets."
+            )
         if dataset_id not in self.available_datasets:
             raise ValueError(
                 f"The dataset_id {dataset_id} isn't part of the accepted datasets. "