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