feat(robot): 增强 UMI SLAM 稳定性，支持相机画面缩放

将插件猴子补丁改为通过 UF_LEROBOT_PATCH_TYPE 环境变量按需触发；
SLAM 初始化增加就绪等待、重试及位姿置信度/时间戳校验；
新增 cameras_args 配置支持相机画面 resize；
修复遥操作 config parser 导入及多臂相机清理问题。

src/lerobot_robot_ufactory/__init__.py

@@ -1,19 +1,5 @@
-from lerobot_robot_ufactory.cameras.utils import make_cameras_from_configs as _uf_make_cameras_from_configs
+import os
-from lerobot_robot_ufactory.robots.utils import make_robot_from_config as _uf_make_robot_from_config
+
-from lerobot_robot_ufactory.teleoperators.utils import make_teleoperator_from_config as _uf_make_teleoperator_from_config
-import lerobot.cameras as _lerobot_cameras
-import lerobot.robots as _lerobot_robot
-import lerobot.teleoperators as _lerobot_teleoperators
-import lerobot.cameras.utils as _lerobot_cameras_utils
-import lerobot.robots.utils as _lerobot_robot_utils
-import lerobot.teleoperators.utils as _lerobot_teleoperators_utils
-# patch
-_lerobot_cameras.make_cameras_from_configs = _uf_make_cameras_from_configs
-_lerobot_robot.make_robot_from_config = _uf_make_robot_from_config
-_lerobot_teleoperators.make_teleoperator_from_config = _uf_make_teleoperator_from_config
-_lerobot_cameras_utils.make_cameras_from_configs = _uf_make_cameras_from_configs
-_lerobot_robot_utils.make_robot_from_config = _uf_make_robot_from_config
-_lerobot_teleoperators_utils.make_teleoperator_from_config = _uf_make_teleoperator_from_config
 # register plugin
 import lerobot_robot_ufactory.cameras.umi_camera
 import lerobot_robot_ufactory.robots.uf_robot
@@ -23,3 +9,62 @@ import lerobot_robot_ufactory.teleoperators.gello_teleop
 import lerobot_robot_ufactory.teleoperators.pika_teleop
 import lerobot_robot_ufactory.teleoperators.space_mouse
 import lerobot_robot_ufactory.teleoperators.umi_teleop
+
+def patch_lerobot_modules():
+    """
+    Patch lerobot modules to use uFactory robot implementations.
+    This function replaces the default implementations of certain functions in the lerobot package
+    with the uFactory-specific implementations. It is intended to be called at the beginning of a script
+    that uses the lerobot package, before any other imports from lerobot are made.
+    """
+    from lerobot_robot_ufactory.cameras.utils import make_cameras_from_configs as _uf_make_cameras_from_configs
+    from lerobot_robot_ufactory.robots.utils import make_robot_from_config as _uf_make_robot_from_config
+    from lerobot_robot_ufactory.teleoperators.utils import make_teleoperator_from_config as _uf_make_teleoperator_from_config
+    from lerobot_robot_ufactory.configs.parser import wrap as _uf_config_parser_wrap
+    import lerobot.cameras as _lerobot_cameras
+    import lerobot.robots as _lerobot_robot
+    import lerobot.teleoperators as _lerobot_teleoperators
+    import lerobot.cameras.utils as _lerobot_cameras_utils
+    import lerobot.robots.utils as _lerobot_robot_utils
+    import lerobot.teleoperators.utils as _lerobot_teleoperators_utils
+    import lerobot.configs.parser as _lerobot_configs_parser
+    # patch
+    _lerobot_cameras.make_cameras_from_configs = _uf_make_cameras_from_configs
+    _lerobot_robot.make_robot_from_config = _uf_make_robot_from_config
+    _lerobot_teleoperators.make_teleoperator_from_config = _uf_make_teleoperator_from_config
+    _lerobot_cameras_utils.make_cameras_from_configs = _uf_make_cameras_from_configs
+    _lerobot_robot_utils.make_robot_from_config = _uf_make_robot_from_config
+    _lerobot_teleoperators_utils.make_teleoperator_from_config = _uf_make_teleoperator_from_config
+    _lerobot_configs_parser.wrap = _uf_config_parser_wrap
+
+_UF_LEROBOT_PATCH_TYPE = os.environ.get('UF_LEROBOT_PATCH_TYPE', '0')
+if _UF_LEROBOT_PATCH_TYPE == '1':
+    print("Using UF_LEROBOT_PATCH_TYPE=1: patching lerobot modules for UFACTORY robot support.")
+    patch_lerobot_modules()
+
+# # ── 重新装饰上游脚本中已用旧 wrap 装饰过的函数 ──
+# # 因为 @parser.wrap() 在模块 import 时就执行了，
+# # register_third_party_plugins() 触发本文件时 record() 已经是旧装饰器包好的。
+# # 取出原始函数，用 _uf_config_wrap 重新包一次。
+# import sys as _sys
+# # 执行 python xxx.py 时模块名是 __main__，console_scripts 入口时是完整路径
+# _SCRIPTS_TO_REWRAP = (
+#     "__main__",
+#     "lerobot.scripts.lerobot_record",
+#     "lerobot.scripts.lerobot_teleoperate",
+#     "lerobot.scripts.lerobot_eval",
+# )
+# for _mod_name in _SCRIPTS_TO_REWRAP:
+#     _mod = _sys.modules.get(_mod_name)
+#     if _mod is None:
+#         continue
+#     _record_fn = getattr(_mod, "record", getattr(_mod, "teleoperate", getattr(_mod, "eval_main", None)))
+#     if _record_fn is None:
+#         continue
+#     _original = getattr(_record_fn, "__wrapped__", None)
+#     if _original is not None:
+#         if _mod_name == "__main__":
+#             _sys.modules[_mod_name] = _uf_config_wrap()(_original)
+#         else:
+#             setattr(_mod, _record_fn.__name__, _uf_config_wrap()(_original))
+#         break  # 找到并处理后就退出，避免重复

