feat(robot): Make SO follower P coefficient configurable (#4142)

* Make SO follower P coefficient configurable

* chore(test): minimize tests

* feat(robots): expose PID coeff in SO arms

---------

Co-authored-by: taivu1998 <46636857+taivu1998@users.noreply.github.com>

src/lerobot/robots/bi_so_follower/bi_so_follower.py

@@ -58,6 +58,9 @@ class BiSOFollower(BimanualMixin, Robot):
             port=config.left_arm_config.port,
             disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect,
             max_relative_target=config.left_arm_config.max_relative_target,
+            position_p_coefficient=config.left_arm_config.position_p_coefficient,
+            position_i_coefficient=config.left_arm_config.position_i_coefficient,
+            position_d_coefficient=config.left_arm_config.position_d_coefficient,
             use_degrees=config.left_arm_config.use_degrees,
             cameras=left_arm_cameras,
         )
@@ -68,6 +71,9 @@ class BiSOFollower(BimanualMixin, Robot):
             port=config.right_arm_config.port,
             disable_torque_on_disconnect=config.right_arm_config.disable_torque_on_disconnect,
             max_relative_target=config.right_arm_config.max_relative_target,
+            position_p_coefficient=config.right_arm_config.position_p_coefficient,
+            position_i_coefficient=config.right_arm_config.position_i_coefficient,
+            position_d_coefficient=config.right_arm_config.position_d_coefficient,
             use_degrees=config.right_arm_config.use_degrees,
             cameras=config.right_arm_config.cameras,
         )

src/lerobot/robots/so_follower/config_so_follower.py

@@ -41,6 +41,11 @@ class SOFollowerConfig:
     # Set to `True` for backward compatibility with previous policies/dataset
     use_degrees: bool = True
 
+    # Position-mode PID gains written to Feetech STS3215 motors at connect time.
+    position_p_coefficient: int = 16
+    position_i_coefficient: int = 0
+    position_d_coefficient: int = 32
+
 
 @RobotConfig.register_subclass("so101_follower")
 @RobotConfig.register_subclass("so100_follower")

src/lerobot/robots/so_follower/so_follower.py

@@ -161,11 +161,9 @@ class SOFollower(Robot):
             self.bus.configure_motors()
             for motor in self.bus.motors:
                 self.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value)
-                # Set P_Coefficient to lower value to avoid shakiness (Default is 32)
-                self.bus.write("P_Coefficient", motor, 16)
-                # Set I_Coefficient and D_Coefficient to default value 0 and 32
-                self.bus.write("I_Coefficient", motor, 0)
-                self.bus.write("D_Coefficient", motor, 32)
+                self.bus.write("P_Coefficient", motor, self.config.position_p_coefficient)
+                self.bus.write("I_Coefficient", motor, self.config.position_i_coefficient)
+                self.bus.write("D_Coefficient", motor, self.config.position_d_coefficient)
 
                 if motor == "gripper":
                     self.bus.write("Max_Torque_Limit", motor, 500)  # 50% of max torque to avoid burnout

tests/robots/test_so100_follower.py

@@ -109,3 +109,22 @@ def test_send_action(follower):
 
     goal_pos = {m: (i + 1) * 10 for i, m in enumerate(follower.bus.motors)}
     follower.bus.sync_write.assert_called_once_with("Goal_Position", goal_pos)
+
+
+def test_configure_writes_position_pid_coefficients():
+    bus_mock = _make_bus_mock()
+    bus_mock.motors = ["shoulder_pan"]
+    robot = MagicMock()
+    robot.bus = bus_mock
+    robot.config = SO100FollowerConfig(
+        port="/dev/null",
+        position_p_coefficient=32,
+        position_i_coefficient=1,
+        position_d_coefficient=16,
+    )
+
+    SO100Follower.configure(robot)
+
+    bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32)
+    bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1)
+    bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16)