Add support for dataset download from huggingface

Author: zhuyifengzju

README.md

@@ -72,6 +72,20 @@ python benchmark_scripts/download_libero_datasets.py --datasets DATASET
 ```
 where ```DATASET``` is chosen from `[libero_spatial, libero_object, libero_100, libero_goal`.
 
+**NEW!!!**
+
+Alternatively, you can download the dataset from HuggingFace by using:
+```python
+python benchmark_scripts/download_libero_datasets.py --use-huggingface
+```
+
+This option can also be combined with the specific dataset selection:
+```python
+python benchmark_scripts/download_libero_datasets.py --datasets DATASET --use-huggingface
+```
+
+The datasets hosted on HuggingFace are available at [here](https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets).
+
 
 # Getting Started
 

benchmark_scripts/download_libero_datasets.py

@@ -1,6 +1,7 @@
 import init_path
 import argparse
 import os
+import time
 
 import libero.libero.utils.download_utils as download_utils
 from libero.libero import get_libero_path
@@ -19,6 +20,11 @@ def parse_args():
         choices=["all", "libero_goal", "libero_spatial", "libero_object", "libero_100"],
         default="all",
     )
+    parser.add_argument(
+        "--use-huggingface",
+        action="store_true",
+        help="Use Hugging Face instead of original download links"
+    )
     return parser.parse_args()
 
 
@@ -31,12 +37,26 @@ def main():
     print(f"Datasets downloaded to {args.download_dir}")
     print(f"Downloading {args.datasets} datasets")
 
+    if args.use_huggingface:
+        print("Using Hugging Face as the download source")
+    else:
+        print("Using original download links (note: these may expire soon)")
+        input_str = input("Download from original links may lead to failures. Do you want to continue? (y/n): ")
+        if input_str.lower() != 'y':
+            print("Switching to Hugging Face as the download source...")
+            args.use_huggingface = True
+
     # If not, download
     download_utils.libero_dataset_download(
-        download_dir=args.download_dir, datasets=args.datasets
+        download_dir=args.download_dir, 
+        datasets=args.datasets,
+        use_huggingface=args.use_huggingface
     )
 
-    # (TODO) If datasets exist, check if datasets are the same as benchmark
+
+    # wait for 1 second
+    time.sleep(1)
+    print("\n\n\n")
 
     # Check if datasets exist first
     download_utils.check_libero_dataset(download_dir=args.download_dir)

libero/libero/utils/download_utils.py

@@ -13,14 +13,15 @@
 
 from libero.libero import get_libero_path
 
-DIR = os.path.dirname(__file__)
-
-DATASET_LINKS = {
-    "libero_object": "https://utexas.box.com/shared/static/avkklgeq0e1dgzxz52x488whpu8mgspk.zip",
-    "libero_goal": "https://utexas.box.com/shared/static/iv5e4dos8yy2b212pkzkpxu9wbdgjfeg.zip",
-    "libero_spatial": "https://utexas.box.com/shared/static/04k94hyizn4huhbv5sz4ev9p2h1p6s7f.zip",
-    "libero_100": "https://utexas.box.com/shared/static/cv73j8zschq8auh9npzt876fdc1akvmk.zip",
-}
+try:
+    from huggingface_hub import snapshot_download
+    import shutil
+    HUGGINGFACE_AVAILABLE = True
+except ImportError:
+    HUGGINGFACE_AVAILABLE = False
+
+import libero.libero.utils.download_utils as download_utils
+from libero.libero import get_libero_path
 
 
 class DownloadProgressBar(tqdm):
@@ -97,44 +98,109 @@ def download_url(url, download_dir, check_overwrite=True, is_zipfile=True):
             os.remove(file_to_write)
 
 
-def libero_dataset_download(datasets="all", download_dir=None, check_overwrite=True):
+DATASET_LINKS = {
+    "libero_object": "https://utexas.box.com/shared/static/avkklgeq0e1dgzxz52x488whpu8mgspk.zip",
+    "libero_goal": "https://utexas.box.com/shared/static/iv5e4dos8yy2b212pkzkpxu9wbdgjfeg.zip",
+    "libero_spatial": "https://utexas.box.com/shared/static/04k94hyizn4huhbv5sz4ev9p2h1p6s7f.zip",
+    "libero_100": "https://utexas.box.com/shared/static/cv73j8zschq8auh9npzt876fdc1akvmk.zip",
+}
+
+HF_REPO_ID = "yifengzhu-hf/LIBERO-datasets"
+
+
+def download_from_huggingface(dataset_name, download_dir, check_overwrite=True):
+    """
+    Download a specific LIBERO dataset from Hugging Face.
+    
+    Args:
+        dataset_name (str): Name of the dataset to download (e.g., 'libero_spatial')
+        download_dir (str): Directory where the dataset should be downloaded
+        check_overwrite (bool): If True, will check if dataset already exists
+    """
+    if not HUGGINGFACE_AVAILABLE:
+        raise ImportError(
+            "Hugging Face Hub is not available. Install it with 'pip install huggingface_hub'"
+        )
+    
+    # Create the destination folder
+    os.makedirs(download_dir, exist_ok=True)
+    
+    # Check if dataset already exists
+    dataset_dir = os.path.join(download_dir, dataset_name)
+    if check_overwrite and os.path.exists(dataset_dir):
+        user_response = input(
+            f"Warning: dataset {dataset_name} already exists at {dataset_dir}. Overwrite? y/n\n"
+        )
+        if user_response.lower() not in {"yes", "y"}:
+            print(f"Skipping download of {dataset_name}")
+            return
+        
+        # Remove existing directory
+        print(f"Removing existing folder: {dataset_dir}")
+        shutil.rmtree(dataset_dir)
+    
+    # Download the dataset
+    print(f"Downloading {dataset_name} from Hugging Face...")
+    folder_path = snapshot_download(
+        repo_id=HF_REPO_ID,
+        repo_type="dataset",
+        local_dir=download_dir,
+        allow_patterns=f"{dataset_name}/*",
+        local_dir_use_symlinks=False,  # Prevents using symlinks to cached files
+        force_download=True  # Forces re-downloading files
+    )
+    
+    # Verify downloaded files
+    file_count = sum([len(files) for _, _, files in os.walk(os.path.join(download_dir, dataset_name))])
+    print(f"Downloaded {file_count} files for {dataset_name}")
+
+
+def libero_dataset_download(datasets="all", download_dir=None, check_overwrite=True, use_huggingface=False):
     """Download libero datasets
 
     Args:
         datasets (str, optional): Specify which datasets to save. Defaults to "all", downloading all the datasets.
         download_dir (str, optional): Target location for storing datasets. Defaults to None, using the default path.
         check_overwrite (bool, optional): Check if overwriting datasets. Defaults to True.
+        use_huggingface (bool, optional): Use Hugging Face instead of the original download links. Defaults to False.
     """
-
     if download_dir is None:
         download_dir = get_libero_path("datasets")
     if not os.path.exists(download_dir):
         os.makedirs(download_dir)
 
-        assert datasets in [
-            "all",
-            "libero_object",
-            "libero_goal",
-            "libero_spatial",
-            "libero_100",
-        ]
+    assert datasets in [
+        "all",
+        "libero_object",
+        "libero_goal",
+        "libero_spatial",
+        "libero_100",
+    ]
 
-    for dataset_name in [
+    datasets_to_download = [
         "libero_object",
         "libero_goal",
         "libero_spatial",
         "libero_100",
-    ]:
-        if datasets == dataset_name or datasets == "all":
-            print(f"Downloading {dataset_name}")
+    ] if datasets == "all" else [datasets]
+
+    for dataset_name in datasets_to_download:
+        print(f"Downloading {dataset_name}")
+        
+        if use_huggingface:
+            download_from_huggingface(
+                dataset_name=dataset_name,
+                download_dir=download_dir,
+                check_overwrite=check_overwrite
+            )
+        else:
+            print("Using original download links (these may expire soon)")
             download_url(
                 DATASET_LINKS[dataset_name],
                 download_dir=download_dir,
                 check_overwrite=check_overwrite,
             )
 
-            # (TODO): unzip the files
-
 
 def check_libero_dataset(download_dir=None):
     """Check the integrity of the downloaded datasets.