src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot.py

@@ -23,6 +23,7 @@ class MultipleUFRobot(Robot):
         self.action_queues = {}
         self.action_threads = {}
         for key, robot_config in self.config.robots.items():
+            robot_config.cameras_args = self.config.cameras_args
             robot = UFRobot(robot_config, prefix=key)
             self.robots[key] = robot
             if self._is_async_action:

src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot_config.py

@@ -11,6 +11,7 @@ class MultipleUFRobotConfig(RobotConfig):
     async_connect: bool = True
     async_configure: bool = True
     async_action: bool = False
+    cameras_args: dict = None
 
     def __post_init__(self):
         super().__post_init__()

src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py

@@ -75,6 +75,9 @@ class UFRobot(Robot, Thread):
         self._control_space = self.config.control_space
 
         self.real_arm = None
+        cameras_args = self.config.cameras_args or {}
+        self.camera_width = cameras_args.get('w', 0)
+        self.camera_height = cameras_args.get('h', 0)
         self.cameras = make_cameras_from_configs(config.cameras)
 
         self._is_connected = False
@@ -154,7 +157,10 @@ class UFRobot(Robot, Thread):
     def _cam_features(self) -> dict:
         cam_ft = {}
         for cam_key, cam in self.cameras.items():
-            cam_ft[f"{self.prefix}{cam_key}"] = (cam.height, cam.width, 3)
+            camera_width = self.camera_width if self.camera_width != 0 else cam.width
+            camera_height = self.camera_height if self.camera_height != 0 else cam.height
+            cam_ft[f"{self.prefix}{cam_key}"] = (camera_height, camera_width, 3)
+            # cam_ft[f"{self.prefix}{cam_key}"] = (cam.height, cam.width, 3)
         return cam_ft
 
     @property
@@ -332,7 +338,14 @@ class UFRobot(Robot, Thread):
         # Capture images from cameras
         for cam_key, cam in self.cameras.items():
             before_camread_t = time.perf_counter()
-            obs_dict[f"{self.prefix}{cam_key}"] = cam.async_read()
+            frame = cam.async_read()
+            shape = frame.shape
+            if (self.camera_height > 0 and self.camera_height != shape[0]) or (self.camera_width > 0 and self.camera_width != shape[1]):
+                camera_width = self.camera_width if self.camera_width != 0 else shape[1]
+                camera_height = self.camera_height if self.camera_height != 0 else shape[0]
+                import cv2
+                frame = cv2.resize(frame, (camera_height, camera_width), interpolation=cv2.INTER_AREA)
+            obs_dict[f"{self.prefix}{cam_key}"] = frame
             self.logs[f"async_read_camera_{cam_key}_dt_s"] = time.perf_counter() - before_camread_t
 
         return obs_dict

