Merge (No verify)

tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.

tests/artifacts/datasets/save_dataset_to_safetensors.py

@@ -23,7 +23,7 @@ If you know that your change will break backward compatibility, you should write
 doesnt need to be merged into the `main` branch. Then you need to run this script and update the tests artifacts.
 
 Example usage:
-    `python tests/scripts/save_dataset_to_safetensors.py`
+    `python tests/artifacts/datasets/save_dataset_to_safetensors.py`
 """
 
 import shutil
@@ -94,4 +94,4 @@ if __name__ == "__main__":
         "lerobot/nyu_franka_play_dataset",
         "lerobot/cmu_stretch",
     ]:
-        save_dataset_to_safetensors("tests/data/save_dataset_to_safetensors", repo_id=dataset)
+        save_dataset_to_safetensors("tests/artifacts/datasets", repo_id=dataset)

tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py

@@ -27,7 +27,7 @@ from lerobot.common.datasets.transforms import (
 )
 from lerobot.common.utils.random_utils import seeded_context
 
-ARTIFACT_DIR = Path("tests/data/save_image_transforms_to_safetensors")
+ARTIFACT_DIR = Path("tests/artifacts/image_transforms")
 DATASET_REPO_ID = "lerobot/aloha_mobile_shrimp"
 
 

tests/artifacts/policies/save_policy_to_safetensors.py

@@ -33,12 +33,11 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id=ds_repo_id, episodes=[0]),
         policy=make_policy_config(policy_name, **policy_kwargs),
-        device="cpu",
     )
     train_cfg.validate()  # Needed for auto-setting some parameters
 
     dataset = make_dataset(train_cfg)
-    policy = make_policy(train_cfg.policy, ds_meta=dataset.meta, device=train_cfg.device)
+    policy = make_policy(train_cfg.policy, ds_meta=dataset.meta)
     policy.train()
 
     optimizer, _ = make_optimizer_and_scheduler(train_cfg, policy)
@@ -142,5 +141,5 @@ if __name__ == "__main__":
         raise RuntimeError("No policies were provided!")
     for ds_repo_id, policy, policy_kwargs, file_name_extra in artifacts_cfg:
         ds_name = ds_repo_id.split("/")[-1]
-        output_dir = Path("tests/data/save_policy_to_safetensors") / f"{ds_name}_{policy}_{file_name_extra}"
+        output_dir = Path("tests/artifacts/policies") / f"{ds_name}_{policy}_{file_name_extra}"
         save_policy_to_safetensors(output_dir, ds_repo_id, policy, policy_kwargs)

tests/cameras/mock_cv2.py

@@ -1,7 +1,25 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from functools import cache
 
 import numpy as np
 
+CAP_V4L2 = 200
+CAP_DSHOW = 700
+CAP_AVFOUNDATION = 1200
+CAP_ANY = -1
+
 CAP_PROP_FPS = 5
 CAP_PROP_FRAME_WIDTH = 3
 CAP_PROP_FRAME_HEIGHT = 4

tests/cameras/mock_pyrealsense2.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 enum
 
 import numpy as np

tests/cameras/test_cameras.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """
 Tests for physical cameras and their mocked versions.
 If the physical camera is not connected to the computer, or not working,
@@ -72,8 +85,8 @@ def test_camera(request, camera_type, mock):
     camera.connect()
     assert camera.is_connected
     assert camera.fps is not None
-    assert camera.width is not None
-    assert camera.height is not None
+    assert camera.capture_width is not None
+    assert camera.capture_height is not None
 
     # Test connecting twice raises an error
     with pytest.raises(RobotDeviceAlreadyConnectedError):
@@ -133,7 +146,7 @@ def test_camera(request, camera_type, mock):
         camera.connect()
 
         if mock:
-            import tests.mock_cv2 as cv2
+            import tests.cameras.mock_cv2 as cv2
         else:
             import cv2
 
@@ -191,3 +204,49 @@ def test_save_images_from_cameras(tmp_path, request, camera_type, mock):
 
     # Small `record_time_s` to speedup unit tests
     save_images_from_cameras(tmp_path, record_time_s=0.02, mock=mock)
