forked from tangger/lerobot
334 lines
12 KiB
Python
334 lines
12 KiB
Python
import time
|
||
import threading
|
||
from typing import Dict
|
||
from dataclasses import dataclass
|
||
from contextlib import contextmanager
|
||
from lerobot.common.robot_devices.motors.configs import RealmanDualMotorsBusConfig
|
||
from Robotic_Arm.rm_robot_interface import *
|
||
|
||
|
||
def compare_joint_difference(master_joints, follow_joints, threshold=30.0):
|
||
"""
|
||
比较主臂和从臂关节数据的差异
|
||
|
||
Args:
|
||
master_joints (list): 主臂关节数据 [joint1, joint2, ..., joint6]
|
||
follow_joints (list): 从臂关节数据 [joint1, joint2, ..., joint6]
|
||
threshold (float): 差异阈值（度），默认5度
|
||
|
||
Returns:
|
||
bool: True表示差异在阈值内，False表示超过阈值
|
||
"""
|
||
# 检查数据长度
|
||
if len(master_joints) != len(follow_joints):
|
||
return False
|
||
|
||
# 计算每个关节的绝对差异
|
||
for i in range(len(master_joints)):
|
||
diff = abs(master_joints[i] - follow_joints[i])
|
||
if diff > threshold:
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
@dataclass
|
||
class ArmState:
|
||
"""机械臂状态数据类"""
|
||
joint_positions: list
|
||
gripper_position: int
|
||
pose: list
|
||
|
||
|
||
class RealmanDualMotorsBus:
|
||
"""
|
||
对Realman SDK的二次封装
|
||
"""
|
||
def __init__(self, config: RealmanDualMotorsBusConfig):
|
||
self.config = config
|
||
self._initialize_arms()
|
||
self._initialize_parameters()
|
||
self._initialize_positions()
|
||
self._initialize_threading()
|
||
|
||
def _initialize_arms(self):
|
||
"""初始化机械臂连接"""
|
||
self.left_rmarm = RoboticArm(rm_thread_mode_e.RM_TRIPLE_MODE_E)
|
||
self.right_rmarm = RoboticArm(rm_thread_mode_e.RM_TRIPLE_MODE_E)
|
||
self.handle_left = self.left_rmarm.rm_create_robot_arm(
|
||
self.config.left_ip, self.config.left_port
|
||
)
|
||
self.handle_right = self.right_rmarm.rm_create_robot_arm(
|
||
self.config.right_ip, self.config.right_port
|
||
)
|
||
|
||
def _initialize_parameters(self):
|
||
"""初始化参数"""
|
||
self.motors = self.config.motors
|
||
self.axis = self.config.axis
|
||
self.joint_count = sum(self.axis.values())
|
||
self.left_offset = self.axis['left_joint']
|
||
|
||
def _initialize_positions(self):
|
||
"""初始化位置"""
|
||
self.init_joint_position = self.config.init_joint['joint']
|
||
self.safe_disable_position = self.config.init_joint['joint']
|
||
|
||
# 移动到初始位置
|
||
self._move_to_initial_position()
|
||
|
||
# 获取初始姿态
|
||
time.sleep(3)
|
||
self.init_pose = self._get_initial_pose()
|
||
|
||
def _initialize_threading(self):
|
||
"""初始化线程控制"""
|
||
self.left_slow_busy = False
|
||
self.right_slow_busy = False
|
||
self._thread_lock = threading.Lock()
|
||
|
||
# 添加读取相关的线程控制
|
||
self._state_cache = {"joint": {}, "pose": {}}
|
||
self._cache_lock = threading.Lock()
|
||
self._keep_reading = True
|
||
|
||
# 启动后台读取线程
|
||
self._start_background_readers()
|
||
|
||
def _start_background_readers(self):
|
||
"""启动后台读取线程"""
|
||
# 读取线程
|
||
threading.Thread(
|
||
target=self._read_task,
|
||
daemon=True,
|
||
name="arm_reader"
|
||
).start()
|
||
|
||
@property
|
||
def motor_names(self) -> list[str]:
|
||
return list(self.motors.keys())
|
||
|
||
@property
|
||
def motor_models(self) -> list[str]:
|
||
return [model for _, model in self.motors.values()]
|
||
|
||
@property
|
||
def motor_indices(self) -> list[int]:
|
||
return [idx for idx, _ in self.motors.values()]
|
||
|
||
@contextmanager
|
||
def _timeout_context(self, timeout: float = 5.0):
|
||
"""超时上下文管理器"""
|
||
start_time = time.time()
|
||
try:
|
||
yield lambda: time.time() - start_time < timeout
|
||
except Exception as e:
|
||
raise TimeoutError(f"操作超时: {e}")
|
||
|
||
def _read_task(self):
|
||
"""左臂后台读取任务 - 模仿_left_slow_task的风格"""
|
||
while self._keep_reading:
|
||
try:
|
||
left_state = self._read_arm_state(self.left_rmarm, "left")
|
||
with self._cache_lock:
|
||
self._state_cache["joint"].update(left_state["joint"])
|
||
self._state_cache["pose"].update(left_state["pose"])
|
||
except Exception as e:
|
||
print(f"左臂读取失败: {e}")
|
||
|
||
try:
|
||
right_state = self._read_arm_state(self.right_rmarm, "right")
|
||
with self._cache_lock:
|
||
self._state_cache["joint"].update(right_state["joint"])
|
||
self._state_cache["pose"].update(right_state["pose"])
|
||
except Exception as e:
|
||
print(f"右臂读取失败: {e}")
|
||
|
||
def _read_arm_state(self, arm: RoboticArm, prefix: str) -> dict:
|
||
"""读取单臂状态 - 保持原有逻辑"""
|
||
joint_msg = arm.rm_get_current_arm_state()[1]
|
||
gripper_msg = arm.rm_get_gripper_state()[1]
|
||
|
||
joint_state = joint_msg['joint']
|
||
gripper_state = gripper_msg['actpos']
|
||
pose_state = joint_msg['pose']
|
||
|
||
joint_state_dict = {}
|
||
for i in range(len(joint_state)):
|
||
joint_state_dict[f"{prefix}_joint_{i+1}"] = joint_state[i]
|
||
joint_state_dict[f"{prefix}_gripper"] = gripper_state
|
||
|
||
pose_state_dict = {
|
||
f"{prefix}_x": pose_state[0],
|
||
f"{prefix}_y": pose_state[1],
|
||
f"{prefix}_z": pose_state[2],
|
||
f"{prefix}_rx": pose_state[3],
|
||
f"{prefix}_ry": pose_state[4],
|
||
f"{prefix}_rz": pose_state[5],
|
||
}
|
||
|
||
return {"joint": joint_state_dict, 'pose': pose_state_dict}
|
||
|
||
def _move_to_initial_position(self):
|
||
"""移动到初始位置"""
|
||
left_joints = self.init_joint_position[:self.left_offset]
|
||
right_joints = self.init_joint_position[self.left_offset+1:-1]
|
||
|
||
self.left_rmarm.rm_movej(left_joints, 5, 0, 0, 1)
|
||
self.right_rmarm.rm_movej(right_joints, 5, 0, 0, 1)
|
||
|
||
def _get_initial_pose(self) -> list:
|
||
"""获取初始姿态"""
|
||
left_ret = self.left_rmarm.rm_get_current_arm_state()
|
||
right_ret = self.right_rmarm.rm_get_current_arm_state()
|
||
return left_ret[1]['pose'] + right_ret[1]['pose']
|
||
|
||
def _validate_joint_count(self, joints: list, expected_count: int):
|
||
"""验证关节数量"""
|
||
if len(joints) != expected_count:
|
||
raise ValueError(f"关节数量不匹配: 期望 {expected_count}, 实际 {len(joints)}")
|
||
|
||
def _execute_slow_movement(self, arm: str, joint_data: list):
|
||
"""执行慢速运动"""
|
||
busy_flag = f"{arm}_slow_busy"
|
||
|
||
if not getattr(self, busy_flag):
|
||
setattr(self, busy_flag, True)
|
||
|
||
target_method = getattr(self, f"_{arm}_slow_task")
|
||
threading.Thread(
|
||
target=target_method,
|
||
args=(joint_data.copy(),),
|
||
daemon=True
|
||
).start()
|
||
|
||
def _left_slow_task(self, joint_data: list):
|
||
"""左臂慢速任务"""
|
||
try:
|
||
self.write_left_joint_slow(joint_data)
|
||
finally:
|
||
self.left_slow_busy = False
|
||
|
||
def _right_slow_task(self, joint_data: list):
|
||
"""右臂慢速任务"""
|
||
try:
|
||
self.write_right_joint_slow(joint_data)
|
||
finally:
|
||
self.right_slow_busy = False
|
||
|
||
def _execute_arm_action(self, arm: str, action: dict, master_joint: list, follow_joint: list):
|
||
"""执行单臂动作"""
|
||
controller_status = action['master_controller_status'][arm]
|
||
|
||
if controller_status['infrared'] == 1:
|
||
if compare_joint_difference(master_joint, follow_joint):
|
||
if arm == 'left':
|
||
self.write_left_joint_canfd(master_joint)
|
||
else:
|
||
self.write_right_joint_canfd(master_joint)
|
||
else:
|
||
self._execute_slow_movement(arm, master_joint)
|
||
|
||
def write_endpose(self, target_endpose: list):
|
||
assert target_endpose == 12, "the length of target pose is not equal 12"
|
||
self.left_rmarm.rm_movej_p(target_endpose[:6], 50, 0, 0, 1)
|
||
self.right_rmarm.rm_movej_p(target_endpose[6:], 50, 0, 0, 1)
|
||
|
||
def write_left_joint_slow(self, left_joint: list):
|
||
assert len(left_joint) == self.left_offset, "len of left master joint is not equal the count of left joint"
|
||
self.left_rmarm.rm_movej(left_joint, 5, 0, 0, 1)
|
||
|
||
def write_right_joint_slow(self, right_joint: list):
|
||
assert len(right_joint) == self.left_offset, "len of right master joint is not equal the count of right joint"
|
||
self.right_rmarm.rm_movej(right_joint, 5, 0, 0, 1)
|
||
|
||
def write_left_joint_canfd(self, left_joint: list):
|
||
assert len(left_joint) == self.left_offset, "len of left master joint is not equal the count of left joint"
|
||
self.left_rmarm.rm_movej_canfd(left_joint, False)
|
||
|
||
def write_right_joint_canfd(self, right_joint: list):
|
||
assert len(right_joint) == self.left_offset, "len of right master joint is not equal the count of right joint"
|
||
self.right_rmarm.rm_movej_canfd(right_joint, False)
|
||
|
||
def write_endpose_canfd(self, target_endpose: list):
|
||
assert len(target_endpose) == 12, "the length of target pose is not equal 12"
|
||
self.left_rmarm.rm_movep_canfd(target_endpose[:6], False)
|
||
self.right_rmarm.rm_movep_canfd(target_endpose[6:], False)
|
||
|
||
def write_dual_gripper(self, left_gripper: int, right_gripper: int):
|
||
self.left_rmarm.rm_set_gripper_position(left_gripper, False, 2)
|
||
self.right_rmarm.rm_set_gripper_position(right_gripper, False, 2)
|
||
|
||
def read_current_arm_joint_state(self):
|
||
return self.left_rmarm.rm_get_current_arm_state()[1]['joint'] + self.right_rmarm.rm_get_current_arm_state()[1]['joint']
|
||
|
||
def read_current_arm_endpose_state(self):
|
||
return self.left_rmarm.rm_get_current_arm_state()[1]['pose'] + self.right_rmarm.rm_get_current_arm_state()[1]['pose']
|
||
|
||
########################## lerobot function ##########################
|
||
|
||
def connect(self, enable: bool = True) -> bool:
|
||
"""使能机械臂并检测使能状态"""
|
||
with self._timeout_context() as is_timeout_valid:
|
||
while is_timeout_valid():
|
||
try:
|
||
if enable:
|
||
left_ret = self.left_rmarm.rm_get_current_arm_state()
|
||
right_ret = self.right_rmarm.rm_get_current_arm_state()
|
||
if left_ret[0] == 0 and right_ret[0] == 0:
|
||
print("机械臂使能成功")
|
||
return True
|
||
else:
|
||
RoboticArm.rm_destory()
|
||
print("机械臂断开连接")
|
||
return True
|
||
except Exception as e:
|
||
print(f"连接异常: {e}")
|
||
time.sleep(1)
|
||
print("连接超时")
|
||
return False
|
||
|
||
def set_calibration(self):
|
||
raise NotImplementedError
|
||
|
||
def revert_calibration(self):
|
||
raise NotImplementedError
|
||
|
||
def apply_calibration(self):
|
||
"""
|
||
移动到初始位置
|
||
"""
|
||
self.write(target_joint=self.init_joint_position)
|
||
|
||
def write(self, target_joint: list):
|
||
"""写入关节位置"""
|
||
self._validate_joint_count(target_joint, self.joint_count)
|
||
|
||
left_joints = target_joint[:self.left_offset]
|
||
left_gripper = target_joint[self.left_offset]
|
||
right_joints = target_joint[self.left_offset+1:-1]
|
||
right_gripper = target_joint[-1]
|
||
|
||
self.left_rmarm.rm_movej_canfd(left_joints, follow=False)
|
||
# self.left_rmarm.rm_movej_follow(left_joints)
|
||
# self.left_rmarm.rm_set_gripper_position(left_gripper, block=False, timeout=2)
|
||
self.right_rmarm.rm_movej_canfd(right_joints, follow=False)
|
||
# self.right_rmarm.rm_movej_follow(right_joints)
|
||
# self.right_rmarm.rm_set_gripper_position(right_gripper, block=False, timeout=2)
|
||
|
||
def read(self) -> Dict:
|
||
"""读取机械臂状态 - 直接从缓存获取"""
|
||
with self._cache_lock:
|
||
return self._state_cache.copy()
|
||
|
||
def safe_disconnect(self):
|
||
"""安全断开连接"""
|
||
try:
|
||
self.write(target_joint=self.safe_disable_position)
|
||
time.sleep(2) # 等待移动完成
|
||
except Exception as e:
|
||
print(f"移动到安全位置失败: {e}")
|
||
finally:
|
||
RoboticArm.rm_destory()
|
||
|
||
########################## lerobot function ########################## |