419 lines
14 KiB
Python
419 lines
14 KiB
Python
import pygame
|
||
import threading
|
||
import time
|
||
import logging
|
||
from typing import Dict
|
||
from dataclasses import dataclass
|
||
from .find_gamepad import find_controller_index
|
||
from .servo_server import ServoArmServer
|
||
|
||
|
||
class RealmanAlohaMaster:
|
||
def __init__(self, config):
|
||
self.config = config
|
||
self._initialize_master_arm()
|
||
|
||
def _initialize_master_arm(self):
|
||
"""初始化主控臂"""
|
||
try:
|
||
self.master_dual_arm = ServoArmServer(self.config.config_file)
|
||
except Exception as e:
|
||
logging.error(f"初始化主控臂失败: {e}")
|
||
raise
|
||
|
||
def get_action(self) -> Dict:
|
||
"""获取控制动作"""
|
||
try:
|
||
master_joint_actions = self.master_dual_arm.get_joint_data()
|
||
return self._format_action(master_joint_actions)
|
||
except Exception as e:
|
||
logging.error(f"获取动作失败: {e}")
|
||
|
||
def _format_action(self, master_joint_actions: dict) -> dict:
|
||
"""格式化动作数据"""
|
||
master_controller_status = {
|
||
'left': master_joint_actions['left_controller_status'],
|
||
'right': master_joint_actions['right_controller_status']
|
||
}
|
||
|
||
return {
|
||
'control_mode': 'joint',
|
||
'master_joint_actions': master_joint_actions['dual_joint_actions'],
|
||
'left_joint_actions': master_joint_actions['left_joint_actions'][:-1],
|
||
'right_joint_actions': master_joint_actions['right_joint_actions'][:-1],
|
||
'left_gripper_actions': master_joint_actions['left_joint_actions'][-1], # 修复bug
|
||
'right_gripper_actions': master_joint_actions['right_joint_actions'][-1],
|
||
'master_controller_status': master_controller_status
|
||
}
|
||
|
||
def stop(self):
|
||
"""停止控制器"""
|
||
try:
|
||
if hasattr(self, 'master_dual_arm') and self.master_dual_arm:
|
||
self.master_dual_arm.shutdown()
|
||
print("混合控制器已退出")
|
||
except Exception as e:
|
||
logging.error(f"停止控制器失败: {e}")
|
||
|
||
|
||
|
||
class DummyEndposeMaster:
|
||
def __init__(self, config):
|
||
# 初始化pygame
|
||
pygame.init()
|
||
pygame.joystick.init()
|
||
# 获取所有 USB 游戏控制器的信息
|
||
self.joysticks = find_controller_index()
|
||
print(self.joysticks)
|
||
self.control_info = config.end_control_info
|
||
left_stick = self._init_stick('left')
|
||
right_stick = self._init_stick('right')
|
||
self.controllers = [left_stick, right_stick]
|
||
|
||
def _init_stick(self, arm_name:str):
|
||
stick_info = {}
|
||
stick_info['index'] = self.joysticks[self.control_info[arm_name]]['index']
|
||
stick_info['guid'] = self.control_info[arm_name]
|
||
stick_info['name'] = f'{arm_name}'
|
||
device_name = self.joysticks[self.control_info[arm_name]]['device_name']
|
||
stick = XboxStick(stick_info) if "Xbox" in device_name else FlightStick(stick_info)
|
||
stick.start_polling()
|
||
return stick
|
||
|
||
def get_action(self, state) -> Dict:
|
||
from copy import deepcopy
|
||
|
||
new_state = deepcopy(state)
|
||
gamepad_action = {}
|
||
xyz = []
|
||
rxryrz = []
|
||
gripper = []
|
||
"""获取控制动作"""
|
||
try:
|
||
for i, controller in enumerate(self.controllers):
|
||
# states = controller.get_raw_states()
|
||
gamepad_action.update(controller.get_control_signal(controller.name))
|
||
xyz += [f"{controller.name}_x", f"{controller.name}_y", f"{controller.name}_z"]
|
||
rxryrz += [f"{controller.name}_joint_4", f"{controller.name}_joint_5", f"{controller.name}_joint_6"]
|
||
gripper += [f"{controller.name}_gripper"]
|
||
|
||
for name in xyz:
|
||
new_state['pose'][name] += (gamepad_action[name] * gamepad_action['xyz_vel'] * gamepad_action[name.split('_')[0]+'_ratio'])
|
||
|
||
for name in gripper:
|
||
new_state['joint'][name] += int(gamepad_action[name] * gamepad_action['gripper_vel'] * gamepad_action[name.split('_')[0]+'_ratio'])
|
||
new_state['joint'][name] = min(990, max(0, new_state['joint'][name]))
|
||
|
||
for name in rxryrz:
|
||
new_state['joint'][name] += (gamepad_action[name] * gamepad_action['rxyz_vel'] * gamepad_action[name.split('_')[0]+'_ratio'])
|
||
|
||
new_state['control_mode'] = 'endpose'
|
||
return new_state
|
||
|
||
except Exception as e:
|
||
logging.error(f"获取动作失败: {e}")
|
||
|
||
def stop(self):
|
||
"""停止控制器"""
|
||
try:
|
||
# 停止轮询线程
|
||
for controller in self.controllers:
|
||
controller.stop_polling()
|
||
except Exception as e:
|
||
logging.error(f"停止控制器失败: {e}")
|
||
|
||
|
||
|
||
class ControllerBase:
|
||
def __init__(self, joystick_info: dict):
|
||
# 初始化手柄对象
|
||
self.joystick = pygame.joystick.Joystick(joystick_info['index'])
|
||
self.joystick.init()
|
||
self.name = joystick_info['name']
|
||
self.guid = joystick_info['guid']
|
||
|
||
# 存储所有控制器状态的字典
|
||
self.states = {
|
||
'buttons': [False] * self.joystick.get_numbuttons(), # 按钮状态
|
||
'axes': [0.0] * self.joystick.get_numaxes(), # 摇杆和轴状态
|
||
'hats': [(0, 0)] * self.joystick.get_numhats() # 舵状态
|
||
}
|
||
|
||
# deadzone
|
||
self.deadzone = 0.15
|
||
# validzone
|
||
self.validzone = 0.05
|
||
self.ratio = 1
|
||
self.gripper_vel = 50
|
||
self.rxyz_vel = 5
|
||
self.xyz_vel = 0.02
|
||
|
||
# 线程控制标志
|
||
self.running = False
|
||
|
||
def start_polling(self):
|
||
"""启动线程以轮询控制器状态"""
|
||
if not self.running:
|
||
self.running = True
|
||
self.thread = threading.Thread(target=self._poll_controller)
|
||
self.thread.start()
|
||
|
||
def stop_polling(self):
|
||
"""停止线程"""
|
||
if self.running:
|
||
self.running = False
|
||
self.thread.join()
|
||
|
||
def _poll_controller(self):
|
||
"""后台线程函数，用于轮询控制器状态"""
|
||
while self.running:
|
||
# 处理pygame事件
|
||
pygame.event.pump()
|
||
|
||
# 获取按钮状态
|
||
for i in range(self.joystick.get_numbuttons()):
|
||
self.states['buttons'][i] = self.joystick.get_button(i)
|
||
|
||
# 获取摇杆和轴状态（通常范围是 -1.0 到 1.0）
|
||
for i in range(self.joystick.get_numaxes()):
|
||
self.states['axes'][i] = self.joystick.get_axis(i)
|
||
|
||
# 获取舵状态（通常返回一个元组 (x, y)，值范围为 -1, 0, 1）
|
||
for i in range(self.joystick.get_numhats()):
|
||
self.states['hats'][i] = self.joystick.get_hat(i)
|
||
|
||
# 控制轮询频率
|
||
time.sleep(0.01)
|
||
|
||
def get_raw_states(self):
|
||
"""获取当前控制器状态"""
|
||
return self.states
|
||
|
||
class FlightStick(ControllerBase):
|
||
def __init__(self, joystick_info):
|
||
super().__init__(joystick_info)
|
||
|
||
def get_x_control_signal(self):
|
||
x = 0
|
||
if self.states['axes'][0] > self.validzone:
|
||
x = 1
|
||
elif self.states['axes'][0] < -self.validzone:
|
||
x = -1
|
||
return x
|
||
|
||
def get_y_control_signal(self):
|
||
y = 0
|
||
if self.states['axes'][1] > self.validzone:
|
||
y = -1
|
||
elif self.states['axes'][1] < -self.validzone:
|
||
y = 1
|
||
return y
|
||
|
||
def get_z_control_signal(self):
|
||
z = 0
|
||
if self.states['buttons'][0]:
|
||
z = 1
|
||
elif self.states['buttons'][1]:
|
||
z = -1
|
||
return z
|
||
|
||
def get_gripper_control_signal(self):
|
||
gripper = 0
|
||
if self.states['buttons'][2] == 1:
|
||
gripper = 1
|
||
elif self.states['buttons'][3] == 1:
|
||
gripper = -1
|
||
return gripper
|
||
|
||
def get_ratio_control_signal(self):
|
||
ratio = self.ratio
|
||
if self.states['axes'][2] > 0.8:
|
||
ratio = self.ratio / 5
|
||
elif self.states['axes'][2] < -0.8:
|
||
ratio = self.ratio * 2
|
||
return ratio
|
||
|
||
def get_rx_control_signal(self):
|
||
rx = 0
|
||
if self.states['hats'][0][0] == -1:
|
||
rx = 1
|
||
elif self.states['hats'][0][0] == 1:
|
||
rx = -1
|
||
else:
|
||
rx = 0
|
||
return rx
|
||
|
||
def get_ry_control_signal(self):
|
||
ry = 0
|
||
if self.states['hats'][0][1] == 1:
|
||
ry = -1
|
||
elif self.states['hats'][0][1] == -1:
|
||
ry = 1
|
||
else:
|
||
ry = 0
|
||
return ry
|
||
|
||
def get_rz_control_signal(self):
|
||
rz = 0
|
||
if self.states['axes'][3] < -self.validzone:
|
||
rz = -1
|
||
elif self.states['axes'][3] > self.validzone:
|
||
rz = 1
|
||
else:
|
||
rz = 0
|
||
return rz
|
||
|
||
def get_control_signal(self, prefix: str = ""):
|
||
"""获取所有控制信号"""
|
||
return {
|
||
f'{prefix}_x': self.get_x_control_signal(),
|
||
f'{prefix}_y': self.get_y_control_signal(),
|
||
f'{prefix}_z': self.get_z_control_signal(),
|
||
f'{prefix}_joint_4': self.get_rx_control_signal(),
|
||
f'{prefix}_joint_5': self.get_ry_control_signal(),
|
||
f'{prefix}_joint_6': self.get_rz_control_signal(),
|
||
f'{prefix}_gripper': self.get_gripper_control_signal(),
|
||
f'{prefix}_ratio': self.get_ratio_control_signal(),
|
||
'gripper_vel': self.gripper_vel,
|
||
'rxyz_vel': self.rxyz_vel,
|
||
'xyz_vel': self.xyz_vel
|
||
}
|
||
|
||
|
||
|
||
class XboxStick(ControllerBase):
|
||
def __init__(self, joystick_info: dict):
|
||
super().__init__(joystick_info)
|
||
|
||
def get_x_control_signal(self):
|
||
"""获取 X 轴控制信号"""
|
||
x = 0
|
||
if self.states['hats'][0][0] == -1:
|
||
x = 1
|
||
elif self.states['hats'][0][0] == 1:
|
||
x = -1
|
||
return x
|
||
|
||
def get_y_control_signal(self):
|
||
"""获取 Y 轴控制信号"""
|
||
y = 0
|
||
if self.states['hats'][0][1] == 1:
|
||
y = -1
|
||
elif self.states['hats'][0][1] == -1:
|
||
y = 1
|
||
return y
|
||
|
||
def get_z_control_signal(self):
|
||
"""获取 Z 轴控制信号"""
|
||
z = 0
|
||
if self.states['axes'][4] > self.deadzone: # A 按钮
|
||
z = -1
|
||
elif self.states['axes'][4] < -self.deadzone: # B 按钮
|
||
z = 1
|
||
return z
|
||
|
||
def get_ratio_control_signal(self):
|
||
"""获取速度控制信号"""
|
||
ratio = self.ratio
|
||
if self.states['axes'][2] > 0.8: # LT 按钮
|
||
ratio = self.ratio * 2
|
||
elif self.states['axes'][5] > 0.8: # RT 按钮
|
||
ratio = self.ratio / 5
|
||
return ratio
|
||
|
||
def get_gripper_control_signal(self):
|
||
gripper = 0
|
||
if self.states['buttons'][0] == 1:
|
||
gripper = 1
|
||
elif self.states['buttons'][1] == 1:
|
||
gripper = -1
|
||
return gripper
|
||
|
||
def get_rx_control_signal(self):
|
||
"""获取 RX 轴控制信号"""
|
||
rx = 0
|
||
if self.states['axes'][0] > self.deadzone: # 左舵
|
||
rx = -1
|
||
elif self.states['axes'][0] < -self.deadzone: # 右舵
|
||
rx = 1
|
||
return rx
|
||
|
||
def get_ry_control_signal(self):
|
||
"""获取 RY 轴控制信号"""
|
||
ry = 0
|
||
if self.states['axes'][1] > self.deadzone: # 上舵
|
||
ry = 1
|
||
elif self.states['axes'][1] < -self.deadzone: # 下舵
|
||
ry = -1
|
||
return ry
|
||
|
||
def get_rz_control_signal(self):
|
||
"""获取 RZ 轴控制信号"""
|
||
rz = 0
|
||
if self.states['buttons'][4] == 1: # 左摇杆
|
||
rz = 1
|
||
elif self.states['buttons'][5] == 1: # 右摇杆
|
||
rz = -1
|
||
return rz
|
||
|
||
def get_control_signal(self, prefix: str = ""):
|
||
"""获取所有控制信号"""
|
||
return {
|
||
f'{prefix}_x': self.get_x_control_signal(),
|
||
f'{prefix}_y': self.get_y_control_signal(),
|
||
f'{prefix}_z': self.get_z_control_signal(),
|
||
f'{prefix}_joint_4': self.get_rx_control_signal(),
|
||
f'{prefix}_joint_5': self.get_ry_control_signal(),
|
||
f'{prefix}_joint_6': self.get_rz_control_signal(),
|
||
f'{prefix}_gripper': self.get_gripper_control_signal(),
|
||
f'{prefix}_ratio': self.get_ratio_control_signal(),
|
||
'gripper_vel': self.gripper_vel,
|
||
'rxyz_vel': self.rxyz_vel,
|
||
'xyz_vel': self.xyz_vel
|
||
}
|
||
|
||
|
||
@dataclass
|
||
class ControllerConfig:
|
||
"""控制器配置"""
|
||
init_joint: list
|
||
init_pose: list
|
||
max_gripper: int
|
||
min_gripper: int
|
||
config_file: str
|
||
end_control_info: dict
|
||
|
||
|
||
def parse_init_info(init_info: dict) -> ControllerConfig:
|
||
"""解析初始化信息"""
|
||
return ControllerConfig(
|
||
init_joint=init_info['init_joint'],
|
||
init_pose=init_info.get('init_pose', [0]*12),
|
||
max_gripper=init_info['max_gripper'],
|
||
min_gripper=init_info['min_gripper'],
|
||
config_file=init_info['servo_config_file'],
|
||
end_control_info=init_info['end_control_info']
|
||
)
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
config = {
|
||
'init_joint': {'joint': [-170, 90, 0, 90, 120, 0, 10, 170, 90, 0, -90, 120, 0, 10]},
|
||
'init_pose': {},
|
||
'max_gripper': {},
|
||
'min_gripper': {},
|
||
'servo_config_file': {},
|
||
'end_control_info': {'left': "0300b14bff1100003708000010010000" , 'right': '0300509d5e040000120b000009050000'}
|
||
}
|
||
config = parse_init_info(config)
|
||
endpose_arm = DummyEndposeMaster(config)
|
||
while True:
|
||
gamepad_action = {}
|
||
xyz = []
|
||
for i, controller in enumerate(endpose_arm.controllers):
|
||
# states = controller.get_raw_states()
|
||
gamepad_action.update(controller.get_control_signal(controller.name))
|
||
xyz += [f"{controller.name}_x", f"{controller.name}_y", f"{controller.name}_z"]
|
||
time.sleep(1)
|
||
print(gamepad_action) |