+
+
+@pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
+@require_camera
+def test_camera_rotation(request, camera_type, mock):
+    config_kwargs = {"camera_type": camera_type, "mock": mock, "width": 640, "height": 480, "fps": 30}
+
+    # No rotation.
+    camera = make_camera(**config_kwargs, rotation=None)
+    camera.connect()
+    assert camera.capture_width == 640
+    assert camera.capture_height == 480
+    assert camera.width == 640
+    assert camera.height == 480
+    no_rot_img = camera.read()
+    h, w, c = no_rot_img.shape
+    assert h == 480 and w == 640 and c == 3
+    camera.disconnect()
+
+    # Rotation = 90 (clockwise).
+    camera = make_camera(**config_kwargs, rotation=90)
+    camera.connect()
+    # With a 90° rotation, we expect the metadata dimensions to be swapped.
+    assert camera.capture_width == 640
+    assert camera.capture_height == 480
+    assert camera.width == 480
+    assert camera.height == 640
+    import cv2
+
+    assert camera.rotation == cv2.ROTATE_90_CLOCKWISE
+    rot_img = camera.read()
+    h, w, c = rot_img.shape
+    assert h == 640 and w == 480 and c == 3
+    camera.disconnect()
+
+    # Rotation = 180.
+    camera = make_camera(**config_kwargs, rotation=None)
+    camera.connect()
+    assert camera.capture_width == 640
+    assert camera.capture_height == 480
+    assert camera.width == 640
+    assert camera.height == 480
+    no_rot_img = camera.read()
+    h, w, c = no_rot_img.shape
+    assert h == 480 and w == 640 and c == 3
+    camera.disconnect()

tests/configs/test_plugin_loading.py

@@ -0,0 +1,89 @@
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Generator
+
+import pytest
+
+from lerobot.common.envs.configs import EnvConfig
+from lerobot.configs.parser import PluginLoadError, load_plugin, parse_plugin_args, wrap
+
+
+def create_plugin_code(*, base_class: str = "EnvConfig", plugin_name: str = "test_env") -> str:
+    """Creates a dummy plugin module that implements its own EnvConfig subclass."""
+    return f"""
+from dataclasses import dataclass
+from lerobot.common.envs.configs import {base_class}
+
+@{base_class}.register_subclass("{plugin_name}")
+@dataclass
+class TestPluginConfig:
+    value: int = 42
+    """
+
+
+@pytest.fixture
+def plugin_dir(tmp_path: Path) -> Generator[Path, None, None]:
+    """Creates a temporary plugin package structure."""
+    plugin_pkg = tmp_path / "test_plugin"
+    plugin_pkg.mkdir()
+    (plugin_pkg / "__init__.py").touch()
+
+    with open(plugin_pkg / "my_plugin.py", "w") as f:
+        f.write(create_plugin_code())
+
+    # Add tmp_path to Python path so we can import from it
+    sys.path.insert(0, str(tmp_path))
+    yield plugin_pkg
+    sys.path.pop(0)
+
+
+def test_parse_plugin_args():
+    cli_args = [
+        "--env.type=test",
+        "--model.discover_packages_path=some.package",
+        "--env.discover_packages_path=other.package",
+    ]
+    plugin_args = parse_plugin_args("discover_packages_path", cli_args)
+    assert plugin_args == {
+        "model.discover_packages_path": "some.package",
+        "env.discover_packages_path": "other.package",
+    }
+
+
+def test_load_plugin_success(plugin_dir: Path):
+    # Import should work and register the plugin with the real EnvConfig
+    load_plugin("test_plugin")
+
+    assert "test_env" in EnvConfig.get_known_choices()
+    plugin_cls = EnvConfig.get_choice_class("test_env")
+    plugin_instance = plugin_cls()
+    assert plugin_instance.value == 42
+
+
+def test_load_plugin_failure():
+    with pytest.raises(PluginLoadError) as exc_info:
+        load_plugin("nonexistent_plugin")
+    assert "Failed to load plugin 'nonexistent_plugin'" in str(exc_info.value)
+
+
+def test_wrap_with_plugin(plugin_dir: Path):
+    @dataclass
+    class Config:
+        env: EnvConfig
+
+    @wrap()
+    def dummy_func(cfg: Config):
+        return cfg
+
+    # Test loading plugin via CLI args
+    sys.argv = [
+        "dummy_script.py",
+        "--env.discover_packages_path=test_plugin",
+        "--env.type=test_env",
+    ]
+
+    cfg = dummy_func()
+    assert isinstance(cfg, Config)
+    assert isinstance(cfg.env, EnvConfig.get_choice_class("test_env"))
+    assert cfg.env.value == 42

