diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1d0fb555..4a28292e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,4 @@ -exclude: ^(data/|tests/) +exclude: ^(data/|tests/data) default_language_version: python: python3.10 repos: diff --git a/tests/test_available.py b/tests/test_available.py index 675e68d2..ed3b22bf 100644 --- a/tests/test_available.py +++ b/tests/test_available.py @@ -1,9 +1,9 @@ import importlib -import pytest -import lerobot -import gymnasium as gym -from lerobot.common.utils.import_utils import is_package_available +import gymnasium as gym +import pytest + +import lerobot from lerobot.common.policies.act.modeling_act import ActionChunkingTransformerPolicy from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy from lerobot.common.policies.tdmpc.policy import TDMPCPolicy @@ -21,7 +21,7 @@ def test_available_env_task(env_name: str, task_name: list): package_name = f"gym_{env_name}" importlib.import_module(package_name) gym_handle = f"{package_name}/{task_name}" - assert gym_handle in gym.envs.registry.keys(), gym_handle + assert gym_handle in gym.envs.registry, gym_handle def test_available_policies(): diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 9ec5c3fc..e488c30b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -1,24 +1,35 @@ +import logging import os from pathlib import Path + import einops import pytest import torch +from datasets import Dataset import lerobot -from lerobot.common.datasets.utils import compute_stats, get_stats_einops_patterns, load_previous_and_future_frames +from lerobot.common.datasets.factory import make_dataset +from lerobot.common.datasets.utils import ( + compute_stats, + get_stats_einops_patterns, + load_previous_and_future_frames, +) from lerobot.common.transforms import Prod from lerobot.common.utils.utils import init_hydra_config -import logging -from lerobot.common.datasets.factory import make_dataset -from datasets import Dataset -from .utils import DEVICE, DEFAULT_CONFIG_PATH + +from .utils import DEFAULT_CONFIG_PATH, DEVICE @pytest.mark.parametrize("env_name, dataset_id, policy_name", lerobot.env_dataset_policy_triplets) def test_factory(env_name, dataset_id, policy_name): cfg = init_hydra_config( DEFAULT_CONFIG_PATH, - overrides=[f"env={env_name}", f"dataset_id={dataset_id}", f"policy={policy_name}", f"device={DEVICE}"] + overrides=[ + f"env={env_name}", + f"dataset_id={dataset_id}", + f"policy={policy_name}", + f"device={DEVICE}", + ], ) dataset = make_dataset(cfg) delta_timestamps = dataset.delta_timestamps @@ -42,7 +53,7 @@ def test_factory(env_name, dataset_id, policy_name): (key, 3, True), ) assert dataset.hf_dataset[key].dtype == torch.uint8, f"{key}" - + # test number of dimensions for key, ndim, required in keys_ndim_required: if key not in item: @@ -51,13 +62,13 @@ def test_factory(env_name, dataset_id, policy_name): else: logging.warning(f'Missing key in dataset: "{key}" not in {dataset}.') continue - + if delta_timestamps is not None and key in delta_timestamps: assert item[key].ndim == ndim + 1, f"{key}" assert item[key].shape[0] == len(delta_timestamps[key]), f"{key}" else: assert item[key].ndim == ndim, f"{key}" - + if key in image_keys: assert item[key].dtype == torch.float32, f"{key}" # TODO(rcadene): we assume for now that image normalization takes place in the model @@ -68,10 +79,9 @@ def test_factory(env_name, dataset_id, policy_name): # test t,c,h,w assert item[key].shape[1] == 3, f"{key}" else: - # test c,h,w + # test c,h,w assert item[key].shape[0] == 3, f"{key}" - if delta_timestamps is not None: # test missing keys in delta_timestamps for key in delta_timestamps: @@ -86,20 +96,20 @@ def test_compute_stats_on_xarm(): """ from lerobot.common.datasets.xarm import XarmDataset - DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None + data_dir = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None # get transform to convert images from uint8 [0,255] to float32 [0,1] transform = Prod(in_keys=XarmDataset.image_keys, prod=1 / 255.0) dataset = XarmDataset( dataset_id="xarm_lift_medium", - root=DATA_DIR, + root=data_dir, transform=transform, ) # Note: we set the batch size to be smaller than the whole dataset to make sure we are testing batched # computation of the statistics. While doing this, we also make sure it works when we don't divide the - # dataset into even batches. + # dataset into even batches. computed_stats = compute_stats(dataset, batch_size=int(len(dataset) * 0.25)) # get einops patterns to aggregate batches and compute statistics @@ -119,7 +129,9 @@ def test_compute_stats_on_xarm(): for k, pattern in stats_patterns.items(): expected_stats[k] = {} expected_stats[k]["mean"] = einops.reduce(hf_dataset[k], pattern, "mean") - expected_stats[k]["std"] = torch.sqrt(einops.reduce((hf_dataset[k] - expected_stats[k]["mean"]) ** 2, pattern, "mean")) + expected_stats[k]["std"] = torch.sqrt( + einops.reduce((hf_dataset[k] - expected_stats[k]["mean"]) ** 2, pattern, "mean") + ) expected_stats[k]["min"] = einops.reduce(hf_dataset[k], pattern, "min") expected_stats[k]["max"] = einops.reduce(hf_dataset[k], pattern, "max") @@ -144,12 +156,14 @@ def test_compute_stats_on_xarm(): def test_load_previous_and_future_frames_within_tolerance(): - hf_dataset = Dataset.from_dict({ - "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], - "index": [0, 1, 2, 3, 4], - "episode_data_index_from": [0, 0, 0, 0, 0], - "episode_data_index_to": [5, 5, 5, 5, 5], - }) + hf_dataset = Dataset.from_dict( + { + "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], + "index": [0, 1, 2, 3, 4], + "episode_data_index_from": [0, 0, 0, 0, 0], + "episode_data_index_to": [5, 5, 5, 5, 5], + } + ) hf_dataset = hf_dataset.with_format("torch") item = hf_dataset[2] delta_timestamps = {"index": [-0.2, 0, 0.139]} @@ -161,12 +175,14 @@ def test_load_previous_and_future_frames_within_tolerance(): def test_load_previous_and_future_frames_outside_tolerance_inside_episode_range(): - hf_dataset = Dataset.from_dict({ - "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], - "index": [0, 1, 2, 3, 4], - "episode_data_index_from": [0, 0, 0, 0, 0], - "episode_data_index_to": [5, 5, 5, 5, 5], - }) + hf_dataset = Dataset.from_dict( + { + "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], + "index": [0, 1, 2, 3, 4], + "episode_data_index_from": [0, 0, 0, 0, 0], + "episode_data_index_to": [5, 5, 5, 5, 5], + } + ) hf_dataset = hf_dataset.with_format("torch") item = hf_dataset[2] delta_timestamps = {"index": [-0.2, 0, 0.141]} @@ -176,12 +192,14 @@ def test_load_previous_and_future_frames_outside_tolerance_inside_episode_range( def test_load_previous_and_future_frames_outside_tolerance_outside_episode_range(): - hf_dataset = Dataset.from_dict({ - "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], - "index": [0, 1, 2, 3, 4], - "episode_data_index_from": [0, 0, 0, 0, 0], - "episode_data_index_to": [5, 5, 5, 5, 5], - }) + hf_dataset = Dataset.from_dict( + { + "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], + "index": [0, 1, 2, 3, 4], + "episode_data_index_from": [0, 0, 0, 0, 0], + "episode_data_index_to": [5, 5, 5, 5, 5], + } + ) hf_dataset = hf_dataset.with_format("torch") item = hf_dataset[2] delta_timestamps = {"index": [-0.3, -0.24, 0, 0.26, 0.3]} @@ -189,6 +207,6 @@ def test_load_previous_and_future_frames_outside_tolerance_outside_episode_range item = load_previous_and_future_frames(item, hf_dataset, delta_timestamps, tol) data, is_pad = item["index"], item["index_is_pad"] assert torch.equal(data, torch.tensor([0, 0, 2, 4, 4])), "Data does not match expected values" - assert torch.equal(is_pad, torch.tensor([True, False, False, True, True])), "Padding does not match expected values" - - + assert torch.equal( + is_pad, torch.tensor([True, False, False, True, True]) + ), "Padding does not match expected values" diff --git a/tests/test_envs.py b/tests/test_envs.py index 75d86274..33928a62 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -1,18 +1,17 @@ import importlib + +import gymnasium as gym import pytest import torch -from lerobot.common.datasets.factory import make_dataset -import gymnasium as gym from gymnasium.utils.env_checker import check_env +import lerobot +from lerobot.common.datasets.factory import make_dataset from lerobot.common.envs.factory import make_env -from lerobot.common.utils.import_utils import is_package_available +from lerobot.common.envs.utils import preprocess_observation from lerobot.common.utils.utils import init_hydra_config -import lerobot -from lerobot.common.envs.utils import preprocess_observation - -from .utils import DEVICE, DEFAULT_CONFIG_PATH, require_env +from .utils import DEFAULT_CONFIG_PATH, DEVICE, require_env OBS_TYPES = ["state", "pixels", "pixels_agent_pos"] @@ -23,7 +22,7 @@ OBS_TYPES = ["state", "pixels", "pixels_agent_pos"] def test_env(env_name, env_task, obs_type): if env_name == "aloha" and obs_type == "state": pytest.skip("`state` observations not available for aloha") - + package_name = f"gym_{env_name}" importlib.import_module(package_name) env = gym.make(f"{package_name}/{env_task}", obs_type=obs_type) diff --git a/tests/test_examples.py b/tests/test_examples.py index 0ca1f212..a3f90cf3 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1,5 +1,5 @@ -from pathlib import Path import subprocess +from pathlib import Path def _find_and_replace(text: str, finds_and_replaces: list[tuple[str, str]]) -> str: @@ -10,7 +10,7 @@ def _find_and_replace(text: str, finds_and_replaces: list[tuple[str, str]]) -> s def _run_script(path): - subprocess.run(['python', path], check=True) + subprocess.run(["python", path], check=True) def test_example_1(): @@ -33,7 +33,7 @@ def test_examples_4_and_3(): path = "examples/4_train_policy.py" - with open(path, "r") as file: + with open(path) as file: file_contents = file.read() # Do less steps, use smaller batch, use CPU, and don't complicate things with dataloader workers. @@ -55,7 +55,7 @@ def test_examples_4_and_3(): path = "examples/3_evaluate_pretrained_policy.py" - with open(path, "r") as file: + with open(path) as file: file_contents = file.read() # Do less evals, use CPU, and use the local model. @@ -74,4 +74,4 @@ def test_examples_4_and_3(): ], ) - assert Path(f"outputs/train/example_pusht_diffusion").exists() + assert Path("outputs/train/example_pusht_diffusion").exists() diff --git a/tests/test_policies.py b/tests/test_policies.py index 24b30a45..ab679fcb 100644 --- a/tests/test_policies.py +++ b/tests/test_policies.py @@ -1,14 +1,15 @@ import pytest import torch +from lerobot.common.datasets.factory import make_dataset from lerobot.common.datasets.utils import cycle +from lerobot.common.envs.factory import make_env from lerobot.common.envs.utils import postprocess_action, preprocess_observation from lerobot.common.policies.factory import make_policy from lerobot.common.policies.policy_protocol import Policy -from lerobot.common.envs.factory import make_env -from lerobot.common.datasets.factory import make_dataset from lerobot.common.utils.utils import init_hydra_config -from .utils import DEVICE, DEFAULT_CONFIG_PATH, require_env + +from .utils import DEFAULT_CONFIG_PATH, DEVICE, require_env # TODO(aliberts): refactor using lerobot/__init__.py variables diff --git a/tests/utils.py b/tests/utils.py index 6709cde1..f3fe5790 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -8,6 +8,7 @@ DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + def require_env(func): """ Decorator that skips the test if the required environment package is not installed. @@ -18,11 +19,11 @@ def require_env(func): @wraps(func) def wrapper(*args, **kwargs): # Determine if 'env_name' is provided and extract its value - arg_names = func.__code__.co_varnames[:func.__code__.co_argcount] - if 'env_name' in arg_names: + arg_names = func.__code__.co_varnames[: func.__code__.co_argcount] + if "env_name" in arg_names: # Get the index of 'env_name' and retrieve the value from args - index = arg_names.index('env_name') - env_name = args[index] if len(args) > index else kwargs.get('env_name') + index = arg_names.index("env_name") + env_name = args[index] if len(args) > index else kwargs.get("env_name") else: raise ValueError("Function does not have 'env_name' as an argument.") @@ -33,4 +34,4 @@ def require_env(func): return func(*args, **kwargs) - return wrapper \ No newline at end of file + return wrapper