opentau 0.1.0__py3-none-any.whl

opentau/policies/value/siglip_gemma.py

@@ -0,0 +1,221 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SigLip + Gemma Model for Value Function Estimation.
+
+This module defines the configuration and model classes for a value function estimator
+that combines a SigLIP vision encoder and a Gemma language model.
+"""
+
+import torch
+from einops import rearrange
+from torch import nn
+from transformers import (
+    AutoConfig,
+    Gemma3ForCausalLM,
+    PretrainedConfig,
+    PreTrainedModel,
+    SiglipVisionModel,
+)
+from transformers.models.auto import CONFIG_MAPPING
+
+
+class SiglipGemmaValueConfig(PretrainedConfig):
+    """Configuration class for SiglipGemmaValueModel.
+
+    Args:
+        siglip_config: Configuration for the SigLIP vision model.
+        gemma_config: Configuration for the Gemma language model.
+        num_value_bins: Number of bins for value discretization. Defaults to 201.
+        **kwargs: Additional keyword arguments passed to `PretrainedConfig`.
+    """
+
+    model_type = "SiglipGemmaValueModel"
+    sub_configs = {"siglip_config": AutoConfig, "gemma_config": AutoConfig}
+
+    def __init__(
+        self,
+        siglip_config: dict | None = None,
+        gemma_config: dict | None = None,
+        num_value_bins: int = 201,
+        **kwargs,
+    ):
+        self.num_value_bins = num_value_bins
+
+        if siglip_config is None:
+            # Default SIGLIP config similar to PaliGemma vision config
+            self.siglip_config = CONFIG_MAPPING["siglip_vision_model"](
+                hidden_size=1152,
+                intermediate_size=4304,
+                model_type="siglip_vision_model",
+                num_attention_heads=16,
+                num_hidden_layers=27,
+                num_image_tokens=256,
+                patch_size=14,
+                projector_hidden_act="gelu_fast",
+                torch_dtype="float32",
+                vision_use_head=False,
+            )
+        elif isinstance(siglip_config, dict):
+            if "model_type" not in siglip_config:
+                siglip_config["model_type"] = "siglip_vision_model"
+
+            cfg_cls = CONFIG_MAPPING[siglip_config["model_type"]]
+            self.siglip_config = cfg_cls(**siglip_config)
+        else:
+            self.siglip_config = siglip_config
+
+        if gemma_config is None:
+            # Default config for Gemma 3 270M
+            # Based on typical scaling: smaller than 1B model
+            self.gemma_config = CONFIG_MAPPING["gemma"](
+                attention_bias=False,
+                attention_dropout=0.0,
+                bos_token_id=2,
+                eos_token_id=1,
+                head_dim=128,
+                hidden_act="gelu_pytorch_tanh",
+                hidden_activation="gelu_pytorch_tanh",
+                hidden_size=640,
+                initializer_range=0.02,
+                intermediate_size=2048,
+                max_position_embeddings=8192,
+                model_type="gemma",
+                num_attention_heads=8,
+                num_hidden_layers=18,
+                num_key_value_heads=1,
+                pad_token_id=0,
+                rms_norm_eps=1e-06,
+                rope_theta=10000.0,
+                torch_dtype="float32",
+                transformers_version="4.48.1",
+                use_cache=True,
+                vocab_size=257152,
+            )
+        elif isinstance(gemma_config, dict):
+            if "model_type" not in gemma_config:
+                gemma_config["model_type"] = "gemma"
+
+            cfg_cls = CONFIG_MAPPING[gemma_config["model_type"]]
+            self.gemma_config = cfg_cls(**gemma_config)
+        else:
+            self.gemma_config = gemma_config
+
+        super().__init__(**kwargs)
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.attention_implementation not in ["eager", "fa2"]:
+            raise ValueError(
+                f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager' or 'fa2'."
+            )
+
+
+class SiglipGemmaValueModel(PreTrainedModel):
+    """SigLIP + Gemma Model for Value Function Estimation.
+
+    This model combines a SigLIP vision encoder and a Gemma language model to estimate
+    state values. It projects the final hidden state to a set of discretized value bins.
+
+    Args:
+        config: Configuration object of type `SiglipGemmaValueConfig`.
+    """
+
+    config_class = SiglipGemmaValueConfig
+
+    def __init__(self, config: SiglipGemmaValueConfig):
+        """Initializes the SiglipGemmaValueModel.
+
+        Args:
+            config: Configuration object of type `SiglipGemmaValueConfig`.
+        """
+        super().__init__(config=config)
+
+        self.vision_encoder = SiglipVisionModel.from_pretrained("google/siglip2-so400m-patch14-224")
+
+        # Initialize language model (Gemma 3 270M)
+        self.gemma = Gemma3ForCausalLM.from_pretrained("google/gemma-3-270m")
+        self.gemma = self.gemma.model  # we do not want the LM head
+
+        # Value head: projects final hidden state to discretized value bins
+        self.value_head = nn.Linear(640, config.num_value_bins)
+
+    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Embeds images using the SIGLIP vision encoder.
+
+        Args:
+            image: Tensor containing image pixel values.
+
+        Returns:
+            torch.Tensor: The embedded image features.
+        """
+        # Handle different transformers versions
+        if hasattr(self.vision_encoder, "get_image_features"):
+            return self.vision_encoder.get_image_features(image)
+        else:
+            outputs = self.vision_encoder(pixel_values=image)
+            return outputs.last_hidden_state
+
+    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Embeds language tokens using the Gemma embedding layer.
+
+        Args:
+            tokens: Tensor containing language token IDs.
+
+        Returns:
+            torch.Tensor: The embedded language tokens.
+        """
+        return self.gemma.embed_tokens(tokens)
+
+    def forward(
+        self,
+        inputs_embeds: torch.FloatTensor,
+        attention_mask: torch.Tensor,
+        position_ids: torch.LongTensor,
+    ) -> torch.Tensor:
+        """Forward pass that processes vision and language inputs and outputs a value.
+
+        Args:
+            inputs_embeds: Tensor of shape [batch_size, sequence_length, embedding_dim]
+                containing the combined embeddings of images and text.
+            attention_mask: Attention mask for the sequence.
+            position_ids: Position IDs for RoPE.
+
+        Returns:
+            torch.Tensor: Logits for discretized values of shape [batch_size, num_value_bins].
+        """
+
+        attention_mask = rearrange(attention_mask, "b n1 n2 -> b 1 n1 n2")  # support multihead attention
+        # HACK: use full attention for sliding attention as well since our context length is almost the same size as the sliding window
+        mask_mapping = {
+            "full_attention": attention_mask,
+            "sliding_attention": attention_mask,
+        }
+        outputs = self.gemma(
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            attention_mask=mask_mapping,
+        )
+        hidden_states = outputs.last_hidden_state
+
+        # Extract the last token's hidden state for value prediction
+        # Use the last token (which should be the last language token)
+        final_hidden = hidden_states[:, -1]
+
+        # Project to logits for discretized values
+        logits = self.value_head(final_hidden)
+
+        return logits

opentau/scripts/actions_mse_loss.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from dataclasses import asdict
+from pprint import pformat
+
+import numpy as np
+import torch
+from sklearn.metrics import r2_score
+from torch.utils.data import DataLoader
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.factory import make_dataset_mixture
+from opentau.policies.factory import get_policy_class
+from opentau.utils.random_utils import set_seed
+from opentau.utils.utils import (
+    attempt_torch_compile,
+    auto_torch_device,
+    init_logging,
+)
+
+
+@parser.wrap()
+def inference_main(cfg: TrainPipelineConfig):
+    logging.info(pformat(asdict(cfg)))
+    # build lerobot dataset and dataloader
+    datasets = make_dataset_mixture(cfg)
+
+    # load trained or finetunned model. Change the batch size to 1 in the config
+
+    device = auto_torch_device()
+    if cfg.seed is not None:
+        set_seed(cfg.seed)
+
+    logging.info("Creating policy")
+    policy_class = get_policy_class(cfg.policy.type)
+    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=cfg.policy)
+    policy = policy.to(device=device, dtype=torch.bfloat16)
+    policy.eval()
+    policy_sample_actions = attempt_torch_compile(policy.sample_actions, device_hint=device)
+
+    # Always reset policy before episode to clear out action cache.
+    policy.reset()
+
+    for dataset in datasets.datasets:
+        robot_dof = dataset.meta.info["features"]["actions"]["shape"][0]
+        assert cfg.max_action_dim >= robot_dof
+        print(f"The batch size is {cfg.batch_size}")
+        dataloader = DataLoader(dataset, batch_size=cfg.batch_size)
+
+        pred = []
+        truth = []
+        with torch.inference_mode():
+            for batch in dataloader:
+                for key, value in batch.items():
+                    if isinstance(value, torch.Tensor):
+                        batch[key] = batch[key].to(device)
+                action = policy_sample_actions(batch)
+                predicted_action = action.to("cpu", torch.float32).numpy()
+                pred.append(predicted_action[0, :, :robot_dof].squeeze(0))
+                truth.append(batch["actions"][:, 0, :].squeeze(0)[:robot_dof].to(torch.float32).cpu().numpy())
+
+        pred = np.stack(pred, axis=0)
+        truth = np.stack(truth, axis=0)
+
+        print(f"the mean squared error loss per dimension is {np.mean((pred - truth) ** 2, axis=0)}")
+
+        print(f"the r2 score per dimension is {r2_score(pred, truth, multioutput='raw_values')}")
+    logging.info("End of inference")
+
+
+if __name__ == "__main__":
+    init_logging()
+    inference_main()

opentau/scripts/bin_to_safetensors.py

@@ -0,0 +1,116 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import copy  # Import copy module for deepcopy
+import os
+import sys
+
+import torch
+from safetensors.torch import save_file  # Removed load_file as it's not used here
+
+
+def convert_bin_to_safetensors(input_path: str, output_path: str, map_location: str = "cpu"):
+    """
+    Converts a PyTorch .bin checkpoint (state_dict) to .safetensors format.
+
+    This script attempts to handle cases where tensors might share memory by
+    creating a deep copy of the state_dict before saving to .safetensors.
+    This ensures that each tensor has its own memory, which is a requirement
+    for `safetensors.torch.save_file` when shared tensors are detected.
+
+    Args:
+        input_path (str): Path to the input .bin file.
+        output_path (str): Path where the .safetensors file will be saved.
+        map_location (str): Device to map the tensors to when loading the .bin file.
+                            Defaults to 'cpu' to avoid GPU memory issues during conversion.
+                            Can be 'cuda', 'cuda:0', etc.
+    """
+    if not os.path.exists(input_path):
+        print(f"Error: Input file not found at '{input_path}'", file=sys.stderr)
+        sys.exit(1)
+
+    if not output_path.lower().endswith(".safetensors"):
+        print(
+            f"Warning: Output path '{output_path}' does not end with '.safetensors'. Appending it.",
+            file=sys.stderr,
+        )
+        output_path += ".safetensors"
+
+    print(f"Attempting to load state_dict from '{input_path}'...")
+    try:
+        # Load the state_dict from the .bin file
+        state_dict = torch.load(input_path, map_location=torch.device(map_location))  # nosec B614
+        print("State_dict loaded successfully.")
+
+        # Handle shared memory tensors for safetensors compatibility
+        # Create a deep copy of the state_dict to ensure all tensors have unique memory locations.
+        # This resolves the "Some tensors share memory" error from safetensors.
+        # Note: This might increase the file size if many tensors were originally shared.
+        unique_state_dict = {}
+        for key, value in state_dict.items():
+            if isinstance(value, torch.Tensor):
+                unique_state_dict[key] = value.clone().detach()
+            else:
+                unique_state_dict[key] = copy.deepcopy(value)  # Handle non-tensor items (e.g., lists, dicts)
+
+        # Save the state_dict to .safetensors format
+        print(f"Saving state_dict to '{output_path}' in .safetensors format...")
+        save_file(unique_state_dict, output_path)
+        print(f"Conversion successful! Output saved to '{output_path}'")
+
+    except Exception as e:
+        print(f"An error occurred during conversion: {e}", file=sys.stderr)
+        # Provide more specific guidance if it's the known safetensors shared memory error
+        if "Some tensors share memory" in str(e):
+            print(
+                "\nThis error typically occurs when the PyTorch state_dict contains tensors that share the same underlying memory (e.g., `lm_head.weight` and `embed_tokens.weight`)."
+            )
+            print("The script attempted to resolve this by deep-copying the state_dict before saving.")
+            print(
+                "If the issue persists, ensure your `safetensors` library is up-to-date (`pip install --upgrade safetensors`)."
+            )
+            print(
+                "For models with complex shared weight patterns, manually loading the model architecture and using `safetensors.torch.save_model(model, output_path)` might be necessary, as it handles shared weights more robustly."
+            )
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Convert a PyTorch .bin checkpoint (state_dict) to .safetensors format."
+    )
+    parser.add_argument("input_file", type=str, help="Path to the input .bin model weights file.")
+    parser.add_argument(
+        "--output_file",
+        type=str,
+        help="Path to save the output .safetensors file. If not provided, "
+        "it will be inferred from the input filename (e.g., model.bin -> model.safetensors).",
+    )
+    parser.add_argument(
+        "--map_location",
+        type=str,
+        default="cpu",
+        help="Device to map the tensors to when loading the .bin file. "
+        "Defaults to 'cpu'. Can be 'cuda', 'cuda:0', etc.",
+    )
+
+    args = parser.parse_args()
+
+    # If output_file is not provided, infer it
+    if args.output_file is None:
+        base_name = os.path.splitext(args.input_file)[0]
+        args.output_file = base_name + ".safetensors"
+
+    convert_bin_to_safetensors(args.input_file, args.output_file, args.map_location)

opentau/scripts/compute_max_token_length.py

@@ -0,0 +1,111 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import math
+from collections import Counter
+from dataclasses import dataclass
+from functools import partial
+from itertools import accumulate
+from multiprocessing import Pool, cpu_count
+from pathlib import Path
+
+from tqdm import tqdm
+from transformers import AutoTokenizer, PreTrainedTokenizer
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.factory import make_dataset_mixture
+from opentau.datasets.lerobot_dataset import BaseDataset
+from opentau.policies.factory import get_policy_class
+from opentau.policies.pi0.modeling_pi0 import PI0Policy
+from opentau.policies.pi05.modeling_pi05 import PI05Policy
+
+
+@dataclass
+class Args:
+    target_cfg: str  # path to the training configuration file
+    keys: tuple[str] = (
+        "response",
+        "prompt",
+    )  # keys to compute max token length for, e.g. ["response", "prompt"]
+    num_workers: int | None = None
+    chunk_size: int = 1000
+    output_path: str | None = None
+
+
+def get_tokenizer(cfg: TrainPipelineConfig) -> callable:
+    r"""Returns a tokenizer function based on the policy type in the configuration."""
+    policy_class = get_policy_class(cfg.policy.type)
+
+    # TODO: Add `elif` for other policy types if needed
+    if issubclass(policy_class, PI0Policy) or issubclass(policy_class, PI05Policy):
+        return AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")
+
+    raise ValueError(f"Unsupported policy type: {cfg.policy.type}")
+
+
+def chunked(dataset: BaseDataset, key: str, chunk_size: int):
+    n = len(dataset)
+    for start in range(0, n, chunk_size):
+        end = min(n, start + chunk_size)
+        yield [dataset[i][key] for i in range(start, end)]
+
+
+def worker_fn(chunk, tokenizer: PreTrainedTokenizer):
+    return Counter(len(tokenizer(s)["input_ids"]) for s in chunk)
+
+
+def to_percentile(counter: Counter) -> dict[int, float]:
+    r"""Convert counter to a dictionary with token lengths as keys and their percentile as values."""
+    total = counter.total()
+    sorted_keys = sorted(counter.keys())
+    values = accumulate(counter[k] for k in sorted_keys)
+    return {k: v / total for k, v in reversed(list(zip(sorted_keys, values, strict=False)))}
+
+
+@parser.wrap()
+def main(args: Args):
+    cfg = TrainPipelineConfig.from_pretrained(args.target_cfg)
+    datasets = make_dataset_mixture(cfg).datasets
+    tokenizer = get_tokenizer(cfg)
+    worker = partial(worker_fn, tokenizer=tokenizer)
+
+    output = {}
+    for key in args.keys:
+        counter = Counter()
+        for ds in datasets:
+            tasks = tqdm(
+                chunked(ds, key, args.chunk_size),
+                total=math.ceil(len(ds) / args.chunk_size),
+                desc=f"Processing {key} in {ds._get_feature_mapping_key()}",
+            )
+            # TODO: multiprocessing doesn't seem to speed things up. debug why.
+            with Pool(args.num_workers or cpu_count()) as pool:
+                parts = pool.imap_unordered(worker, tasks)
+                counter = sum(parts, start=counter)
+
+        output[key] = to_percentile(counter)
+
+    if args.output_path:
+        path = Path(args.output_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            json.dump(output, f, indent=2)
+    else:
+        print(json.dumps(output, indent=2))
+
+
+if __name__ == "__main__":
+    main()

opentau/scripts/display_sys_info.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Use this script to get a quick summary of your system config.
+It should be able to run without any of OpenTau's dependencies or OpenTau itself installed.
+"""
+
+import platform
+
+HAS_HF_HUB = True
+HAS_HF_DATASETS = True
+HAS_NP = True
+HAS_TORCH = True
+HAS_OPENTAU = True
+
+try:
+    import huggingface_hub
+except ImportError:
+    HAS_HF_HUB = False
+
+try:
+    import datasets
+except ImportError:
+    HAS_HF_DATASETS = False
+
+try:
+    import numpy as np
+except ImportError:
+    HAS_NP = False
+
+try:
+    import torch
+except ImportError:
+    HAS_TORCH = False
+
+try:
+    import opentau
+except ImportError:
+    HAS_OPENTAU = False
+
+
+opentau_version = opentau.__version__ if HAS_OPENTAU else "N/A"
+hf_hub_version = huggingface_hub.__version__ if HAS_HF_HUB else "N/A"
+hf_datasets_version = datasets.__version__ if HAS_HF_DATASETS else "N/A"
+np_version = np.__version__ if HAS_NP else "N/A"
+
+torch_version = torch.__version__ if HAS_TORCH else "N/A"
+torch_cuda_available = torch.cuda.is_available() if HAS_TORCH else "N/A"
+cuda_version = torch._C._cuda_getCompiledVersion() if HAS_TORCH and torch.version.cuda is not None else "N/A"
+
+
+def display_sys_info() -> dict:
+    """Run this to get basic system info to help for tracking issues & bugs."""
+    info = {
+        "`opentau` version": opentau_version,
+        "Platform": platform.platform(),
+        "Python version": platform.python_version(),
+        "Huggingface_hub version": hf_hub_version,
+        "Dataset version": hf_datasets_version,
+        "Numpy version": np_version,
+        "PyTorch version (GPU?)": f"{torch_version} ({torch_cuda_available})",
+        "Cuda version": cuda_version,
+        "Using GPU in script?": "<fill in>",
+        # "Using distributed or parallel set-up in script?": "<fill in>",
+    }
+    print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the last point.\n")
+    print(format_dict(info))
+    return info
+
+
+def format_dict(d: dict) -> str:
+    return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
+
+
+if __name__ == "__main__":
+    display_sys_info()