tests/conftest.py

@@ -36,51 +36,27 @@ def pytest_collection_finish():
     print(f"\nTesting with {DEVICE=}")
 
 
-@pytest.fixture
-def is_robot_available(robot_type):
-    if robot_type not in available_robots:
+def _check_component_availability(component_type, available_components, make_component):
+    """Generic helper to check if a hardware component is available"""
+    if component_type not in available_components:
         raise ValueError(
-            f"The robot type '{robot_type}' is not valid. Expected one of these '{available_robots}"
+            f"The {component_type} type is not valid. Expected one of these '{available_components}'"
         )
 
     try:
-        robot = make_robot(robot_type)
-        robot.connect()
-        del robot
+        component = make_component(component_type)
+        component.connect()
+        del component
         return True
 
     except Exception as e:
-        print(f"\nA {robot_type} robot is not available.")
+        print(f"\nA {component_type} is not available.")
 
         if isinstance(e, ModuleNotFoundError):
             print(f"\nInstall module '{e.name}'")
         elif isinstance(e, SerialException):
-            print("\nNo physical motors bus detected.")
-        else:
-            traceback.print_exc()
-
-        return False
-
-
-@pytest.fixture
-def is_camera_available(camera_type):
-    if camera_type not in available_cameras:
-        raise ValueError(
-            f"The camera type '{camera_type}' is not valid. Expected one of these '{available_cameras}"
-        )
-
-    try:
-        camera = make_camera(camera_type)
-        camera.connect()
-        del camera
-        return True
-
-    except Exception as e:
-        print(f"\nA {camera_type} camera is not available.")
-
-        if isinstance(e, ModuleNotFoundError):
-            print(f"\nInstall module '{e.name}'")
-        elif isinstance(e, ValueError) and "camera_index" in e.args[0]:
+            print("\nNo physical device detected.")
+        elif isinstance(e, ValueError) and "camera_index" in str(e):
             print("\nNo physical camera detected.")
         else:
             traceback.print_exc()
@@ -88,30 +64,19 @@ def is_camera_available(camera_type):
         return False
 
 
+@pytest.fixture
+def is_robot_available(robot_type):
+    return _check_component_availability(robot_type, available_robots, make_robot)
+
+
+@pytest.fixture
+def is_camera_available(camera_type):
+    return _check_component_availability(camera_type, available_cameras, make_camera)
+
+
 @pytest.fixture
 def is_motor_available(motor_type):
-    if motor_type not in available_motors:
-        raise ValueError(
-            f"The motor type '{motor_type}' is not valid. Expected one of these '{available_motors}"
-        )
-
-    try:
-        motors_bus = make_motors_bus(motor_type)
-        motors_bus.connect()
-        del motors_bus
-        return True
-
-    except Exception as e:
-        print(f"\nA {motor_type} motor is not available.")
-
-        if isinstance(e, ModuleNotFoundError):
-            print(f"\nInstall module '{e.name}'")
-        elif isinstance(e, SerialException):
-            print("\nNo physical motors bus detected.")
-        else:
-            traceback.print_exc()
-
-        return False
+    return _check_component_availability(motor_type, available_motors, make_motors_bus)
 
 
 @pytest.fixture

tests/datasets/test_datasets.py

@@ -45,7 +45,7 @@ from lerobot.common.robot_devices.robots.utils import make_robot
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
-from tests.utils import DEVICE, require_x86_64_kernel
+from tests.utils import require_x86_64_kernel
 
 
 @pytest.fixture
