feat(robot): 增强 UMI SLAM 稳定性，支持相机画面缩放
将插件猴子补丁改为通过 UF_LEROBOT_PATCH_TYPE 环境变量按需触发； SLAM 初始化增加就绪等待、重试及位姿置信度/时间戳校验； 新增 cameras_args 配置支持相机画面 resize； 修复遥操作 config parser 导入及多臂相机清理问题。
This commit is contained in:
Vinman
parent
e2632889b3
commit
8c7cea1cb6
@ -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 # 找到并处理后就退出，避免重复
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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__()
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user