Add torque_disabled context

lerobot/common/motors/motors_bus.py

@@ -21,6 +21,7 @@
 
 import abc
 import logging
+from contextlib import contextmanager
 from dataclasses import dataclass
 from enum import Enum
 from functools import cached_property
@@ -463,6 +464,14 @@ class MotorsBus(abc.ABC):
     def enable_torque(self, motors: str | list[str] | None = None, num_retry: int = 0) -> None:
         pass
 
+    @contextmanager
+    def torque_disabled(self):
+        self.disable_torque()
+        try:
+            yield
+        finally:
+            self.enable_torque()
+
     def set_timeout(self, timeout_ms: int | None = None):
         timeout_ms = timeout_ms if timeout_ms is not None else self.default_timeout
         self.port_handler.setPacketTimeoutMillis(timeout_ms)

lerobot/common/robots/koch/koch_follower.py

@@ -146,21 +146,21 @@ class KochFollower(Robot):
         logger.info(f"Calibration saved to {self.calibration_fpath}")
 
     def configure(self) -> None:
-        self.arm.disable_torque()
+        with self.arm.torque_disabled():
             self.arm.configure_motors()
             # Use 'extended position mode' for all motors except gripper, because in joint mode the servos
-        # can't rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while
+            # can't rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling
-        # assembling the arm, you could end up with a servo with a position 0 or 4095 at a crucial
+            # the arm, you could end up with a servo with a position 0 or 4095 at a crucial point
-        # point
             for name in self.arm.names:
                 if name != "gripper":
                     self.arm.write("Operating_Mode", name, OperatingMode.EXTENDED_POSITION.value)
 
-        # Use 'position control current based' for gripper to be limited by the limit of the current.
+            # Use 'position control current based' for gripper to be limited by the limit of the current. For
-        # For the follower gripper, it means it can grasp an object without forcing too much even tho,
+            # the follower gripper, it means it can grasp an object without forcing too much even tho, its
-        # its goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
+            # goal position is a complete grasp (both gripper fingers are ordered to join and reach a touch).
-        # For the leader gripper, it means we can use it as a physical trigger, since we can force with our finger
+            # For the leader gripper, it means we can use it as a physical trigger, since we can force with
-        # to make it move, and it will move back to its original target position when we release the force.
+            # our finger to make it move, and it will move back to its original target position when we
+            # release the force.
             self.arm.write("Operating_Mode", "gripper", OperatingMode.CURRENT_POSITION.value)
 
             # Set better PID values to close the gap between recorded states and actions
@@ -168,7 +168,6 @@ class KochFollower(Robot):
             self.arm.write("Position_P_Gain", "elbow_flex", 1500)
             self.arm.write("Position_I_Gain", "elbow_flex", 0)
             self.arm.write("Position_D_Gain", "elbow_flex", 600)
-        self.arm.enable_torque()
 
     def get_observation(self) -> dict[str, Any]:
         if not self.is_connected:

lerobot/common/robots/so100/so100_follower.py

@@ -144,7 +144,7 @@ class SO100Follower(Robot):
         print("Calibration saved to", self.calibration_fpath)
 
     def configure(self) -> None:
-        self.arm.disable_torque()
+        with self.arm.torque_disabled():
             self.arm.configure_motors()
             for name in self.arm.names:
                 self.arm.write("Operating_Mode", name, OperatingMode.POSITION.value)
@@ -154,8 +154,6 @@ class SO100Follower(Robot):
                 self.arm.write("I_Coefficient", name, 0)
                 self.arm.write("D_Coefficient", name, 32)
 
-        self.arm.enable_torque()
-
     def get_observation(self) -> dict[str, Any]:
         if not self.is_connected:
             raise DeviceNotConnectedError(f"{self} is not connected.")

lerobot/common/robots/viperx/viperx.py

@@ -141,7 +141,7 @@ class ViperX(Robot):
         logger.info(f"Calibration saved to {self.calibration_fpath}")
 
     def configure(self) -> None:
-        self.arm.disable_torque()
+        with self.arm.torque_disabled():
             self.arm.configure_motors()
 
             # Set secondary/shadow ID for shoulder and elbow. These joints have two motors.
@@ -154,19 +154,18 @@ class ViperX(Robot):
             # TODO(aliberts): remove as it's actually useless in position control
             self.arm.write("Velocity_Limit", 131)
 
-        # Use 'extended position mode' for all motors except gripper, because in joint mode the servos can't
+            # Use 'extended position mode' for all motors except gripper, because in joint mode the servos
-        # rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling the arm,
+            # can't rotate more than 360 degrees (from 0 to 4095) And some mistake can happen while assembling
-        # you could end up with a servo with a position 0 or 4095 at a crucial point. See:
+            # the arm, you could end up with a servo with a position 0 or 4095 at a crucial point.
-        # https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11
+            # See: https://emanual.robotis.com/docs/en/dxl/x/x_series/#operating-mode11
             for name in self.arm.names:
                 if name != "gripper":
                     self.arm.write("Operating_Mode", name, OperatingMode.EXTENDED_POSITION.value)
 
-        # Use 'position control current based' for follower gripper to be limited by the limit of the current.
+            # Use 'position control current based' for follower gripper to be limited by the limit of the
-        # It can grasp an object without forcing too much even tho, it's goal position is a complete grasp
+            # current. It can grasp an object without forcing too much even tho, it's goal position is a
-        # (both gripper fingers are ordered to join and reach a touch).
+            # complete grasp (both gripper fingers are ordered to join and reach a touch).
             self.arm.write("Operating_Mode", "gripper", OperatingMode.CURRENT_POSITION.value)
-        self.arm.enable_torque()
 
     def get_observation(self) -> dict[str, Any]:
         """The returned observations do not have a batch dimension."""