src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py

@@ -9,6 +9,7 @@ class UFRobotConfig(RobotConfig):
     cameras: dict[str, CameraConfig] = field(
         default_factory=lambda: {}
     )
+    cameras_args: dict = None
     robot_ip: str = "192.168.1.127"
     robot_dof: int | None = None  # Set it correctly if controlling in joint space!
     control_space: str = "joint"

src/lerobot_robot_ufactory/scripts/uf_robot_teleop.py

@@ -23,7 +23,7 @@ from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import (
     init_logging,
 )
-from lerobot.configs import parser
+from lerobot_robot_ufactory.configs import parser
 from lerobot_robot_ufactory.utils.utils import is_headless, init_keyboard_listener
 from lerobot_robot_ufactory.teleoperators.base_teleop import UFBaseTeleop
 
@@ -35,6 +35,10 @@ class TeleopConfig:
     fps: int = 30
 
     def __post_init__(self):
+        if hasattr(self.robot, 'robots'):
+            for _, robot in self.robot.robots.items():
+                robot.cameras = {}
+        else:
             self.robot.cameras = {}
 
 

src/lerobot_robot_ufactory/teleoperators/umi_teleop/umi_teleop.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python
 
 import math
+import time
 from typing import Any
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.errors import DeviceNotConnectedError
 from lerobot_robot_ufactory.devices.umi.vive_tracker.transformations import Transformations
 from lerobot_robot_ufactory.devices.umi.vive_tracker import ViveTracker
 from ..base_teleop import UFBaseTeleop
@@ -37,6 +38,7 @@ class UmiTeleop(UFBaseTeleop):
         self.begin_tracker_robot_matrix = None
         self._last_robot_pose = Transformations.rotation_matrix_to_xyzrxryrz(self.robot_base_matrix)
         self._last_gripper_pos = 0.0
+        self._last_timestamp = 0
 
     @property
     def action_features(self) -> dict:
@@ -83,17 +85,66 @@ class UmiTeleop(UFBaseTeleop):
     def configure(self) -> None:
         pass
 
+    def wait_slam_ready(self, repeat_times=3, timeout=8.0, min_confidence=0.5, stable_frames=5) -> bool:
+        def _wait_slam_ready():
+            deadline = time.monotonic() + timeout
+            ok_count = 0
+            last_ts = None
+            cnt = 0
+            while time.monotonic() < deadline and self.is_connected:
+                time.sleep(0.04)
+                ret, pose = self.xvlib.xv_get_slam_data()
+                if cnt % 25 == 0:
+                    print(f'[{self.prefix}UMI{self.config.serial_number}] Waiting for SLAM ready ... confidence: {pose.confidence}, hostTimestamp: {pose.hostTimestamp}, edgeTimestampUs: {pose.edgeTimestampUs}')
+                cnt += 1
+                if ret == 0:
+                    if pose.confidence < min_confidence or pose.hostTimestamp == last_ts:
+                        ok_count = 0
+                        last_ts = pose.hostTimestamp
+                        continue
+                    ok_count += 1
+                    last_ts = pose.hostTimestamp
+                    if ok_count >= stable_frames:
+                        return True
+            return False
+
+        print(f'[{self.prefix}UMI{self.config.serial_number}] Waiting for SLAM ready ...')
+        ready = _wait_slam_ready()
+        if ready:
+            print(f'[{self.prefix}UMI{self.config.serial_number}] ******* SLAM is ready! ******')
+        else:
+            for i in range(repeat_times):
+                if not self.is_connected:
+                    break
+                print(f'[{self.prefix}UMI{self.config.serial_number}] ******* SLAM is not ready! ******, try {i+1}/{repeat_times}')
+                self.xvlib.xv_slam_uninit()
+                time.sleep(0.5)
+                self.xvlib.xv_slam_init()
+                time.sleep(0.5)
+                ready = _wait_slam_ready()
+                if ready:
+                    print(f'[{self.prefix}UMI{self.config.serial_number}] ******* SLAM is ready! ******')
+                    return
+            print(f'[{self.prefix}UMI{self.config.serial_number}] ******* SLAM is not ready after {repeat_times} tries! ******')
+            self._is_connected = False
+            raise DeviceNotConnectedError(f'[{self.prefix}UMI{self.config.serial_number}] SLAM is not ready, please check the device and try again.')
+
     def connect(self, calibrate: bool = False) -> None:
         from lerobot_robot_ufactory.devices.umi.xvlib import XVLib
         self.tracker = ViveTracker() if self.config.use_vive_tracker else None