@@ -349,7 +349,6 @@ def test_factory(env_name, repo_id, policy_name):
         dataset=DatasetConfig(repo_id=repo_id, episodes=[0]),
         env=make_env_config(env_name),
         policy=make_policy_config(policy_name),
-        device=DEVICE,
     )
 
     dataset = make_dataset(cfg)
@@ -474,12 +473,12 @@ def test_flatten_unflatten_dict():
 )
 @require_x86_64_kernel
 def test_backward_compatibility(repo_id):
-    """The artifacts for this test have been generated by `tests/scripts/save_dataset_to_safetensors.py`."""
+    """The artifacts for this test have been generated by `tests/artifacts/datasets/save_dataset_to_safetensors.py`."""
 
     # TODO(rcadene, aliberts): remove dataset download
     dataset = LeRobotDataset(repo_id, episodes=[0])
 
-    test_dir = Path("tests/data/save_dataset_to_safetensors") / repo_id
+    test_dir = Path("tests/artifacts/datasets") / repo_id
 
     def load_and_compare(i):
         new_frame = dataset[i]  # noqa: B023

tests/datasets/test_delta_timestamps.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from itertools import accumulate
 
 import datasets

tests/datasets/test_image_transforms.py

@@ -33,7 +33,7 @@ from lerobot.scripts.visualize_image_transforms import (
     save_all_transforms,
     save_each_transform,
 )
-from tests.scripts.save_image_transforms_to_safetensors import ARTIFACT_DIR
+from tests.artifacts.image_transforms.save_image_transforms_to_safetensors import ARTIFACT_DIR
 from tests.utils import require_x86_64_kernel
 
 

tests/datasets/test_image_writer.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 queue
 import time
 from multiprocessing import queues

tests/datasets/test_utils.py

@@ -14,9 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
+from datasets import Dataset
 from huggingface_hub import DatasetCard
 
-from lerobot.common.datasets.utils import create_lerobot_dataset_card
+from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
+from lerobot.common.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch
 
 
 def test_default_parameters():
@@ -36,3 +39,17 @@ def test_with_tags():
     tags = ["tag1", "tag2"]
     card = create_lerobot_dataset_card(tags=tags)
     assert card.data.tags == ["LeRobot", "tag1", "tag2"]
+
+
+def test_calculate_episode_data_index():
+    dataset = Dataset.from_dict(
+        {
+            "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
+            "index": [0, 1, 2, 3, 4, 5],
+            "episode_index": [0, 0, 1, 2, 2, 2],
+        },
+    )
+    dataset.set_transform(hf_transform_to_torch)
+    episode_data_index = calculate_episode_data_index(dataset)
+    assert torch.equal(episode_data_index["from"], torch.tensor([0, 2, 3]))
+    assert torch.equal(episode_data_index["to"], torch.tensor([2, 3, 6]))

tests/envs/test_envs.py

@@ -23,8 +23,7 @@ from gymnasium.utils.env_checker import check_env
 import lerobot
 from lerobot.common.envs.factory import make_env, make_env_config
 from lerobot.common.envs.utils import preprocess_observation
-
-from .utils import require_env
+from tests.utils import require_env
 
 OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
 

tests/fixtures/constants.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from lerobot.common.constants import HF_LEROBOT_HOME
 
 LEROBOT_TEST_DIR = HF_LEROBOT_HOME / "_testing"

tests/fixtures/dataset_factories.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 random
 from functools import partial
 from pathlib import Path

tests/fixtures/files.py

@@ -1,3 +1,17 @@
+# Copyright 2024 The HuggingFace Inc. team. 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
 from pathlib import Path
 
 import datasets

tests/fixtures/hub.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from pathlib import Path
 
 import datasets

tests/fixtures/optimizers.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 pytest
 import torch
 

tests/motors/mock_dynamixel_sdk.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """Mocked classes and functions from dynamixel_sdk to allow for continuous integration
 and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
 

tests/motors/mock_scservo_sdk.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """Mocked classes and functions from dynamixel_sdk to allow for continuous integration
 and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
 

