Tests cleaning & simplification (#81)

lerobot/__init__.py

@@ -7,16 +7,22 @@ Example:
         import lerobot
         print(lerobot.available_envs)
         print(lerobot.available_tasks_per_env)
-        print(lerobot.available_datasets_per_env)
         print(lerobot.available_datasets)
         print(lerobot.available_policies)
+        print(lerobot.available_policies_per_env)
     ```
 
-When implementing a new dataset (e.g. `AlohaDataset`), policy (e.g. `DiffusionPolicy`), or environment, follow these steps:
-- Set the required class attributes: `available_datasets`.
-- Set the required class attributes: `name`.
-- Update variables in `lerobot/__init__.py` (e.g. `available_envs`, `available_datasets_per_envs`, `available_policies`)
-- Update variables in `tests/test_available.py` by importing your new class
+When implementing a new dataset class (e.g. `AlohaDataset`) follow these steps:
+- Update `available_datasets` in `lerobot/__init__.py`
+- Set the required `available_datasets` class attribute using the previously updated `lerobot.available_datasets`
+
+When implementing a new environment (e.g. `gym_aloha`), follow these steps:
+- Update `available_envs`, `available_tasks_per_env` and `available_datasets` in `lerobot/__init__.py`
+
+When implementing a new policy class (e.g. `DiffusionPolicy`) follow these steps:
+- Update `available_policies` in `lerobot/__init__.py`
+- Set the required `name` class attribute.
+- Update variables in `tests/test_available.py` by importing your new Policy class
 """
 
 from lerobot.__version__ import __version__  # noqa: F401
@@ -36,7 +42,7 @@ available_tasks_per_env = {
     "xarm": ["XarmLift-v0"],
 }
 
-available_datasets_per_env = {
+available_datasets = {
     "aloha": [
         "aloha_sim_insertion_human",
         "aloha_sim_insertion_scripted",
@@ -47,10 +53,23 @@ available_datasets_per_env = {
     "xarm": ["xarm_lift_medium"],
 }
 
-available_datasets = [dataset for env in available_envs for dataset in available_datasets_per_env[env]]
-
 available_policies = [
     "act",
     "diffusion",
     "tdmpc",
 ]
+
+available_policies_per_env = {
+    "aloha": ["act"],
+    "pusht": ["diffusion"],
+    "xarm": ["tdmpc"],
+}
+
+env_task_pairs = [(env, task) for env, tasks in available_tasks_per_env.items() for task in tasks]
+env_dataset_pairs = [(env, dataset) for env, datasets in available_datasets.items() for dataset in datasets]
+env_dataset_policy_triplets = [
+    (env, dataset, policy)
+    for env, datasets in available_datasets.items()
+    for dataset in datasets
+    for policy in available_policies_per_env[env]
+]

lerobot/common/datasets/aloha.py

@@ -14,6 +14,7 @@ class AlohaDataset(torch.utils.data.Dataset):
     https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_scripted
     """
 
+    # Copied from lerobot/__init__.py
     available_datasets = [
         "aloha_sim_insertion_human",
         "aloha_sim_insertion_scripted",

lerobot/common/datasets/pusht.py

@@ -17,6 +17,7 @@ class PushtDataset(torch.utils.data.Dataset):
         If `None`, no shift is applied to current timestamp and the data from the current frame is loaded.
     """
 
+    # Copied from lerobot/__init__.py
     available_datasets = ["pusht"]
     fps = 10
     image_keys = ["observation.image"]

lerobot/common/datasets/xarm.py

@@ -11,9 +11,8 @@ class XarmDataset(torch.utils.data.Dataset):
     https://huggingface.co/datasets/lerobot/xarm_lift_medium
     """
 
-    available_datasets = [
-        "xarm_lift_medium",
-    ]
+    # Copied from lerobot/__init__.py
+    available_datasets = ["xarm_lift_medium"]
     fps = 15
     image_keys = ["observation.image"]
 

lerobot/common/policies/factory.py

@@ -2,7 +2,7 @@ import inspect
 
 from omegaconf import DictConfig, OmegaConf
 
-from lerobot.common.utils import get_safe_torch_device
+from lerobot.common.utils.utils import get_safe_torch_device
 
 
 def _policy_cfg_from_hydra_cfg(policy_cfg_class, hydra_cfg):

lerobot/common/policies/tdmpc/policy.py

@@ -11,7 +11,7 @@ import torch.nn as nn
 
 import lerobot.common.policies.tdmpc.helper as h
 from lerobot.common.policies.utils import populate_queues
-from lerobot.common.utils import get_safe_torch_device
+from lerobot.common.utils.utils import get_safe_torch_device
 
 FIRST_FRAME = 0
 

lerobot/common/utils/import_utils.py

@@ -0,0 +1,44 @@
+import importlib
+import logging
+
+
+def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[bool, str] | bool:
+    """Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/utils/import_utils.py
+    Check if the package spec exists and grab its version to avoid importing a local directory.
+    **Note:** this doesn't work for all packages.
+    """
+    package_exists = importlib.util.find_spec(pkg_name) is not None
+    package_version = "N/A"
+    if package_exists:
+        try:
+            # Primary method to get the package version
+            package_version = importlib.metadata.version(pkg_name)
+        except importlib.metadata.PackageNotFoundError:
+            # Fallback method: Only for "torch" and versions containing "dev"
+            if pkg_name == "torch":
+                try:
+                    package = importlib.import_module(pkg_name)
+                    temp_version = getattr(package, "__version__", "N/A")
+                    # Check if the version contains "dev"
+                    if "dev" in temp_version:
+                        package_version = temp_version
+                        package_exists = True
+                    else:
+                        package_exists = False
+                except ImportError:
+                    # If the package can't be imported, it's not available
+                    package_exists = False
+            else:
+                # For packages other than "torch", don't attempt the fallback and set as not available
+                package_exists = False
+        logging.debug(f"Detected {pkg_name} version: {package_version}")
+    if return_version:
+        return package_exists, package_version
+    else:
+        return package_exists
+
+
+_torch_available, _torch_version = is_package_available("torch", return_version=True)
+_gym_xarm_available = is_package_available("gym_xarm")
+_gym_aloha_available = is_package_available("gym_aloha")
+_gym_pusht_available = is_package_available("gym_pusht")

lerobot/scripts/display_sys_info.py

@@ -15,7 +15,7 @@ cuda_version = torch._C._cuda_getCompiledVersion() if torch.version.cuda is not
 
 
 # TODO(aliberts): refactor into an actual command `lerobot env`
-def get_env_info() -> dict:
+def display_sys_info() -> dict:
     """Run this to get basic system info to help for tracking issues & bugs."""
     info = {
         "`lerobot` version": version,
@@ -40,4 +40,4 @@ def format_dict(d: dict) -> str:
 
 
 if __name__ == "__main__":
-    get_env_info()
+    display_sys_info()

lerobot/scripts/eval.py

@@ -50,7 +50,7 @@ from lerobot.common.envs.factory import make_env
 from lerobot.common.envs.utils import postprocess_action, preprocess_observation
 from lerobot.common.logger import log_output_dir
 from lerobot.common.policies.factory import make_policy
-from lerobot.common.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
+from lerobot.common.utils.utils import get_safe_torch_device, init_hydra_config, init_logging, set_global_seed
 
 
 def write_video(video_path, stacked_frames, fps):

lerobot/scripts/train.py

@@ -13,7 +13,7 @@ from lerobot.common.datasets.utils import cycle
 from lerobot.common.envs.factory import make_env
 from lerobot.common.logger import Logger, log_output_dir
 from lerobot.common.policies.factory import make_policy
-from lerobot.common.utils import (
+from lerobot.common.utils.utils import (
     format_big_number,
     get_safe_torch_device,
     init_logging,

lerobot/scripts/visualize_dataset.py

@@ -9,7 +9,7 @@ import torch
 
 from lerobot.common.datasets.factory import make_dataset
 from lerobot.common.logger import log_output_dir
-from lerobot.common.utils import init_logging
+from lerobot.common.utils.utils import init_logging
 
 NUM_EPISODES_TO_RENDER = 50
 MAX_NUM_STEPS = 1000