diff --git a/src/lerobot/scripts/lerobot_setup_motors.py b/src/lerobot/scripts/lerobot_setup_motors.py index f86c31af2..5c1ced2c9 100644 --- a/src/lerobot/scripts/lerobot_setup_motors.py +++ b/src/lerobot/scripts/lerobot_setup_motors.py @@ -51,19 +51,7 @@ from lerobot.teleoperators import ( # noqa: F401 rebot_102_leader, so_leader, ) - -COMPATIBLE_DEVICES = [ - "koch_follower", - "koch_leader", - "omx_follower", - "omx_leader", - "openarm_mini", - "so100_follower", - "so100_leader", - "so101_follower", - "so101_leader", - "lekiwi", -] +from lerobot.utils.import_utils import register_third_party_plugins @dataclass @@ -80,18 +68,19 @@ class SetupConfig: @draccus.wrap() def setup_motors(cfg: SetupConfig): - if cfg.device.type not in COMPATIBLE_DEVICES: - raise NotImplementedError - if isinstance(cfg.device, RobotConfig): device = make_robot_from_config(cfg.device) else: device = make_teleoperator_from_config(cfg.device) - device.setup_motors() + setup = getattr(device, "setup_motors", None) + if not callable(setup): + raise NotImplementedError(f"Device type '{cfg.device.type}' does not support motor setup.") + setup() def main(): + register_third_party_plugins() setup_motors() diff --git a/tests/scripts/test_setup_motors.py b/tests/scripts/test_setup_motors.py new file mode 100644 index 000000000..4d3f4acba --- /dev/null +++ b/tests/scripts/test_setup_motors.py @@ -0,0 +1,49 @@ +# Copyright 2026 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 types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +import lerobot.scripts.lerobot_setup_motors as motors_module + + +def test_main_registers_plugins_before_parsing(monkeypatch): + calls = [] + monkeypatch.setattr(motors_module, "register_third_party_plugins", lambda: calls.append("register")) + monkeypatch.setattr(motors_module, "setup_motors", lambda: calls.append("setup")) + + motors_module.main() + + assert calls == ["register", "setup"] + + +def test_setup_motors_accepts_third_party_device(monkeypatch): + device = MagicMock() + monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device) + cfg = SimpleNamespace(device=SimpleNamespace(type="third_party")) + + motors_module.setup_motors.__wrapped__(cfg) + + device.setup_motors.assert_called_once_with() + + +def test_setup_motors_reports_unsupported_device(monkeypatch): + device = object() + monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device) + cfg = SimpleNamespace(device=SimpleNamespace(type="third_party")) + + with pytest.raises(NotImplementedError, match="third_party"): + motors_module.setup_motors.__wrapped__(cfg)