202 lines
9.0 KiB
Python
202 lines
9.0 KiB
Python
#!/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
|