opentau/scripts/download_libero_benchmarks.py

@@ -0,0 +1,54 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+from libero.libero import get_libero_path
+from libero.libero.utils.download_utils import check_libero_dataset, libero_dataset_download
+
+from opentau.configs import parser
+
+
+@dataclass
+class Args:
+    suite: str | None = None
+    download_dir: str = get_libero_path("datasets")
+
+    def __post_init__(self):
+        if self.suite not in [None, "object", "spatial", "goal", "10", "90"]:
+            raise ValueError(
+                f"Invalid suite: {self.suite}. Available suites are: 'object', 'spatial', 'goal', '10', or '90'."
+            )
+
+
+@parser.wrap()
+def main(args: Args):
+    # Ask users to specify the download directory of datasets
+    download_dir = Path(args.download_dir)
+    download_dir.mkdir(parents=True, exist_ok=True)
+    download_dir = str(download_dir.resolve())
+    print("Datasets will be downloaded to:", download_dir)
+
+    datasets = "all" if args.suite is None else f"libero_{args.suite}"
+    print("Datasets to download:", datasets)
+
+    libero_dataset_download(datasets=datasets, download_dir=download_dir, use_huggingface=True)
+    time.sleep(1)
+    check_libero_dataset(download_dir=args.download_dir)
+
+
+if __name__ == "__main__":
+    main()