From da5a076f4c2a6c4cdcc710bde183adff0f86a864 Mon Sep 17 00:00:00 2001 From: Saberlve Date: Fri, 7 Aug 2026 16:04:08 +0800 Subject: [PATCH] Fix bug --- .../gello_teleop/gello_teleop.py | 50 +++++++++++++------ tests/test_gello_reset.py | 33 ++++++++++++ 2 files changed, 69 insertions(+), 14 deletions(-) diff --git a/src/lerobot_robot_ufactory/teleoperators/gello_teleop/gello_teleop.py b/src/lerobot_robot_ufactory/teleoperators/gello_teleop/gello_teleop.py index f1416af..2f5bc44 100644 --- a/src/lerobot_robot_ufactory/teleoperators/gello_teleop/gello_teleop.py +++ b/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.") diff --git a/tests/test_gello_reset.py b/tests/test_gello_reset.py index 7b160c3..5d4e59d 100644 --- a/tests/test_gello_reset.py +++ b/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 = []