tests/motors/test_motors.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """
 Tests for physical motors and their mocked versions.
 If the physical motors are not connected to the computer, or not working,

tests/optim/test_optimizers.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 pytest
 import torch
 

tests/optim/test_schedulers.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from torch.optim.lr_scheduler import LambdaLR
 
 from lerobot.common.constants import SCHEDULER_STATE

tests/policies/test_policies.py

@@ -40,7 +40,7 @@ from lerobot.common.utils.random_utils import seeded_context
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
-from tests.scripts.save_policy_to_safetensors import get_policy_stats
+from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats
 from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
 
 
@@ -147,12 +147,11 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
         dataset=DatasetConfig(repo_id=ds_repo_id, episodes=[0]),
         policy=make_policy_config(policy_name, **policy_kwargs),
         env=make_env_config(env_name, **env_kwargs),
-        device=DEVICE,
     )
 
     # Check that we can make the policy object.
     dataset = make_dataset(train_cfg)
-    policy = make_policy(train_cfg.policy, ds_meta=dataset.meta, device=DEVICE)
+    policy = make_policy(train_cfg.policy, ds_meta=dataset.meta)
     assert isinstance(policy, PreTrainedPolicy)
 
     # Check that we run select_actions and get the appropriate output.
@@ -218,7 +217,6 @@ def test_act_backbone_lr():
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id="lerobot/aloha_sim_insertion_scripted", episodes=[0]),
         policy=make_policy_config("act", optimizer_lr=0.01, optimizer_lr_backbone=0.001),
-        device=DEVICE,
     )
     cfg.validate()  # Needed for auto-setting some parameters
 
@@ -226,7 +224,7 @@ def test_act_backbone_lr():
     assert cfg.policy.optimizer_lr_backbone == 0.001
 
     dataset = make_dataset(cfg)
-    policy = make_policy(cfg.policy, device=DEVICE, ds_meta=dataset.meta)
+    policy = make_policy(cfg.policy, ds_meta=dataset.meta)
     optimizer, _ = make_optimizer_and_scheduler(cfg, policy)
     assert len(optimizer.param_groups) == 2
     assert optimizer.param_groups[0]["lr"] == cfg.policy.optimizer_lr
@@ -258,10 +256,11 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name:
         key: ft for key, ft in features.items() if key not in policy_cfg.output_features
     }
     policy = policy_cls(policy_cfg)
+    policy.to(policy_cfg.device)
     save_dir = tmp_path / f"test_save_and_load_pretrained_{policy_cls.__name__}"
     policy.save_pretrained(save_dir)
-    policy_ = policy_cls.from_pretrained(save_dir, config=policy_cfg)
-    assert all(torch.equal(p, p_) for p, p_ in zip(policy.parameters(), policy_.parameters(), strict=True))
+    loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg)
+    torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)
 
 
 @pytest.mark.parametrize("insert_temporal_dim", [False, True])
@@ -373,7 +372,7 @@ def test_normalize(insert_temporal_dim):
         # was changed to true. For some reason, tests would pass locally, but not in CI. So here we override
         # to test with `policy.use_mpc=false`.
         ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": False}, "use_policy"),
-        ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": True}, "use_mpc"),
+        # ("lerobot/xarm_lift_medium", "tdmpc", {"use_mpc": True}, "use_mpc"),
         # TODO(rcadene): the diffusion model was normalizing the image in mean=0.5 std=0.5 which is a hack supposed to
         # to normalize the image at all. In our current codebase we dont normalize at all. But there is still a minor difference
         # that fails the test. However, by testing to normalize the image with 0.5 0.5 in the current codebase, the test pass.
@@ -412,12 +411,10 @@ def test_backward_compatibility(ds_repo_id: str, policy_name: str, policy_kwargs
            should be updated.
         4. Check that this test now passes.
         5. Remember to restore `tests/scripts/save_policy_to_safetensors.py` to its original state.
-        6. Remember to stage and commit the resulting changes to `tests/data`.
+        6. Remember to stage and commit the resulting changes to `tests/artifacts`.
     """
     ds_name = ds_repo_id.split("/")[-1]
