Fix bug

src/lerobot_robot_ufactory/teleoperators/gello_teleop/gello_teleop.py

@@ -28,6 +28,7 @@ class GelloTeleop(UFBaseTeleop):
         self.config = config
         self._is_connected = False
         self._teleop_enabled = False
+        self._needs_alignment = True
         self._is_calibrated = True # CHECK!!
 
         from gello.dynamixel.driver import DynamixelDriver
@@ -95,17 +96,22 @@ class GelloTeleop(UFBaseTeleop):
             raise DeviceAlreadyConnectedError(f"{self} already connected")
         from gello.agents.gello_agent import GelloAgent
 
-        self.gello_agent = GelloAgent(port=self.config.port, dynamixel_config=self._dynamixel_robo_config)
-        self.gello_agent._robot.set_torque_mode(False)
-        if not self._is_calibrated and calibrate:
-            logger.info(
-                "Mismatch between calibration values in the motor and the calibration file or no calibration file found"
-            )
-            self.calibrate()
+        try:
+            self.gello_agent = GelloAgent(port=self.config.port, dynamixel_config=self._dynamixel_robo_config)
+            self.gello_agent._robot.set_torque_mode(False)
+            if not self._is_calibrated and calibrate:
+                logger.info(
+                    "Mismatch between calibration values in the motor and the calibration file or no calibration file found"
+                )
+                self.calibrate()
 
-        self.configure()
-        self._is_connected = True
-        super().connect(calibrate)
+            self.configure()
+            self._is_connected = True
+            super().connect(calibrate)
+        except BaseException:
+            self._is_connected = False
+            self._close_gello_driver()
+            raise
         logger.info(f"{self} connected.")
 
     @property
@@ -175,11 +181,16 @@ class GelloTeleop(UFBaseTeleop):
             gello_robot.set_torque_mode(False)
             gello_robot._last_pos = None
 
+        self._needs_alignment = False
+
     def set_teleop_enabled(self, enabled: bool, obs=None):
         if enabled and not self._is_connected:
             raise DeviceNotConnectedError("Gello teleop is not connected")
+        if enabled and self._needs_alignment and obs is not None:
+            self.reset_to_robot_observation(obs)
         if not enabled and self._is_connected and hasattr(self, "gello_agent"):
             self.gello_agent._robot.set_torque_mode(False)
+            self._needs_alignment = True
         self._teleop_enabled = enabled
         logger.info("Gello teleoperation %s", "enabled" if enabled else "disabled")
 
@@ -201,9 +212,20 @@ class GelloTeleop(UFBaseTeleop):
     def send_feedback(self, feedback: dict[str, float]) -> None:
         raise NotImplementedError
 
+    def _close_gello_driver(self) -> None:
+        if not hasattr(self, "gello_agent"):
+            return
+        gello_robot = self.gello_agent._robot
+        try:
+            gello_robot.set_torque_mode(False)
+        finally:
+            gello_robot._driver.close()
+
     def disconnect(self) -> None:
-        if hasattr(self, "gello_agent"):
-            self.gello_agent._robot.set_torque_mode(False)
-        self._is_connected = False
-        self._teleop_enabled = False
+        try:
+            self._close_gello_driver()
+        finally:
+            self._is_connected = False
+            self._teleop_enabled = False
+            self._needs_alignment = True
         logger.info(f"{self} disconnected.")

tests/test_gello_reset.py

@@ -19,6 +19,9 @@ class FakeDriver:
         if self.follow_commands:
             self.positions = self.commands[-1].copy()
 
+    def close(self):
+        pass
+
 
 class FakeGelloRobot:
     def __init__(self, follow_commands=True):
@@ -35,8 +38,10 @@ class FakeGelloRobot:
 
 def make_teleop(robot):
     teleop = gello_module.GelloTeleop.__new__(gello_module.GelloTeleop)
+    teleop.id = "test_gello"
     teleop._is_connected = True
     teleop._teleop_enabled = False
+    teleop._needs_alignment = True
     teleop.dof = 2
     teleop.gello_agent = type("FakeAgent", (), {"_robot": robot})()
     return teleop
@@ -65,6 +70,7 @@ def test_gello_reset_moves_to_robot_observation_and_disables_torque(monkeypatch)
     assert np.allclose(robot._driver.positions, [0.4, 0.6, 0.5])
     assert robot._last_pos is None
     assert teleop._teleop_enabled is False
+    assert teleop._needs_alignment is False
 
 
 def test_gello_reset_failure_leaves_torque_off_and_teleop_disabled(monkeypatch):
@@ -81,6 +87,33 @@ def test_gello_reset_failure_leaves_torque_off_and_teleop_disabled(monkeypatch):
     assert teleop._teleop_enabled is False
 
 
+def test_gello_enable_after_pause_realigns_before_output(monkeypatch):
+    patch_clock(monkeypatch)
+    robot = FakeGelloRobot()
+    teleop = make_teleop(robot)
+
+    teleop.set_teleop_enabled(
+        True,
+        {"J1.pos": 0.3, "J2.pos": -0.4, "gripper.pos": 0.5},
+    )
+
+    assert robot.torque_calls == [True, False]
+    assert np.allclose(robot._driver.positions, [0.4, 0.6, 0.5])
+    assert teleop._teleop_enabled is True
+
+
+def test_gello_disconnect_closes_driver():
+    robot = FakeGelloRobot()
+    closed = []
+    robot._driver.close = lambda: closed.append(True)
+    teleop = make_teleop(robot)
+    teleop.disconnect()
+
+    assert closed == [True]
+    assert robot.torque_calls == [False]
+    assert teleop._is_connected is False
+
+
 def test_recording_reset_disables_before_robot_and_enables_after_alignment():
     calls = []