From 2e705ff554d668299e6d085c728f34161cc6af8e Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Thu, 8 May 2025 13:10:52 +0200
Subject: [PATCH] Add setup_motors

---
 lerobot/setup_motors.py | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 lerobot/setup_motors.py

diff --git a/lerobot/setup_motors.py b/lerobot/setup_motors.py
new file mode 100644
index 000000000..7e004c324
--- /dev/null
+++ b/lerobot/setup_motors.py
@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+
+import draccus
+
+from .common.robots import RobotConfig, koch_follower, make_robot_from_config, so100_follower  # noqa: F401
+from .common.teleoperators import (  # noqa: F401
+    TeleoperatorConfig,
+    koch_leader,
+    make_teleoperator_from_config,
+    so100_leader,
+)
+
+COMPATIBLE_DEVICES = [
+    "koch_follower",
+    "koch_leader",
+    "so100_follower",
+    "so100_leader",
+]
+
+
+@dataclass
+class SetupConfig:
+    device: RobotConfig | TeleoperatorConfig
+
+
+@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()
+
+
+if __name__ == "__main__":
+    setup_motors()