-        self.xvlib = XVLib(self.config.serial_number, not self.config.use_vive_tracker, self.config.use_gripper)
+        self.xvlib = XVLib(self.config.serial_number, False, self.config.use_gripper)
         if not self.config.use_vive_tracker:
+            time.sleep(1) # wait xvlib init
             self.xvlib.xv_slam_init()
+            time.sleep(1) # wait slam init
         if self.config.use_gripper:
             self.xvlib.xv_clamp_stream_init()
         self._is_connected = True
         super().connect(calibrate)
 
+        if not self.config.use_vive_tracker:
+            self.wait_slam_ready()
+
     def disconnect(self):
         super().disconnect()
         if self.xvlib:
@@ -121,7 +172,7 @@ class UmiTeleop(UFBaseTeleop):
             self.begin_tracker_robot_matrix = None
             self._last_action = None
             self._teleop_enabled = True
-            print(f'[{self.prefix}UMI] Teleoperation is start')
+            print(f'[{self.prefix}UMI{self.config.serial_number}] Teleoperation is start')
         else:
             obs = self._last_action
             if obs:
@@ -130,7 +181,7 @@ class UmiTeleop(UFBaseTeleop):
                     self._last_gripper_pos = obs[f"{self.prefix}gripper.pos"]
             self._teleop_enabled = False
             self._last_action = None
-            print(f'[{self.prefix}UMI] Teleoperation has paused')
+            print(f'[{self.prefix}UMI{self.config.serial_number}] Teleoperation has paused')
 
     # delta action
     def get_action(self) -> dict[str, Any]:
@@ -156,10 +207,31 @@ class UmiTeleop(UFBaseTeleop):
         if self.tracker is not None:
             pose_data = self.tracker.get_pose(self.config.vive_tracker_id)
             if pose_data is None:
-                print('cant not get pose from vive tracker')
+                print(f'[{self.prefix}UMI{self.config.serial_number}] cant not get pose from vive tracker')
-                _, pose_data = self.xvlib.xv_get_slam_data()
+                ret, pose_data = self.xvlib.xv_get_slam_data()
+                if ret != 0:
+                    print(f'[{self.prefix}UMI{self.config.serial_number}] cant not get pose from xvlib, ret: {ret}')
+                    return self._last_action
+                elif pose_data.hostTimestamp == self._last_timestamp:
+                    print(f'[{self.prefix}UMI{self.config.serial_number}] pose hostTimestamp is the same as last time, use last action')
+                    return self._last_action
+                self._last_timestamp = pose_data.hostTimestamp
         else:
-            _, pose_data = self.xvlib.xv_get_slam_data()
+            ret, pose_data = self.xvlib.xv_get_slam_data()
+            if ret != 0:
+                print(f'[{self.prefix}UMI{self.config.serial_number}] cant not get pose from xvlib, ret: {ret}')
+                return self._last_action
+            elif pose_data.confidence < 0.3:
+                print(f'[{self.prefix}UMI{self.config.serial_number}] pose confidence is too low: {pose_data.confidence}, use last action')
+                return self._last_action
+            # elif pose_data.hostTimestamp == self._last_timestamp:
+            #     print(f'[{self.prefix}UMI{self.config.serial_number}] pose hostTimestamp({pose_data.hostTimestamp} {self._last_timestamp}) is the same as last time, use last action')
+            #     return self._last_action
+            self._last_timestamp = pose_data.hostTimestamp
+            # print(f'[{self.prefix}UMI{self.config.serial_number}] pose11: {pose_data.position.to_list(6)}, confidence: {pose_data.confidence}, hostTimestamp: {pose_data.hostTimestamp}, edgeTimestampUs: {pose_data.edgeTimestampUs}')
+            # self.xvlib.xv_get_slam_pose(0)
+            # print(f'[{self.prefix}UMI{self.config.serial_number}] pose22: {pose_data.position.to_list(6)}, confidence: {pose_data.confidence}, hostTimestamp: {pose_data.hostTimestamp}, edgeTimestampUs: {pose_data.edgeTimestampUs}')
+
         position = pose_data.position.to_list(6)
         quaternion = pose_data.quaternion.to_list(6)