From e047074825eb6ab81e35dc48416ed96f53d40e3d Mon Sep 17 00:00:00 2001
From: Simon Alibert <simon.alibert@huggingface.co>
Date: Sun, 23 Mar 2025 10:20:08 +0100
Subject: [PATCH] Add CalibrationMode

---
 lerobot/common/motors/motors_bus.py | 19 ++++++++++---------
 tests/motors/test_dynamixel.py      |  8 ++++----
 tests/motors/test_feetech.py        |  8 ++++----
 3 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/lerobot/common/motors/motors_bus.py b/lerobot/common/motors/motors_bus.py
index 2a64a887..6da8b36d 100644
--- a/lerobot/common/motors/motors_bus.py
+++ b/lerobot/common/motors/motors_bus.py
@@ -74,10 +74,17 @@ class DriveMode(Enum):
 
 
 class CalibrationMode(Enum):
-    # Joints with rotational motions are expressed in degrees in nominal range of [-180, 180]
     DEGREE = 0
-    # Joints with liner motions (like gripper of Aloha) are expressed in nominal range of [0, 100]
-    LINEAR = 1
+    RANGE_0_100 = 1
+    RANGE_M100_100 = 2
+    VELOCITY = 3
+
+
+@dataclass
+class Motor:
+    id: int
+    model: str
+    calibration: CalibrationMode
 
 
 class JointOutOfRangeError(Exception):
@@ -190,12 +197,6 @@ class GroupSyncWrite(Protocol):
     def txPacket(self): ...
 
 
-@dataclass
-class Motor:
-    id: int
-    model: str
-
-
 class MotorsBus(abc.ABC):
     """The main LeRobot class for implementing motors buses.
 
diff --git a/tests/motors/test_dynamixel.py b/tests/motors/test_dynamixel.py
index 5f70ba6a..688a7367 100644
--- a/tests/motors/test_dynamixel.py
+++ b/tests/motors/test_dynamixel.py
@@ -5,7 +5,7 @@ from unittest.mock import patch
 import dynamixel_sdk as dxl
 import pytest
 
-from lerobot.common.motors import Motor
+from lerobot.common.motors import CalibrationMode, Motor
 from lerobot.common.motors.dynamixel import DynamixelMotorsBus
 from tests.mocks.mock_dynamixel import MockMotors, MockPortHandler
 
@@ -30,9 +30,9 @@ def mock_motors() -> Generator[MockMotors, None, None]:
 @pytest.fixture
 def dummy_motors() -> dict[str, Motor]:
     return {
-        "dummy_1": Motor(id=1, model="xl430-w250"),
-        "dummy_2": Motor(id=2, model="xm540-w270"),
-        "dummy_3": Motor(id=3, model="xl330-m077"),
+        "dummy_1": Motor(1, "xl430-w250", CalibrationMode.RANGE_M100_100),
+        "dummy_2": Motor(2, "xm540-w270", CalibrationMode.RANGE_M100_100),
+        "dummy_3": Motor(3, "xl330-m077", CalibrationMode.RANGE_M100_100),
     }
 
 
diff --git a/tests/motors/test_feetech.py b/tests/motors/test_feetech.py
index e0d63c40..cb9dcda8 100644
--- a/tests/motors/test_feetech.py
+++ b/tests/motors/test_feetech.py
@@ -5,7 +5,7 @@ from unittest.mock import patch
 import pytest
 import scservo_sdk as scs
 
-from lerobot.common.motors import Motor
+from lerobot.common.motors import CalibrationMode, Motor
 from lerobot.common.motors.feetech import FeetechMotorsBus
 from tests.mocks.mock_feetech import MockMotors, MockPortHandler
 
@@ -30,9 +30,9 @@ def mock_motors() -> Generator[MockMotors, None, None]:
 @pytest.fixture
 def dummy_motors() -> dict[str, Motor]:
     return {
-        "dummy_1": Motor(id=1, model="sts3215"),
-        "dummy_2": Motor(id=2, model="sts3215"),
-        "dummy_3": Motor(id=3, model="sts3215"),
+        "dummy_1": Motor(1, "sts3215", CalibrationMode.RANGE_M100_100),
+        "dummy_2": Motor(2, "sts3215", CalibrationMode.RANGE_M100_100),
+        "dummy_3": Motor(3, "sts3215", CalibrationMode.RANGE_M100_100),
     }