From 4467e193227c785383c40b4153e7e9733e369949 Mon Sep 17 00:00:00 2001 From: Vinman Date: Thu, 11 Jun 2026 11:55:23 +0800 Subject: [PATCH] =?UTF-8?q?feat(robot):=20=E6=B7=BB=E5=8A=A0=20uFactory=20?= =?UTF-8?q?=E6=9C=BA=E6=A2=B0=E8=87=82=E5=AE=8C=E6=95=B4=E5=8A=9F=E8=83=BD?= =?UTF-8?q?=E5=8C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 包含机器人控制(uf_robot)、遥操作(teleoperators)、 摄像头(cameras)、设备驱动(devices)和执行脚本(scripts)等模块。 --- .gitignore | 3 +- pyproject.toml | 2 +- ufactory_lerobot/__init__.py | 16 + ufactory_lerobot/cameras/__init__.py | 0 .../cameras/umi_camera/__init__.py | 16 + .../cameras/umi_camera/camera_umi.py | 130 ++++ .../cameras/umi_camera/configuration_umi.py | 48 ++ ufactory_lerobot/cameras/utils.py | 38 ++ ufactory_lerobot/devices/__init__.py | 0 ufactory_lerobot/devices/pika/__init__.py | 1 + ufactory_lerobot/devices/pika/pika_device.py | 227 +++++++ ufactory_lerobot/devices/umi/__init__.py | 2 + .../devices/umi/vive_tracker/__init__.py | 2 + .../umi/vive_tracker/transformations.py | 248 +++++++ .../devices/umi/vive_tracker/vive_tracker.py | 234 +++++++ .../devices/umi/xvlib/__init__.py | 1 + ufactory_lerobot/devices/umi/xvlib/xvlib.py | 611 ++++++++++++++++++ ufactory_lerobot/robots/__init__.py | 0 .../robots/uf_mock_robot/__init__.py | 20 + .../uf_mock_robot/multiple_uf_mock_robot.py | 74 +++ .../multiple_uf_mock_robot_config.py | 8 + .../robots/uf_mock_robot/uf_mock_robot.py | 133 ++++ .../uf_mock_robot/uf_mock_robot_config.py | 28 + ufactory_lerobot/robots/uf_robot/__init__.py | 20 + .../robots/uf_robot/multiple_uf_robot.py | 122 ++++ .../uf_robot/multiple_uf_robot_config.py | 11 + ufactory_lerobot/robots/uf_robot/uf_robot.py | 465 +++++++++++++ .../robots/uf_robot/uf_robot_config.py | 44 ++ ufactory_lerobot/robots/utils.py | 34 + ufactory_lerobot/scripts/uf_lerobot_eval.py | 443 +++++++++++++ ufactory_lerobot/scripts/uf_lerobot_record.py | 364 +++++++++++ ufactory_lerobot/scripts/uf_robot_teleop.py | 118 ++++ ufactory_lerobot/teleoperators/__init__.py | 0 .../teleoperators/base_teleop/__init__.py | 1 + .../teleoperators/base_teleop/base_teleop.py | 13 + .../teleoperators/gello_teleop/__init__.py | 4 + .../gello_teleop/gello_teleop.py | 136 ++++ .../gello_teleop/gello_teleop_config.py | 20 + .../teleoperators/pika_teleop/__init__.py | 4 + .../teleoperators/pika_teleop/pika_teleop.py | 263 ++++++++ .../pika_teleop/pika_teleop_config.py | 32 + .../teleoperators/space_mouse/__init__.py | 4 + .../teleoperators/space_mouse/space_mouse.py | 187 ++++++ .../space_mouse/space_mouse_config.py | 31 + .../teleoperators/uf_mock_teleop/__init__.py | 4 + .../uf_mock_teleop/uf_mock_teleop.py | 479 ++++++++++++++ .../uf_mock_teleop/uf_mock_teleop_config.py | 48 ++ .../teleoperators/umi_teleop/__init__.py | 6 + .../umi_teleop/multiple_umi_teleop.py | 71 ++ .../umi_teleop/multiple_umi_teleop_config.py | 25 + .../teleoperators/umi_teleop/umi_teleop.py | 201 ++++++ .../umi_teleop/umi_teleop_config.py | 35 + ufactory_lerobot/teleoperators/utils.py | 40 ++ ufactory_lerobot/utils/__init__.py | 1 + ufactory_lerobot/utils/utils.py | 18 + 55 files changed, 5084 insertions(+), 2 deletions(-) create mode 100644 ufactory_lerobot/__init__.py create mode 100644 ufactory_lerobot/cameras/__init__.py create mode 100644 ufactory_lerobot/cameras/umi_camera/__init__.py create mode 100644 ufactory_lerobot/cameras/umi_camera/camera_umi.py create mode 100644 ufactory_lerobot/cameras/umi_camera/configuration_umi.py create mode 100644 ufactory_lerobot/cameras/utils.py create mode 100644 ufactory_lerobot/devices/__init__.py create mode 100644 ufactory_lerobot/devices/pika/__init__.py create mode 100644 ufactory_lerobot/devices/pika/pika_device.py create mode 100644 ufactory_lerobot/devices/umi/__init__.py create mode 100644 ufactory_lerobot/devices/umi/vive_tracker/__init__.py create mode 100644 ufactory_lerobot/devices/umi/vive_tracker/transformations.py create mode 100644 ufactory_lerobot/devices/umi/vive_tracker/vive_tracker.py create mode 100644 ufactory_lerobot/devices/umi/xvlib/__init__.py create mode 100644 ufactory_lerobot/devices/umi/xvlib/xvlib.py create mode 100644 ufactory_lerobot/robots/__init__.py create mode 100644 ufactory_lerobot/robots/uf_mock_robot/__init__.py create mode 100644 ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot.py create mode 100644 ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py create mode 100644 ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py create mode 100644 ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py create mode 100644 ufactory_lerobot/robots/uf_robot/__init__.py create mode 100644 ufactory_lerobot/robots/uf_robot/multiple_uf_robot.py create mode 100644 ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py create mode 100644 ufactory_lerobot/robots/uf_robot/uf_robot.py create mode 100644 ufactory_lerobot/robots/uf_robot/uf_robot_config.py create mode 100644 ufactory_lerobot/robots/utils.py create mode 100644 ufactory_lerobot/scripts/uf_lerobot_eval.py create mode 100644 ufactory_lerobot/scripts/uf_lerobot_record.py create mode 100644 ufactory_lerobot/scripts/uf_robot_teleop.py create mode 100644 ufactory_lerobot/teleoperators/__init__.py create mode 100644 ufactory_lerobot/teleoperators/base_teleop/__init__.py create mode 100644 ufactory_lerobot/teleoperators/base_teleop/base_teleop.py create mode 100644 ufactory_lerobot/teleoperators/gello_teleop/__init__.py create mode 100644 ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py create mode 100644 ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py create mode 100644 ufactory_lerobot/teleoperators/pika_teleop/__init__.py create mode 100644 ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py create mode 100644 ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py create mode 100644 ufactory_lerobot/teleoperators/space_mouse/__init__.py create mode 100644 ufactory_lerobot/teleoperators/space_mouse/space_mouse.py create mode 100644 ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py create mode 100644 ufactory_lerobot/teleoperators/uf_mock_teleop/__init__.py create mode 100644 ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop.py create mode 100644 ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py create mode 100644 ufactory_lerobot/teleoperators/umi_teleop/__init__.py create mode 100644 ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py create mode 100644 ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py create mode 100644 ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py create mode 100644 ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py create mode 100644 ufactory_lerobot/teleoperators/utils.py create mode 100644 ufactory_lerobot/utils/__init__.py create mode 100644 ufactory_lerobot/utils/utils.py diff --git a/.gitignore b/.gitignore index 13ba2f7..37e2e71 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,7 @@ __pycache__/ *.py[cod] *$py.class -*.so +# *.so *.egg *.egg-info/ dist/ @@ -83,3 +83,4 @@ models/ *.onnx *.engine *.trt +third_party/ diff --git a/pyproject.toml b/pyproject.toml index 71b6e62..fc1449c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ train = [ ] [tool.setuptools.packages.find] -where = ["src"] +where = ["."] [tool.ruff] line-length = 100 diff --git a/ufactory_lerobot/__init__.py b/ufactory_lerobot/__init__.py new file mode 100644 index 0000000..5308395 --- /dev/null +++ b/ufactory_lerobot/__init__.py @@ -0,0 +1,16 @@ +from ufactory_lerobot.cameras.utils import make_cameras_from_configs as _uf_make_cameras_from_configs +from ufactory_lerobot.robots.utils import make_robot_from_config as _uf_make_robot_from_config +from ufactory_lerobot.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 diff --git a/ufactory_lerobot/cameras/__init__.py b/ufactory_lerobot/cameras/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ufactory_lerobot/cameras/umi_camera/__init__.py b/ufactory_lerobot/cameras/umi_camera/__init__.py new file mode 100644 index 0000000..b0f36d3 --- /dev/null +++ b/ufactory_lerobot/cameras/umi_camera/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .camera_umi import UmiCamera +from .configuration_umi import UmiCameraConfig diff --git a/ufactory_lerobot/cameras/umi_camera/camera_umi.py b/ufactory_lerobot/cameras/umi_camera/camera_umi.py new file mode 100644 index 0000000..157d0c2 --- /dev/null +++ b/ufactory_lerobot/cameras/umi_camera/camera_umi.py @@ -0,0 +1,130 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Provides the RealSenseCamera class for capturing frames from Intel RealSense cameras. +""" + +import logging +import time +from threading import Event, Lock, Thread +from typing import Any + +import cv2 # type: ignore # TODO: add type stubs for OpenCV +import numpy as np # type: ignore # TODO: add type stubs for numpy +from numpy.typing import NDArray # type: ignore # TODO: add type stubs for numpy.typing +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from lerobot.cameras.camera import Camera +from ufactory_lerobot.devices.umi.xvlib import XVLib +from .configuration_umi import UmiCameraConfig +from lerobot.cameras.configs import ColorMode +from lerobot.cameras.utils import get_cv2_rotation + + +logger = logging.getLogger(__name__) + + +class UmiCamera(Camera): + def __init__(self, config: UmiCameraConfig): + """ + Initializes the RealSenseCamera instance. + + Args: + config: The configuration settings for the camera. + """ + + super().__init__(config) + + self.config = config + self.serial_number = self.config.serial_number + + self.fps = config.fps if config.fps else 30 + self.width = config.width if config.width else 1280 + self.height = config.height if config.height else 1280 + self.color_mode = config.color_mode + self.use_depth = config.use_depth + self.warmup_s = config.warmup_s + self.rotation: int | None = get_cv2_rotation(config.rotation) + + self.last_frame = None + self.xvlib = XVLib(self.serial_number) + self.xvlib.xv_color_camera_init() + self.xvlib.xv_set_color_camera_framerate(self.config.fps) + self.frame_lock = Lock() + self.new_frame_event: Event = Event() + self.thread = Thread(target=self._read_loop, daemon=True) + self.thread.start() + + def __str__(self) -> str: + return f"{self.__class__.__name__}({self.serial_number})" + + @property + def is_connected(self) -> bool: + return True + + @staticmethod + def find_cameras() -> list[dict[str, Any]]: + return [] + + def connect(self, warmup: bool = True) -> None: + self.xvlib.xv_color_camera_init() + if warmup: + start_time = time.monotonic() + while time.monotonic() - start_time < self.warmup_s: + time.sleep(0.1) + + def _read_loop(self): + while True: + try: + frame = self.read() + if frame is not None: + with self.frame_lock: + self.last_frame = frame + self.new_frame_event.set() + except Exception as e: + print('Read Frame Ex: {}'.format(e)) + time.sleep(0.01) + + def read(self, color_mode = None): + ret, img_data = self.xvlib.xv_get_color_image_rgb_data() + if ret <= 0: + return None + requested_color_mode = self.color_mode if color_mode is None else color_mode + if requested_color_mode not in (ColorMode.RGB, ColorMode.BGR): + raise ValueError( + f"Invalid color mode '{requested_color_mode}'. Expected {ColorMode.RGB} or {ColorMode.BGR}." + ) + if requested_color_mode == ColorMode.RGB: + frame = img_data.frame(rgb=True) + else: + frame = img_data.frame(rgb=False) + if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE, cv2.ROTATE_180]: + frame = cv2.rotate(frame, self.rotation) + return frame + + def async_read(self, timeout_ms: float = 200): + if not self.new_frame_event.wait(timeout=timeout_ms / 1000.0): + thread_alive = self.thread is not None and self.thread.is_alive() + raise TimeoutError( + f"Timed out waiting for frame from camera {self} after {timeout_ms} ms. " + f"Read thread alive: {thread_alive}." + ) + + with self.frame_lock: + frame = self.last_frame + self.new_frame_event.clear() + return frame + + def disconnect(self) -> None: + pass diff --git a/ufactory_lerobot/cameras/umi_camera/configuration_umi.py b/ufactory_lerobot/cameras/umi_camera/configuration_umi.py new file mode 100644 index 0000000..2ff25e2 --- /dev/null +++ b/ufactory_lerobot/cameras/umi_camera/configuration_umi.py @@ -0,0 +1,48 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from lerobot.cameras.configs import CameraConfig, ColorMode, Cv2Rotation + + +@CameraConfig.register_subclass("uf::umi_camera") +@dataclass +class UmiCameraConfig(CameraConfig): + serial_number: str + color_mode: ColorMode = ColorMode.RGB + use_depth: bool = False + rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION + warmup_s: int = 1 + + def __post_init__(self) -> None: + if self.color_mode not in (ColorMode.RGB, ColorMode.BGR): + raise ValueError( + f"`color_mode` is expected to be {ColorMode.RGB.value} or {ColorMode.BGR.value}, but {self.color_mode} is provided." + ) + + if self.rotation not in ( + Cv2Rotation.NO_ROTATION, + Cv2Rotation.ROTATE_90, + Cv2Rotation.ROTATE_180, + Cv2Rotation.ROTATE_270, + ): + raise ValueError( + f"`rotation` is expected to be in {(Cv2Rotation.NO_ROTATION, Cv2Rotation.ROTATE_90, Cv2Rotation.ROTATE_180, Cv2Rotation.ROTATE_270)}, but {self.rotation} is provided." + ) + + values = (self.fps, self.width, self.height) + if any(v is not None for v in values) and any(v is None for v in values): + raise ValueError( + "For `fps`, `width` and `height`, either all of them need to be set, or none of them." + ) diff --git a/ufactory_lerobot/cameras/utils.py b/ufactory_lerobot/cameras/utils.py new file mode 100644 index 0000000..03920fd --- /dev/null +++ b/ufactory_lerobot/cameras/utils.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import cast +from lerobot.utils.import_utils import make_device_from_device_class +from lerobot.cameras.utils import make_cameras_from_configs as lerobot_make_cameras_from_configs +from lerobot.cameras.camera import Camera +from lerobot.cameras.configs import CameraConfig + + +def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[str, Camera]: + lerobot_camera_configs = {key: cfg for key, cfg in camera_configs.items() if not cfg.type.startswith("uf::")} + uf_camera_configs = {key: cfg for key, cfg in camera_configs.items() if cfg.type.startswith("uf::")} + cameras = lerobot_make_cameras_from_configs(lerobot_camera_configs) + + for key, cfg in uf_camera_configs.items(): + if cfg.type == "uf::umi_camera": + from .umi_camera import UmiCamera + cameras[key] = UmiCamera(cfg) + else: + try: + cameras[key] = cast(Camera, make_device_from_device_class(cfg)) + except Exception as e: + raise ValueError(f"Error creating camera with config {cfg}: {e}") from e + return cameras diff --git a/ufactory_lerobot/devices/__init__.py b/ufactory_lerobot/devices/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ufactory_lerobot/devices/pika/__init__.py b/ufactory_lerobot/devices/pika/__init__.py new file mode 100644 index 0000000..73edced --- /dev/null +++ b/ufactory_lerobot/devices/pika/__init__.py @@ -0,0 +1 @@ +from .pika_device import PikaDevice \ No newline at end of file diff --git a/ufactory_lerobot/devices/pika/pika_device.py b/ufactory_lerobot/devices/pika/pika_device.py new file mode 100644 index 0000000..1b31e17 --- /dev/null +++ b/ufactory_lerobot/devices/pika/pika_device.py @@ -0,0 +1,227 @@ +import time +import logging +import threading +import serial +from serial.tools import list_ports + +# 配置日志 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger('pika_device') + + +def get_serial_ports(vidpid='1a86:7522'): + """ + 搜索所有指定vidpid的串口 + vidpid: 指定设备的VID:PID字符串, 默认值为'1a86:7522' + 返回找到的所有符合的串口号列表 + """ + ports = list_ports.comports() + pika_ports = [] + for port in ports: + if port.vid is not None and port.pid is not None: + if '{:04x}:{:04x}'.format(port.vid, port.pid) == vidpid: + pika_ports.append(port.device) + # else: + # print('pidvid:', '{:04x}:{:04x}'.format(port.vid, port.pid)) + return pika_ports + +def check_pika_device(port): + """ + 检测串口对应的Pika设备类型 + 返回值: + -1: 无法打开串口 + 0: 不是Pika设备 + 1: Pika Sense设备 + 2: Pika Gripper设备 + """ + try: + ser = serial.Serial( + port=port, + baudrate=460800, + bytesize=serial.EIGHTBITS, + parity=serial.PARITY_NONE, + stopbits=serial.STOPBITS_ONE, + timeout=1.0 + ) + time.sleep(0.5) # 等待串口稳定 + data = b'' + expired_time = time.monotonic() + 1.0 # 最多等待1秒 + while time.monotonic() < expired_time: + if ser.in_waiting > 0: + data += ser.read(ser.in_waiting) + if len(data) > 200: # 足够的数据来判断 + break + time.sleep(0.05) + ser.close() + data_str = data.decode('utf-8', errors='ignore') + if '"Command"' in data_str or '"AS5047"' in data_str or '"IMU"' in data_str: + # logger.info('✓ 检测到 Pika Sense 设备: {}'.format(port)) + return 1 + elif '"motor"' in data_str or '"motorstatus"' in data_str: + # logger.info('✓ 检测到 Pika Gripper 设备: {}'.format(port)) + return 2 + else: + # logger.info('✗ 未检测到 Pika 设备: {}, 数据长度: {}'.format(port, len(data))) + return 0 + except: + pass + return -1 + + +class PikaDevice(object): + # _instance = None + # _pika_sense_port = None + # _pika_gripper_port = None + # _lock = threading.Lock() + + def __init__(self, dev_type=1, **kwargs): + """ + port: serial port + dev_type: 1: sense, 2: gripper + """ + if dev_type not in [1, 2, 3]: + raise ValueError('不支持dev_type={}'.format(dev_type)) + + self._dev_type = dev_type + self._pika_sense_port = kwargs.get('pika_sense_port', None) + self._pika_gripper_port = kwargs.get('pika_gripper_port', None) + + use_pika_sense = self._dev_type in [1, 3] + use_pika_gripper = self._dev_type in [2, 3] + + self._pika_sense = None + self._pika_gripper = None + + if (use_pika_sense and self._pika_sense_port is None) or (use_pika_gripper and self._pika_gripper_port is None): + pika_ports = get_serial_ports() + if not pika_ports: + logger.error('未找到Pika设备, 请检查连接') + exit(1) + + for port in pika_ports: + device_type = check_pika_device(port) + if device_type == 1 and use_pika_sense and self._pika_sense_port is None: + self._pika_sense_port = port + logger.info('✓ 检测到 Pika Sense 设备: {}'.format(port)) + if not use_pika_gripper: + break + if device_type == 2 and use_pika_gripper and self._pika_gripper_port is None: + self._pika_gripper_port = port + logger.info('✓ 检测到 Pika Gripper 设备: {}'.format(port)) + if not use_pika_sense: + break + + if use_pika_sense and self._pika_sense_port is None: + logger.error('未找到Pika Sense设备, 请检查连接') + exit(1) + + if use_pika_gripper and self._pika_gripper_port is None: + logger.error('未找到Pika Gripper设备, 请检查连接') + exit(1) + + if use_pika_sense: + print('Pika Sense设备:', self._pika_sense_port) + if use_pika_gripper: + print('Pika Gripper 设备:', self._pika_gripper_port) + + self.pika_tracker_device = None + + # def __new__(cls, *args, **kwargs): + # if not cls._instance: + # with cls._lock: + # if not cls._instance: + # cls._instance = super().__new__(cls) + # cls._instance.init(*args, *kwargs) + # return cls._instance + + def __del__(self): + if self._pika_sense: + self._pika_sense.disconnect() + if self._pika_gripper: + self._pika_gripper.disconnect() + + @property + def pika_sense(self): + if self._dev_type not in [1, 3]: + return None + if self._pika_sense is None: + from pika.sense import Sense + # 初始化Sense对象 + self._pika_sense = Sense(port=self._pika_sense_port) + # 连接设备 + if not self._pika_sense.connect(): + logger.error('连接Pika Sense设备失败') + exit(1) + logger.info('Pika Sense设备连接成功') + + # 配置Vive Tracker（可选） + # sense.set_vive_tracker_config(config_path='path/to/config', lh_config='lighthouse_config') + + tracker = self._pika_sense.get_vive_tracker() + if not tracker: + logger.error('Vive Tracker初始化失败') + self._pika_sense.disconnect() + exit(1) + logger.info('Vive Tracker初始化成功') + time.sleep(2) + + devices = self._pika_sense.get_tracker_devices() + if not devices: + logger.error('未检测到Vive Tracker设备') + self._pika_sense.disconnect() + exit(1) + logger.info('检测到Vive Tracker设备: {}'.format(devices)) + + self.pika_tracker_device = None + for device in devices: + if device.startswith('WM'): + self.pika_tracker_device = device + break + else: + self.pika_tracker_device = devices[0] + logger.info('开始跟踪设备: {}\n'.format(self.pika_tracker_device)) + return self._pika_sense + + @property + def pika_gripper(self): + if self._dev_type not in [2, 3]: + return None + if self._pika_gripper is None: + if self._dev_type in [2, 3]: + from pika.gripper import Gripper + self._pika_gripper = Gripper(port=self._pika_gripper_port) + # 连接设备 + if not self._pika_gripper.connect(): + logger.error('连接Pika Gripper设备失败') + if self._dev_type in [1, 3]: + self.pika_sense.disconnect() + exit(1) + logger.info('Pika Gripper设备连接成功') + return self._pika_gripper + + +if __name__ == '__main__': + pika_device1 = PikaDevice(1) + pika_device1.pika_sense + pika_device1.pika_gripper + time.sleep(3) + + # input('=================') + + pika_device2 = PikaDevice(2) + pika_device2.pika_sense + pika_device2.pika_gripper + + input('=================') + + print(pika_device1) + print(pika_device1.pika_sense) + print(pika_device1.pika_gripper) + + print(pika_device2) + print(pika_device2.pika_sense) + print(pika_device2.pika_gripper) + + input('=================') + + \ No newline at end of file diff --git a/ufactory_lerobot/devices/umi/__init__.py b/ufactory_lerobot/devices/umi/__init__.py new file mode 100644 index 0000000..a5d5610 --- /dev/null +++ b/ufactory_lerobot/devices/umi/__init__.py @@ -0,0 +1,2 @@ +from .xvlib import XVLib +from .vive_tracker import Transformations, ViveTracker \ No newline at end of file diff --git a/ufactory_lerobot/devices/umi/vive_tracker/__init__.py b/ufactory_lerobot/devices/umi/vive_tracker/__init__.py new file mode 100644 index 0000000..79f23ff --- /dev/null +++ b/ufactory_lerobot/devices/umi/vive_tracker/__init__.py @@ -0,0 +1,2 @@ +from .transformations import Transformations +from .vive_tracker import ViveTracker \ No newline at end of file diff --git a/ufactory_lerobot/devices/umi/vive_tracker/transformations.py b/ufactory_lerobot/devices/umi/vive_tracker/transformations.py new file mode 100644 index 0000000..f7ed23f --- /dev/null +++ b/ufactory_lerobot/devices/umi/vive_tracker/transformations.py @@ -0,0 +1,248 @@ +import numpy as np + + +class Transformations: + @staticmethod + def quaternion_to_rotation_matrix(q): + """ + 将四元数转换为旋转矩阵 + + 注: 四元素顺序为xyzw + """ + norm = np.linalg.norm(q) + if norm < 1e-6: + raise ValueError('零四元数无法归一化') + + x, y, z, w = q / norm # 归一化 + xx, yy, zz = x * x, y * y, z * z + xy, xz, yz = x * y, x * z, y * z + wx, wy, wz = w * x, w * y, w * z + + R = np.array([ + [1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy)], + [ 2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx)], + [ 2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy)] + ]) + return R + + @staticmethod + def rotation_matrix_to_quaternion(R): + """ + 将3x3变换矩阵转换四元数 + 注: 四元素顺序为xyzw + """ + # 提取旋转矩阵部分 + rot_matrix = R[:3, :3] + + # 计算四元数 + trace = rot_matrix[0, 0] + rot_matrix[1, 1] + rot_matrix[2, 2] + + if trace > 0: + s = 0.5 / np.sqrt(trace + 1.0) + qw = 0.25 / s + qx = (rot_matrix[2, 1] - rot_matrix[1, 2]) * s + qy = (rot_matrix[0, 2] - rot_matrix[2, 0]) * s + qz = (rot_matrix[1, 0] - rot_matrix[0, 1]) * s + elif rot_matrix[0, 0] > rot_matrix[1, 1] and rot_matrix[0, 0] > rot_matrix[2, 2]: + s = 2.0 * np.sqrt(1.0 + rot_matrix[0, 0] - rot_matrix[1, 1] - rot_matrix[2, 2]) + qw = (rot_matrix[2, 1] - rot_matrix[1, 2]) / s + qx = 0.25 * s + qy = (rot_matrix[0, 1] + rot_matrix[1, 0]) / s + qz = (rot_matrix[0, 2] + rot_matrix[2, 0]) / s + elif rot_matrix[1, 1] > rot_matrix[2, 2]: + s = 2.0 * np.sqrt(1.0 + rot_matrix[1, 1] - rot_matrix[0, 0] - rot_matrix[2, 2]) + qw = (rot_matrix[0, 2] - rot_matrix[2, 0]) / s + qx = (rot_matrix[0, 1] + rot_matrix[1, 0]) / s + qy = 0.25 * s + qz = (rot_matrix[1, 2] + rot_matrix[2, 1]) / s + else: + s = 2.0 * np.sqrt(1.0 + rot_matrix[2, 2] - rot_matrix[0, 0] - rot_matrix[1, 1]) + qw = (rot_matrix[1, 0] - rot_matrix[0, 1]) / s + qx = (rot_matrix[0, 2] + rot_matrix[2, 0]) / s + qy = (rot_matrix[1, 2] + rot_matrix[2, 1]) / s + qz = 0.25 * s + + return [qx, qy, qz, qw] + + @staticmethod + def rpy_to_rotation_matrix(roll, pitch, yaw): + """RPY角到旋转矩阵的转换""" + cr, sr = np.cos(roll), np.sin(roll) + cp, sp = np.cos(pitch), np.sin(pitch) + cy, sy = np.cos(yaw), np.sin(yaw) + + R = np.array([ + [cp*cy, -cr*sy + sr*sp*cy, sr*sy + cr*sp*cy], + [cp*sy, cr*cy + sr*sp*sy, -sr*cy + cr*sp*sy], + [ -sp, sr*cp, cr*cp], + ]) + + return R + + @staticmethod + def rotation_matrix_to_rpy(R, yaw_zero=True): + """ + 旋转矩阵到RPY角的转换 + + yaw_zero: 万向节锁情况下, True就把yaw置0, False就把roll置0 + 返回: roll, pitch, yaw + """ + epsilon = 1e-6 + if abs(R[2, 0]) > 1 - epsilon: # 万向节锁(pitch=±90°) + pitch = np.arcsin(-R[2, 0]) + roll_yaw = np.arctan2(-R[0, 1], R[1, 1]) + if yaw_zero: + # 保留roll, 把yaw置0 + roll, yaw = roll_yaw, 0 + else: + # 保留yaw, 把roll置0 + roll, yaw = 0, roll_yaw + else: + roll = np.arctan2(R[2, 1], R[2, 2]) + pitch = np.arcsin(-R[2, 0]) + yaw = np.arctan2(R[1, 0], R[0, 0]) + + return roll, pitch, yaw + + @staticmethod + def rxryrz_to_matrix(axis_angle): + """ + 将轴角向量 (rx, ry, rz) 转换为 3x3 旋转矩阵。 + 输入: np.array([rx, ry, rz]) + - 方向: 旋转轴 + - 模长: 旋转角度 (弧度) + """ + theta = np.linalg.norm(axis_angle) + + # 如果角度接近0，返回单位矩阵 + if theta < 1e-8: + return np.eye(3) + + # 归一化旋转轴 + axis = axis_angle / theta + x, y, z = axis + + c = np.cos(theta) + s = np.sin(theta) + t = 1 - c + + # 罗德里格斯旋转公式 + # R = I + sin(theta)*K + (1-cos(theta))*K^2 + # 展开为矩阵形式： + R = np.array([ + [t*x*x + c, t*x*y - s*z, t*x*z + s*y], + [t*x*y + s*z, t*y*y + c, t*y*z - s*x], + [t*x*z - s*y, t*y*z + s*x, t*z*z + c] + ]) + + return R + + @staticmethod + def rotation_matrix_to_rxryrz(R): + """ + 旋转矩阵到轴角的转换 (rx, ry, rz = aixs * angle) + 返回: rx, ry, rz + """ + R = np.asarray(R) + if R.shape[-2:] != (3, 3): + raise ValueError("Input must be (..., 3, 3)") + + # 计算旋转角度 theta + trace = np.trace(R) + cos_theta = (trace - 1) / 2.0 + cos_theta = np.clip(cos_theta, -1.0, 1.0) # 防止数值误差导致 arccos 越界 + theta = np.arccos(cos_theta) + eps = 1e-8 + + # 情况 1: 无旋转 (theta ≈ 0) + if theta < eps: + axis = np.array([1.0, 0.0, 0.0]) + return axis * 0.0 + + # 情况 2: 旋转角度接近 pi (180 度) + if np.pi - theta < eps: + # 此时 sin(theta) ≈ 0，不能用反对称公式 + # 从 R 对角线提取轴：R = I + 2 * (uu^T - I) => uu^T = (R + I)/2 + # 所以 u_i^2 = (R_ii + 1)/2 + diag = np.diag(R) + axis = np.sqrt(np.maximum(diag + 1, 0)) # 取非负根 + + # 确定符号：利用非对角元素，例如 R[0,1] = 2*u0*u1 + if axis[0] > eps: + if R[0, 1] < 0: + axis[1] *= -1 + if R[0, 2] < 0: + axis[2] *= -1 + elif axis[1] > eps: + if R[1, 2] < 0: + axis[2] *= -1 + # 注意：可能存在符号歧义，但旋转效果相同 + + axis = axis / np.linalg.norm(axis) + return axis * theta + + # 情况 3: 一般情况 (0 < theta < pi) + sin_theta = np.sin(theta) + axis = np.array([ + R[2, 1] - R[1, 2], + R[0, 2] - R[2, 0], + R[1, 0] - R[0, 1] + ]) / (2 * sin_theta) + + axis = axis / np.linalg.norm(axis) # 确保单位长度（数值误差可能破坏） + return axis * theta + + @classmethod + def xyzq_to_rotation_matrix(cls, x, y, z, q): + T = np.eye(4) + T[:3, :3] = cls.quaternion_to_rotation_matrix(q) + T[:3, 3] = [x, y, z] + return T + + @classmethod + def xyzrpy_to_rotation_matrix(cls, x, y, z, roll, pitch, yaw): + """构造4x4齐次变换矩阵""" + T = np.eye(4) + T[:3, :3] = cls.rpy_to_rotation_matrix(roll, pitch, yaw) + T[:3, 3] = [x, y, z] + return T + + @classmethod + def rotation_matrix_to_xyzq(cls, rotation_matrix): + """从4x4齐次变换矩阵到xyzq的转换""" + x, y, z = rotation_matrix[0, 3], rotation_matrix[1, 3], rotation_matrix[2, 3] + q = cls.rotation_matrix_to_quaternion(rotation_matrix[:3, :3]) + return [x, y, z, q] + + @classmethod + def rotation_matrix_to_xyzrpy(cls, rotation_matrix): + """从4x4齐次变换矩阵到xyzrpy的转换""" + x, y, z = rotation_matrix[0, 3], rotation_matrix[1, 3], rotation_matrix[2, 3] + roll, pitch, yaw = cls.rotation_matrix_to_rpy(rotation_matrix) + return [x, y, z, roll, pitch, yaw] + + @classmethod + def rotation_matrix_to_xyzrxryrz(cls, rotation_matrix): + """从4x4齐次变换矩阵到xyzrxryrz的转换""" + x, y, z = rotation_matrix[0, 3], rotation_matrix[1, 3], rotation_matrix[2, 3] + rx, ry, rz = cls.rotation_matrix_to_rxryrz(rotation_matrix[:3,:3]) + return [x, y, z, rx, ry, rz] + + @classmethod + def tracker_pose_to_robot_matrix(cls, x, y, z, q, tracker_to_robot_matrix): + # Tracker位置对应的变换矩阵 + tracker_matrix = cls.xyzq_to_rotation_matrix(x, y, z, q) + # Tracker位置转换到机械臂坐标系后对应的变换矩阵 + robot_matrix = np.dot(tracker_matrix, tracker_to_robot_matrix) + return robot_matrix + + @classmethod + def tracker_robot_matrix_to_robot_pose(cls, begin_tracker_robot_matrix, end_tracker_robot_matrix, robot_base_matrix, is_axis_angle=False): + # 机械臂目标位置对应的变换矩阵 + # 机械臂目标 = 机械臂初始位置 + (当前手姿 - 初始手姿) + delta_matrix = np.dot(np.linalg.inv(begin_tracker_robot_matrix), end_tracker_robot_matrix) + robot_martix = np.dot(robot_base_matrix, delta_matrix) + if is_axis_angle: + return cls.rotation_matrix_to_xyzrxryrz(robot_martix) + else: + return cls.rotation_matrix_to_xyzrpy(robot_martix) diff --git a/ufactory_lerobot/devices/umi/vive_tracker/vive_tracker.py b/ufactory_lerobot/devices/umi/vive_tracker/vive_tracker.py new file mode 100644 index 0000000..0454e0f --- /dev/null +++ b/ufactory_lerobot/devices/umi/vive_tracker/vive_tracker.py @@ -0,0 +1,234 @@ +import sys +import ctypes +import logging +import threading +import pysurvive +import numpy as np +from .transformations import Transformations + +# 配置日志 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger('uf.vive_tracker') + + +class Vector(ctypes.Structure): + def __getitem__(self, index): + # 获取字段名列表 + field_name = self._fields_[index][0] + # 使用 getattr 获取对应属性的值 + return getattr(self, field_name) + + def __setitem__(self, index, value): + # 获取字段名列表 + field_name = self._fields_[index][0] + # 使用 setattr 设置对应属性的值 + setattr(self, field_name, value) + + def __str__(self): + return f'{self.to_list(6)}' + + def to_list(self, ndigits=6): + return [round(getattr(self, item[0]), ndigits=ndigits) for item in self._fields_] + + +class Vector3D(Vector): + _fields_ = [ + ("x", ctypes.c_double), + ("y", ctypes.c_double), + ("z", ctypes.c_double) + ] + + +class Vector4D(Vector): + _fields_ = [ + ("x", ctypes.c_double), + ("y", ctypes.c_double), + ("z", ctypes.c_double), + ("w", ctypes.c_double) + ] + + +class PoseData(ctypes.Structure): + _fields_ = [ + ("position", Vector3D), + # ("orientation", Vector3D), + ("quaternion", Vector4D), + ("hostTimestamp", ctypes.c_double), + # ("edgeTimestampUs", ctypes.c_longlong), + # ("confidence", ctypes.c_double) + ] + + +class SingletonMeta(type): + _instances = {} + + def __call__(cls, *args, **kwargs): + if cls not in cls._instances: + cls._instances[cls] = super().__call__(*args, **kwargs) + return cls._instances[cls] + + +class ViveTracker(metaclass=SingletonMeta): + # _instance = None + # _initialized = False + def __init__(self, config_path=None, lh_config=None, args=None): + # if self._initialized: + # return + # self._initialized = True + self.config_path = config_path + self.lh_config = lh_config + self.args = args if args else [] + self.running = False + self.context = None + self.collector_thread = None + self.data_lock = threading.Lock() + self.latest_poses = {} + self.latest_raw_poses = {} + self.init() + + # def __new__(cls, *args, **kwargs): + # if cls._instance is None: + # cls._instance = super().__new__(cls) + # return cls._instance + + def __del__(self): + logger.info("正在停止Vive Tracker位姿追踪...") + self.running = False + # 等待线程结束 + if self.collector_thread: + self.collector_thread.join(timeout=2.0) + # 清理资源 + self.context = None + logger.info("Vive Tracker已断开连接") + + @staticmethod + def to_str(v): + return v.decode("utf-8") if isinstance(v, bytes) else str(v) + + def list_devices(self): + # import pysurvive + # for obj in self.context.Objects(): + # name = self.to_str(obj.Name()) + # serial_number = None + # if hasattr(pysurvive, "simple_serial_number"): + # serial_number = self.to_str(pysurvive.simple_serial_number(obj.ptr)) + # print("object:", name, "serial:", serial_number) + return [key for key in self.latest_poses.keys() if not key.startswith('WM')] + + def init(self): + # 构建pysurvive参数 + survive_args = sys.argv[:1] # 保留程序名 + + # 添加配置文件参数 + if self.config_path: + survive_args.extend(['--config', self.config_path]) + + # 添加灯塔配置参数 + if self.lh_config: + survive_args.extend(['--lh', self.lh_config]) + + # 添加其他参数 + survive_args.extend(self.args) + try: + logger.info("正在初始化pysurvive...") + self.context = pysurvive.SimpleContext(survive_args) + if not self.context: + logger.error("错误: 无法初始化pysurvive上下文") + return False + + logger.info("pysurvive初始化成功") + # 标记为运行状态 + self.running = True + + # 创建并启动位姿收集线程 + self.collector_thread = threading.Thread(target=self._pose_collector) + self.collector_thread.daemon = True + self.collector_thread.start() + except Exception as e: + logger.error(f"连接Vive Tracker时发生错误: {e}") + self.running = False + return False + + def _pose_collector(self): + initial_rotation = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, -30 / 180.0 * np.pi, 0, 0) + + # alignment_rotation = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, -np.pi / 2, -np.pi / 2, 0) + alignment_rotation = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, -np.pi / 2, np.pi, np.pi) + + rotate_matrix = np.dot(initial_rotation, alignment_rotation) + # 应用平移变换 - 将采集到的pose数据变换到夹爪中心 + # transform_matrix = Transformations.xyzrpy_to_rotation_matrix(0.172, 0, -0.076, 0, 0, 0) + # transform_matrix = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, 0, 0, 0) + + # tracker_to_robot_matrix = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, np.pi, 0, np.pi) + tracker_to_robot_matrix = Transformations.xyzrpy_to_rotation_matrix(0, 0, 0, 0, 0, -np.pi / 2) + + robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*[0, 0, 0, np.pi, -np.pi / 2, 0]) + begin_tracker_robot_matrix = None + + cnt = 0 + + # 持续获取最新位姿 + while self.running and self.context.Running(): + updated = self.context.NextUpdated() + if not updated: + continue + if cnt < 100: + cnt += 1 + continue + # 获取设备名称 + device_name = str(updated.Name(), 'utf-8') + serial_number = None + if hasattr(pysurvive, "simple_serial_number"): + serial_number = self.to_str(pysurvive.simple_serial_number(updated.ptr)) + # 获取位姿数据 + pose_obj = updated.Pose() + pose_data = pose_obj[0] # 位姿数据 + timestamp = pose_obj[1] # 时间戳 + position = [pose_data.Pos[0], pose_data.Pos[1], pose_data.Pos[2]] + quaternion = [pose_data.Rot[1], pose_data.Rot[2], pose_data.Rot[3], pose_data.Rot[0]] + origin_mat = Transformations.xyzq_to_rotation_matrix(*position, quaternion) + # tracker_matrix = np.dot(origin_mat, rotate_matrix) + tracker_matrix = np.matmul(origin_mat, rotate_matrix) + # tracker_matrix = np.matmul(np.matmul(origin_mat, rotate_matrix), transform_matrix) + + x, y, z, q = Transformations.rotation_matrix_to_xyzq(tracker_matrix) + pose_data = PoseData(position=Vector3D(x, y, z), quaternion=Vector4D(*q), hostTimestamp=timestamp) + pose_raw_data = PoseData(position=Vector3D(*position), quaternion=Vector4D(*quaternion), hostTimestamp=timestamp) + with self.data_lock: + self.latest_poses[device_name] = pose_data + self.latest_raw_poses[device_name] = pose_raw_data + if serial_number: + self.latest_poses[serial_number] = pose_data + self.latest_raw_poses[serial_number] = pose_raw_data + + # tracker_robot_matrix = np.dot(tracker_matrix, tracker_to_robot_matrix) + # if begin_tracker_robot_matrix is None: + # begin_tracker_robot_matrix = tracker_robot_matrix + # pose = Transformations.tracker_robot_matrix_to_robot_pose(begin_tracker_robot_matrix, tracker_robot_matrix, robot_base_matrix, is_axis_angle=True) + + # pose_data = PoseData(position=Vector3D(*pose[:3]), orientation=Vector3D(*pose[3:]), quaternion=Vector4D(*quaternion), hostTimestamp=timestamp) + # with self.data_lock: + # self.latest_poses[device_name] = pose_data + + def get_pose(self, device_name=None): + if device_name: + with self.data_lock: + if device_name in self.latest_poses: + return self.latest_poses[device_name] + else: + return None + else: + with self.data_lock: + return self.latest_poses.copy() + + def get_raw_pose(self, device_name=None): + if device_name: + with self.data_lock: + if device_name in self.latest_raw_poses: + return self.latest_raw_poses[device_name] + else: + return None + else: + with self.data_lock: + return self.latest_raw_poses.copy() diff --git a/ufactory_lerobot/devices/umi/xvlib/__init__.py b/ufactory_lerobot/devices/umi/xvlib/__init__.py new file mode 100644 index 0000000..0400e92 --- /dev/null +++ b/ufactory_lerobot/devices/umi/xvlib/__init__.py @@ -0,0 +1 @@ +from .xvlib import XVLib \ No newline at end of file diff --git a/ufactory_lerobot/devices/umi/xvlib/xvlib.py b/ufactory_lerobot/devices/umi/xvlib/xvlib.py new file mode 100644 index 0000000..12c3291 --- /dev/null +++ b/ufactory_lerobot/devices/umi/xvlib/xvlib.py @@ -0,0 +1,611 @@ +import os +import ctypes +import cv2 +import time +import logging +import numpy as np + +# 配置日志 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger('uf.xvlib') + + +class DeviceStruct(ctypes.Structure): + _fields_ = [ + ("uuid", ctypes.c_char * 100) + ] + + @property + def serial_number(self): + return self.uuid.decode('utf-8') + + +class Vector(ctypes.Structure): + def __getitem__(self, index): + # 获取字段名列表 + field_name = self._fields_[index][0] + # 使用 getattr 获取对应属性的值 + return getattr(self, field_name) + + def __setitem__(self, index, value): + # 获取字段名列表 + field_name = self._fields_[index][0] + # 使用 setattr 设置对应属性的值 + setattr(self, field_name, value) + + def __str__(self): + return f'{self.to_list(6)}' + + def to_list(self, ndigits=6): + return [round(getattr(self, item[0]), ndigits=ndigits) for item in self._fields_] + + +class Vector3B(Vector): + _fields_ = [ + ("x", ctypes.c_bool), + ("y", ctypes.c_bool), + ("z", ctypes.c_bool) + ] + + +class Vector3D(Vector): + _fields_ = [ + ("x", ctypes.c_double), + ("y", ctypes.c_double), + ("z", ctypes.c_double) + ] + + +class Vector4D(Vector): + _fields_ = [ + ("x", ctypes.c_double), + ("y", ctypes.c_double), + ("z", ctypes.c_double), + ("w", ctypes.c_double) + ] + + +class ClampData(ctypes.Structure): + _fields_ = [ + ("timestamp", ctypes.c_double), + ("data", ctypes.c_double) + ] + + +class ColorImageData(ctypes.Structure): + _fields_ = [ + ("codec", ctypes.c_int), + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("data", ctypes.c_uint8 * int(1280*1280*3)), + ("dataSize", ctypes.c_uint), + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong) + ] + def frame(self, rgb=False): + np_array = np.frombuffer(bytes(self.data[:self.dataSize]), dtype=np.uint8) + if self.codec == 0: # YUYV 格式, 重塑为 (h, w, 2)，因为每两个字节包含 Y 和 UV 信息 + yuv_mat = np_array.reshape((self.height, self.width, 2)) + if rgb: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2RGB_YUYV) + else: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2BGR_YUYV) + elif self.codec == 1: # YU12 (即 I420) 格式 (UV 平面) + yuv_mat = np_array.reshape((int(self.height * 1.5), self.width)) + if rgb: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2RGB_I420) + else: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2BGR_I420) + elif self.codec == 2: # JPEG 格式, 直接解码，不需要知道宽高（宽高包含在 JPEG 头中，但可以用 w,h 校验） + frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR) + elif self.codec == 3: # NV12 格式 (UV 交错) + yuv_mat = np_array.reshape((int(self.height * 1.5), self.width)) + if rgb: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2RGB_NV12) + else: + frame = cv2.cvtColor(yuv_mat, cv2.COLOR_YUV2BGR_NV12) + elif self.codec == 4: # BITSTREAM (H.264/H.265) 格式, 同样使用 imdecode，OpenCV 会自动处理常见的视频流头 + frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR) + else: + frame = np_array + return frame + + +class DepthImageData(ctypes.Structure): + _fields_ = [ + ("type", ctypes.c_int), + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("confidence", ctypes.c_double), + ("data", ctypes.c_uint8 * int(1280*1280*3)), + ("dataSize", ctypes.c_uint), + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong) + ] + def frame(self): + np_array = np.frombuffer(bytes(self.data[:self.dataSize]), dtype=np.uint8) + if self.type == 0: # Depth_16, 数据大小应为 w * h * 2 + # 1. 转换为 uint16 类型 + depth_uint16 = np_array.view(dtype=np.uint16).reshape((self.height, self.width)) + + # 2. 归一化用于显示 (0-255) + # 深度值通常在 0-65535 (mm)，直接显示是全黑的 + # cv2.normalize 将数据拉伸到 0-255 范围 + depth_norm = cv2.normalize(depth_uint16, None, 0, 255, cv2.NORM_MINMAX) + depth_norm = np.uint8(depth_norm) # 转为 8位灰度图 + + # 3. 可选：转为伪彩色以便观察 + depth_color = cv2.applyColorMap(depth_norm, cv2.COLORMAP_JET) + frame = depth_color + elif self.type == 1: # Depth_32, 数据大小应为 w * h * 4 + depth_float = np_array.view(dtype=np.float32).reshape((self.height, self.width)) + + # 显示处理：截取有效范围 (例如 0-5米) 并归一化 + # 注意：这里假设最大值是 5000mm 或 5.0m，根据实际情况调整 + # max_depth = 5000.0 if np.max(depth_float) > 100 else 5.0 + depth_norm = cv2.normalize(depth_float, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U) + depth_color = cv2.applyColorMap(depth_norm, cv2.COLORMAP_JET) + frame = depth_color + elif self.type == 2: # IR, 通常是 8位或16位灰度图, 数据大小可能是 w*h (8bit) 或 w*h*2 (16bit) + # 尝试根据数据长度判断位深 + if len(np_array) == self.width * self.height: + ir_img = np_array.reshape((self.height, self.width)) # 8位 + elif len(np_array) == self.width * self.height * 2: + ir_img = np_array.view(dtype=np.uint16).reshape((self.height, self.width)) # 16位 + # 16位转8位显示 + ir_img = cv2.normalize(ir_img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) + else: + raise ValueError("IR 数据长度不匹配") + # 转伪彩色 + ir_color = cv2.applyColorMap(ir_img, cv2.COLORMAP_JET) + frame = ir_color + elif self.type == 3: # Cloud, 这不是图像，是 xyz 坐标集合, 数据大小应为 w * h * 3 * 4 (float32) 或类似 + # 假设是 float32 格式 + cloud_data = np_array.view(dtype=np.float32).reshape((-1, 3)) + # cloud_data 现在是一个 N x 3 的数组，每一行是 (x, y, z) + # 这里不返回图像，返回点云数据供 PCL 或 Open3D 处理 + frame = cloud_data + elif self.type in [4, 5, 6]: # 4: Raw, 5: Eeprom, 6: IQ, 非图像数据，无法直接显示 + frame = None + else: + frame = np_array + return frame + + +class RgbImageData(ctypes.Structure): + _fields_ = [ + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("data", ctypes.c_uint8 * (1280*1280*3)), + ("dataSize", ctypes.c_uint), + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong) + ] + def frame(self, rgb=False): + rgb_frame = np.frombuffer(bytes(self.data[:self.dataSize]), dtype=np.uint8).reshape((self.height, self.width, 3)) + if rgb: + return rgb_frame.copy() + # RGB => BGR + return cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR) # 转换颜色空间 + + +class GrayScaleImage(ctypes.Structure): + _fields_ = [ + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("data", ctypes.c_uint8 * (640*480)) + ] + def frame(self): + buf_size = 640 * 480 + count = min(self.width * self.height, buf_size) + return np.array(self.data[:count], dtype=np.uint8).reshape((self.height, self.width, 1)) + + +class FisheyeImagesData(ctypes.Structure): + _fields_ = [ + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong), + ("images", GrayScaleImage * 4), + ("id", ctypes.c_longlong) + ] + + def frame(self, inx=-1): + inx_vaild = inx >= 0 and inx < 4 + if inx_vaild: + return self.images[inx].frame() + else: + frame0 = cv2.resize(self.images[0].frame(), (480, 360)) + frame1 = cv2.resize(self.images[1].frame(), (480, 360)) + frame2 = cv2.resize(self.images[2].frame(), (480, 360)) + frame3 = cv2.resize(self.images[3].frame(), (480, 360)) + + up_frame = cv2.hconcat([frame0, frame1]) + down_frame = cv2.hconcat([frame2, frame3]) + return cv2.vconcat([up_frame, down_frame]) + + +class EyetrackingImageData(ctypes.Structure): + _fields_ = [ + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong), + ("images", GrayScaleImage * 4), + ] + + +class PoseData(ctypes.Structure): + _fields_ = [ + ("position", Vector3D), + ("orientation", Vector3D), + ("quaternion", Vector4D), + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong), + ("confidence", ctypes.c_double) + ] + + +class ImuData(ctypes.Structure): + _fields_ = [ + ("gyro", Vector3D), + ("accel", Vector3D), + ("accelSaturation", Vector3B), + ("magneto", Vector3D), + ("temperature", ctypes.c_double), + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong) + ] + + +class EventData(ctypes.Structure): + _fields_ = [ + ("hostTimestamp", ctypes.c_double), + ("edgeTimestampUs", ctypes.c_longlong), + ("type", ctypes.c_int), + ("state", ctypes.c_int) + ] + + +class XVLib: + _xvlib = None + + def __init__(self, serial_number, init_slam=False, init_clamp_stream=False, init_color_camera=False, init_fisheye_cameras=False): + self.instance_id = -1 + + self._clamp_data = ClampData() + self._color_image_data = ColorImageData() + self._color_image_rgb_data = RgbImageData() + self._depth_image_data = DepthImageData() + self._fisheye_images_data = FisheyeImagesData() + self._slam_data = PoseData() + self._external_stream_data = PoseData() + self._spheretrack_stream_data = PoseData() + + self.xv_get_devices() + time.sleep(1) + + serial_number = ctypes.c_char_p(serial_number.encode('utf-8')) + self.instance_id = self.xv_init(serial_number, init_slam, init_clamp_stream, init_color_camera, init_fisheye_cameras) + if self.instance_id > 0: + logger.info('Device initialized successfully.') + else: + raise Exception('Device initialized failure.') + + def __del__(self): + self.xv_uninit() + + @classmethod + def __load_library(cls): + if cls._xvlib is None: + # 加载动态库 + lib_dir = os.path.dirname(__file__) + if os.path.exists(os.path.join(lib_dir, 'libopencv_core.so.4.2')): + ctypes.CDLL(os.path.join(lib_dir, 'libopencv_core.so.4.2'), mode=ctypes.RTLD_GLOBAL) + if os.path.exists(os.path.join(lib_dir, 'libopencv_imgproc.so.4.2')): + ctypes.CDLL(os.path.join(lib_dir, 'libopencv_imgproc.so.4.2'), mode=ctypes.RTLD_GLOBAL) + lib_path = os.path.abspath(os.path.join(lib_dir, 'libxvlib.so')) + logger.info(f"Loading library from: {lib_path}") + cls._xvlib = ctypes.CDLL(lib_path) + logger.info('Library initialized successfully.') + + @classmethod + def xv_get_devices(cls, max_devices=16): + cls.__load_library() + devices = (DeviceStruct * max_devices)() + device_count = ctypes.c_int(0) + cls._xvlib.xv_get_devices( + ctypes.byref(devices), + ctypes.byref(device_count), + ctypes.c_int(max_devices) + ) + return device_count.value, list(devices[:device_count.value]) + + def xv_init(self, serial_number, init_slam, init_clamp_stream, init_color_camera, init_fisheye_cameras): + return self._xvlib.xv_init(serial_number, init_slam, init_clamp_stream, init_color_camera, init_fisheye_cameras) + + def xv_uninit(self): + if self._xvlib is not None and self.instance_id > 0: + return self._xvlib.xv_uninit(self.instance_id) + else: + return -1 + + def xv_sleep(self, level = 0): + return self._xvlib.xv_sleep(self.instance_id, level) + + def xv_wakeup(self): + return self._xvlib.xv_wakeup(self.instance_id) + + def xv_slam_init(self): + return self._xvlib.xv_slam_init(self.instance_id) + + def xv_slam_uninit(self): + return self._xvlib.xv_slam_uninit(self.instance_id) + + def xv_imu_sensor_init(self): + return self._xvlib.xv_imu_sensor_init(self.instance_id) + + def xv_imu_sensor_uninit(self): + return self._xvlib.xv_imu_sensor_uninit(self.instance_id) + + def xv_event_stream_init(self): + return self._xvlib.xv_event_stream_init(self.instance_id) + + def xv_event_stream_uninit(self): + return self._xvlib.xv_event_stream_uninit(self.instance_id) + + def xv_orientation_stream_init(self): + return self._xvlib.xv_orientation_stream_init(self.instance_id) + + def xv_orientation_stream_uninit(self): + return self._xvlib.xv_orientation_stream_uninit(self.instance_id) + + def xv_fisheye_cameras_init(self): + return self._xvlib.xv_fisheye_cameras_init(self.instance_id) + + def xv_fisheye_cameras_uninit(self): + return self._xvlib.xv_fisheye_cameras_uninit(self.instance_id) + + def xv_color_camera_init(self): + return self._xvlib.xv_color_camera_init(self.instance_id) + + def xv_color_camera_uninit(self): + return self._xvlib.xv_color_camera_uninit(self.instance_id) + + def xv_tof_camera_init(self): + return self._xvlib.xv_tof_camera_init(self.instance_id) + + def xv_tof_camera_uninit(self): + return self._xvlib.xv_tof_camera_uninit(self.instance_id) + + def xv_sgbm_camera_init(self, config): + return self._xvlib.xv_sgbm_camera_init(self.instance_id, ctypes.c_char_p(config.encode('utf-8'))) + + def xv_sgbm_camera_uninit(self): + return self._xvlib.xv_sgbm_camera_uninit(self.instance_id) + + def xv_eyetracking_camera_init(self): + return self._xvlib.xv_eyetracking_camera_init(self.instance_id) + + def xv_eyetracking_camera_uninit(self): + return self._xvlib.xv_eyetracking_camera_uninit(self.instance_id) + + def xv_gaze_stream_init(self): + return self._xvlib.xv_gaze_stream_init(self.instance_id) + + def xv_gaze_stream_uninit(self): + return self._xvlib.xv_gaze_stream_uninit(self.instance_id) + + def xv_iris_stream_init(self): + return self._xvlib.xv_iris_stream_init(self.instance_id) + + def xv_iris_stream_uninit(self): + return self._xvlib.xv_iris_stream_uninit(self.instance_id) + + def xv_gesture_stream_init(self): + return self._xvlib.xv_gesture_stream_init(self.instance_id) + + def xv_gesture_stream_uninit(self): + return self._xvlib.xv_gesture_stream_uninit(self.instance_id) + + def xv_gps_stream_init(self): + return self._xvlib.xv_gps_stream_init(self.instance_id) + + def xv_gps_stream_uninit(self): + return self._xvlib.xv_gps_stream_uninit(self.instance_id) + + def xv_gps_distance_stream_init(self): + return self._xvlib.xv_gps_distance_stream_init(self.instance_id) + + def xv_gps_distance_stream_uninit(self): + return self._xvlib.xv_gps_distance_stream_uninit(self.instance_id) + + def xv_terrestrial_magnetism_stream_init(self): + return self._xvlib.xv_terrestrial_magnetism_stream_init(self.instance_id) + + def xv_terrestrial_magnetism_stream_uninit(self): + return self._xvlib.xv_terrestrial_magnetism_stream_uninit(self.instance_id) + + def xv_external_stream_init(self): + return self._xvlib.xv_external_stream_init(self.instance_id) + + def xv_external_stream_uninit(self): + return self._xvlib.xv_external_stream_uninit(self.instance_id) + + def xv_mic_stream_init(self): + return self._xvlib.xv_mic_stream_init(self.instance_id) + + def xv_mic_stream_uninit(self): + return self._xvlib.xv_mic_stream_uninit(self.instance_id) + + def xv_object_detector_init(self): + return self._xvlib.xv_object_detector_init(self.instance_id) + + def xv_object_detector_uninit(self): + return self._xvlib.xv_object_detector_uninit(self.instance_id) + + def xv_object_detector_RKNN3588_init(self): + return self._xvlib.xv_object_detector_RKNN3588_init(self.instance_id) + + def xv_object_detector_RKNN3588_uninit(self): + return self._xvlib.xv_object_detector_RKNN3588_uninit(self.instance_id) + + def xv_device_status_stream_init(self): + return self._xvlib.xv_device_status_stream_init(self.instance_id) + + def xv_device_status_stream_uninit(self): + return self._xvlib.xv_device_status_stream_uninit(self.instance_id) + + def xv_clamp_stream_init(self): + return self._xvlib.xv_clamp_stream_init(self.instance_id) + + def xv_clamp_stream_uninit(self): + return self._xvlib.xv_clamp_stream_uninit(self.instance_id) + + def xv_spheretrack_stream_init(self): + return self._xvlib.xv_spheretrack_stream_init(self.instance_id) + + def xv_spheretrack_stream_uninit(self): + return self._xvlib.xv_spheretrack_stream_uninit(self.instance_id) + + def xv_get_clamp_stream_data(self): + ret = self._xvlib.xv_get_clamp_stream_data(self.instance_id, ctypes.byref(self._clamp_data)) + return ret, self._clamp_data + + def xv_get_color_image_data(self): + ret = self._xvlib.xv_get_color_image_data(self.instance_id, ctypes.byref(self._color_image_data)) + return ret, self._color_image_data + + def xv_get_color_image_rgb_data(self): + ret = self._xvlib.xv_get_color_image_rgb_data(self.instance_id, ctypes.byref(self._color_image_rgb_data)) + return ret, self._color_image_rgb_data + + def xv_get_depth_image_data(self): + ret = self._xvlib.xv_get_depth_image_data(self.instance_id, ctypes.byref(self._depth_image_data)) + return ret, self._depth_image_data + + def xv_get_fisheye_images_data(self, index=5): + ret = self._xvlib.xv_get_fisheye_images_data(self.instance_id, ctypes.byref(self._fisheye_images_data), ctypes.c_size_t(index)) + return ret, self._fisheye_images_data + + def xv_get_slam_data(self): + ret = self._xvlib.xv_get_slam_data(self.instance_id, ctypes.byref(self._slam_data)) + return ret, self._slam_data + + def xv_get_slam_pose(self, prediction): + ret = self._xvlib.xv_get_slam_pose(self.instance_id, ctypes.byref(self._slam_data), ctypes.c_double(prediction)) + return ret, self._slam_data + + def xv_get_slam_pose_at(self, timestamp): + ret = self._xvlib.xv_get_slam_pose_at(self.instance_id, ctypes.byref(self._slam_data), ctypes.c_double(timestamp)) + return ret, self._slam_data + + def xv_get_external_stream_data(self): + ret = self._xvlib.xv_get_external_stream_data(self.instance_id, ctypes.byref(self._external_stream_data)) + return ret, self._external_stream_data + + def xv_get_spheretrack_stream_data(self): + ret = self._xvlib.xv_get_spheretrack_stream_data(self.instance_id, ctypes.byref(self._spheretrack_stream_data)) + return ret, self._spheretrack_stream_data + + def xv_set_color_camera_rgb_mode(self, mode): + """ + Docstring for xv_set_color_camera_rgb_mode + + :param mode: Description + 0: AF + 1: MF + 2: Unknown + """ + return self._xvlib.xv_set_color_camera_rgb_mode(self.instance_id, ctypes.c_int(mode)) + + def xv_set_color_camera_resolution(self, resolution): + """ + Docstring for xv_set_color_camera_resolution + + :param resolution: + 0: RGB_1920x1080 + 1: RGB_1280x720 + 2: RGB_640x480 + 3: RGB_320x240 (not supported now) + 4: RGB_2560x1920 (not supported now) + 5: RGB_3840x2160 (not supported now) + """ + return self._xvlib.xv_set_color_camera_resolution(self.instance_id, ctypes.c_int(resolution)) + + def xv_set_color_camera_framerate(self, framerate): + return self._xvlib.xv_set_color_camera_framerate(self.instance_id, ctypes.c_float(framerate)) + + def xv_set_color_camera_brightness(self, brightness): + return self._xvlib.xv_set_color_camera_brightness(self.instance_id, ctypes.c_int(brightness)) + + def xv_set_tof_camera_mode(self, mode): + return self._xvlib.xv_set_tof_camera_mode(self.instance_id, ctypes.c_int(mode)) + + def xv_set_tof_camera_stream_mode(self, mode): + """ + Docstring for xv_set_tof_camera_stream_mode + + :param mode: Description + 0: DepthOnly + 1: CloudOnly + 2: DepthAndCloud + 3: None + 4: CloudOnLeftHandSlam + """ + return self._xvlib.xv_set_tof_camera_stream_mode(self.instance_id, ctypes.c_int(mode)) + + def xv_set_tof_camera_distance_mode(self, mode): + """ + Docstring for xv_set_tof_camera_distance_mode + + :param mode: Description + 0: Short + 1: Middle + 2: Long + """ + return self._xvlib.xv_set_tof_camera_distance_mode(self.instance_id, ctypes.c_int(mode)) + + def xv_set_tof_camera_resolution(self, resolution): + """ + Docstring for xv_set_tof_camera_resolution + + :param resolution: Description + -1: Unknown + 0: VGA + 1: QVGA + 2: HQVGA + """ + return self._xvlib.xv_set_tof_camera_resolution(self.instance_id, ctypes.c_int(resolution)) + + def xv_set_tof_camera_framerate(self, framerate): + """ + Docstring for xv_set_tof_camera_framerate + + :param framerate: Description + 0: FPS_5 + 1: FPS_10 + 2: FPS_15 + 3: FPS_20 + 4: FPS_25 + 5: FPS_30 + """ + return self._xvlib.xv_set_tof_camera_framerate(self.instance_id, ctypes.c_float(framerate)) + + def xv_set_tof_camera_brightness(self, brightness): + return self._xvlib.xv_set_tof_camera_brightness(self.instance_id, ctypes.c_int(brightness)) + + def xv_set_fisheye_cameras_resolution(self, resolution): + return self._xvlib.xv_set_fisheye_cameras_resolution(self.instance_id, ctypes.c_int(resolution)) + + def xv_set_fisheye_cameras_framerate(self, framerate): + return self._xvlib.xv_set_fisheye_cameras_framerate(self.instance_id, ctypes.c_float(framerate)) + + def xv_set_fisheye_cameras_brightness(self, brightness): + return self._xvlib.xv_set_fisheye_cameras_brightness(self.instance_id, ctypes.c_int(brightness)) + + def xv_set_eyetracking_camera_resolution(self, resolution): + return self._xvlib.xv_set_eyetracking_camera_resolution(self.instance_id, ctypes.c_int(resolution)) + + def xv_set_eyetracking_camera_framerate(self, framerate): + return self._xvlib.xv_set_eyetracking_camera_framerate(self.instance_id, ctypes.c_float(framerate)) + + def xv_set_eyetracking_camera_brightness(self, brightness): + return self._xvlib.xv_set_eyetracking_camera_brightness(self.instance_id, ctypes.c_int(brightness)) diff --git a/ufactory_lerobot/robots/__init__.py b/ufactory_lerobot/robots/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ufactory_lerobot/robots/uf_mock_robot/__init__.py b/ufactory_lerobot/robots/uf_mock_robot/__init__.py new file mode 100644 index 0000000..fa74844 --- /dev/null +++ b/ufactory_lerobot/robots/uf_mock_robot/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .uf_mock_robot_config import UFMockRobotConfig +from .uf_mock_robot import UFMockRobot +from .multiple_uf_mock_robot_config import MultipleUFMockRobotConfig +from .multiple_uf_mock_robot import MultipleUFMockRobot diff --git a/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot.py b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot.py new file mode 100644 index 0000000..23a6046 --- /dev/null +++ b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python + +from lerobot.processor.core import RobotAction, RobotObservation +from lerobot.processor.core import RobotObservation +from lerobot.robots import Robot +from .multiple_uf_mock_robot_config import MultipleUFMockRobotConfig +from .uf_mock_robot import UFMockRobot + + +class MultipleUFMockRobot(Robot): + + config_class = MultipleUFMockRobotConfig + name = "UFACTORY Multiple Mock Robot" + + def __init__(self, config: MultipleUFMockRobotConfig): + super().__init__(config) + self.config = config + self.keys = [] + self.robots = [] + for key, robot_config in self.config.robots.items(): + self.keys.append(key) + self.robots.append(UFMockRobot(robot_config, prefix=key)) + + @property + def observation_features(self) -> dict: + observation_features = {} + for robot in self.robots: + observation_features.update(robot.observation_features) + return observation_features + + @property + def action_features(self) -> dict: + action_features = {} + for robot in self.robots: + action_features.update(robot.action_features) + return action_features + + @property + def is_connected(self) -> bool: + return all(robot.is_connected for robot in self.robots) + + @property + def is_calibrated(self) -> bool: + return all(robot.is_calibrated for robot in self.robots) + + def connect(self, calibrate: bool = True) -> None: + for robot in self.robots: + robot.connect(calibrate=calibrate) + + def calibrate(self) -> None: + for robot in self.robots: + robot.calibrate() + + def configure(self) -> None: + for robot in self.robots: + robot.configure() + + def disconnect(self) -> None: + for robot in self.robots: + robot.disconnect() + + def get_observation(self) -> RobotObservation: + observations = [robot.get_observation() for robot in self.robots] + combined_observation = RobotObservation() + for obs in observations: + combined_observation.update(obs) + return combined_observation + + def send_action(self, action: RobotAction) -> RobotAction: + for i in range(len(self.keys)): + key = self.keys[i] + action_subset = {k: v for k, v in action.items() if k.startswith(f"{key}.")} + self.robots[i].send_action(action_subset) + return action diff --git a/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py new file mode 100644 index 0000000..d2d8eca --- /dev/null +++ b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass +from lerobot.robots import RobotConfig +from .uf_mock_robot_config import UFMockRobotConfig + +@RobotConfig.register_subclass("uf::multiple_mock_robot") +@dataclass +class MultipleUFMockRobotConfig(RobotConfig): + robots: dict[str, UFMockRobotConfig] diff --git a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py new file mode 100644 index 0000000..f6161c5 --- /dev/null +++ b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py @@ -0,0 +1,133 @@ +import numpy as np +from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.robots import Robot +from .uf_mock_robot_config import UFMockRobotConfig + + +CARTESIAN_OBS_KEYS = [ + "pose.x", "pose.y", "pose.z", "pose.rx", "pose.ry", "pose.rz", + # un-comment if you need more features below: + # "velo.x", "velo.y", "velo.z", "velo.rx", "velo.ry", "velo.rz", +] + +CARTESIAN_ACTION_KEYS = [ + "pose.x", "pose.y", "pose.z", "pose.rx", "pose.ry", "pose.rz", +] + + +class UFMockRobot(Robot): + + config_class = UFMockRobotConfig + name = "UFACTORY Mock Robot" + + def __init__(self, config: UFMockRobotConfig, prefix=''): + super().__init__(config) + self.prefix = '' if not prefix else f"{prefix}." + self.config = config + self._dof = config.robot_dof + if self._dof == None or (not self._dof in (5, 6, 7)): + raise ValueError(f"Please specify the correct DOF uf_robot!, got {self._dof}") + + self._control_space = self.config.control_space + self._jnt_obs_has_vel = config.observe_joint_vel if self._control_space == "joint" else False + self._is_connected = False + self._is_calibrated =True + self._teleop = self.config.teleop + self._cache_num = self.config.state_offset_action + self._teleop_actions = [] + + self.cameras = make_cameras_from_configs(config.cameras) + + self._gripper_type = 0 if not self.config.gripper_control else self.config.gripper_type if self.config.gripper_type > 0 else 1 + + @property + def _robot_state_features(self)-> dict: + if self._control_space == "joint": + state_features = {f"{self.prefix}J{motor}.pos": float for motor in range(1, self._dof+1)} + if self._jnt_obs_has_vel: + state_features.update({f"{self.prefix}J{motor}.vel": float for motor in range(1, self._dof+1)}) + if self._gripper_type > 0: + state_features.update({f"{self.prefix}gripper.pos": float}) + elif self._control_space == "cartesian": + state_features = {f"{self.prefix}{key}": float for key in CARTESIAN_OBS_KEYS} + if self._gripper_type > 0: + state_features.update({f"{self.prefix}gripper.pos": float}) + else: + raise ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + return state_features + + @property + # CHECK!! channel first or last? + 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) + return cam_ft + + @property + def observation_features(self) -> dict[str, type | tuple]: + return {**self._robot_state_features, **self._cam_features} + + @property + def action_features(self)-> dict: + if self._control_space == "joint": + action_ft = {f"{self.prefix}J{motor}.pos": float for motor in range(1, self._dof+1)} + elif self._control_space == "cartesian": + action_ft = {f"{self.prefix}{key}": float for key in CARTESIAN_ACTION_KEYS} + else: + raise ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + # Consider adding velocity configuration ?? + if self._gripper_type > 0: + action_ft.update({f"{self.prefix}gripper.pos": float}) + return action_ft + + def connect(self, calibrate: bool = True) -> None: + self._is_connected = True + for cam in self.cameras.values(): + cam.connect() + self._is_connected = self._is_connected and cam.is_connected + + self.configure() + if calibrate: + self.calibrate() + + def configure(self) -> None: + pass + + def calibrate(self) -> None: + self._is_calibrated = True + pass # CHECK! currently No-op + + def get_observation(self) -> dict[str, np.ndarray]: + new_act = self._teleop.get_action() + # Capture images from cameras + for cam_key, cam in self.cameras.items(): + new_act[f"{self.prefix}{cam_key}"] = cam.async_read() + + # if len(self._teleop_actions) >= self._cache_num: + # act = self._teleop_actions.pop(0) + # else: + # act = new_act + # self._teleop_actions.append(new_act) + + return new_act + + def send_action(self, action: dict) -> np.ndarray: + return action + + def disconnect(self) -> None: + for cam in self.cameras.values(): + cam.disconnect() + + self._is_connected = False + + def is_calibrated(self) -> bool: + """Whether the robot is currently calibrated or not. Should be always `True` if not applicable""" + return self._is_calibrated + + def is_connected(self) -> bool: + """Whether the robot is currently calibrated or not. Should be always `True` if not applicable""" + return self._is_connected + + + diff --git a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py new file mode 100644 index 0000000..737706c --- /dev/null +++ b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass, field +from lerobot.cameras import CameraConfig +from lerobot.cameras.opencv import OpenCVCameraConfig +from lerobot.robots import RobotConfig + +@RobotConfig.register_subclass("uf::mock_robot") +@dataclass +class UFMockRobotConfig(RobotConfig): + # cameras + cameras: dict[str, CameraConfig] = field( + default_factory=lambda: { + "fisheye": OpenCVCameraConfig( + index_or_path=6, + width=640, + height=480, + fps=30, + fourcc="MJPG" + ) + } + ) + + robot_dof: int | None = None # Set it correctly if controlling in joint space! + control_space: str = "joint" + gripper_control: bool = True + gripper_type: int = 1 # 1: xArm Gripper, 10: Pika Gripper + observe_joint_vel: bool = False # only effective in joint control mode + teleop: None = None # from lerobot.teleoperators import Teleoperator + state_offset_action: int = 3 # the number of previous teleop actions to be included in the observation diff --git a/ufactory_lerobot/robots/uf_robot/__init__.py b/ufactory_lerobot/robots/uf_robot/__init__.py new file mode 100644 index 0000000..a89d438 --- /dev/null +++ b/ufactory_lerobot/robots/uf_robot/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .uf_robot_config import UFRobotConfig +from .uf_robot import UFRobot +from .multiple_uf_robot_config import MultipleUFRobotConfig +from .multiple_uf_robot import MultipleUFRobot diff --git a/ufactory_lerobot/robots/uf_robot/multiple_uf_robot.py b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot.py new file mode 100644 index 0000000..17ad7ac --- /dev/null +++ b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python + +import queue +import threading +from typing import Any +from lerobot.robots import Robot +from .multiple_uf_robot_config import MultipleUFRobotConfig +from .uf_robot import UFRobot + + +class MultipleUFRobot(Robot): + + config_class = MultipleUFRobotConfig + name = "UFACTORY Multiple Robot" + + def __init__(self, config: MultipleUFRobotConfig): + super().__init__(config) + self.config = config + self._is_async_connect = config.async_connect + self._is_async_configure = config.async_configure + self._is_async_action = config.async_action + self.robots = {} + self.action_queues = {} + self.action_threads = {} + for key, robot_config in self.config.robots.items(): + robot = UFRobot(robot_config, prefix=key) + self.robots[key] = robot + if self._is_async_action: + action_queue = queue.Queue() + self.action_queues[key] = action_queue + action_thread = threading.Thread(target=self.run_action_loop, args=(action_queue, robot), daemon=True) + self.action_threads[key] = action_thread + # action_thread.start() + self.cameras = {} + for robot in self.robots.values(): + self.cameras.update(robot.cameras) + + def run_action_loop(self, action_queue: queue.Queue, robot: Robot): + while robot.is_connected: + try: + action = action_queue.get(timeout=0.5) + if action is None: + break + robot.send_action(action) + except queue.Empty: + continue + + @property + def observation_features(self) -> dict: + observation_features = {} + for robot in self.robots.values(): + observation_features.update(robot.observation_features) + return observation_features + + @property + def action_features(self) -> dict: + action_features = {} + for robot in self.robots.values(): + action_features.update(robot.action_features) + return action_features + + @property + def is_connected(self) -> bool: + return all(robot.is_connected for robot in self.robots.values()) + + @property + def is_calibrated(self) -> bool: + return all(robot.is_calibrated for robot in self.robots.values()) + + def connect(self, calibrate: bool = True) -> None: + if self._is_async_connect: + threads = [] + for robot in self.robots.values(): + thread = threading.Thread(target=robot.connect, kwargs={"calibrate": calibrate}, daemon=True) + threads.append(thread) + thread.start() + for thread in threads: + thread.join() + else: + for robot in self.robots.values(): + robot.connect(calibrate=calibrate) + if self._is_async_action: + for thread in self.action_threads.values(): + thread.start() + + def calibrate(self) -> None: + for robot in self.robots.values(): + robot.calibrate() + + def configure(self) -> None: + if self._is_async_configure: + threads = [] + for robot in self.robots.values(): + thread = threading.Thread(target=robot.configure, daemon=True) + threads.append(thread) + thread.start() + for thread in threads: + thread.join() + else: + for robot in self.robots.values(): + robot.configure() + + def disconnect(self) -> None: + for robot in self.robots.values(): + robot.disconnect() + + def get_observation(self) -> dict[str, Any]: + observations = {} + for robot in self.robots.values(): + observations.update(robot.get_observation()) + return observations + + def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + if self._is_async_action: + for key, robot in self.robots.items(): + action_subset = {k: v for k, v in action.items() if k.startswith(f"{key}.")} + self.action_queues[key].put(action_subset) + else: + for key, robot in self.robots.items(): + action_subset = {k: v for k, v in action.items() if k.startswith(f"{key}.")} + robot.send_action(action_subset) + return action diff --git a/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py new file mode 100644 index 0000000..95f5d1b --- /dev/null +++ b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass +from lerobot.robots import RobotConfig +from .uf_robot_config import UFRobotConfig + +@RobotConfig.register_subclass("uf::multiple_robot") +@dataclass +class MultipleUFRobotConfig(RobotConfig): + robots: dict[str, UFRobotConfig] + async_connect: bool = True + async_configure: bool = True + async_action: bool = False diff --git a/ufactory_lerobot/robots/uf_robot/uf_robot.py b/ufactory_lerobot/robots/uf_robot/uf_robot.py new file mode 100644 index 0000000..a4db6f4 --- /dev/null +++ b/ufactory_lerobot/robots/uf_robot/uf_robot.py @@ -0,0 +1,465 @@ +#!/usr/bin/env python + +import time +import math +import logging +import struct +import numpy as np +from enum import IntEnum +from dataclasses import dataclass +from threading import Thread, Event, Lock +from lerobot.robots import Robot +from lerobot.cameras.utils import make_cameras_from_configs +from ufactory_lerobot.devices.pika import PikaDevice +from .uf_robot_config import UFRobotConfig +from xarm.wrapper import XArmAPI +from xarm.core.utils import convert + +## Configurations: +INIT_SYNC_JOINT_VELOCITY_RAD = 0.2 + +CARTESIAN_OBS_KEYS = [ + "pose.x", "pose.y", "pose.z", "pose.rx", "pose.ry", "pose.rz", + # un-comment if you need more features below: + # "velo.x", "velo.y", "velo.z", "velo.rx", "velo.ry", "velo.rz", +] + +CARTESIAN_ACTION_KEYS = [ + "pose.x", "pose.y", "pose.z", "pose.rx", "pose.ry", "pose.rz", +] + +class GripperType(IntEnum): + NoGripper = 0 + xArmGripper = 1 + xArmGripperG2 = 2 + BioGripperG2 = 3 + PikaGripper = 10 + RobotiqGripper = 11 + + +@dataclass +class GripperParam: + name: str + open_pos: int + close_pos: int + speed: int = 0 + force: int = 0 + gripper_norm: float = 0 + + def get_grippos(self, gripper_norm): + pos = self.open_pos + gripper_norm * (self.close_pos - self.open_pos) + min_pos, max_pos = min(self.open_pos, self.close_pos), max(self.open_pos, self.close_pos) + return int(min(max(min_pos, pos), max_pos)) + + def get_gripper_norm(self, grippos): + if grippos is None: + return self.gripper_norm + self.gripper_norm = (self.open_pos - grippos) / (self.open_pos - self.close_pos) + return self.gripper_norm + + +class UFRobot(Robot, Thread): + + config_class = UFRobotConfig + name = "UFACTORY Robot" + + def __init__(self, config: UFRobotConfig, prefix=''): + super().__init__(config) + Thread.__init__(self) + self.prefix = '' if not prefix else f"{prefix}." + self.config = config + self._dof = config.robot_dof + if self._dof == None or (not self._dof in (5,6,7)): + raise ValueError(f"Please specify the correct DOF uf_robot!, got {self._dof}") + + self._control_space = self.config.control_space + + self.real_arm = None + self.cameras = make_cameras_from_configs(config.cameras) + + self._is_connected = False + self._is_calibrated =True + + self.logs = {} + + self._cmd_cnt = 0 + + self._max_joint_velocity = math.radians(self.config.max_joint_velocity) + self._max_linear_velocity = self.config.max_linear_velocity + + self._start_tcp_pose = self.config.start_tcp_pose + self._start_joints = self.config.start_joints + + self.report_stop_event = Event() + self._rt_report_normal = False + self._update_lock = Lock() + self._use_rt_report = (self._control_space == "cartesian") # Cartesian observations must utilize rt_report + self._cart_obs_has_vel = any('velo.' in key for key in CARTESIAN_OBS_KEYS) + self._jnt_obs_has_vel = self.config.observe_joint_vel + + self._gripper_type = 0 if not self.config.gripper_control else self.config.gripper_type + if self._gripper_type == GripperType.xArmGripper: + gripper_speed = 5000 if self.config.gripper_speed < 0 else min(max(50, self.config.gripper_speed), 5000) + gripper_force = 50 if self.config.gripper_force < 0 else self.config.gripper_force # # not support + self._gripper_param = GripperParam('xArmGripper', open_pos=800, close_pos=0, speed=gripper_speed, force=gripper_force) + elif self._gripper_type == GripperType.xArmGripperG2: + speed = 225 if self.config.gripper_speed < 0 else min(max(15, self.config.gripper_speed), 225) + gripper_speed = int(((speed * 60) / 9.88235 + 140) / 0.4) + gripper_force = 50 if self.config.gripper_force < 0 else min(max(1, self.config.gripper_force), 100) + self._gripper_param = GripperParam('xArmGripperG2', open_pos=84, close_pos=0, speed=gripper_speed, force=gripper_force) + elif self._gripper_type == GripperType.BioGripperG2: + gripper_speed = 2000 if self.config.gripper_speed < 0 else min(max(500, self.config.gripper_speed), 4500) + gripper_force = 100 if self.config.gripper_force < 0 else min(max(1, self.config.gripper_force), 100) + self._gripper_param = GripperParam('BioGripperG2', open_pos=150, close_pos=71, speed=gripper_speed, force=gripper_force) + elif self._gripper_type == GripperType.PikaGripper: + self.pika_device = PikaDevice(2, pika_gripper_port=self.config.gripper_port) + self.pika_gripper = self.pika_device.pika_gripper + logger = logging.getLogger('pika.gripper') + logger.setLevel(logging.WARNING) + gripper_speed = 0 if self.config.gripper_speed < 0 else self.config.gripper_speed # not support + gripper_force = 0 if self.config.gripper_force < 0 else self.config.gripper_force # not support + self._gripper_param = GripperParam('PikaGripper', open_pos=100, close_pos=0, speed=gripper_speed, force=gripper_force) + elif self._gripper_type == GripperType.RobotiqGripper: + gripper_speed = 255 if self.config.gripper_speed < 0 else min(max(1, self.config.gripper_speed), 255) + gripper_force = 255 if self.config.gripper_force < 0 else min(max(1, self.config.gripper_force), 255) + self._gripper_param = GripperParam('RobotiqGripper', open_pos=0, close_pos=0xFF, speed=gripper_speed, force=gripper_force) + else: # no gripper or not support + self._gripper_type = 0 + self._gripper_param = GripperParam('NoGripper', open_pos=0, close_pos=0, speed=0, force=0) + + @property + def _robot_state_features(self)-> dict: + if self._control_space == "joint": + state_features = {f"{self.prefix}J{motor}.pos": float for motor in range(1, self._dof+1)} + if self._jnt_obs_has_vel: + state_features.update({f"{self.prefix}J{motor}.vel": float for motor in range(1, self._dof+1)}) + if self._gripper_type > GripperType.NoGripper: + state_features.update({f"{self.prefix}gripper.pos": float}) + elif self._control_space == "cartesian": + state_features = {f"{self.prefix}{key}": float for key in CARTESIAN_OBS_KEYS} + if self._gripper_type > GripperType.NoGripper: + state_features.update({f"{self.prefix}gripper.pos": float}) + else: + raise ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + return state_features + + @property + # CHECK!! channel first or last? + 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) + return cam_ft + + @property + def observation_features(self) -> dict[str, type | tuple]: + return {**self._robot_state_features, **self._cam_features} + + @property + def action_features(self)-> dict: + if self._control_space == "joint": + action_ft = {f"{self.prefix}J{motor}.pos": float for motor in range(1, self._dof+1)} + elif self._control_space == "cartesian": + action_ft = {f"{self.prefix}{key}": float for key in CARTESIAN_ACTION_KEYS} + else: + raise ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + # Consider adding velocity configuration ?? + if self._gripper_type > GripperType.NoGripper: + action_ft.update({f"{self.prefix}gripper.pos": float}) + return action_ft + + def connect(self, calibrate: bool = True) -> None: + self.real_arm = XArmAPI(self.config.robot_ip) + time.sleep(0.2) + self._is_connected = self.real_arm.connected + if not self._is_connected: + print(f"UF Robot connection Failed, please check the hardware availability at ip: {self.config.robot_ip}") + raise ConnectionError() + + if not self._dof == self.real_arm.axis: + print(f"[ERROR: ] Real Robot DOF({self.real_arm.axis}) does not match configuration ({self._dof})!") + self._is_connected = False + raise ConnectionError() + + for cam in self.cameras.values(): + cam.connect() + self._is_connected = self._is_connected and cam.is_connected + + if not self._is_connected: + print("Could not connect to the cameras, check that all cameras are plugged-in.") + raise ConnectionError() + + # if self._gripper_type == GripperType.PikaGripper: + # if not self.pika_gripper.connect(): + # print('Could not connect to pika gripper.') + # raise ConnectionError() + + self.configure() + if calibrate: + self.calibrate() + + self._is_connected = True + + def configure(self) -> None: + self.real_arm.motion_enable() + self.real_arm.clean_error() + self.real_arm.set_mode(0) # set to idle mode + self.real_arm.set_state(0) # set to start state + time.sleep(0.5) + if self._start_tcp_pose is None: + self.real_arm.set_servo_angle(angle=self._start_joints, is_radian=True, wait=True) + else: + self.real_arm.set_servo_angle(angle=self._start_joints, is_radian=True, wait=True) + self.real_arm.set_position(*self._start_tcp_pose, speed=100, is_radian=True, wait=True) + _, self._start_joints = self.real_arm.get_servo_angle(is_radian=True) + self._start_tcp_pose = None + + if self._control_space == "joint": + self.real_arm.set_mode(6) + elif self._control_space == "cartesian": + self.real_arm.set_mode(7) + else: + raise ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + + self.real_arm.set_state(0) + + _, err_warn = self.real_arm.get_err_warn_code() + if err_warn[0] != 0: + raise RuntimeError(f"Failed to set correct state to UF robot! Controller Error code: {err_warn[0]} !") + + if self._gripper_type > GripperType.NoGripper: + self.real_arm._arm._baud_checkset = True + if self._gripper_type == GripperType.xArmGripper: + self.real_arm.set_gripper_enable(True) + self.real_arm.set_gripper_mode(0) + self.real_arm.set_gripper_speed(self._gripper_param.speed) + self.real_arm.set_gripper_position(self._gripper_param.open_pos) + elif self._gripper_type == GripperType.xArmGripperG2: + self.real_arm.set_gripper_enable(True) + self.real_arm.set_gripper_mode(0) + self.real_arm.set_gripper_g2_position(self._gripper_param.open_pos) + elif self._gripper_type == GripperType.BioGripperG2: + _, mode = self.real_arm.get_bio_gripper_control_mode() + if mode != 1: + self.real_arm.set_bio_gripper_control_mode(1) + self.real_arm.set_bio_gripper_enable(True) + self.real_arm.open_bio_gripper() + elif self._gripper_type == GripperType.PikaGripper: + self.pika_gripper.enable() + time.sleep(0.5) + self.pika_gripper.set_gripper_distance(self._gripper_param.open_pos) + elif self._gripper_type == GripperType.RobotiqGripper: + self.real_arm.robotiq_reset() + self.real_arm.robotiq_set_activate(wait=True) + self.real_arm.robotiq_set_position(self._gripper_param.open_pos, wait=True) + self._gripper_param.grippos = self._gripper_param.open_pos + self._gripper_param.gripper_norm = self._gripper_param.open_pos + self.real_arm._arm._baud_checkset = False + _, err_warn = self.real_arm.get_err_warn_code() + if err_warn[0] != 0: + raise RuntimeError(f"Failed to set correct state to Gripper! Controller Error code: {err_warn[0]} !") + + if self._use_rt_report and not self._rt_report_normal: + self.start() + time.sleep(0.2) + + def calibrate(self) -> None: + self._is_calibrated = True + pass # CHECK! currently No-op + + def get_observation(self) -> dict[str, np.ndarray]: + obs_dict = {} + + # Read Stretch state + before_read_t = time.perf_counter() + if self._control_space == "joint": + code, states = self.real_arm.get_joint_states(is_radian=True, num=3) + pos_list = states[0].copy() + obs_dict = {f"{self.prefix}J{k+1}.pos": pos_list[k] for k in range(self._dof)} + if self._jnt_obs_has_vel: + vel_list = states[1].copy() + obs_dict.update({f"{self.prefix}J{k+1}.vel": vel_list[k] for k in range(self._dof)}) + elif self._control_space == "cartesian": + if not self._rt_report_normal: + raise ConnectionError("RT Report for target robot NOT READY! ") + + with self._update_lock: + pos_list = self.rt_actual_tcp_pose.copy() + vel_list = self.rt_actual_tcp_speed.copy() + # pos_cmd_list = self.rt_cmd_tcp_pose.copy() + # vel_cmd_list = self.rt_cmd_tcp_vel.copy() + # jpos_fbk_list = self.rt_actual_joint_pos.copy() + # jvel_fbk_list = self.rt_actual_joint_speed.copy() + + obs_dict = {f"{self.prefix}pose.x": pos_list[0], f"{self.prefix}pose.y": pos_list[1], f"{self.prefix}pose.z": pos_list[2], f"{self.prefix}pose.rx": pos_list[3], f"{self.prefix}pose.ry": pos_list[4], f"{self.prefix}pose.rz": pos_list[5]} + if self._cart_obs_has_vel: + obs_dict.update({f"{self.prefix}velo.x": vel_list[0], f"{self.prefix}velo.y": vel_list[1], f"{self.prefix}velo.z": vel_list[2], f"{self.prefix}velo.rx": vel_list[3], f"{self.prefix}velo.ry": vel_list[4], f"{self.prefix}velo.rz": vel_list[5]}) + else: + ValueError(f"Please check the given control space of uf_robot! got {self._control_space}") + + if self._gripper_type > GripperType.NoGripper: + if self._gripper_type == GripperType.xArmGripper: + code, grippos = self.real_arm.get_gripper_position() + grippos_norm = self._gripper_param.get_gripper_norm(grippos) + elif self._gripper_type == GripperType.xArmGripperG2: + code, grippos = self.real_arm.get_gripper_g2_position() + grippos_norm = self._gripper_param.get_gripper_norm(grippos) + elif self._gripper_type == GripperType.BioGripperG2: + code, grippos = self.real_arm.get_bio_gripper_g2_position() + grippos_norm = self._gripper_param.get_gripper_norm(grippos) + elif self._gripper_type == GripperType.PikaGripper: + grippos = self.pika_gripper.get_gripper_distance() + grippos_norm = self._gripper_param.get_gripper_norm(grippos) + elif self._gripper_type == GripperType.RobotiqGripper: + self.real_arm.robotiq_get_status(number_of_registers=3) + grippos = self.real_arm.robotiq_status['gPO'] # 0..255 + grippos_norm = self._gripper_param.get_gripper_norm(grippos) # 0=open, 1=closed + self.logs["read_pos_dt_s"] = time.perf_counter() - before_read_t + obs_dict[f"{self.prefix}gripper.pos"] = grippos_norm + + # 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() + self.logs[f"async_read_camera_{cam_key}_dt_s"] = time.perf_counter() - before_camread_t + + return obs_dict + + def send_action(self, action: dict) -> np.ndarray: + if not self._is_connected: + raise ConnectionError() + if self.real_arm.error_code != 0: + return action + if self.config.no_action: + return action + + before_write_t = time.perf_counter() + if self._control_space == "joint": + # first sync with gello or other control device SLOWLY! + jnt_spd = INIT_SYNC_JOINT_VELOCITY_RAD if self._cmd_cnt < 20 else self._max_joint_velocity + wait_ = True if self._cmd_cnt == 0 else False + + cmd_list = [0]*(self._dof) + for i in range(self._dof): + cmd_list[i] = action[f"{self.prefix}J{i+1}.pos"] + + # TODO: make mode 6 compatible with wait=True + if wait_== False and self.real_arm.mode != 6: + self.real_arm.set_mode(6) + self.real_arm.set_state(0) + time.sleep(0.1) + elif wait_ and self.real_arm.mode != 0: + self.real_arm.set_mode(0) + self.real_arm.set_state(0) + time.sleep(0.1) + + self.real_arm.set_servo_angle(angle=cmd_list[:self._dof], speed=jnt_spd, is_radian=True, wait=wait_) + elif self._control_space == "cartesian": # unit: mm? + lin_spd = self._max_linear_velocity + + if not self._rt_report_normal: + raise ConnectionError("RT Report for target robot NOT READY! ") + cmd_list = [action[f"{self.prefix}pose.x"], action[f"{self.prefix}pose.y"], action[f"{self.prefix}pose.z"], action[f"{self.prefix}pose.rx"], action[f"{self.prefix}pose.ry"], action[f"{self.prefix}pose.rz"]] + self.real_arm.set_position_aa(axis_angle_pose=cmd_list, speed=lin_spd, is_radian=True, wait=False) + # self.real_arm.set_position(*cmd_list, radius=0, speed=lin_spd, is_radian=True, wait=False) + + if self._cmd_cnt < 99999: + self._cmd_cnt += 1 # CHECK!! possibility of overflow? + if self._gripper_type > GripperType.NoGripper: + gripper_norm = action[f"{self.prefix}gripper.pos"] + if self._gripper_type == GripperType.xArmGripper: + grippos = self._gripper_param.get_grippos(gripper_norm) + modbus_datas = [0x08, 0x10, 0x07, 0x00, 0x00, 0x02, 0x04] + modbus_datas.extend(list(struct.pack('>i', grippos))) + self.real_arm.getset_tgpio_modbus_data(modbus_datas) + # self.real_arm.set_gripper_position(grippos, wait=False, wait_motion=False) # CHECK! the command unit + elif self._gripper_type == GripperType.xArmGripperG2: + grippos = self._gripper_param.get_grippos(gripper_norm) + grippos = int((math.degrees(math.asin((grippos - 16) / 110)) + 8.33) * 18.28) + modbus_datas = [0x08, 0x10, 0x0C, 0x00, 0x00, 0x05, 0x0A, 0x00, 0x01] + modbus_datas.extend(list(struct.pack('>h', self._gripper_param.speed))) + modbus_datas.extend(list(struct.pack('>h', self._gripper_param.force))) + modbus_datas.extend(list(struct.pack('>i', grippos))) + self.real_arm.getset_tgpio_modbus_data(modbus_datas) + elif self._gripper_type == GripperType.BioGripperG2: + grippos = self._gripper_param.get_grippos(gripper_norm) + grippos = int(grippos * 3.7342 - 265.13) + modbus_datas = [0x08, 0x10, 0x0C, 0x00, 0x00, 0x05, 0x0A, 0x00, 0x01] + modbus_datas.extend(list(struct.pack('>h', self._gripper_param.speed))) + modbus_datas.extend(list(struct.pack('>h', self._gripper_param.force))) + modbus_datas.extend(list(struct.pack('>i', grippos))) + self.real_arm.getset_tgpio_modbus_data(modbus_datas) + elif self._gripper_type == GripperType.PikaGripper: + grippos = self._gripper_param.get_grippos(gripper_norm) + self.pika_gripper.set_gripper_distance(grippos) + elif self._gripper_type == GripperType.RobotiqGripper: + grippos = self._gripper_param.get_grippos(gripper_norm) + modbus_datas = [0x09, 0x10, 0x03, 0xE8, 0x00, 0x03, 0x06, 0x09, 0x00, 0x00, grippos, self._gripper_param.speed, self._gripper_param.force] + self.real_arm.getset_tgpio_modbus_data(modbus_datas) + # self.real_arm.robotiq_set_position( + # grippos, speed=self._gripper_param.speed, force=self._gripper_param.force, + # wait=False, wait_motion=False, + # ) + + self.logs["write_pos_dt_s"] = time.perf_counter() - before_write_t + return action + + def print_logs(self) -> None: + pass + + def disconnect(self) -> None: + self.real_arm.set_state(4) # stop + self.real_arm.set_mode(0) + if self._use_rt_report: + self.report_stop_event.set() + self.join() + self.real_arm.disconnect() + # CHECK!! how about gripper? + + for cam in self.cameras.values(): + cam.disconnect() + + self._is_connected = False + + def is_calibrated(self) -> bool: + """Whether the robot is currently calibrated or not. Should be always `True` if not applicable""" + return self._is_calibrated + + def is_connected(self) -> bool: + """Whether the robot is currently calibrated or not. Should be always `True` if not applicable""" + return self._is_connected + + def run(self): + import socket + + robot_port = 30000 # DO NOT CHANGE + # create socket connection + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.setblocking(True) + sock.settimeout(1) + sock.connect((self.config.robot_ip, robot_port)) + + buffer = sock.recv(4) + print(buffer) + while len(buffer) < 4: + buffer += sock.recv(4 - len(buffer)) + size = convert.bytes_to_u32(buffer[:4]) + print(f"UFACTORY Robot ({self.config.robot_ip}) RT Report Thread starts!! =======") + while not self.report_stop_event.is_set(): + buffer += sock.recv(size - len(buffer)) + if len(buffer) < size: + continue + data = buffer[:size] + buffer = buffer[size:] + with self._update_lock: + self.rt_actual_joint_pos = convert.bytes_to_fp32s(data[116:144], 7) + self.rt_actual_joint_speed = convert.bytes_to_fp32s(data[144:172], 7) + self.rt_cmd_tcp_pose = convert.bytes_to_fp32s(data[424:448], 6) + self.rt_cmd_tcp_vel = convert.bytes_to_fp32s(data[448:472], 6) + self.rt_actual_tcp_pose = convert.bytes_to_fp32s(data[472:496], 6) + self.rt_actual_tcp_speed = convert.bytes_to_fp32s(data[496:520], 6) + self._rt_report_normal = True + + self._rt_report_normal = False + print(f"UFACTORY Robot ({self.config.robot_ip}) RT Report Thread Exit!! =======") diff --git a/ufactory_lerobot/robots/uf_robot/uf_robot_config.py b/ufactory_lerobot/robots/uf_robot/uf_robot_config.py new file mode 100644 index 0000000..3093997 --- /dev/null +++ b/ufactory_lerobot/robots/uf_robot/uf_robot_config.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass, field +from typing import Tuple +import numpy as np +from lerobot.cameras import CameraConfig +from lerobot.cameras.realsense import RealSenseCameraConfig +from lerobot.robots import RobotConfig + +@RobotConfig.register_subclass("uf::robot") +@dataclass +class UFRobotConfig(RobotConfig): + # cameras + cameras: dict[str, CameraConfig] = field( + default_factory=lambda: { + "overhead": RealSenseCameraConfig( + serial_number_or_name="Intel RealSense D435I", + fps=30, + width=640, # 1280 + height=480, # 720 + # rotation=90, + ), + "tool": RealSenseCameraConfig( + serial_number_or_name="Intel RealSense D435", + fps=30, + width=640, # 1280 + height=480, # 720 + ), + } + ) + + robot_ip: str = "192.168.1.127" + robot_dof: int | None = None # Set it correctly if controlling in joint space! + control_space: str = "joint" + gripper_control: bool = True + gripper_type: int = 1 # 1: xArm Gripper, 2: xArm Gripper G2, 10: Pika Gripper, 11: Robotiq 2F-85 + gripper_port: str = None # only used by pika gripper (gripper_type=10) + gripper_speed: int = -1 # auto + gripper_force: int = -1 # auto + observe_joint_vel: bool = False # only effective in joint control mode + start_joints: Tuple[float, ...] = (0, 0, 0, np.pi/2, 0, np.pi/2, 0) + start_tcp_pose: Tuple[float, ...] = None # xyzrpy + max_joint_velocity: int = 90 # °/s, only effective in joint control mode + max_linear_velocity: int = 200 # mm/s, only effective in cartesian control mode + rx_continuous: bool = False + no_action: bool = False diff --git a/ufactory_lerobot/robots/utils.py b/ufactory_lerobot/robots/utils.py new file mode 100644 index 0000000..7f29e7b --- /dev/null +++ b/ufactory_lerobot/robots/utils.py @@ -0,0 +1,34 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from lerobot.robots.utils import make_robot_from_config as lerobot_make_robot_from_config +from lerobot.robots.config import RobotConfig +from lerobot.robots.robot import Robot + + +def make_robot_from_config(config: RobotConfig) -> Robot: + if config.type == "uf::robot": + from .uf_robot import UFRobot + return UFRobot(config) + elif config.type == "uf::multiple_robot": + from .uf_robot import MultipleUFRobot + return MultipleUFRobot(config) + elif config.type == "uf::mock_robot": + from .uf_mock_robot import UFMockRobot + return UFMockRobot(config) + elif config.type == "uf::multiple_mock_robot": + from .uf_mock_robot import MultipleUFMockRobot + return MultipleUFMockRobot(config) + else: + return lerobot_make_robot_from_config(config) diff --git a/ufactory_lerobot/scripts/uf_lerobot_eval.py b/ufactory_lerobot/scripts/uf_lerobot_eval.py new file mode 100644 index 0000000..3c01225 --- /dev/null +++ b/ufactory_lerobot/scripts/uf_lerobot_eval.py @@ -0,0 +1,443 @@ +import yaml +import argparse +import logging +import time +import math +import torch +from pathlib import Path +from dataclasses import asdict, dataclass +from pprint import pformat +from contextlib import nullcontext +import numpy as np +import ufactory_lerobot # patch +from lerobot.scripts.lerobot_record import register_third_party_plugins +from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features +from lerobot.datasets.utils import build_dataset_frame, combine_feature_dicts +from lerobot.policies.utils import make_robot_action +from lerobot.policies.factory import make_policy, make_pre_post_processors +from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata +from lerobot.utils.constants import OBS_STR +from lerobot.processor import ( + make_default_processors, +) +from lerobot.robots import ( # noqa: F401 + Robot, + RobotConfig, + make_robot_from_config, +) +from lerobot.utils.control_utils import ( + is_headless, + init_keyboard_listener, + predict_action, +) +from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.robot_utils import precise_sleep +from lerobot.utils.utils import ( + get_safe_torch_device, + init_logging, +) +from lerobot.configs import parser +from lerobot.configs.policies import PreTrainedConfig +from lerobot.scripts.lerobot_record import DatasetRecordConfig +from lerobot.datasets.lerobot_dataset import LeRobotDataset +from ufactory_lerobot.utils.utils import instantiate_from_dict +from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations + + +def continuous_rotvec(new_rv, prev_rv): + """Keep rotvec in the same sign-hemisphere as prev to avoid ±π flips. + When accumulated rotation crosses π, as_rotvec() can flip the axis sign + (e.g. rx jumps 3.14 → -3.13), causing the robot to make a large motion. + This re-maps the equivalent rotation to stay consistent with prev.""" + new_rv = np.asarray(new_rv, dtype=np.float64) + prev_rv = np.asarray(prev_rv, dtype=np.float64) + if np.dot(new_rv, prev_rv) < 0: + angle = np.linalg.norm(new_rv) + if angle > 1e-6: + axis = new_rv / angle + new_rv = -(2 * np.pi - angle) * axis + return new_rv + +def blend_poses(pose_a, pose_b, alpha): + """位姿混合: (1-alpha)*A + alpha*B, 旋转用 SO(3) 插值。 + 先将 pose_b 的 rotvec 归一化到与 pose_a 同符号半球，避免 ±π 跳变破坏线性混合。""" + blended_pos = (1 - alpha) * np.array(pose_a[:3]) + alpha * np.array(pose_b[:3]) + # 旋转用线性混合 rotvec (delta 很小时近似 SLERP) + rot_b = continuous_rotvec(np.array(pose_b[3:6]), np.array(pose_a[3:6])) + blended_rot = (1 - alpha) * np.array(pose_a[3:6]) + alpha * rot_b + return np.concatenate([blended_pos, blended_rot]).tolist() + +def compute_relative_axis_angle(rot_prev, rot_curr): + """ + 计算两个轴角之间的相对旋转。 + 逻辑: R_diff = R_prev.T @ R_curr + 返回: 相对轴角向量 + """ + # 1. 转为矩阵 + R_prev = Transformations.rxryrz_to_matrix(rot_prev) + R_curr = Transformations.rxryrz_to_matrix(rot_curr) + + # 2. 计算相对旋转矩阵 + # R_delta 表示从 prev 坐标系到 curr 坐标系的旋转 + R_delta = R_prev.T @ R_curr + + # 3. 转回轴角 + return Transformations.rotation_matrix_to_rxryrz(R_delta) + +def compute_target_axis_angle(rot_prev, rot_delta): + """ + 根据起始轴角和相对轴角计算目标轴角 + """ + R_prev = Transformations.rxryrz_to_matrix(rot_prev) + R_delta = Transformations.rxryrz_to_matrix(rot_delta) + R_curr = R_prev @ R_delta + # R_curr = R_prev.apply(R_delta) + return Transformations.rotation_matrix_to_rxryrz(R_curr) + + +@dataclass +class EvalConfig: + robot: RobotConfig + dataset: DatasetRecordConfig + # Whether to control the robot with a policy + policy: PreTrainedConfig | None = None + n_episodes: int = 50 + single_task: str | None = "pick_place" + + def __post_init__(self): + # HACK: We parse again the cli args here to get the pretrained path if there was one. + policy_path = parser.get_path_arg("policy") + if policy_path: + cli_overrides = parser.get_cli_overrides("policy") + policy_path = Path(policy_path).expanduser() + self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides) + self.policy.pretrained_path = policy_path + + if self.policy is None: + raise ValueError("Choose a policy to control the robot") + + @classmethod + def __get_path_fields__(cls) -> list[str]: + """This enables the parser to load config from the policy using `--policy.path=local/dir`""" + return ["policy"] + + +def eval_loop(cfg: EvalConfig, relative=False): + init_logging() + logging.info(pformat(asdict(cfg))) + + # print(cfg.robot) + print(type(cfg.robot)) + if hasattr(cfg.robot, 'robots'): + print(cfg.robot.robots.keys()) + exit(1) + + robot = make_robot_from_config(cfg.robot) + + teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + + try: + dataset_metadata = LeRobotDatasetMetadata(repo_id=cfg.dataset.repo_id, root=cfg.dataset.root) + dataset_features = dataset_metadata.features + print("Loaded dataset metadata successfully.") + except Exception: + dataset_features = combine_feature_dicts( + aggregate_pipeline_dataset_features( + pipeline=teleop_action_processor, + initial_features=create_initial_features( + action=robot.action_features + ), # TODO(steven, pepijn): in future this should be come from teleop or policy + use_videos=cfg.dataset.video, + ), + aggregate_pipeline_dataset_features( + pipeline=robot_observation_processor, + initial_features=create_initial_features(observation=robot.observation_features), + use_videos=cfg.dataset.video, + ), + ) + # Create empty dataset or load existing saved episodes + dataset = LeRobotDataset.create( + cfg.dataset.repo_id, + cfg.dataset.fps, + root=cfg.dataset.root, + robot_type=robot.name, + features=dataset_features, + use_videos=cfg.dataset.video, + image_writer_processes=cfg.dataset.num_image_writer_processes, + image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), + batch_encoding_size=cfg.dataset.video_encoding_batch_size, + ) + dataset_metadata = dataset.meta + print("Created new dataset metadata successfully.") + + if cfg.dataset.fps != dataset_metadata.fps: + raise ValueError(f"The dataset fps should be equal to requested fps ({dataset_metadata.fps} != {cfg.dataset.fps}).") + + + policy = make_policy(cfg=cfg.policy, ds_meta=dataset_metadata) + # policy.eval() + + # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. + preprocessor_overrides = { + "device_processor": {"device": str(policy.config.device)}, + "rename_observations_processor": {"rename_map": cfg.dataset.rename_map}, + } + + preprocessor, postprocessor = make_pre_post_processors( + policy_cfg=cfg.policy, + pretrained_path=cfg.policy.pretrained_path, + preprocessor_overrides=preprocessor_overrides, + dataset_stats=dataset_metadata.stats + ) + + robot.connect() + + events = {"reset": False, "exit": False} + listener = None + + if not is_headless(): + from pynput import keyboard + + def on_press(key): + try: + if key == keyboard.Key.right: + print("Right arrow key pressed. Resetting...") + events["reset"] = True + elif key == keyboard.Key.left: + print("Left arrow key pressed. Resetting....") + events["reset"] = True + elif key == keyboard.Key.esc: + print("Escape key pressed. Stopping ...") + events["exit"] = True + except Exception as e: + print(f"Error handling key press: {e}") + + listener, events = init_keyboard_listener(events=events, on_press=on_press) + + device = get_safe_torch_device(policy.config.device, log=True) + sleep_time_s = 1 / dataset_metadata.fps + + # Gripper look-ahead: denormalization stats for peeking into the action queue + _gripper_mean = dataset_metadata.stats['action']['mean'][-1].item() + _gripper_std = dataset_metadata.stats['action']['std'][-1].item() + _gripper_min = dataset_metadata.stats['action']['min'][-1].item() + _gripper_max = dataset_metadata.stats['action']['max'][-1].item() + is_act_policy = hasattr(policy.config, 'chunk_size') + # ACT: lookahead 30 (~1s) 补偿 chunk 慢启动; DP: 队列仅 8 步，lookahead 4 + GRIPPER_LOOKAHEAD = 0 if is_act_policy else 4 + + # ===================================================== + # Chunk boundary smoothing: only damp large discontinuities at action chunk + # boundaries (mean ~6mm jump) while preserving smooth within-chunk motion (~1mm). + # When step-to-step cmd change exceeds SMOOTH_THRESHOLD, clamp it to that limit. + # ===================================================== + SMOOTH_THRESHOLD = 0 # mm: smooth chunk boundary jumps while preserving trajectory + SMOOTH_ROT_THRESHOLD = 0.05 # rad: max allowed rotation jump per step + prev_smoothed_pose = None + + print("\n********** Policy Eval Episode Loop Start **********") + print(f'relative: {relative}') + + rx_continuous = getattr(cfg.robot, 'rx_continuous', False) + + # with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(): + while True: + robot.configure() + policy.reset() + preprocessor.reset() + postprocessor.reset() + + obs = robot.get_observation() + + prev_robot_dict = {} + prev_action_dict = {} + + is_multiple_robot = False + if hasattr(cfg.robot, 'robots'): + keys = cfg.robot.robots.keys() + is_multiple_robot = True + else: + keys = [''] + for key in keys: + prefix = f'.{key}' if key else '' + is_tcp = f'{prefix}pose.x' in obs and f'{prefix}pose.y' in obs and f'{prefix}pose.z' in obs and f'{prefix}pose.rx' in obs and f'{prefix}pose.ry' in obs and f'{prefix}pose.rz' in obs + if is_tcp: + pose = [obs[f'{prefix}pose.x'], obs[f'{prefix}pose.y'], obs[f'{prefix}pose.z'], obs[f'{prefix}pose.rx'], obs[f'{prefix}pose.ry'], obs[f'{prefix}pose.rz']] + else: + pose = [] + prev_robot_dict[key] = {'type': 1 if is_tcp else 0, 'pose': np.array(pose)} + prev_action_dict[key] = {'type': 1 if is_tcp else 0, 'pose': np.array(pose)} + + while True: + start_loop_t = time.perf_counter() + + if events["reset"] or events["exit"]: + events["reset"] = False + print("\n********** Policy Eval Episode (Reset) **********") + break + + # Get robot observation + obs = robot.get_observation() + # if rx_continuous and not relative and 'pose.rx' in obs and obs['pose.rx'] < 0: + # obs['pose.rx'] += 2 * math.pi + curr_robot_dict = {} + curr_action_dict = {} + for key in keys: + prefix = f'.{key}' if key else '' + is_tcp = f'{prefix}pose.x' in obs and f'{prefix}pose.y' in obs and f'{prefix}pose.z' in obs and f'{prefix}pose.rx' in obs and f'{prefix}pose.ry' in obs and f'{prefix}pose.rz' in obs + if not is_tcp or prev_robot_dict[key]['type'] != 1: + curr_robot_dict[key] = {'type': 0, 'pose': []} + continue + if rx_continuous and not relative and f'{prefix}pose.rx' in obs and obs[f'{prefix}pose.rx'] < 0: + obs[f'{prefix}pose.rx'] += 2 * np.pi + prev_robot_pose = prev_robot_dict[key]['pose'] + curr_robot_pose = np.array([obs[f'{prefix}pose.x'], obs[f'{prefix}pose.y'], obs[f'{prefix}pose.z'], obs[f'{prefix}pose.rx'], obs[f'{prefix}pose.ry'], obs[f'{prefix}pose.rz']]) + curr_rot_normalized = continuous_rotvec(curr_robot_pose[3:6], prev_robot_pose[3:6]) + curr_robot_pose[3] = float(curr_rot_normalized[0]) + curr_robot_pose[4] = float(curr_rot_normalized[1]) + curr_robot_pose[5] = float(curr_rot_normalized[2]) + curr_robot_dict[key] = {'type': 1, 'pose': curr_robot_pose} + + if relative: + delta = compute_relative_axis_angle(prev_robot_pose[3:6], curr_robot_pose[3:6]) + obs[f'{prefix}pose.x'] = curr_robot_pose[0] - prev_robot_pose[0] + obs[f'{prefix}pose.y'] = curr_robot_pose[1] - prev_robot_pose[1] + obs[f'{prefix}pose.z'] = curr_robot_pose[2] - prev_robot_pose[2] + obs[f'{prefix}pose.rx'] = delta[0] + obs[f'{prefix}pose.ry'] = delta[1] + obs[f'{prefix}pose.rz'] = delta[2] + prev_robot_dict[key]['pose'] = curr_robot_pose + + # Applies a pipeline to the raw robot observation, default is IdentityProcessor + obs_processed = robot_observation_processor(obs) + + observation_frame = build_dataset_frame(dataset_features, obs_processed, prefix=OBS_STR) + + action_values = predict_action( + observation=observation_frame, + policy=policy, + device=device, + preprocessor=preprocessor, + postprocessor=postprocessor, + use_amp=policy.config.use_amp, + task=cfg.single_task, + robot_type=robot.robot_type, + ) + act_processed_policy = make_robot_action(action_values, dataset_features) + robot_action_to_send = robot_action_processor((act_processed_policy, obs)) + + for key in keys: + prefix = f'.{key}' if key else '' + if not curr_robot_dict[key]['type'] != 1 or prev_action_dict[key]['type'] != 1: + continue + if relative: + rot_delta = np.array([robot_action_to_send[f'{prefix}pose.rx'], robot_action_to_send[f'{prefix}pose.ry'], robot_action_to_send[f'{prefix}pose.rz']]) + prev_action_pose = prev_action_dict[key]['pose'] + rot_curr = compute_target_axis_angle(prev_action_pose[3:6], rot_delta) + robot_action_to_send[f'{prefix}pose.x'] = prev_action_pose[0] + robot_action_to_send[f'{prefix}pose.x'] + robot_action_to_send[f'{prefix}pose.y'] = prev_action_pose[1] + robot_action_to_send[f'{prefix}pose.y'] + robot_action_to_send[f'{prefix}pose.z'] = prev_action_pose[2] + robot_action_to_send[f'{prefix}pose.z'] + robot_action_to_send[f'{prefix}pose.rx'] = rot_curr[0] + robot_action_to_send[f'{prefix}pose.ry'] = rot_curr[1] + robot_action_to_send[f'{prefix}pose.rz'] = rot_curr[2] + elif rx_continuous and f'{prefix}pose.rx' in robot_action_to_send and robot_action_to_send[f'{prefix}pose.rx'] > math.pi: + robot_action_to_send[f'{prefix}pose.rx'] -= 2 * np.pi + + # robot_action_to_send[f'{prefix}pose.z'] = max(robot_action_to_send[f'{prefix}pose.z'], 199) + + # Rate-limited smoothing: cap position velocity to reduce chunk boundary jerks + # Uses vector-norm clamping to preserve motion direction + if SMOOTH_THRESHOLD > 0: + pos_keys = [f'{prefix}pose.x', f'{prefix}pose.y', f'{prefix}pose.z'] + rot_keys = [f'{prefix}pose.rx', f'{prefix}pose.ry', f'{prefix}pose.rz'] + if prev_smoothed_pose is None: + prev_smoothed_pose = {k: robot_action_to_send[k] for k in pos_keys + rot_keys} + else: + # Vector-norm clamp on position (preserves direction) + delta_pos = np.array([robot_action_to_send[k] - prev_smoothed_pose[k] for k in pos_keys]) + norm = np.linalg.norm(delta_pos) + if norm > SMOOTH_THRESHOLD: + delta_pos = delta_pos * (SMOOTH_THRESHOLD / norm) + for i, k in enumerate(pos_keys): + prev_smoothed_pose[k] = prev_smoothed_pose[k] + delta_pos[i] + robot_action_to_send[k] = prev_smoothed_pose[k] + # Per-axis clamp on rotation + for k in rot_keys: + delta = robot_action_to_send[k] - prev_smoothed_pose[k] + if abs(delta) > SMOOTH_ROT_THRESHOLD: + delta = SMOOTH_ROT_THRESHOLD * (1 if delta > 0 else -1) + prev_smoothed_pose[k] = prev_smoothed_pose[k] + delta + robot_action_to_send[k] = prev_smoothed_pose[k] + + if relative: + curr_action_pose = np.array([ + robot_action_to_send[f'{prefix}pose.x'], robot_action_to_send[f'{prefix}pose.y'], robot_action_to_send[f'{prefix}pose.z'], + robot_action_to_send[f'{prefix}pose.rx'], robot_action_to_send[f'{prefix}pose.ry'], robot_action_to_send[f'{prefix}pose.rz'] + ]) + # 相对增量模式: 漂移修正，将指令位姿温和拉回实际位姿 + # prev_action_pose = blend_poses(curr_action_pose, curr_robot_pose, 0.05) + prev_action_dict[key]['pose'] = curr_action_pose + + # # Gripper look-ahead: peek ahead in the action queue to compensate + # # for the slow ramp in the ACT chunk (eliminates 1-2s gripper delay) + # gripper_raw = robot_action_to_send.get('left.gripper.pos', 0) + # if hasattr(policy, '_action_queue') and len(policy._action_queue) > 0: + # # ACT: 队列为 deque of tensors, 归一化方式 MEAN_STD + # lookahead_idx = min(GRIPPER_LOOKAHEAD, len(policy._action_queue) - 1) + # future_gripper_norm = policy._action_queue[lookahead_idx][0, -1].item() + # gripper_raw = future_gripper_norm * _gripper_std + _gripper_mean + # elif hasattr(policy, '_queues') and 'action' in policy._queues and len(policy._queues['action']) > 0: + # # DP: 队列结构不同, 归一化方式 MIN_MAX → [-1,1] → [min,max] + # lookahead_idx = min(GRIPPER_LOOKAHEAD, len(policy._queues['action']) - 1) + # future_gripper_norm = policy._queues['action'][lookahead_idx][0, -1].item() + # gripper_raw = (future_gripper_norm + 1) / 2 * (_gripper_max - _gripper_min) + _gripper_min + # robot_action_to_send['left.gripper.pos'] = 1.0 if gripper_raw > 0.4 else 0.0 + + # gripper_raw = robot_action_to_send.get('right.gripper.pos', 0) + # if hasattr(policy, '_action_queue') and len(policy._action_queue) > 0: + # # ACT: 队列为 deque of tensors, 归一化方式 MEAN_STD + # lookahead_idx = min(GRIPPER_LOOKAHEAD, len(policy._action_queue) - 1) + # future_gripper_norm = policy._action_queue[lookahead_idx][0, -1].item() + # gripper_raw = future_gripper_norm * _gripper_std + _gripper_mean + # elif hasattr(policy, '_queues') and 'action' in policy._queues and len(policy._queues['action']) > 0: + # # DP: 队列结构不同, 归一化方式 MIN_MAX → [-1,1] → [min,max] + # lookahead_idx = min(GRIPPER_LOOKAHEAD, len(policy._queues['action']) - 1) + # future_gripper_norm = policy._queues['action'][lookahead_idx][0, -1].item() + # gripper_raw = (future_gripper_norm + 1) / 2 * (_gripper_max - _gripper_min) + _gripper_min + # robot_action_to_send['right.gripper.pos'] = 1.0 if gripper_raw > 0.4 else 0.0 + + robot.send_action(robot_action_to_send) + + dt_s = time.perf_counter() - start_loop_t + precise_sleep(sleep_time_s - dt_s) + + if events["exit"]: + break + + print("\n********** Policy Eval Loop Exit **********") + if not is_headless() and listener is not None: + listener.stop() + +def main(): + parser = argparse.ArgumentParser(description='configuration args') + parser.add_argument('-c', '--config', type=str, required=True, + help='configuration file path, e.g.my_config.yaml') + parser.add_argument('--policy.path', type=str, required=True, + help='configuration file path, e.g.my_config.yaml') + parser.add_argument('--relative', action='store_true', help='is relative motion or not') + args = parser.parse_args() + try: + with open(args.config, 'r') as f: + cfg = yaml.safe_load(f) + except Exception as e: + print(f"Error loading config yaml file: {e}") + else: + register_third_party_plugins() + config = instantiate_from_dict(cfg) + + eval_cfg = EvalConfig(robot=config["RobotConfig"], dataset=config["DatasetRecordConfig"]) + eval_loop(eval_cfg, args.relative) + + +if __name__ == "__main__": + main() diff --git a/ufactory_lerobot/scripts/uf_lerobot_record.py b/ufactory_lerobot/scripts/uf_lerobot_record.py new file mode 100644 index 0000000..c98c939 --- /dev/null +++ b/ufactory_lerobot/scripts/uf_lerobot_record.py @@ -0,0 +1,364 @@ +import yaml +import argparse +from pathlib import Path +import ufactory_lerobot # patch +from lerobot.scripts.lerobot_record import * +from ufactory_lerobot.teleoperators.uf_mock_teleop import UFMockTeleop +from ufactory_lerobot.teleoperators.base_teleop import UFBaseTeleop +from ufactory_lerobot.utils.utils import instantiate_from_dict + + +@safe_stop_image_writer +def record_loop( + robot: Robot, + events: dict, + fps: int, + teleop_action_processor: RobotProcessorPipeline[ + tuple[RobotAction, RobotObservation], RobotAction + ], # runs after teleop + robot_action_processor: RobotProcessorPipeline[ + tuple[RobotAction, RobotObservation], RobotAction + ], # runs before robot + robot_observation_processor: RobotProcessorPipeline[ + RobotObservation, RobotObservation + ], # runs after robot + dataset: LeRobotDataset | None = None, + teleop: Teleoperator | list[Teleoperator] | None = None, + policy: PreTrainedPolicy | None = None, + preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None, + postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None, + control_time_s: int | None = None, + single_task: str | None = None, + display_data: bool = False, + display_compressed_images: bool = False, + frame_callback: callable = None, +): + if dataset is not None and dataset.fps != fps: + raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).") + + teleop_arm = teleop_keyboard = None + if isinstance(teleop, list): + teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None) + teleop_arm = next( + ( + t + for t in teleop + if isinstance( + t, + ( + so_leader.SO100Leader + | so_leader.SO101Leader + | koch_leader.KochLeader + | omx_leader.OmxLeader + ), + ) + ), + None, + ) + + if not (teleop_arm and teleop_keyboard and len(teleop) == 2 and robot.name == "lekiwi_client"): + raise ValueError( + "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot." + ) + + # Reset policy and processor if they are provided + if policy is not None and preprocessor is not None and postprocessor is not None: + policy.reset() + preprocessor.reset() + postprocessor.reset() + + last_robot_cmd = robot.get_observation() + # only positional cmd for now: Remove velo from observation for cmd if needed! + last_robot_cmd = { k: v for k,v in last_robot_cmd.items() if not "vel" in k } + + timestamp = 0 + start_episode_t = time.perf_counter() + while timestamp < control_time_s: + start_loop_t = time.perf_counter() + + if events["exit_early"]: + events["exit_early"] = False + break + + # Get robot observation + obs = robot.get_observation() + + # Applies a pipeline to the raw robot observation, default is IdentityProcessor + obs_processed = robot_observation_processor(obs) + + if policy is not None or dataset is not None: + observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR) + + # Get action from either policy or teleop + if policy is not None and preprocessor is not None and postprocessor is not None: + action_values = predict_action( + observation=observation_frame, + policy=policy, + device=get_safe_torch_device(policy.config.device), + preprocessor=preprocessor, + postprocessor=postprocessor, + use_amp=policy.config.use_amp, + task=single_task, + robot_type=robot.robot_type, + ) + + act_processed_policy: RobotAction = make_robot_action(action_values, dataset.features) + + elif policy is None and isinstance(teleop, Teleoperator): + act = teleop.get_action() + + # (space mouse) from delta Cartesian cmd to absolute command + if "pose.dx" in act: + last_robot_cmd.update({"pose.x": last_robot_cmd["pose.x"] + act["pose.dx"], "pose.y": last_robot_cmd["pose.y"] + act["pose.dy"], "pose.z": last_robot_cmd["pose.z"] + act["pose.dz"]}) + act = last_robot_cmd.copy() # watch out this is shallow copy, not for nested dict + + # Applies a pipeline to the raw teleop action, default is IdentityProcessor + act_processed_teleop = teleop_action_processor((act, obs)) + + elif policy is None and isinstance(teleop, list): + arm_action = teleop_arm.get_action() + arm_action = {f"arm_{k}": v for k, v in arm_action.items()} + keyboard_action = teleop_keyboard.get_action() + base_action = robot._from_keyboard_to_base_action(keyboard_action) + act = {**arm_action, **base_action} if len(base_action) > 0 else arm_action + act_processed_teleop = teleop_action_processor((act, obs)) + else: + logging.info( + "No policy or teleoperator provided, skipping action generation." + "This is likely to happen when resetting the environment without a teleop device." + "The robot won't be at its rest position at the start of the next episode." + ) + continue + + # Applies a pipeline to the action, default is IdentityProcessor + if policy is not None and act_processed_policy is not None: + action_values = act_processed_policy + robot_action_to_send = robot_action_processor((act_processed_policy, obs)) + else: + action_values = act_processed_teleop + robot_action_to_send = robot_action_processor((act_processed_teleop, obs)) + + # Send action to robot + # Action can eventually be clipped using `max_relative_target`, + # so action actually sent is saved in the dataset. action = postprocessor.process(action) + # TODO(steven, pepijn, adil): we should use a pipeline step to clip the action, so the sent action is the action that we input to the robot. + _sent_action = robot.send_action(robot_action_to_send) + + # Write to dataset + if dataset is not None: + action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION) + frame = {**observation_frame, **action_frame, "task": single_task} + if frame_callback is not None: + frame = frame_callback(frame) + dataset.add_frame(frame) + + if display_data: + log_rerun_data( + observation=obs_processed, action=action_values, compress_images=display_compressed_images + ) + + dt_s = time.perf_counter() - start_loop_t + precise_sleep(max(1 / fps - dt_s, 0.0)) + + timestamp = time.perf_counter() - start_episode_t + + +@parser.wrap() +def record(cfg: RecordConfig) -> LeRobotDataset: + init_logging() + logging.info(pformat(asdict(cfg))) + if cfg.display_data: + init_rerun(session_name="recording") + + teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None + if hasattr(cfg.robot, "teleop"): + cfg.robot.teleop = teleop + robot = make_robot_from_config(cfg.robot) + + teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + + dataset_features = combine_feature_dicts( + aggregate_pipeline_dataset_features( + pipeline=teleop_action_processor, + initial_features=create_initial_features( + action=robot.action_features + ), # TODO(steven, pepijn): in future this should be come from teleop or policy + use_videos=cfg.dataset.video, + ), + aggregate_pipeline_dataset_features( + pipeline=robot_observation_processor, + initial_features=create_initial_features(observation=robot.observation_features), + use_videos=cfg.dataset.video, + ), + ) + + if cfg.resume: + dataset = LeRobotDataset( + cfg.dataset.repo_id, + root=cfg.dataset.root, + batch_encoding_size=cfg.dataset.video_encoding_batch_size, + ) + + if hasattr(robot, "cameras") and len(robot.cameras) > 0: + dataset.start_image_writer( + num_processes=cfg.dataset.num_image_writer_processes, + num_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), + ) + sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features) + else: + # Create empty dataset or load existing saved episodes + sanity_check_dataset_name(cfg.dataset.repo_id, cfg.policy) + dataset = LeRobotDataset.create( + cfg.dataset.repo_id, + cfg.dataset.fps, + root=cfg.dataset.root, + robot_type=robot.name, + features=dataset_features, + use_videos=cfg.dataset.video, + image_writer_processes=cfg.dataset.num_image_writer_processes, + image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), + batch_encoding_size=cfg.dataset.video_encoding_batch_size, + ) + + # Load pretrained policy + policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta) + preprocessor = None + postprocessor = None + if cfg.policy is not None: + preprocessor, postprocessor = make_pre_post_processors( + policy_cfg=cfg.policy, + pretrained_path=cfg.policy.pretrained_path, + dataset_stats=rename_stats(dataset.meta.stats, cfg.dataset.rename_map), + preprocessor_overrides={ + "device_processor": {"device": cfg.policy.device}, + "rename_observations_processor": {"rename_map": cfg.dataset.rename_map}, + }, + ) + + robot.connect() + if teleop is not None: + teleop.connect() + + listener, events = init_keyboard_listener() + + print("\n********** Episode Record Loop Start **********") + + if isinstance(teleop, UFBaseTeleop): + # if getattr(cfg.robot, 'rx_continuous', False): + # def frame_callback(frame): + # if frame['action'][3] < 0: + # frame['action'][3] += 2 * math.pi + # if frame['observation.state'][3] < 0: + # frame['observation.state'][3] += 2 * math.pi + # return frame + # else: + # frame_callback = None + frame_callback = None + input('\nPress Enter to record this episode >>>>> ') + time.sleep(0.5) + teleop.set_ctrl_status(True) + time.sleep(0.5) + else: + frame_callback = None + + with VideoEncodingManager(dataset): + recorded_episodes = 0 + while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]: + if teleop is not None and isinstance(teleop, UFMockTeleop): + if events["stop_recording"]: + continue + teleop.configure(events=events) + if events["rerecord_episode"]: + events["rerecord_episode"] = False + events["exit_early"] = False + input('\nPress Enter to regenerate random target location >>>>> ') + continue + if events["stop_recording"]: + continue + + log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds) + record_loop( + robot=robot, + events=events, + fps=cfg.dataset.fps, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, + teleop=teleop, + policy=policy, + preprocessor=preprocessor, + postprocessor=postprocessor, + dataset=dataset, + control_time_s=cfg.dataset.episode_time_s, + single_task=cfg.dataset.single_task, + display_data=cfg.display_data, + frame_callback=frame_callback, + ) + + if events["rerecord_episode"]: + log_say("Re-record episode", cfg.play_sounds) + events["rerecord_episode"] = False + events["exit_early"] = False + if isinstance(teleop, UFBaseTeleop): + teleop.set_ctrl_status(False) + dataset.clear_episode_buffer() + input('\nPress Enter to rerecord this episode >>>>> ') + + if isinstance(teleop, UFBaseTeleop): + robot.configure() + time.sleep(0.5) + teleop.set_ctrl_status(True) + time.sleep(0.5) + continue + + if not events['stop_recording']: + if isinstance(teleop, UFBaseTeleop): + teleop.set_ctrl_status(False) + dataset.save_episode() + recorded_episodes += 1 + input('Press Enter to record at the next episode >>>>> ') + if isinstance(teleop, UFBaseTeleop): + robot.configure() + time.sleep(1) + teleop.set_ctrl_status(True) + + print("\n********** Episode Record Loop Exit **********") + + robot.disconnect() + if teleop is not None: + teleop.disconnect() + + if not is_headless() and listener is not None: + listener.stop() + + if cfg.dataset.push_to_hub: + dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private) + + log_say("Exiting", cfg.play_sounds) + return dataset + + +def main(): + parser = argparse.ArgumentParser(description='configuration args') + parser.add_argument('-c', '--config', type=str, required=True, + help='configuration file path, e.g.my_config.yaml') + parser.add_argument('-r', '--resume', + action='store_true', # specify --resume if resume needs to be True + default=False, + help='Whether contitue recording on existing dataset (default: False)') + args = parser.parse_args() + try: + with open(Path(args.config).expanduser(), 'r') as f: + cfg = yaml.safe_load(f) + except Exception as e: + print(f"Error loading config yaml file: {e}") + else: + register_third_party_plugins() + config = instantiate_from_dict(cfg) + + record_cfg = RecordConfig(resume=args.resume, play_sounds=False, robot=config["RobotConfig"], dataset=config["DatasetRecordConfig"], teleop=config["TeleoperatorConfig"]) + record(record_cfg) + + +if __name__ == "__main__": + main() diff --git a/ufactory_lerobot/scripts/uf_robot_teleop.py b/ufactory_lerobot/scripts/uf_robot_teleop.py new file mode 100644 index 0000000..2f1b41d --- /dev/null +++ b/ufactory_lerobot/scripts/uf_robot_teleop.py @@ -0,0 +1,118 @@ +import yaml +import argparse +import logging +import time +from pathlib import Path +from dataclasses import asdict, dataclass +from pprint import pformat +import ufactory_lerobot # patch +from lerobot.scripts.lerobot_record import register_third_party_plugins +from lerobot.processor import ( + make_default_processors, +) +from lerobot.robots import ( # noqa: F401 + RobotConfig, + make_robot_from_config, +) +from lerobot.teleoperators import ( # noqa: F401 + TeleoperatorConfig, + make_teleoperator_from_config, +) +from lerobot.utils.control_utils import ( + is_headless, + init_keyboard_listener +) +from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.robot_utils import precise_sleep +from lerobot.utils.utils import ( + init_logging, +) +from ufactory_lerobot.utils.utils import instantiate_from_dict + +@dataclass +class TeleopConfig: + robot: RobotConfig + teleop: TeleoperatorConfig + fps: int = 30 + + +def teleop_loop(cfg: TeleopConfig): + init_logging() + logging.info(pformat(asdict(cfg))) + + teleop = make_teleoperator_from_config(cfg.teleop) + if hasattr(cfg.robot, "teleop"): + cfg.robot.teleop = teleop + robot = make_robot_from_config(cfg.robot) + + teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + + robot.connect() + teleop.connect() + + events = {"exit": False} + listener = None + + if not is_headless(): + from pynput import keyboard + + def on_press(key): + try: + if key == keyboard.Key.esc: + print("Escape key pressed. Stopping ...") + events["exit"] = True + except Exception as e: + print(f"Error handling key press: {e}") + + listener, events = init_keyboard_listener(events=events, on_press=on_press) + + sleep_time_s = 1 / cfg.fps + + print("\n********** Test Teleop With Robot **********") + input('Enter to control robot with teleop >>> ') + + print("\n********** Teleop Control Loop Start **********") + + while not events["exit"]: + start_loop_t = time.perf_counter() + + # Get robot observation + obs = robot.get_observation() + + act = teleop.get_action() + act_processed_teleop = teleop_action_processor((act, obs)) + + robot_action_to_send = robot_action_processor((act_processed_teleop, obs)) + robot.send_action(robot_action_to_send) + + dt_s = time.perf_counter() - start_loop_t + precise_sleep(sleep_time_s - dt_s) + + print("\n********** Teleop Control Loop Exit **********") + robot.disconnect() + teleop.disconnect() + if not is_headless() and listener is not None: + listener.stop() + +def main(): + parser = argparse.ArgumentParser(description='configuration args') + parser.add_argument('-c', '--config', type=str, required=True, + help='configuration file path, e.g.my_config.yaml') + parser.add_argument('-f', '--fps', type=int, default=30, + help='control loop frequency in Hz (default: 30)') + args = parser.parse_args() + try: + with open(Path(args.config).expanduser(), 'r') as f: + cfg = yaml.safe_load(f) + except Exception as e: + print(f"Error loading config yaml file: {e}") + else: + register_third_party_plugins() + config = instantiate_from_dict(cfg, ignore_cameras=True) + + teleop_cfg = TeleopConfig(robot=config["RobotConfig"], teleop=config["TeleoperatorConfig"], fps=args.fps) + teleop_loop(teleop_cfg) + + +if __name__ == "__main__": + main() diff --git a/ufactory_lerobot/teleoperators/__init__.py b/ufactory_lerobot/teleoperators/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ufactory_lerobot/teleoperators/base_teleop/__init__.py b/ufactory_lerobot/teleoperators/base_teleop/__init__.py new file mode 100644 index 0000000..ee59fc4 --- /dev/null +++ b/ufactory_lerobot/teleoperators/base_teleop/__init__.py @@ -0,0 +1 @@ +from .base_teleop import UFBaseTeleop \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py b/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py new file mode 100644 index 0000000..309495f --- /dev/null +++ b/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py @@ -0,0 +1,13 @@ +from lerobot.teleoperators import Teleoperator + + +class UFBaseTeleop(Teleoperator): + config_class = None + name = "Base Teleop For UFACTORY" + + def __init__(self, config): + super().__init__(config) + self.config = config + + def set_ctrl_status(self, status): + pass diff --git a/ufactory_lerobot/teleoperators/gello_teleop/__init__.py b/ufactory_lerobot/teleoperators/gello_teleop/__init__.py new file mode 100644 index 0000000..1fce0a2 --- /dev/null +++ b/ufactory_lerobot/teleoperators/gello_teleop/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python + +from .gello_teleop_config import GelloTeleopConfig +from .gello_teleop import GelloTeleop diff --git a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py new file mode 100644 index 0000000..f4e70c7 --- /dev/null +++ b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python +import logging +import time +import numpy as np +from gello.dynamixel.driver import DynamixelDriver +from gello.agents.gello_agent import GelloAgent, DynamixelRobotConfig +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from ..base_teleop import UFBaseTeleop +from .gello_teleop_config import GelloTeleopConfig + + +logger = logging.getLogger(__name__) + +class GelloTeleop(UFBaseTeleop): + """ + GELLO for xArm tele-op, ref: https://wuphilipp.github.io/gello_site/ + """ + + config_class = GelloTeleopConfig + name = "Gello Teleop For xArm" + + def __init__(self, config: GelloTeleopConfig): + super().__init__(config) + self.config = config + self._is_connected = False + self._is_calibrated = True # CHECK!! + + # auto get joint offset from gello + joint_ids = [] + joint_ids.extend(self.config.joint_ids) + if self.config.gripper_id >= 0: + joint_ids.append(self.config.gripper_id) + driver = DynamixelDriver(joint_ids, port=self.config.port, baudrate=57600) + for _ in range(10): + driver.get_joints() # warmup + curr_joints = driver.get_joints() + driver.close() + joint_offsets = [] + for i in range(len(self.config.start_joints)): + offset = curr_joints[i] - self.config.start_joints[i] / self.config.joint_signs[i] + joint_offsets.append(offset) + if self.config.gripper_id >= 0: + gripper_config = [self.config.gripper_id, np.rad2deg(curr_joints[-1]) - 0.2, np.rad2deg(curr_joints[-1]) - 42] + else: + gripper_config = None + + param_dict = { + "joint_ids": self.config.joint_ids, + "joint_signs": self.config.joint_signs, + "joint_offsets": joint_offsets, + "gripper_config": gripper_config + } + self._dynamixel_robo_config = DynamixelRobotConfig(**param_dict) + print(self._dynamixel_robo_config) + self.dof = len(self.config.start_joints) + + if self.config.torque_joint_ids: + driver = DynamixelDriver(self.config.torque_joint_ids, port=self.config.port, baudrate=57600) + driver.set_torque_mode(True) + driver.close() + + @property + def action_features(self) -> dict: + # Add one more dof for gripper + # act_ft = { + # "joint_position": { + # "dtype": "float", + # "shape": (self.dof+1,) + # } + # } + act_ft = { f"J{i+1}.pos": float for i in range(self.dof) } | {"gripper.pos": float} + return act_ft + + @property + def feedback_features(self) -> dict: + # fbk_ft = { + # "joint_position": { + # "dtype": "float", + # "shape": (self.dof+1,) + # } + # } + fbk_ft = { f"J{i+1}.pos": float for i in range(self.dof) } | {"gripper.pos": float} + return fbk_ft + + @property + def is_connected(self) -> bool: + return self._is_connected + + def connect(self, calibrate: bool = True) -> None: + if self._is_connected: + raise DeviceAlreadyConnectedError(f"{self} already connected") + + self.gello_agent = GelloAgent(port=self.config.port, dynamixel_config=self._dynamixel_robo_config) + 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 + logger.info(f"{self} connected.") + + @property + def is_calibrated(self) -> bool: + return self._is_calibrated + + def calibrate(self) -> None: + pass + + def configure(self) -> None: + # TODO: Go to sync position slowly? Can not + pass + + def get_action(self) -> dict[str, np.ndarray]: + start = time.perf_counter() + fake_obs = dict({"joint_state": np.array([0.0]*(self.dof+1))}) # for agent.act() argument, actually no use + action_array = self.gello_agent.act(fake_obs) # current gello joint pos as np.ndarray + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read action: {dt_ms:.1f}ms") + + action = {} + for i in range(self.dof): + action.update({f"J{i+1}.pos": action_array[i]}) + action.update({"gripper.pos": action_array[self.dof]}) + return action + + def send_feedback(self, feedback: dict[str, float]) -> None: + raise NotImplementedError + + def disconnect(self) -> None: + if not self._is_connected: + DeviceNotConnectedError(f"{self} is not connected.") + + self._is_connected = False + logger.info(f"{self} disconnected.") diff --git a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py new file mode 100644 index 0000000..4c283ca --- /dev/null +++ b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +import numpy as np +from dataclasses import dataclass +from typing import Tuple +from lerobot.teleoperators import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("uf::gello_teleop") +@dataclass +class GelloTeleopConfig(TeleoperatorConfig): + # Port to connect to the gello dummy arm + port: str = "/dev/serial/by-id/usb-FTDI_USB__-__Serial_Converter_FTAJZYC7-if00-port0" + + # Others: Calibration angles, joint directions etc + joint_ids: Tuple[int, ...] = (1, 2, 3, 4, 5, 6, 7) + joint_signs: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1) # if follow the original open-sourced gello xarm7 setup + start_joints: Tuple[float, ...] = (0, 0, 0, np.pi/2, 0, np.pi/2, 0) + gripper_id: int = 8 # -1: no gripper + torque_joint_ids: Tuple[int, ...] = None # the joints will activate torque mode. diff --git a/ufactory_lerobot/teleoperators/pika_teleop/__init__.py b/ufactory_lerobot/teleoperators/pika_teleop/__init__.py new file mode 100644 index 0000000..6e0ead2 --- /dev/null +++ b/ufactory_lerobot/teleoperators/pika_teleop/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python + +from .pika_teleop_config import PikaTeleopConfig +from .pika_teleop import PikaTeleop diff --git a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py new file mode 100644 index 0000000..5515d59 --- /dev/null +++ b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python + +import time +import math +import numpy as np +from typing import Any +from threading import Thread, Event, Lock +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from xarm.wrapper import XArmAPI +from ufactory_lerobot.devices.pika import PikaDevice +from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations +from ..base_teleop import UFBaseTeleop +from .pika_teleop_config import PikaTeleopConfig + + +class PikaTeleop(UFBaseTeleop, Thread): + + config_class = PikaTeleopConfig + name = "Pika Teleop For xArm" + + def __init__(self, config: PikaTeleopConfig): + + super().__init__(config) + Thread.__init__(self) # Do NOT REMOVE! + self.stop_event = Event() + self.config = config + self._is_connected = False + self._is_calibrated = True + self._data_lock = Lock() + self._ctrl_flag = False + self._need_initial = False + + self.pika_device = PikaDevice(1, pika_sense_port=self.config.port) + self.pika_sense = self.pika_device.pika_sense + + if self.config.robot_ip: + self.arm = XArmAPI(self.config.robot_ip, is_radian=True) + else: + self.arm = None + + self._robot_target_pose = None + self._gripper_target_pos = None + + @property + def action_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (7,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5, "gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5}, + } + + @property + def feedback_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (7,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5, "gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5}, + } + + @property + def is_connected(self) -> bool: + return self._is_connected + + @property + def is_calibrated(self) -> bool: + return self._is_calibrated + + def calibrate(self) -> None: + # CHECK!! + pass + + def configure(self) -> None: + pass + + def connect(self, calibrate: bool = False) -> None: + self.start() + + def disconnect(self): + if not self._is_connected: + DeviceNotConnectedError(f"{self} is not connected.") + + self.stop_event.set() + self._is_connected = False + self.join() + + def set_ctrl_status(self, status): + if status: + if not self._ctrl_flag: + print('开始遥操作') + self._ctrl_flag = True + self._need_initial = True + else: + self._ctrl_flag = False + self._need_initial = False + print('停止遥操作') + + def run(self): + self._is_connected = True + init_state = self.pika_sense.get_command_state() + curr_state = init_state + + last_gripper_distance = 0 + + self._ctrl_flag = False # 是否开启遥操作 + self._need_initial = False + + sleep_time = 1 / self.config.frequency + + if self.arm: + self.arm.set_linear_spd_limit_factor(2.0) + + pika_to_robot_eef = [0, 0, 0, math.pi, -math.pi / 2, 0] # rpy + # pika_to_robot_eef = [0, 0, 0, math.pi, 0, 0] + + # pika坐标系到机械臂坐标系的变换关系对应的变换矩阵 + pika_to_robot_matrix = Transformations.xyzrpy_to_rotation_matrix(*pika_to_robot_eef) + # 机械臂初始位置对应的变换矩阵 + # robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*[0, 0, 190, -np.pi, -np.radians(41), 0]) + robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*[300, 0, 365, np.pi, 0, 0]) + # pika初始位置转换到机械臂坐标系后对应的变换矩阵 + pika_begin_robot_matrix = None + # pika目标位置转换到机械臂坐标系后对应的变换矩阵 + pika_end_robot_matrix = None + + scale_xyz = self.config.scale_xyz + + while not self.stop_event.is_set(): + time.sleep(sleep_time) + + if not self.arm and pika_begin_robot_matrix is None: + pose = self.pika_sense.get_pose(self.pika_device.pika_tracker_device) + if not pose: + continue + x, y, z = pose.position[0] * 1000 * scale_xyz, pose.position[1] * 1000 * scale_xyz, pose.position[2] * 1000 * scale_xyz + pika_begin_robot_matrix = Transformations.tracker_pose_to_robot_matrix(x, y, z, pose.rotation, pika_to_robot_matrix) + robot_target_pose = Transformations.tracker_robot_matrix_to_robot_pose(pika_begin_robot_matrix, pika_begin_robot_matrix, robot_base_matrix, is_axis_angle=True) + print('初始绑定, 当前Pika位置对应的机械臂目标位置: x={:.6f}, y={:.6f}, z={:.6f}, rx={:.6f}, ry={:.6f}, rz={:.6f}'.format(robot_target_pose[0], robot_target_pose[1], robot_target_pose[2], math.degrees(robot_target_pose[3]), math.degrees(robot_target_pose[4]), math.degrees(robot_target_pose[5]))) + continue + + state = self.pika_sense.get_command_state() + if state != curr_state: + curr_state = state + if not self._ctrl_flag and curr_state != init_state: + self._ctrl_flag = True + self._need_initial = True + # self.robot_init() + print('开始遥操作') + time.sleep(1) + elif self._ctrl_flag and curr_state == init_state: + self._ctrl_flag = False + print('停止遥操作') + continue + + if self._ctrl_flag and self.arm and (not self.arm.connected or self.arm.error_code != 0 or self.arm.state >= 4): + print('机械臂原因, 遥操作自动停止') + init_state = state + curr_state = state + self._ctrl_flag = False + continue + + if not self._ctrl_flag: + continue + + if self.config.use_gripper: + distance = min(max(self.pika_sense.get_gripper_distance(), 0), 100) + + if abs(last_gripper_distance - distance) > 2: + last_gripper_distance = distance + with self._data_lock: + self._gripper_target_pos = last_gripper_distance + + pose = self.pika_sense.get_pose(self.pika_device.pika_tracker_device) + if not pose: + continue + x, y, z = pose.position[0] * 1000 * scale_xyz, pose.position[1] * 1000 * scale_xyz, pose.position[2] * 1000 * scale_xyz + + if not self.arm: + # 只有PIKA设备, 没有机械臂 + pika_end_robot_matrix = Transformations.tracker_pose_to_robot_matrix(x, y, z, pose.rotation, pika_to_robot_matrix) + robot_target_pose = Transformations.tracker_robot_matrix_to_robot_pose(pika_begin_robot_matrix, pika_end_robot_matrix, robot_base_matrix, is_axis_angle=True) + + if self._need_initial: + self._need_initial = False + print('[初始位置] x={:.6f}, y={:.6f}, z={:.6f}, rx={:.6f}, ry={:.6f}, rz={:.6f}'.format(robot_target_pose[0], robot_target_pose[1], robot_target_pose[2], math.degrees(robot_target_pose[3]), math.degrees(robot_target_pose[4]), math.degrees(robot_target_pose[5]))) + else: + if self._need_initial: + self._need_initial = False + # _, robot_pos = self.arm.get_position() + _, robot_pos = self.arm.get_position(is_radian=True) + robot_base_pose = robot_pos + print('[初始] 机械臂位置: {}'.format(robot_pos)) + + # 机械臂初始位置对应的变换矩阵 + robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*robot_pos) + + # pika初始位置转换到机械臂坐标系后对应的变换矩阵 + pika_begin_robot_matrix = Transformations.tracker_pose_to_robot_matrix(x, y, z, pose.rotation, pika_to_robot_matrix) + pika_end_robot_matrix = pika_begin_robot_matrix + else: + # pika目标位置转换到机械臂坐标系后对应的变换矩阵 + pika_end_robot_matrix = Transformations.tracker_pose_to_robot_matrix(x, y, z, pose.rotation, pika_to_robot_matrix) + + robot_target_pose = Transformations.tracker_robot_matrix_to_robot_pose(pika_begin_robot_matrix, pika_end_robot_matrix, robot_base_matrix, is_axis_angle=True) + + with self._data_lock: + self._robot_target_pose = robot_target_pose + + # delta action + def get_action(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError( + "PikaTeleop is not connected. You need to run `connect()` before `get_action()`." + ) + + with self._data_lock: + if self._robot_target_pose is not None: + robot_target_pose = self._robot_target_pose.copy() + else: + robot_target_pose = None + if self._gripper_target_pos is not None: + gripper_target_pos = (100 - self._gripper_target_pos) / (100 - 0) + else: + gripper_target_pos = 0.0 + + if robot_target_pose is None: + if self.arm: + _, robot_target_pose = self.arm.get_position_aa(is_radian=True) + else: + # robot_target_pose = [0, 0, 190, -np.pi, -np.radians(41), 0] + robot_target_pose = [300, 0, 365, np.pi, 0, 0] + # print(self._robot_target_pose, robot_target_pose) + + # output is delta change of the robot pose + action_dict = { + "pose.x": robot_target_pose[0], + "pose.y": robot_target_pose[1], + "pose.z": robot_target_pose[2], + "pose.rx": robot_target_pose[3], + "pose.ry": robot_target_pose[4], + "pose.rz": robot_target_pose[5], + } + + if self.config.use_gripper: + action_dict.update({"gripper.pos": gripper_target_pos}) + + return action_dict + + def send_feedback(self, feedback: dict[str, float]) -> None: + raise NotImplementedError \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py new file mode 100644 index 0000000..453adae --- /dev/null +++ b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +from lerobot.teleoperators import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("uf::pika_teleop") +@dataclass +class PikaTeleopConfig(TeleoperatorConfig): + # robot_ip to connect to the arm + robot_ip: str = None + # Port to connect to the pika + port: str = None + frequency: int = 100 # hz + use_gripper: bool = True + scale_xyz: float = 1.0 # + rx_continuous: bool = False diff --git a/ufactory_lerobot/teleoperators/space_mouse/__init__.py b/ufactory_lerobot/teleoperators/space_mouse/__init__.py new file mode 100644 index 0000000..cc0bcab --- /dev/null +++ b/ufactory_lerobot/teleoperators/space_mouse/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python + +from .space_mouse_config import SpaceMouseTeleopConfig +from .space_mouse import SpaceMouseTeleop diff --git a/ufactory_lerobot/teleoperators/space_mouse/space_mouse.py b/ufactory_lerobot/teleoperators/space_mouse/space_mouse.py new file mode 100644 index 0000000..5028f5c --- /dev/null +++ b/ufactory_lerobot/teleoperators/space_mouse/space_mouse.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python + +import time +import numpy as np +from typing import Any +from threading import Thread, Event +from spnav import spnav_open, spnav_poll_event, spnav_close, SpnavMotionEvent, SpnavButtonEvent +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError +from collections import defaultdict + +from lerobot.teleoperators import Teleoperator +from .space_mouse_config import SpaceMouseTeleopConfig + +class SpaceMouseTeleop(Teleoperator, Thread): + + config_class = SpaceMouseTeleopConfig + name = "Space Mouse Teleop For xArm" + + def __init__(self, config: SpaceMouseTeleopConfig): + + super().__init__(config) + Thread.__init__(self) # Do NOT REMOVE! + self.stop_event = Event() + self.config = config + self.max_value = config.max_value + self.frequency = config.frequency + self.max_pos_speed = config.max_pos_speed + deadzone = config.deadzone + self._is_connected = False + self.dtype = np.float32 # CHECK! make it configurable ??? + + if np.issubdtype(type(deadzone), np.number): + self.deadzone = np.full(6, fill_value=deadzone, dtype=self.dtype) + else: + self.deadzone = np.array(deadzone, dtype=self.dtype) + assert (self.deadzone >= 0).all() + + self.motion_event = SpnavMotionEvent([0,0,0], [0,0,0], 0) + self.button_state = defaultdict(lambda: False) + self.tx_zup_spnav = np.array([ + [0,0,-1], + [1,0,0], + [0,1,0] + ], dtype=np.float32) + + @property + def action_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (4,), + "names": {"pose.dx": 0, "pose.dy": 1, "pose.dz": 2, "gripper.pos": 3}, + } + else: + return { + "dtype": "float32", + "shape": (3,), + "names": {"pose.dx": 0, "pose.dy": 1, "pose.dz": 2}, + } + + @property + def feedback_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (4,), + "names": {"pose.dx": 0, "pose.dy": 1, "pose.dz": 2, "gripper.pos": 3}, + } + else: + return { + "dtype": "float32", + "shape": (3,), + "names": {"pose.dx": 0, "pose.dy": 1, "pose.dz": 2}, + } + + @property + def is_connected(self) -> bool: + return self._is_connected + + @property + def is_calibrated(self) -> bool: + return self._is_calibrated + + def calibrate(self) -> None: + # CHECK!! + pass + + def configure(self) -> None: + pass + + def connect(self, calibrate: bool = False) -> None: + self.start() + + def disconnect(self): + if not self._is_connected: + DeviceNotConnectedError(f"{self} is not connected.") + + self.stop_event.set() + self._is_connected = False + self.join() + + def get_motion_state(self): + me = self.motion_event + state = np.array(me.translation + me.rotation, + dtype=self.dtype) / self.max_value + is_dead = (-self.deadzone < state) & (state < self.deadzone) + state[is_dead] = 0 + return state + + def get_motion_state_transformed(self): + """ + Return in right-handed coordinate + z + *------>y right + | _ + | (O) space mouse + v + x + back + + """ + state = self.get_motion_state() + tf_state = np.zeros_like(state) + tf_state[:3] = self.tx_zup_spnav @ state[:3] + tf_state[3:] = self.tx_zup_spnav @ state[3:] + return tf_state + + def is_button_pressed(self, button_id): + return self.button_state[button_id] + + def run(self): + spnav_open() + self._is_connected = True + try: + while not self.stop_event.is_set(): + event = spnav_poll_event() + if isinstance(event, SpnavMotionEvent): + self.motion_event = event + elif isinstance(event, SpnavButtonEvent): + self.button_state[event.bnum] = event.press + else: + time.sleep(1/200) + finally: + self._is_connected = False + spnav_close() + + # delta action + def get_action(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError( + "SpaceMouseTeleop is not connected. You need to run `connect()` before `get_action()`." + ) + + # self._drain_pressed_keys() + sm_state = self.get_motion_state_transformed() + + dpos = sm_state[:3] * self.max_pos_speed / self.frequency + + # Currently No rotation operation + # drot_xyz = sm_state[3:] * (max_rot_speed / frequency) + + # if not self.is_button_pressed(0): + # # translation mode + # drot_xyz[:] = 0 + # else: + # dpos[:] = 0 + # if not self.is_button_pressed(1): + + # X-Y 2D translation mode, no gripper control. Modify the code if you need more DOF control + dpos[2] = 0 + + gripper_action = 1.0 + + # output is delta change of the robot pose + action_dict = { + "pose.dx": dpos[0], + "pose.dy": dpos[1], + "pose.dz": dpos[2], + } + + if self.config.use_gripper: + action_dict.update({"gripper.pos": gripper_action}) + + return action_dict + + def send_feedback(self, feedback: dict[str, float]) -> None: + raise NotImplementedError \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py b/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py new file mode 100644 index 0000000..0138c39 --- /dev/null +++ b/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass + +from lerobot.teleoperators import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("uf::spacemouse_teleop") +@dataclass +class SpaceMouseTeleopConfig(TeleoperatorConfig): + # Port to connect to the arm + max_value: int = 300 + deadzone: tuple = (0,0,0,0,0,0) + use_gripper: bool = False + frequency: int = 10 # hz + max_pos_speed: int = 250 # mm/s + # Others: Calibration angles, joint directions etc. diff --git a/ufactory_lerobot/teleoperators/uf_mock_teleop/__init__.py b/ufactory_lerobot/teleoperators/uf_mock_teleop/__init__.py new file mode 100644 index 0000000..5f57960 --- /dev/null +++ b/ufactory_lerobot/teleoperators/uf_mock_teleop/__init__.py @@ -0,0 +1,4 @@ +#!/usr/bin/env python + +from .uf_mock_teleop_config import UFMockTeleopConfig +from .uf_mock_teleop import UFMockTeleop diff --git a/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop.py b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop.py new file mode 100644 index 0000000..a182bce --- /dev/null +++ b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop.py @@ -0,0 +1,479 @@ +#!/usr/bin/env python + +import time +import socket +import random +from typing import Any +import threading +from xarm.wrapper import XArmAPI +from xarm.core.utils import convert +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError + +from lerobot.teleoperators import Teleoperator +from .uf_mock_teleop_config import UFMockTeleopConfig + +RECORD_NONE = 0 +RECORD_POSE = RECORD_NO_GRIPPER = 1 +RECORD_GRIPPER = 2 +RECORD_POSE_GRIPPER = RECORD_POSE | RECORD_GRIPPER + + +class UFMockTeleop(Teleoperator, threading.Thread): + + config_class = UFMockTeleopConfig + name = "Mock Teleop For xArm" + + def __init__(self, config: UFMockTeleopConfig): + + super().__init__(config) + threading.Thread.__init__(self) # Do NOT REMOVE! + self.config = config + self._is_connected = False + self._is_calibrated = True # CHECK!! + self._is_joint_space = config.control_space == 'joint' + self._action_datas = [] + + self.arm = XArmAPI(config.robot_ip, do_not_open=True) + self.rt_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + assert len(self.config.mock_x_range) >= 1, 'The length of the parameter mock_x_range cannot be less than 1.' + assert len(self.config.mock_y_range) >= 1, 'The length of the parameter mock_y_range cannot be less than 1.' + assert len(self.config.mock_z_range) >= 1, 'The length of the parameter mock_z_range cannot be less than 1.' + assert len(self.config.mock_roll_range) >= 1, 'The length of the parameter mock_roll_range cannot be less than 1.' + assert len(self.config.mock_pitch_range) >= 1, 'The length of the parameter mock_pitch_range cannot be less than 1.' + assert len(self.config.mock_yaw_range) >= 1, 'The length of the parameter mock_yaw_range cannot be less than 1.' + + self.mock_x_range = self.config.mock_x_range if len(self.config.mock_x_range) > 1 else (self.config.mock_x_range[0], self.config.mock_x_range[0]) + self.mock_y_range = self.config.mock_y_range if len(self.config.mock_y_range) > 1 else (self.config.mock_y_range[0], self.config.mock_y_range[0]) + self.mock_z_range = self.config.mock_z_range if len(self.config.mock_z_range) > 1 else (self.config.mock_z_range[0], self.config.mock_z_range[0]) + self.mock_roll_range = self.config.mock_roll_range if len(self.config.mock_roll_range) > 1 else (self.config.mock_roll_range[0], self.config.mock_roll_range[0]) + self.mock_pitch_range = self.config.mock_pitch_range if len(self.config.mock_pitch_range) > 1 else (self.config.mock_pitch_range[0], self.config.mock_pitch_range[0]) + self.mock_yaw_range = self.config.mock_yaw_range if len(self.config.mock_yaw_range) > 1 else (self.config.mock_yaw_range[0], self.config.mock_yaw_range[0]) + + self._last_mock_x = self.mock_x_range[0] + self._last_mock_y = self.mock_y_range[0] + self._last_mock_z = self.mock_z_range[0] + self._last_mock_roll = self.mock_roll_range[0] + self._last_mock_pitch = self.mock_pitch_range[0] + self._last_mock_yaw = self.mock_yaw_range[0] + + self._update_lock = threading.Lock() + self._action_inx = 0 + self._record_status = 0 # 0: not recording; 1: recording joint/tcp pose; 3: recording joint/tcp pose and gripper pos + self._report_gripper = False + + def __set_record_status(self, status: int): + self._record_status = status + + def __reset_action_data(self): + self.__set_record_status(RECORD_NONE) + self._action_datas.clear() + + def __init_robot(self): + self.arm.connect(self.config.robot_ip) + self.arm.motion_enable() + self.arm.set_mode(0) + self.arm.set_state(0) + + if self.config.initial_tcp_pose is not None: + self.arm.set_position(*self.config.initial_tcp_pose, speed=self.config.mock_tcp_speed, wait=True) + if self.config.gripper_type > 0: + self.arm.set_gripper_enable(True) + self.arm.set_gripper_speed(3000) + self.arm.set_gripper_position(800, wait=True) + + if self.arm.arm.version_is_ge(2, 7, 100) and hasattr(self.arm, 'set_external_device_monitor_params'): + self.arm.set_external_device_monitor_params(self.config.gripper_type, self.config.gripper_freq) + self._report_gripper = True + + self.rt_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.rt_sock.setblocking(True) + self.rt_sock.settimeout(1) + self.rt_sock.connect((self.config.robot_ip, 30000)) + self._is_connected = True + + def __mock(self, events=None) -> bool: + mock_tcp_speed = self.config.mock_tcp_speed + grasp_tcp_speed = self.config.grasp_tcp_speed + initial_tcp_pose = self.config.initial_tcp_pose + self.arm.set_mode(0) + self.arm.set_state(0) + self.arm.set_position(*initial_tcp_pose, speed=mock_tcp_speed) + if self.config.gripper_type > 0: + self.arm.set_gripper_position(800, wait=True) + + if abs(self._last_mock_x - self.mock_x_range[0]) >= abs(self._last_mock_x - self.mock_x_range[1]): + mock_x_range = [self.mock_x_range[0], self._last_mock_x] + else: + mock_x_range = [self._last_mock_x, self.mock_x_range[1]] + if abs(self._last_mock_y - self.mock_y_range[0]) >= abs(self._last_mock_y - self.mock_y_range[1]): + mock_y_range = [self.mock_y_range[0], self._last_mock_y] + else: + mock_y_range = [self._last_mock_y, self.mock_y_range[1]] + if abs(self._last_mock_z - self.mock_z_range[0]) >= abs(self._last_mock_z - self.mock_z_range[1]): + mock_z_range = [self.mock_z_range[0], self._last_mock_z] + else: + mock_z_range = [self._last_mock_z, self.mock_z_range[1]] + if abs(self._last_mock_roll - self.mock_roll_range[0]) >= abs(self._last_mock_roll - self.mock_roll_range[1]): + mock_roll_range = [self.mock_roll_range[0], self._last_mock_roll] + else: + mock_roll_range = [self._last_mock_roll, self.mock_roll_range[1]] + if abs(self._last_mock_pitch - self.mock_pitch_range[0]) >= abs(self._last_mock_pitch - self.mock_pitch_range[1]): + mock_pitch_range = [self.mock_pitch_range[0], self._last_mock_pitch] + else: + mock_pitch_range = [self._last_mock_pitch, self.mock_pitch_range[1]] + if abs(self._last_mock_yaw - self.mock_yaw_range[0]) >= abs(self._last_mock_yaw - self.mock_yaw_range[1]): + mock_yaw_range = [self.mock_yaw_range[0], self._last_mock_yaw] + else: + mock_yaw_range = [self._last_mock_yaw, self.mock_yaw_range[1]] + + x = random.uniform(mock_x_range[0], mock_x_range[1]) + y = random.uniform(mock_y_range[0], mock_y_range[1]) + z = random.uniform(mock_z_range[0], mock_z_range[1]) + roll = random.uniform(mock_roll_range[0], mock_roll_range[1]) + pitch = random.uniform(mock_pitch_range[0], mock_pitch_range[1]) + yaw = random.uniform(mock_yaw_range[0], mock_yaw_range[1]) + + # x = random.uniform(self.mock_x_range[0], self.mock_x_range[1]) + # y = random.uniform(self.mock_y_range[0], self.mock_y_range[1]) + # z = random.uniform(self.mock_z_range[0], self.mock_z_range[1]) + # roll = random.uniform(self.mock_roll_range[0], self.mock_roll_range[1]) + # pitch = random.uniform(self.mock_pitch_range[0], self.mock_pitch_range[1]) + # yaw = random.uniform(self.mock_yaw_range[0], self.mock_yaw_range[1]) + + print(f'\n[MOCK] x: {x:.1f} mm, y: {y:.1f} mm, z: {z:.1f} mm, roll: {roll:.1f} °, pitch: {pitch:.1f} °, yaw: {yaw:.1f} °') + self._last_mock_x = x + self._last_mock_y = y + self._last_mock_z = z + self._last_mock_roll = roll + self._last_mock_pitch = pitch + self._last_mock_yaw = yaw + + code = self.arm.set_position(x=x, y=y, z=z, roll=roll, pitch=pitch, yaw=yaw, speed=mock_tcp_speed, wait=True) + if code != 0: + print(f'[MOCK ERROR] Failed to move to mock position, error code: {code}') + return False + print('[MOCK] Reached mock position.') + if events and events['exit_early']: # exit early + return False + print('*** Place the target in the correct gripping position, adjust the robotic arm height.') + input('*** Enter to continue >>> ') + if events and events['exit_early']: # exit early + return False + + self.arm.set_mode(0) + self.arm.set_state(0) + + _, tcp_target_pose = self.arm.get_position(is_radian=False) + # 回到抓取目标正上方 + code = self.arm.set_position(z=z + self.config.hover_offset, speed=mock_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + return False + _, tcp_hover_pose = self.arm.get_position(is_radian=False) + + # 回到初始位置 + code = self.arm.set_position(*initial_tcp_pose, speed=mock_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + return False + + ##################### 模拟抓取 ############################ + + self.__reset_action_data() + self._action_inx = 0 + self.__set_record_status(RECORD_POSE_GRIPPER) + time.sleep(1.0) + + case_type = self.config.mock_type if self.config.mock_type in [1, 2, 3] else 1 + x0, y0, yaw0 = tcp_hover_pose[0], tcp_hover_pose[1], tcp_hover_pose[5] + dis = 100 + yaw_dis = 45 + spd = grasp_tcp_speed // 3 * 2 + x_range = [x0 - dis, x0 + dis] + y_range = [y0 - dis, y0 + dis] + yaw_range = [yaw0 - yaw_dis, yaw0 + yaw_dis] + + if self.config.mock_type == 0: + # 按一定概率随机执行方案1/2/3 + val = val = random.random() + case_type = 1 if val < 0.66 else 2 if val < 0.88 else 3 + + cnt = 0 + if case_type == 1: + # 方案1: 不做额外处理 + pass + elif case_type == 2: + # 方案2: 在目标点位上方固定高度一定范围内游走 + cnt = random.randint(3, 8) + for i in range(cnt): + x1 = min(max(self.mock_x_range[0], random.uniform(x_range[0], x_range[1])), self.mock_x_range[1]) + y1 = min(max(self.mock_y_range[0], random.uniform(y_range[0], y_range[1])), self.mock_y_range[1]) + yaw1 = min(max(self.mock_yaw_range[0], random.uniform(yaw_range[0], yaw_range[1])), self.mock_yaw_range[1]) + pose = [x1, y1, tcp_hover_pose[2], tcp_hover_pose[3], tcp_hover_pose[4], yaw1] + self.arm.set_position(*pose, speed=grasp_tcp_speed if i == 0 else spd, wait=True if i == cnt - 1 else False) + + x_range[0] = x1 if x1 < x0 else x0 if x1 == x0 else x_range[0] + x_range[1] = x1 if x1 > x0 else x0 if x1 == x0 else x_range[1] + y_range[0] = y1 if y1 < y0 else y0 if y1 == y0 else y_range[0] + y_range[1] = y1 if y1 > y0 else y0 if y1 == y0 else y_range[1] + yaw_range[0] = yaw1 if yaw1 < yaw0 else yaw0 if yaw1 == yaw0 else yaw_range[0] + yaw_range[1] = yaw1 if yaw1 > yaw0 else yaw0 if yaw1 == yaw0 else yaw_range[1] + elif case_type == 3: + # 方案3: 在目标点位上方随机高度一定范围内游走 + cnt = random.randint(2, 5) + for i in range(cnt): + x1 = min(max(self.mock_x_range[0], random.uniform(x_range[0], x_range[1])), self.mock_x_range[1]) + y1 = min(max(self.mock_y_range[0], random.uniform(y_range[0], y_range[1])), self.mock_y_range[1]) + yaw1 = min(max(self.mock_yaw_range[0], random.uniform(yaw_range[0], yaw_range[1])), self.mock_yaw_range[1]) + z1 = tcp_hover_pose[2] - (random.uniform(0, self.config.hover_offset - 50) if i != 0 else 0) + pose = [x1, y1, z1, tcp_hover_pose[3], tcp_hover_pose[4], yaw] + self.arm.set_position(*pose, speed=grasp_tcp_speed if i == 0 else spd, wait=True) + if tcp_hover_pose[2] - z1 > 30 or i == cnt - 1: + pose = [x1, y1, tcp_hover_pose[2], tcp_hover_pose[3], tcp_hover_pose[4], yaw1] + self.arm.set_position(*pose, speed=grasp_tcp_speed if i == 0 else spd, wait=True) + + x_range[0] = x1 if x1 < x0 else x0 if x1 == x0 else x_range[0] + x_range[1] = x1 if x1 > x0 else x0 if x1 == x0 else x_range[1] + y_range[0] = y1 if y1 < y0 else y0 if y1 == y0 else y_range[0] + y_range[1] = y1 if y1 > y0 else y0 if y1 == y0 else y_range[1] + yaw_range[0] = yaw1 if yaw1 < yaw0 else yaw0 if yaw1 == yaw0 else yaw_range[0] + yaw_range[1] = yaw1 if yaw1 > yaw0 else yaw0 if yaw1 == yaw0 else yaw_range[1] + + print(f'[MOCK] mock_type={self.config.mock_type}, case_type={case_type}, cnt={cnt}') + + # 去到抓取目标正上方 + code = self.arm.set_position(*tcp_hover_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.5) + + # 下移到抓取目标位置 + code = self.arm.set_position(*tcp_target_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.25) + + if self.config.gripper_type > 0: + # 抓取 + self.arm.set_gripper_position(0, wait=True) + time.sleep(0.25) + self.__set_record_status(RECORD_NONE) + # 记录抓取时的夹爪位置 + _, gripper_pos = self.arm.get_gripper_position() + # 松开(不真正抓取) + self.arm.set_gripper_position(800, wait=True) + + self.__set_record_status(RECORD_NO_GRIPPER) + # 回到抓取目标正上方 + code = self.arm.set_position(*tcp_hover_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + if self.config.gripper_type > 0: + # 把机械爪恢复成抓取时位置 + self.arm.set_gripper_position(gripper_pos) + time.sleep(0.25) + + # 去放置位置正上方 + place_tcp_hover_pose = self.config.place_tcp_pose.copy() + place_tcp_hover_pose[2] += self.config.hover_offset + code = self.arm.set_position(*place_tcp_hover_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.25) + + # 下移到放置位置 + place_tcp_target_pose = self.config.place_tcp_pose + code = self.arm.set_position(*place_tcp_target_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.25) + + self.__set_record_status(RECORD_POSE_GRIPPER) + # 放下 + if self.config.gripper_type > 0: + self.arm.set_gripper_position(800, wait=True) + time.sleep(0.25) + # 回到放置位置正上方 + code = self.arm.set_position(*place_tcp_hover_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.25) + # 回到初始位置 + code = self.arm.set_position(*initial_tcp_pose, speed=grasp_tcp_speed, wait=True) + if code or (events and events['exit_early']): # exit early + self.__reset_action_data() + return False + time.sleep(0.25) + self.__set_record_status(RECORD_NONE) + + print(f'[MOCK] Recorded {len(self._action_datas)} steps of action data.') + + self.arm.set_mode(6) + self.arm.set_state(0) + + return True + + @property + def action_features(self) -> dict: + if self._is_joint_space: + # Add one more dof for gripper + return { f"J{i+1}.pos": float for i in range(7) } | {"gripper.pos": float} + else: + if self.config.gripper_type > 0: + return { + "dtype": "float32", + "shape": (7,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5, "gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5}, + } + + @property + def feedback_features(self) -> dict: + if self._is_joint_space: + # Add one more dof for gripper + return { f"J{i+1}.pos": float for i in range(7) } | {"gripper.pos": float} + else: + if self.config.gripper_type > 0: + return { + "dtype": "float32", + "shape": (7,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5, "gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {"pose.x": 0, "pose.y": 1, "pose.z": 2, "pose.rx": 3, "pose.ry": 4, "pose.rz": 5}, + } + + @property + def is_connected(self) -> bool: + return self._is_connected + + @property + def is_calibrated(self) -> bool: + return self._is_calibrated + + def calibrate(self) -> None: + # CHECK!! + pass + + def configure(self, events=None) -> None: + return self.__mock(events=events) + + def connect(self, calibrate: bool = False) -> None: + self.__init_robot() + self.start() + + def disconnect(self): + self._is_connected = False + self.arm.set_state(4) + self.arm.disconnect() + self.join() + + def run(self): + sleep_time_s = 1 / self.config.fps + next_record_time = time.perf_counter() + sleep_time_s + + buffer = self.rt_sock.recv(4) + while len(buffer) < 4: + buffer += self.rt_sock.recv(4 - len(buffer)) + size = convert.bytes_to_u32(buffer[:4]) + print(f"UFACTORY Robot RT Report Thread starts!! size={size}") + gripper_pos = 1.0 + while self.is_connected: + buffer += self.rt_sock.recv(size - len(buffer)) + if len(buffer) < size: + continue + data = buffer[:size] + buffer = buffer[size:] + + with self._update_lock: + if self._record_status: + time_now = time.perf_counter() + if time_now - next_record_time >= -0.001: + if self._is_joint_space: + pose = convert.bytes_to_fp32s(data[116:144], 7) # joint angles + else: + pose = convert.bytes_to_fp32s(data[472:496], 6) # tcp pose + + if self.config.gripper_type > 0: + if self._record_status & RECORD_GRIPPER: + if size >= 744 and self._report_gripper: + external_device_info = convert.bytes_to_16s(data[738:744], 3) # gripper: [pos, speed, current] + grippos = min(external_device_info[0] * 10, self.config.gripper_open) + else: + _, grippos = self.arm.get_gripper_position() + grippos = min(grippos, self.config.gripper_open) + grippos_norm = (self.config.gripper_open - grippos) / (self.config.gripper_open - self.config.gripper_close) + pose.append(grippos_norm) # add gripper pos + gripper_pos = grippos_norm + else: + pose.append(gripper_pos) + else: + if self._is_joint_space: + pose.append(gripper_pos) + + self._action_datas.append(pose) + next_record_time = time_now + sleep_time_s + + # delta action + def get_action(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError( + "UFMockTeleop is not connected. You need to run `connect()` before `get_action()`." + ) + + if not self._action_datas: + if self._is_joint_space: + _, pose = self.arm.get_servo_angle(is_radian=True) + if self.config.gripper_type > 0: + _, grippos = self.arm.get_gripper_position() + grippos = min(grippos, self.config.gripper_open) + grippos_norm = (self.config.gripper_open - grippos) / (self.config.gripper_open - self.config.gripper_close) + pose.append(grippos_norm) + else: + pose.append(1.0) # gripper pos + else: + _, pose = self.arm.get_position_aa(is_radian=True) + if self.config.gripper_type > 0: + _, grippos = self.arm.get_gripper_position() + grippos = min(grippos, self.config.gripper_open) + grippos_norm = (self.config.gripper_open - grippos) / (self.config.gripper_open - self.config.gripper_close) + pose.append(grippos_norm) + else: + if self._action_inx >= len(self._action_datas): + pose = self._action_datas[-1] + else: + pose = self._action_datas[self._action_inx] + self._action_inx += 1 + action = {} + if self._is_joint_space: + for i in range(self.arm.axis): + action.update({f"J{i+1}.pos": pose[i]}) + action.update({"gripper.pos": pose[7]}) + else: + action.update({ + "pose.x": pose[0], + "pose.y": pose[1], + "pose.z": pose[2], + "pose.rx": pose[3], + "pose.ry": pose[4], + "pose.rz": pose[5], + }) + if self.config.gripper_type > 0: + action.update({"gripper.pos": pose[6]}) + return action + + def send_feedback(self, feedback: dict[str, float]) -> None: + raise NotImplementedError diff --git a/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py new file mode 100644 index 0000000..717fc19 --- /dev/null +++ b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Tuple +from dataclasses import dataclass + +from lerobot.teleoperators import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("uf::mock_teleop") +@dataclass +class UFMockTeleopConfig(TeleoperatorConfig): + robot_ip: str = "192.168.1.127" + control_space: str = "joint" # joint/cartesian + mock_type: int = 1 + mock_x_range: Tuple[float, ...] = (180, 600) # mm + mock_y_range: Tuple[float, ...] = (-280, 50) # mm + mock_z_range: Tuple[float, ...] = (180, 180) # mm + mock_roll_range: Tuple[float, ...] = (180, 180) # ° + mock_pitch_range: Tuple[float, ...] = (0, 0) # ° + mock_yaw_range: Tuple[float, ...] = (-90, 90) # ° + mock_tcp_speed: int = 200 # mm/s + hover_offset: int = 200 # mm + grasp_tcp_speed: int = 150 # mm/s + + place_tcp_speed: int = 100 # mm/s + place_tcp_pose: Tuple[float, ...] = (383, 253, 320, 180, 0, 0) # x,y,z in mm; roll,pitch,yaw in ° + initial_tcp_pose: Tuple[float, ...] = (450, 0, 520, 180, 0, 0, 0) # x,y,z in mm; roll,pitch,yaw in ° + + fps: int = 30 # Hz + + gripper_type: int = 1 # 0: no gripper, 1: xArm Gripper, 2: xArm Gripper G2 + gripper_freq: int = 50 # Hz + gripper_open: int = 800 + gripper_close: int = 0 diff --git a/ufactory_lerobot/teleoperators/umi_teleop/__init__.py b/ufactory_lerobot/teleoperators/umi_teleop/__init__.py new file mode 100644 index 0000000..703e249 --- /dev/null +++ b/ufactory_lerobot/teleoperators/umi_teleop/__init__.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python + +from .umi_teleop_config import UmiTeleopConfig +from .umi_teleop import UmiTeleop +from .multiple_umi_teleop_config import MultipleUmiTeleopConfig +from .multiple_umi_teleop import MultipleUmiTeleop diff --git a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py new file mode 100644 index 0000000..608fc23 --- /dev/null +++ b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python + +from typing import Any +from .multiple_umi_teleop_config import MultipleUmiTeleopConfig +from ..base_teleop import UFBaseTeleop +from .umi_teleop import UmiTeleop + + +class MultipleUmiTeleop(UFBaseTeleop): + + config_class = MultipleUmiTeleopConfig + name = "Multiple UMI Teleop For xArm" + + def __init__(self, config: MultipleUmiTeleopConfig): + super().__init__(config) + self.config = config + self.teleops = {} + for key, teleop_config in self.config.teleops.items(): + self.teleops[key] = UmiTeleop(teleop_config, prefix=key) + + def action_features(self) -> dict: + action_features = {} + for teleop in self.teleops.values(): + action_features.update(teleop.action_features) + return action_features + + @property + def feedback_features(self) -> dict: + feedback_features = {} + for teleop in self.teleops.values(): + feedback_features.update(teleop.feedback_features) + return feedback_features + + @property + def is_connected(self) -> bool: + return all(teleop.is_connected for teleop in self.teleops.values()) + + @property + def is_calibrated(self) -> bool: + return all(teleop.is_calibrated for teleop in self.teleops.values()) + + def connect(self, calibrate: bool = True) -> None: + for teleop in self.teleops.values(): + teleop.connect(calibrate=calibrate) + + def calibrate(self) -> None: + for teleop in self.teleops.values(): + teleop.calibrate() + + def configure(self) -> None: + for teleop in self.teleops.values(): + teleop.configure() + + def disconnect(self) -> None: + for teleop in self.teleops.values(): + teleop.disconnect() + + def set_ctrl_status(self, status): + for teleop in self.teleops.values(): + teleop.set_ctrl_status(status) + + def get_action(self) -> dict[str, Any]: + actions = {} + for teleop in self.teleops.values(): + actions.update(teleop.get_action()) + return actions + + def send_feedback(self, feedback: dict[str, Any]) -> None: + for key, teleop in self.teleops.items(): + feedback_subset = {k: v for k, v in feedback.items() if k.startswith(f"{key}.")} + teleop.send_feedback(feedback_subset) diff --git a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py new file mode 100644 index 0000000..9d90c18 --- /dev/null +++ b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from lerobot.teleoperators import TeleoperatorConfig +from .umi_teleop_config import UmiTeleopConfig + + +@TeleoperatorConfig.register_subclass("uf::multiple_umi_teleop") +@dataclass +class MultipleUmiTeleopConfig(TeleoperatorConfig): + teleops: dict[str, UmiTeleopConfig] \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py new file mode 100644 index 0000000..1d1df39 --- /dev/null +++ b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py @@ -0,0 +1,201 @@ +#!/usr/bin/env python + +from typing import Any +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError + +from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations +from ufactory_lerobot.devices.umi.vive_tracker import ViveTracker +from ufactory_lerobot.devices.umi.xvlib import XVLib +from ..base_teleop import UFBaseTeleop +from .umi_teleop_config import UmiTeleopConfig + + +class UmiTeleop(UFBaseTeleop): + + config_class = UmiTeleopConfig + name = "UMI Teleop For xArm" + + def __init__(self, config: UmiTeleopConfig, prefix=''): + super().__init__(config) + self.config = config + self.prefix = '' if not prefix else f'{prefix}.' + self._is_connected = False + self._is_calibrated = True + + if self.config.use_gripper: + self.config.init_clamp_stream = True + else: + self.config.init_clamp_stream = False + + self.tracker = ViveTracker() if self.config.use_vive_tracker else None + self.xvlib = XVLib(self.config.serial_number, self.config.init_slam, self.config.init_clamp_stream, self.config.init_color_camera, self.config.init_fisheye_cameras) + + # tracker_to_robot_eef = [0, 0, 0, math.pi / 2, -math.pi / 2, 0] + # tracker_to_robot_eef = [0, 0, 0, 0, 0, -math.pi/2] + # tracker_to_robot_eef = [0, 0, 0, math.pi, math.pi, 0] # Test1 + # tracker_to_robot_eef = [0, 0, 0, math.pi, math.pi, -math.pi/2] # Dual left + # tracker_to_robot_eef = [0, 0, 0, math.pi, math.pi, math.pi/2] # Dual right + tracker_to_robot_eef = self.config.tracker_to_robot_eef + self.tracker_to_robot_matrix = Transformations.xyzrpy_to_rotation_matrix(*tracker_to_robot_eef) + # robot_base_pose = [300, 0, 300, 0, 0, 0] + # robot_base_pose = [300, 0, 300, math.pi, -math.pi/2, 0] + # robot_base_pose = [220, 0, 385, math.pi, 0, 0] + # robot_base_pose = [250, 0, 150, math.pi, 0, 0] # Test 1 + # robot_base_pose = [250, 0, 150, math.pi, 0, math.pi/2] # Dual left + # robot_base_pose = [250, 0, 150, math.pi, 0, math.pi/2] # Dual right + robot_base_pose = self.config.robot_base_pose + self.robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*robot_base_pose) + self.begin_tracker_robot_matrix = None + + @property + def action_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (7,), + "names": {f"{self.prefix}pose.x": 0, f"{self.prefix}pose.y": 1, f"{self.prefix}pose.z": 2, f"{self.prefix}pose.rx": 3, f"{self.prefix}pose.ry": 4, f"{self.prefix}pose.rz": 5, f"{self.prefix}gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {f"{self.prefix}pose.x": 0, f"{self.prefix}pose.y": 1, f"{self.prefix}pose.z": 2, f"{self.prefix}pose.rx": 3, f"{self.prefix}pose.ry": 4, f"{self.prefix}pose.rz": 5}, + } + + @property + def feedback_features(self) -> dict: + if self.config.use_gripper: + return { + "dtype": "float32", + "shape": (7,), + "names": {f"{self.prefix}pose.x": 0, f"{self.prefix}pose.y": 1, f"{self.prefix}pose.z": 2, f"{self.prefix}pose.rx": 3, f"{self.prefix}pose.ry": 4, f"{self.prefix}pose.rz": 5, f"{self.prefix}gripper.pos": 6}, + } + else: + return { + "dtype": "float32", + "shape": (6,), + "names": {f"{self.prefix}pose.x": 0, f"{self.prefix}pose.y": 1, f"{self.prefix}pose.z": 2, f"{self.prefix}pose.rx": 3, f"{self.prefix}pose.ry": 4, f"{self.prefix}pose.rz": 5}, + } + + @property + def is_connected(self) -> bool: + return self._is_connected + + @property + def is_calibrated(self) -> bool: + return self._is_calibrated + + def calibrate(self) -> None: + # CHECK!! + pass + + def configure(self) -> None: + pass + + def connect(self, calibrate: bool = False) -> None: + self.xvlib.xv_init(self.config.serial_number, self.config.init_slam, self.config.init_clamp_stream, self.config.init_color_camera, self.config.init_fisheye_cameras) + self._is_connected = True + + def disconnect(self): + self.xvlib.xv_uninit() + self._is_connected = False + + @staticmethod + def normalize_angle(angle_deg): + """ + 将角度归一化到 [-180, 180] 区间 + """ + while angle_deg > 180: + angle_deg -= 360 + while angle_deg < -180: + angle_deg += 360 + return angle_deg + + def set_ctrl_status(self, status): + if status: + self.begin_tracker_robot_matrix = None + else: + pass + + # delta action + def get_action(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError( + "UmiTeleop is not connected. You need to run `connect()` before `get_action()`." + ) + + 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() + else: + _, pose_data = self.xvlib.xv_get_slam_data() + position = pose_data.position.to_list(6) + quaternion = pose_data.quaternion.to_list(6) + # orientation = pose_data.orientation.to_list() + + # x, y, z = position[0] * 1000, position[1] * 1000, position[2] * 1000 + # roll, pitch, yaw = math.degrees(orientation[0]), math.degrees(orientation[1]), math.degrees(orientation[2]) + # print(f'[1] x={x:.1f}, y={y:.1f}, z={z:.1f}, roll={roll:.1f}, pitch={pitch:.1f}, yaw={yaw:.1f}') + + # x, y, z = position[2] * 1000, position[0] * 1000, position[1] * 1000 + # roll, pitch, yaw = orientation[0], orientation[1], orientation[2] + # roll, pitch, yaw = math.degrees(roll), math.degrees(pitch), math.degrees(yaw) + # print(f'[1] x={x:.1f}, y={y:.1f}, z={z:.1f}, roll={roll:.1f}, pitch={pitch:.1f}, yaw={yaw:.1f}') + + # x, y, z = position[0] * 1000, position[1] * 1000, position[2] * 1000 + # R_A = Transformations.quaternion_to_rotation_matrix(quaternion) + # roll, pitch, yaw = Transformations.rotation_matrix_to_rpy(R_A) + # roll, pitch, yaw = math.degrees(roll), math.degrees(pitch), math.degrees(yaw) + # print(f'[2] x={x:.1f}, y={y:.1f}, z={z:.1f}, roll={roll:.1f}, pitch={pitch:.1f}, yaw={yaw:.1f}') + # print('*' * 50) + # roll, pitch, yaw = math.degrees(orientation[0]), math.degrees(orientation[1]), math.degrees(orientation[2]) + # print(f'[1] x={x:.1f}, y={y:.1f}, z={z:.1f}, roll={roll:.1f}, pitch={pitch:.1f}, yaw={yaw:.1f}') + + x, y, z = position[0] * 1000, position[1] * 1000, position[2] * 1000 + tracker_robot_matrix = Transformations.tracker_pose_to_robot_matrix(x, y, z, quaternion, self.tracker_to_robot_matrix) + if self.begin_tracker_robot_matrix is None: + self.begin_tracker_robot_matrix = tracker_robot_matrix + + robot_target_pose = Transformations.tracker_robot_matrix_to_robot_pose(self.begin_tracker_robot_matrix, tracker_robot_matrix, self.robot_base_matrix, is_axis_angle=True) + x, y, z = robot_target_pose[0:3] + orientation = robot_target_pose[3:6] + # roll, pitch, yaw = list(map(math.degrees, orientation)) + # print(f'[{self.config.serial_number}] x={x:.3f}, y={y:.3f}, z={z:.3f}, rx={roll:.3f}, ry={pitch:.3f}, rz={yaw:.3f}') + + # R_prev = Transformations.rpy_to_rotation_matrix(math.pi, -math.pi / 2, 0) + # R_delta = Transformations.rxryrz_to_matrix(robot_target_pose[3:6]) + # R_curr = R_prev @ R_delta + # # # R_curr = R_prev.apply(R_delta) + # orientation = Transformations.rotation_matrix_to_rxryrz(R_curr) + + # roll, pitch, yaw = math.degrees(orientation[0]), math.degrees(orientation[1]), math.degrees(orientation[2]) + # print(f'[2] x={x:.1f}, y={y:.1f}, z={z:.1f}, rx={roll:.1f}, ry={pitch:.1f}, rz={yaw:.1f}') + # print('*' * 50) + + # output is delta change of the robot pose + action_dict = { + # "pose.x": z, + # "pose.y": -y, + # "pose.z": x, + # "pose.rx": orientation[2], + # "pose.ry": -orientation[1], + # "pose.rz": orientation[0], + f"{self.prefix}pose.x": x, + f"{self.prefix}pose.y": y, + f"{self.prefix}pose.z": z, + f"{self.prefix}pose.rx": orientation[0], + f"{self.prefix}pose.ry": orientation[1], + f"{self.prefix}pose.rz": orientation[2], + } + + if self.config.use_gripper: + _, clamp_data = self.xvlib.xv_get_clamp_stream_data() + gripper_pos = (87 - clamp_data.data) / (87 - 0) + action_dict.update({f"{self.prefix}gripper.pos": gripper_pos}) + + return action_dict + + def send_feedback(self, feedback: dict[str, float]) -> None: + raise NotImplementedError diff --git a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py new file mode 100644 index 0000000..bdcdb47 --- /dev/null +++ b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python + +# Copyright 2025 UFACTORY Inc. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Tuple +import numpy as np +from lerobot.teleoperators import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("uf::umi_teleop") +@dataclass +class UmiTeleopConfig(TeleoperatorConfig): + serial_number: str + init_slam: bool = True + init_clamp_stream: bool = True + init_color_camera: bool = False + init_fisheye_cameras: bool = False + use_gripper: bool = True + use_vive_tracker: bool = False + vive_tracker_id: str = 'WM0' + tracker_to_robot_eef: Tuple[float, ...] = (0, 0, 0, 0, 0, -np.pi/2) + robot_base_pose: Tuple[float, ...] = (300, 0, 300, np.pi, -np.pi/2, 0) diff --git a/ufactory_lerobot/teleoperators/utils.py b/ufactory_lerobot/teleoperators/utils.py new file mode 100644 index 0000000..30648ce --- /dev/null +++ b/ufactory_lerobot/teleoperators/utils.py @@ -0,0 +1,40 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from lerobot.teleoperators.utils import make_teleoperator_from_config as lerobot_make_teleoperator_from_config +from lerobot.teleoperators.config import TeleoperatorConfig +from lerobot.teleoperators.teleoperator import Teleoperator + + +def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator: + if config.type == "uf::gello_teleop": + from .gello_teleop import GelloTeleop + return GelloTeleop(config) + elif config.type == "uf::pika_teleop": + from .pika_teleop import PikaTeleop + return PikaTeleop(config) + elif config.type == "uf::spacemouse_teleop": + from .space_mouse import SpaceMouseTeleop + return SpaceMouseTeleop(config) + elif config.type == "uf::mock_teleop": + from .uf_mock_teleop import UFMockTeleop + return UFMockTeleop(config) + elif config.type == "uf::umi_teleop": + from .umi_teleop import UmiTeleop + return UmiTeleop(config) + elif config.type == "uf::multiple_umi_teleop": + from .umi_teleop import MultipleUmiTeleop + return MultipleUmiTeleop(config) + else: + return lerobot_make_teleoperator_from_config(config) diff --git a/ufactory_lerobot/utils/__init__.py b/ufactory_lerobot/utils/__init__.py new file mode 100644 index 0000000..8ce58c8 --- /dev/null +++ b/ufactory_lerobot/utils/__init__.py @@ -0,0 +1 @@ +from .utils import instantiate_from_dict \ No newline at end of file diff --git a/ufactory_lerobot/utils/utils.py b/ufactory_lerobot/utils/utils.py new file mode 100644 index 0000000..2289560 --- /dev/null +++ b/ufactory_lerobot/utils/utils.py @@ -0,0 +1,18 @@ +import importlib + +# recursive call, inspired from gello_software: +def instantiate_from_dict(cfg, ignore_cameras=False): + """Instantiate objects from configuration.""" + if isinstance(cfg, dict) and "_target_" in cfg: + module_path, class_name = cfg["_target_"].rsplit(".", 1) + cls = getattr(importlib.import_module(module_path), class_name) + kwargs = {k: v for k, v in cfg.items() if k != "_target_"} + # pp_dict ={k: instantiate_from_dict(v, ignore_cameras) for k, v in kwargs.items()} + # print(pp_dict) + return cls(**{k: {} if ignore_cameras and k == 'cameras' else instantiate_from_dict(v, ignore_cameras) for k, v in kwargs.items()}) + elif isinstance(cfg, dict): + return {k: {} if ignore_cameras and k == 'cameras' else instantiate_from_dict(v, ignore_cameras) for k, v in cfg.items()} + elif isinstance(cfg, list): + return [instantiate_from_dict(v, ignore_cameras) for v in cfg] + else: + return cfg