From 8c7cea1cb6883e726d4e5209d976c41d5f0482de Mon Sep 17 00:00:00 2001 From: Vinman Date: Tue, 21 Jul 2026 18:59:17 +0800 Subject: [PATCH] =?UTF-8?q?feat(robot):=20=E5=A2=9E=E5=BC=BA=20UMI=20SLAM?= =?UTF-8?q?=20=E7=A8=B3=E5=AE=9A=E6=80=A7=EF=BC=8C=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E7=9B=B8=E6=9C=BA=E7=94=BB=E9=9D=A2=E7=BC=A9=E6=94=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将插件猴子补丁改为通过 UF_LEROBOT_PATCH_TYPE 环境变量按需触发； SLAM 初始化增加就绪等待、重试及位姿置信度/时间戳校验； 新增 cameras_args 配置支持相机画面 resize； 修复遥操作 config parser 导入及多臂相机清理问题。 --- src/lerobot_robot_ufactory/__init__.py | 77 +++++++++++++---- .../robots/uf_robot/multiple_uf_robot.py | 1 + .../uf_robot/multiple_uf_robot_config.py | 1 + .../robots/uf_robot/uf_robot.py | 17 +++- .../robots/uf_robot/uf_robot_config.py | 1 + .../scripts/uf_robot_teleop.py | 8 +- .../teleoperators/umi_teleop/umi_teleop.py | 86 +++++++++++++++++-- 7 files changed, 164 insertions(+), 27 deletions(-) diff --git a/src/lerobot_robot_ufactory/__init__.py b/src/lerobot_robot_ufactory/__init__.py index c52e431..9eadbbe 100644 --- a/src/lerobot_robot_ufactory/__init__.py +++ b/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 -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 +import os + # 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 # 找到并处理后就退出，避免重复 diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot.py b/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot.py index 17ad7ac..ac2be97 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot.py +++ b/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: diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot_config.py b/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot_config.py index af6b16d..fda8e42 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/multiple_uf_robot_config.py +++ b/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__() diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py index 6c4bf03..db7fd67 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py +++ b/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 diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py index f0720df..590c11c 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot_config.py +++ b/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" diff --git a/src/lerobot_robot_ufactory/scripts/uf_robot_teleop.py b/src/lerobot_robot_ufactory/scripts/uf_robot_teleop.py index 091bf98..31ccb3d 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_robot_teleop.py +++ b/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,7 +35,11 @@ class TeleopConfig: fps: int = 30 def __post_init__(self): - self.robot.cameras = {} + if hasattr(self.robot, 'robots'): + for _, robot in self.robot.robots.items(): + robot.cameras = {} + else: + self.robot.cameras = {} def teleop_loop(cfg: TeleopConfig): diff --git a/src/lerobot_robot_ufactory/teleoperators/umi_teleop/umi_teleop.py b/src/lerobot_robot_ufactory/teleoperators/umi_teleop/umi_teleop.py index 2c6954d..bd66f3e 100644 --- a/src/lerobot_robot_ufactory/teleoperators/umi_teleop/umi_teleop.py +++ b/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') - _, pose_data = self.xvlib.xv_get_slam_data() + print(f'[{self.prefix}UMI{self.config.serial_number}] cant not get pose from vive tracker') + 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)