-    artifact_dir = (
-        Path("tests/data/save_policy_to_safetensors") / f"{ds_name}_{policy_name}_{file_name_extra}"
-    )
+    artifact_dir = Path("tests/artifacts/policies") / f"{ds_name}_{policy_name}_{file_name_extra}"
     saved_output_dict = load_file(artifact_dir / "output_dict.safetensors")
     saved_grad_stats = load_file(artifact_dir / "grad_stats.safetensors")
     saved_param_stats = load_file(artifact_dir / "param_stats.safetensors")

tests/robots/test_control_robot.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """
 Tests for physical robots and their mocked versions.
 If the physical robots are not connected to the computer, or not working,
@@ -38,8 +51,8 @@ from lerobot.common.robot_devices.control_configs import (
 )
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate
-from tests.test_robots import make_robot
-from tests.utils import DEVICE, TEST_ROBOT_TYPES, mock_calibration_dir, require_robot
+from tests.robots.test_robots import make_robot
+from tests.utils import TEST_ROBOT_TYPES, mock_calibration_dir, require_robot
 
 
 @pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES)
@@ -159,8 +172,7 @@ def test_record_and_replay_and_policy(tmp_path, request, robot_type, mock):
         push_to_hub=False,
         # TODO(rcadene, aliberts): test video=True
         video=False,
-        # TODO(rcadene): display cameras through cv2 sometimes crashes on mac
-        display_cameras=False,
+        display_data=False,
         play_sounds=False,
     )
     dataset = record(robot, rec_cfg)
@@ -171,7 +183,7 @@ def test_record_and_replay_and_policy(tmp_path, request, robot_type, mock):
     replay(robot, replay_cfg)
 
     policy_cfg = ACTConfig()
-    policy = make_policy(policy_cfg, ds_meta=dataset.meta, device=DEVICE)
+    policy = make_policy(policy_cfg, ds_meta=dataset.meta)
 
     out_dir = tmp_path / "logger"
 
@@ -213,11 +225,9 @@ def test_record_and_replay_and_policy(tmp_path, request, robot_type, mock):
         num_episodes=2,
         push_to_hub=False,
         video=False,
-        display_cameras=False,
+        display_data=False,
         play_sounds=False,
         num_image_writer_processes=num_image_writer_processes,
-        device=DEVICE,
-        use_amp=False,
     )
 
     rec_eval_cfg.policy = PreTrainedConfig.from_pretrained(pretrained_policy_path)
@@ -262,7 +272,7 @@ def test_resume_record(tmp_path, request, robot_type, mock):
         episode_time_s=1,
         push_to_hub=False,
         video=False,
-        display_cameras=False,
+        display_data=False,
         play_sounds=False,
         num_episodes=1,
     )
@@ -319,7 +329,7 @@ def test_record_with_event_rerecord_episode(tmp_path, request, robot_type, mock)
             num_episodes=1,
             push_to_hub=False,
             video=False,
-            display_cameras=False,
+            display_data=False,
             play_sounds=False,
         )
         dataset = record(robot, rec_cfg)
@@ -369,7 +379,7 @@ def test_record_with_event_exit_early(tmp_path, request, robot_type, mock):
             num_episodes=1,
             push_to_hub=False,
             video=False,
-            display_cameras=False,
+            display_data=False,
             play_sounds=False,
         )
 
@@ -422,7 +432,7 @@ def test_record_with_event_stop_recording(tmp_path, request, robot_type, mock, n
             num_episodes=2,
             push_to_hub=False,
             video=False,
-            display_cameras=False,
+            display_data=False,
             play_sounds=False,
             num_image_writer_processes=num_image_writer_processes,
         )

tests/robots/test_robots.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 """
 Tests for physical robots and their mocked versions.
 If the physical robots are not connected to the computer, or not working,

tests/utils/test_io_utils.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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
 from pathlib import Path
 from typing import Any

tests/utils/test_logging_utils.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 pytest
 
 from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker

tests/utils/test_random_utils.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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 random
 
 import numpy as np

tests/utils/test_train_utils.py

@@ -1,3 +1,16 @@
+# Copyright 2024 The HuggingFace Inc. team. 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.
 from pathlib import Path
 from unittest.mock import Mock, patch