feat(robot): 添加 uFactory 机械臂完整功能包

包含机器人控制(uf_robot)、遥操作(teleoperators)、
摄像头(cameras)、设备驱动(devices)和执行脚本(scripts)等模块。
This commit is contained in:
Vinman 2026-06-11 11:55:23 +08:00
parent c3028aa7b4
commit 4467e19322
55 changed files with 5084 additions and 2 deletions

3
.gitignore vendored
View File

@ -4,7 +4,7 @@
__pycache__/ __pycache__/
*.py[cod] *.py[cod]
*$py.class *$py.class
*.so # *.so
*.egg *.egg
*.egg-info/ *.egg-info/
dist/ dist/
@ -83,3 +83,4 @@ models/
*.onnx *.onnx
*.engine *.engine
*.trt *.trt
third_party/

View File

@ -47,7 +47,7 @@ train = [
] ]
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["."]
[tool.ruff] [tool.ruff]
line-length = 100 line-length = 100

View File

@ -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

View File

View File

@ -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

View File

@ -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

View File

@ -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."
)

View File

@ -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

View File

View File

@ -0,0 +1 @@
from .pika_device import PikaDevice

View File

@ -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('=================')

View File

@ -0,0 +1,2 @@
from .xvlib import XVLib
from .vive_tracker import Transformations, ViveTracker

View File

@ -0,0 +1,2 @@
from .transformations import Transformations
from .vive_tracker import ViveTracker

View File

@ -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)

View File

@ -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()

View File

@ -0,0 +1 @@
from .xvlib import XVLib

View File

@ -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) 格式, 同样使用 imdecodeOpenCV 会自动处理常见的视频流头
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))

View File

View File

@ -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

View File

@ -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

View File

@ -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]

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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!! =======")

View File

@ -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

View File

@ -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)

View File

@ -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()

View File

@ -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()

View File

@ -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()

View File

@ -0,0 +1 @@
from .base_teleop import UFBaseTeleop

View File

@ -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

View File

@ -0,0 +1,4 @@
#!/usr/bin/env python
from .gello_teleop_config import GelloTeleopConfig
from .gello_teleop import GelloTeleop

View File

@ -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.")

View File

@ -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.

View File

@ -0,0 +1,4 @@
#!/usr/bin/env python
from .pika_teleop_config import PikaTeleopConfig
from .pika_teleop import PikaTeleop

View File

@ -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

View File

@ -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

View File

@ -0,0 +1,4 @@
#!/usr/bin/env python
from .space_mouse_config import SpaceMouseTeleopConfig
from .space_mouse import SpaceMouseTeleop

View File

@ -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

View File

@ -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.

View File

@ -0,0 +1,4 @@
#!/usr/bin/env python
from .uf_mock_teleop_config import UFMockTeleopConfig
from .uf_mock_teleop import UFMockTeleop

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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)

View File

@ -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]

View File

@ -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

View File

@ -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)

View File

@ -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)

View File

@ -0,0 +1 @@
from .utils import instantiate_from_dict

View File

@ -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