diff --git a/ufactory_lerobot/devices/umi/vive_tracker/transformations.py b/ufactory_lerobot/devices/umi/vive_tracker/transformations.py index f7ed23f..0783914 100644 --- a/ufactory_lerobot/devices/umi/vive_tracker/transformations.py +++ b/ufactory_lerobot/devices/umi/vive_tracker/transformations.py @@ -105,7 +105,7 @@ class Transformations: return roll, pitch, yaw @staticmethod - def rxryrz_to_matrix(axis_angle): + def rxryrz_to_rotation_matrix(axis_angle): """ 将轴角向量 (rx, ry, rz) 转换为 3x3 旋转矩阵。 输入: np.array([rx, ry, rz]) @@ -207,6 +207,14 @@ class Transformations: T[:3, 3] = [x, y, z] return T + @classmethod + def xyzrxryrz_to_rotation_matrix(cls, x, y, z, rx, ry, rz): + """构造4x4齐次变换矩阵""" + T = np.eye(4) + T[:3, :3] = cls.rxryrz_to_rotation_matrix([rx, ry, rz]) + T[:3, 3] = [x, y, z] + return T + @classmethod def rotation_matrix_to_xyzq(cls, rotation_matrix): """从4x4齐次变换矩阵到xyzq的转换""" diff --git a/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py index d2d8eca..b3c8409 100644 --- a/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py +++ b/ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py @@ -6,3 +6,7 @@ from .uf_mock_robot_config import UFMockRobotConfig @dataclass class MultipleUFMockRobotConfig(RobotConfig): robots: dict[str, UFMockRobotConfig] + + def __post_init__(self): + super().__post_init__() + self.id = 'multiple_uf_mock_robot' if self.id is None else self.id diff --git a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py index f6161c5..72b2e20 100644 --- a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py +++ b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py @@ -38,7 +38,7 @@ class UFMockRobot(Robot): 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 + self._gripper_type = self.config.gripper_type @property def _robot_state_features(self)-> dict: diff --git a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py index 737706c..a91c410 100644 --- a/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py +++ b/ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py @@ -21,8 +21,11 @@ class UFMockRobotConfig(RobotConfig): 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 + + def __post_init__(self): + super().__post_init__() + self.id = 'uf_mock_robot' if self.id is None else self.id diff --git a/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py index 95f5d1b..ba1bc53 100644 --- a/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py +++ b/ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py @@ -9,3 +9,7 @@ class MultipleUFRobotConfig(RobotConfig): async_connect: bool = True async_configure: bool = True async_action: bool = False + + def __post_init__(self): + super().__post_init__() + self.id = 'multiple_uf_robot' if self.id is None else self.id diff --git a/ufactory_lerobot/robots/uf_robot/uf_robot.py b/ufactory_lerobot/robots/uf_robot/uf_robot.py index a4db6f4..db14656 100644 --- a/ufactory_lerobot/robots/uf_robot/uf_robot.py +++ b/ufactory_lerobot/robots/uf_robot/uf_robot.py @@ -87,8 +87,14 @@ class UFRobot(Robot, Thread): 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 + if self.config.start_tcp_pose and len(self.config.start_tcp_pose) >= 6: + self._start_tcp_pose = list(self.config.start_tcp_pose[:3]) + list(map(math.radians, self.config.start_tcp_pose[3:6])) + else: + self._start_tcp_pose = None + if self.config.start_joints: + self._start_joints = list(map(math.radians, self.config.start_joints)) + else: + self._start_joints = None self.report_stop_event = Event() self._rt_report_normal = False @@ -97,7 +103,7 @@ class UFRobot(Robot, Thread): 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 + self._gripper_type = 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 @@ -198,6 +204,8 @@ class UFRobot(Robot, Thread): if calibrate: self.calibrate() + self.real_arm.set_linear_spd_limit_factor(2.0) + self._is_connected = True def configure(self) -> None: @@ -206,22 +214,6 @@ class UFRobot(Robot, Thread): 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: @@ -259,6 +251,26 @@ class UFRobot(Robot, Thread): if err_warn[0] != 0: raise RuntimeError(f"Failed to set correct state to Gripper! Controller Error code: {err_warn[0]} !") + if self._start_joints is not None: + self.real_arm.set_servo_angle(angle=self._start_joints, is_radian=True, wait=True) + if self._start_tcp_pose is not None: + 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._use_rt_report and not self._rt_report_normal: self.start() time.sleep(0.2) diff --git a/ufactory_lerobot/robots/uf_robot/uf_robot_config.py b/ufactory_lerobot/robots/uf_robot/uf_robot_config.py index 3093997..16bc495 100644 --- a/ufactory_lerobot/robots/uf_robot/uf_robot_config.py +++ b/ufactory_lerobot/robots/uf_robot/uf_robot_config.py @@ -1,6 +1,5 @@ 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 @@ -30,15 +29,17 @@ class UFRobotConfig(RobotConfig): 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 + start_joints: Tuple[float, ...] = (0, 0, 0, 90, 0, 90, 0) # ° + start_tcp_pose: Tuple[float, ...] = None # [x, y, z, roll(°), pitch(°), yaw(°)] 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 + no_action: bool = False # only for debug + + def __post_init__(self): + super().__post_init__() + self.id = 'uf_robot' if self.id is None else self.id diff --git a/ufactory_lerobot/scripts/uf_lerobot_eval.py b/ufactory_lerobot/scripts/uf_lerobot_eval.py index 3c01225..170f6f3 100644 --- a/ufactory_lerobot/scripts/uf_lerobot_eval.py +++ b/ufactory_lerobot/scripts/uf_lerobot_eval.py @@ -27,7 +27,6 @@ from lerobot.robots import ( # noqa: F401 ) from lerobot.utils.control_utils import ( is_headless, - init_keyboard_listener, predict_action, ) from lerobot.utils.import_utils import register_third_party_plugins @@ -40,7 +39,7 @@ 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.utils.utils import instantiate_from_dict, init_keyboard_listener from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations @@ -74,8 +73,8 @@ def compute_relative_axis_angle(rot_prev, rot_curr): 返回: 相对轴角向量 """ # 1. 转为矩阵 - R_prev = Transformations.rxryrz_to_matrix(rot_prev) - R_curr = Transformations.rxryrz_to_matrix(rot_curr) + R_prev = Transformations.rxryrz_to_rotation_matrix(rot_prev) + R_curr = Transformations.rxryrz_to_rotation_matrix(rot_curr) # 2. 计算相对旋转矩阵 # R_delta 表示从 prev 坐标系到 curr 坐标系的旋转 @@ -88,8 +87,8 @@ 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_prev = Transformations.rxryrz_to_rotation_matrix(rot_prev) + R_delta = Transformations.rxryrz_to_rotation_matrix(rot_delta) R_curr = R_prev @ R_delta # R_curr = R_prev.apply(R_delta) return Transformations.rotation_matrix_to_rxryrz(R_curr) @@ -122,7 +121,7 @@ class EvalConfig: return ["policy"] -def eval_loop(cfg: EvalConfig, relative=False): +def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False): init_logging() logging.info(pformat(asdict(cfg))) @@ -238,8 +237,6 @@ def eval_loop(cfg: EvalConfig, relative=False): 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() @@ -278,8 +275,7 @@ def eval_loop(cfg: EvalConfig, relative=False): # 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: @@ -425,6 +421,7 @@ def main(): 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') + parser.add_argument('--rx_continuous', action='store_true', help='rx continuous or not') args = parser.parse_args() try: with open(args.config, 'r') as f: @@ -436,7 +433,7 @@ def main(): config = instantiate_from_dict(cfg) eval_cfg = EvalConfig(robot=config["RobotConfig"], dataset=config["DatasetRecordConfig"]) - eval_loop(eval_cfg, args.relative) + eval_loop(eval_cfg, args.relative, args.rx_continuous) if __name__ == "__main__": diff --git a/ufactory_lerobot/scripts/uf_lerobot_record.py b/ufactory_lerobot/scripts/uf_lerobot_record.py index c98c939..3ec75b0 100644 --- a/ufactory_lerobot/scripts/uf_lerobot_record.py +++ b/ufactory_lerobot/scripts/uf_lerobot_record.py @@ -1,11 +1,12 @@ import yaml +import time 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 +from ufactory_lerobot.utils.utils import instantiate_from_dict, init_keyboard_listener @safe_stop_image_writer @@ -239,31 +240,71 @@ def record(cfg: RecordConfig) -> LeRobotDataset: if teleop is not None: teleop.connect() - listener, events = init_keyboard_listener() + is_evt = not is_headless() + is_uf_teleop = isinstance(teleop, UFBaseTeleop) + is_recorded = False + key_dict = {} + events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False} - print("\n********** Episode Record Loop Start **********") + if is_evt: + from pynput import keyboard - 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) + key_dict = { + keyboard.Key.space: 0, # start + keyboard.Key.enter: 0, # help + } + + def on_press(key): + try: + if key == keyboard.Key.right: + print("Right arrow key pressed. Exiting loop...") + events["exit_early"] = True + elif key == keyboard.Key.left: + print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + events["rerecord_episode"] = True + events["exit_early"] = True + elif key == keyboard.Key.esc: + print("Escape key pressed. Stopping data recording...") + events["stop_recording"] = True + events["exit_early"] = True + except Exception as e: + print(f"Error handling key press: {e}") + if key in key_dict: + key_dict[key] = True + + def on_release(key): + try: + if key == keyboard.Key.enter: + if not is_recorded: + print('[HELP] : EXIT, : START, : RESET, : SAVE') + else: + print('[HELP] : EXIT, : RESET, : SAVE') + # is_recorded = True + except Exception as e: + print(f"Error handling key release: {e}") + if key in key_dict: + key_dict[key] = False + + listener, events = init_keyboard_listener(events=events, on_press=on_press, on_release=on_release) + print("\n********** Episode Record Loop Start **********") + print('[HELP] : EXIT, : START, : RESET, : SAVE') else: - frame_callback = None + input('[HELP] Enter to to start record >>> ') + if is_uf_teleop: + teleop.set_teleop_enabled(True) + is_recorded = True + print("\n********** Episode Record Loop Start **********") + + frame_callback = None with VideoEncodingManager(dataset): recorded_episodes = 0 while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]: + time.sleep(0.01) + if is_evt: + if not is_recorded and key_dict[keyboard.Key.space]: + is_recorded = True + if teleop is not None and isinstance(teleop, UFMockTeleop): if events["stop_recording"]: continue @@ -275,52 +316,65 @@ def record(cfg: RecordConfig) -> LeRobotDataset: continue if events["stop_recording"]: continue + is_recorded = True - 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 is_recorded: + events["rerecord_episode"] = False + events["exit_early"] = False + if is_uf_teleop: + robot.configure() + obs = robot.get_observation() + teleop.set_teleop_enabled(True, obs) + 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, + ) + else: + continue + if events['stop_recording']: + break 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) + if is_uf_teleop: + teleop.set_teleop_enabled(False) + if dataset.episode_buffer: + dataset.clear_episode_buffer() + is_recorded = False + if is_evt: + print('[HELP] : EXIT, : START, : RESET, : SAVE') + else: + input('\nPress Enter to rerecord this episode >>>>> ') + is_recorded = True continue - if not events['stop_recording']: - if isinstance(teleop, UFBaseTeleop): - teleop.set_ctrl_status(False) + if is_recorded and not events['stop_recording']: + log_say(f"Save episode {dataset.num_episodes}", cfg.play_sounds) + if is_uf_teleop: + teleop.set_teleop_enabled(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) + is_recorded = False + if is_evt: + print('[HELP] : EXIT, : START, : RESET, : SAVE') + else: + input('Press Enter to record at the next episode >>>>> ') + is_recorded = True print("\n********** Episode Record Loop Exit **********") @@ -328,7 +382,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset: if teleop is not None: teleop.disconnect() - if not is_headless() and listener is not None: + if is_evt and listener is not None: listener.stop() if cfg.dataset.push_to_hub: diff --git a/ufactory_lerobot/scripts/uf_robot_teleop.py b/ufactory_lerobot/scripts/uf_robot_teleop.py index 2f1b41d..7c760e3 100644 --- a/ufactory_lerobot/scripts/uf_robot_teleop.py +++ b/ufactory_lerobot/scripts/uf_robot_teleop.py @@ -18,16 +18,14 @@ 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 +from ufactory_lerobot.utils.utils import instantiate_from_dict, is_headless, init_keyboard_listener +from ufactory_lerobot.teleoperators.base_teleop import UFBaseTeleop + @dataclass class TeleopConfig: @@ -50,32 +48,107 @@ def teleop_loop(cfg: TeleopConfig): 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 >>> ') + is_evt = not is_headless() + is_uf_teleop = isinstance(teleop, UFBaseTeleop) - print("\n********** Teleop Control Loop Start **********") + is_reset = False + is_paused = True + events = {"exit": False} + listener = None + key_dict = {} + + if is_evt: + from pynput import keyboard + + key_dict = { + keyboard.Key.esc: 0, # exit + keyboard.Key.left: 0, # reset and pause + keyboard.Key.space: 0, # start/pause + keyboard.Key.enter: 0, # help + } + + def on_press(key): + if key_dict.get(key, 1) == 0: + try: + if key == keyboard.Key.esc: + events["exit"] = True + print("\nEscape key pressed. Stopping ...") + except Exception as e: + print(f"Error handling key press: {e}") + if key in key_dict: + key_dict[key] = True + + def on_release(key): + try: + if key == keyboard.Key.enter: + if is_paused: + if is_reset: + print('[HELP] : EXIT, : RESET AND START, : RESET') + else: + print('[HELP] : EXIT, : START, : RESET') + else: + print('[HELP] : EXIT, : PAUSE, : RESET') + except Exception as e: + print(f"Error handling key release: {e}") + if key in key_dict: + key_dict[key] = False + + listener, events = init_keyboard_listener(events=events, on_press=on_press, on_release=on_release) + print("\n********** Teleop Control Loop Start **********") + print('[HELP] : EXIT, : START, : RESET') + else: + input('[HELP] Enter to control robot with teleop >>> ') + if is_uf_teleop: + obs = robot.get_observation() + teleop.set_teleop_enabled(True, obs) + is_paused = False + is_reset = False + print("\n********** Teleop Control Loop Start **********") + + key_space_pressed = False + key_left_pressed = False while not events["exit"]: start_loop_t = time.perf_counter() + if is_evt: + if key_dict[keyboard.Key.left] and not key_left_pressed: + key_left_pressed = True + is_reset = True + if not is_paused: + is_paused = True + if is_uf_teleop: + teleop.set_teleop_enabled(False) + print('[HELP] : EXIT, : RESET AND START, : RESET') + elif not key_dict[keyboard.Key.left] and key_left_pressed: + key_left_pressed = False + + if key_dict[keyboard.Key.space] and not key_space_pressed: + key_space_pressed = True + is_paused = not is_paused + if is_paused: + if is_uf_teleop: + teleop.set_teleop_enabled(False) + # print('========== Teleop is paused ==========') + print('[HELP] : EXIT, : START, : RESET') + else: + if is_reset: + is_reset = False + robot.configure() + # print('========== Teleop is start ==========') + if is_uf_teleop: + obs = robot.get_observation() + teleop.set_teleop_enabled(True, obs) + print('[HELP] : EXIT, : PAUSE, : RESET') + continue + elif not key_dict[keyboard.Key.space] and key_space_pressed: + key_space_pressed = False + + if is_reset or is_paused: + continue + # Get robot observation obs = robot.get_observation() @@ -91,7 +164,7 @@ def teleop_loop(cfg: TeleopConfig): print("\n********** Teleop Control Loop Exit **********") robot.disconnect() teleop.disconnect() - if not is_headless() and listener is not None: + if is_evt and listener is not None: listener.stop() def main(): diff --git a/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py b/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py index 309495f..ef9d612 100644 --- a/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py +++ b/ufactory_lerobot/teleoperators/base_teleop/base_teleop.py @@ -9,5 +9,9 @@ class UFBaseTeleop(Teleoperator): super().__init__(config) self.config = config - def set_ctrl_status(self, status): + def set_teleop_enabled(self, enabled: bool, obs=None): + """ + 启用/停用遥操作 + 当enabled为True且obs不为None时, 顺便设置机械臂初始位置映射 + """ pass diff --git a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py index f4e70c7..7da4855 100644 --- a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py +++ b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import logging import time +import math import numpy as np from gello.dynamixel.driver import DynamixelDriver from gello.agents.gello_agent import GelloAgent, DynamixelRobotConfig @@ -36,8 +37,9 @@ class GelloTeleop(UFBaseTeleop): 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] + start_joints = list(map(math.radians, self.config.start_joints)) + for i in range(len(start_joints)): + offset = curr_joints[i] - 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] @@ -52,7 +54,7 @@ class GelloTeleop(UFBaseTeleop): } self._dynamixel_robo_config = DynamixelRobotConfig(**param_dict) print(self._dynamixel_robo_config) - self.dof = len(self.config.start_joints) + self.dof = len(start_joints) if self.config.torque_joint_ids: driver = DynamixelDriver(self.config.torque_joint_ids, port=self.config.port, baudrate=57600) diff --git a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py index 4c283ca..72525c8 100644 --- a/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py +++ b/ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py @@ -1,6 +1,5 @@ #!/usr/bin/env python -import numpy as np from dataclasses import dataclass from typing import Tuple from lerobot.teleoperators import TeleoperatorConfig @@ -15,6 +14,9 @@ class GelloTeleopConfig(TeleoperatorConfig): # 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) + start_joints: Tuple[float, ...] = (0, 0, 0, 90, 0, 90, 0) # ° gripper_id: int = 8 # -1: no gripper torque_joint_ids: Tuple[int, ...] = None # the joints will activate torque mode. + + def __post_init__(self): + self.id = 'gello_teleop' if self.id is None else self.id diff --git a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py index 5515d59..2b49200 100644 --- a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py +++ b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py @@ -2,11 +2,9 @@ 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 lerobot.utils.errors import DeviceNotConnectedError from ufactory_lerobot.devices.pika import PikaDevice from ufactory_lerobot.devices.umi.vive_tracker.transformations import Transformations from ..base_teleop import UFBaseTeleop @@ -18,7 +16,7 @@ class PikaTeleop(UFBaseTeleop, Thread): config_class = PikaTeleopConfig name = "Pika Teleop For xArm" - def __init__(self, config: PikaTeleopConfig): + def __init__(self, config: PikaTeleopConfig, prefix=''): super().__init__(config) Thread.__init__(self) # Do NOT REMOVE! @@ -27,20 +25,22 @@ class PikaTeleop(UFBaseTeleop, Thread): self._is_connected = False self._is_calibrated = True self._data_lock = Lock() - self._ctrl_flag = False + self._teleop_enabled = False + self._last_action = None self._need_initial = False + self.prefix = '' if not prefix else f'{prefix}.' + + tracker_to_robot_eef = list(self.config.tracker_to_robot_eef[:3]) + list(map(math.radians, self.config.tracker_to_robot_eef[3:6])) + self.tracker_to_robot_matrix = Transformations.xyzrpy_to_rotation_matrix(*tracker_to_robot_eef) + robot_base_pose = list(self.config.robot_base_pose[:3]) + list(map(math.radians, self.config.robot_base_pose[3:6])) + self.robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*robot_base_pose) + self.begin_tracker_robot_matrix = None + self._last_robot_pose = Transformations.rotation_matrix_to_xyzrxryrz(self.robot_base_matrix) + self._last_gripper_pos = 0.0 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: @@ -97,127 +97,45 @@ class PikaTeleop(UFBaseTeleop, Thread): 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 set_teleop_enabled(self, enabled: bool, obs=None): + with self._data_lock: + if enabled: + if obs is not None: + self._last_robot_pose = [obs[f"{self.prefix}pose.x"], obs[f"{self.prefix}pose.y"], obs[f"{self.prefix}pose.z"], obs[f"{self.prefix}pose.rx"], obs[f"{self.prefix}pose.ry"], obs[f"{self.prefix}pose.rz"]] + if self.config.use_gripper: + self._last_gripper_pos = obs[f"{self.prefix}gripper.pos"] + self.robot_base_matrix = Transformations.xyzrxryrz_to_rotation_matrix(*self._last_robot_pose) + self.begin_tracker_robot_matrix = None + self._last_action = None + self._teleop_enabled = True + print(f'[{self.prefix}PIKA] Teleoperation is start') + else: + obs = self._last_action + self._last_robot_pose = [obs[f"{self.prefix}pose.x"], obs[f"{self.prefix}pose.y"], obs[f"{self.prefix}pose.z"], obs[f"{self.prefix}pose.rx"], obs[f"{self.prefix}pose.ry"], obs[f"{self.prefix}pose.rz"]] + if self.config.use_gripper: + self._last_gripper_pos = obs[f"{self.prefix}gripper.pos"] + self._teleop_enabled = False + self._last_action = None + print(f'[{self.prefix}PIKA] Teleoperation has paused') 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('开始遥操作') + if not self._teleop_enabled and curr_state != init_state: + self.set_teleop_enabled(True, self._last_action) time.sleep(1) - elif self._ctrl_flag and curr_state == init_state: - self._ctrl_flag = False - print('停止遥操作') + elif self._teleop_enabled and curr_state == init_state: + self.self.set_teleop_enabled(False) 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]: @@ -225,39 +143,47 @@ class PikaTeleop(UFBaseTeleop, Thread): 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 self._last_action is None: + self._last_action = { + f"{self.prefix}pose.x": self._last_robot_pose[0], + f"{self.prefix}pose.y": self._last_robot_pose[1], + f"{self.prefix}pose.z": self._last_robot_pose[2], + f"{self.prefix}pose.rx": self._last_robot_pose[3], + f"{self.prefix}pose.ry": self._last_robot_pose[4], + f"{self.prefix}pose.rz": self._last_robot_pose[5], + } + if self.config.use_gripper: + self._last_action.update({f"{self.prefix}gripper.pos": self._last_gripper_pos}) + if not self._teleop_enabled: + return self._last_action - 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], - } + pose = self.pika_sense.get_pose(self.pika_device.pika_tracker_device) + if pose: + x, y, z = pose.position[0] * 1000 * self.config.scale_xyz, pose.position[1] * 1000 * self.config.scale_xyz, pose.position[2] * 1000 * self.config.scale_xyz + quaternion = pose.rotation + 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) + # print(['{:.6f}'.format(val) for val in robot_target_pose]) + self._last_action[f"{self.prefix}pose.x"] = robot_target_pose[0] + self._last_action[f"{self.prefix}pose.y"] = robot_target_pose[1] + self._last_action[f"{self.prefix}pose.z"] = robot_target_pose[2] + self._last_action[f"{self.prefix}pose.rx"] = robot_target_pose[3] + self._last_action[f"{self.prefix}pose.ry"] = robot_target_pose[4] + self._last_action[f"{self.prefix}pose.rz"] = robot_target_pose[5] + else: + pass if self.config.use_gripper: - action_dict.update({"gripper.pos": gripper_target_pos}) - - return action_dict + distance = min(max(self.pika_sense.get_gripper_distance(), 0), 100) + if distance is not None: + gripper_pos = (100 - distance) / (100 - 0) + else: + gripper_pos = 0.0 + self._last_action.update({f"{self.prefix}gripper.pos": gripper_pos}) + return self._last_action def send_feedback(self, feedback: dict[str, float]) -> None: raise NotImplementedError \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py index 453adae..30c642e 100644 --- a/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py +++ b/ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py @@ -15,18 +15,20 @@ # limitations under the License. from dataclasses import dataclass - +from typing import Tuple 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 + scale_xyz: float = 1.0 + tracker_to_robot_eef: Tuple[float, ...] = (0, 0, 0, 180, -90, 0) # [x, y, z, roll(°), pitch(°), yaw(°)] + robot_base_pose: Tuple[float, ...] = (400, 0, 400, 180, 0, 0) # [x, y, z, roll(°), pitch(°), yaw(°)] + + def __post_init__(self): + self.id = 'pika_teleop' if self.id is None else self.id diff --git a/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py b/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py index 0138c39..9b39769 100644 --- a/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py +++ b/ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py @@ -29,3 +29,6 @@ class SpaceMouseTeleopConfig(TeleoperatorConfig): frequency: int = 10 # hz max_pos_speed: int = 250 # mm/s # Others: Calibration angles, joint directions etc. + + def __post_init__(self): + self.id = 'spacemouse_teleop' if self.id is None else self.id diff --git a/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py index 717fc19..178aafe 100644 --- a/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py +++ b/ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py @@ -46,3 +46,6 @@ class UFMockTeleopConfig(TeleoperatorConfig): gripper_freq: int = 50 # Hz gripper_open: int = 800 gripper_close: int = 0 + + def __post_init__(self): + self.id = 'uf_mock_teleop' if self.id is None else self.id diff --git a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py index 608fc23..d2dccc0 100644 --- a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py +++ b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py @@ -55,9 +55,13 @@ class MultipleUmiTeleop(UFBaseTeleop): 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 set_teleop_enabled(self, enabled: bool, obs=None): + for key, teleop in self.teleops.items(): + if obs is not None: + teleop_obs = {k: v for k, v in obs.items() if k.startswith(f"{key}.")} + else: + teleop_obs = None + teleop.set_teleop_enabled(enabled, teleop_obs) def get_action(self) -> dict[str, Any]: actions = {} diff --git a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py index 9d90c18..0872dac 100644 --- a/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py +++ b/ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py @@ -22,4 +22,7 @@ from .umi_teleop_config import UmiTeleopConfig @TeleoperatorConfig.register_subclass("uf::multiple_umi_teleop") @dataclass class MultipleUmiTeleopConfig(TeleoperatorConfig): - teleops: dict[str, UmiTeleopConfig] \ No newline at end of file + teleops: dict[str, UmiTeleopConfig] + + def __post_init__(self): + self.id = 'multiple_umi_teleop' if self.id is None else self.id \ No newline at end of file diff --git a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py index 1d1df39..2da100c 100644 --- a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py +++ b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py @@ -1,8 +1,8 @@ #!/usr/bin/env python +import math 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 @@ -21,31 +21,22 @@ class UmiTeleop(UFBaseTeleop): self.prefix = '' if not prefix else f'{prefix}.' self._is_connected = False self._is_calibrated = True + self._teleop_enabled = False + self._last_action = None - if self.config.use_gripper: - self.config.init_clamp_stream = True - else: - self.config.init_clamp_stream = False + self.tracker = None + self.xvlib = None - 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 = ViveTracker() if self.config.use_vive_tracker else None + # self.xvlib = XVLib(self.config.serial_number, not self.config.use_vive_tracker, self.config.use_gripper) + + tracker_to_robot_eef = list(self.config.tracker_to_robot_eef[:3]) + list(map(math.radians, self.config.tracker_to_robot_eef[3:6])) 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 + robot_base_pose = list(self.config.robot_base_pose[:3]) + list(map(math.radians, self.config.robot_base_pose[3:6])) self.robot_base_matrix = Transformations.xyzrpy_to_rotation_matrix(*robot_base_pose) self.begin_tracker_robot_matrix = None + self._last_robot_pose = Transformations.rotation_matrix_to_xyzrxryrz(self.robot_base_matrix) + self._last_gripper_pos = 0.0 @property def action_features(self) -> dict: @@ -93,11 +84,13 @@ class UmiTeleop(UFBaseTeleop): 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.tracker = ViveTracker() if self.config.use_vive_tracker else None + self.xvlib = XVLib(self.config.serial_number, not self.config.use_vive_tracker, self.config.use_gripper) self._is_connected = True def disconnect(self): - self.xvlib.xv_uninit() + if self.xvlib: + self.xvlib.xv_uninit() self._is_connected = False @staticmethod @@ -111,11 +104,26 @@ class UmiTeleop(UFBaseTeleop): angle_deg += 360 return angle_deg - def set_ctrl_status(self, status): - if status: + def set_teleop_enabled(self, enabled: bool, obs=None): + if enabled: + if obs is not None: + self._last_robot_pose = [obs[f"{self.prefix}pose.x"], obs[f"{self.prefix}pose.y"], obs[f"{self.prefix}pose.z"], obs[f"{self.prefix}pose.rx"], obs[f"{self.prefix}pose.ry"], obs[f"{self.prefix}pose.rz"]] + if self.config.use_gripper: + self._last_gripper_pos = obs[f"{self.prefix}gripper.pos"] + self.robot_base_matrix = Transformations.xyzrxryrz_to_rotation_matrix(*self._last_robot_pose) self.begin_tracker_robot_matrix = None + self._last_action = None + self._teleop_enabled = True + print(f'[{self.prefix}UMI] Teleoperation is start') else: - pass + obs = self._last_action + if obs: + self._last_robot_pose = [obs[f"{self.prefix}pose.x"], obs[f"{self.prefix}pose.y"], obs[f"{self.prefix}pose.z"], obs[f"{self.prefix}pose.rx"], obs[f"{self.prefix}pose.ry"], obs[f"{self.prefix}pose.rz"]] + if self.config.use_gripper: + self._last_gripper_pos = obs[f"{self.prefix}gripper.pos"] + self._teleop_enabled = False + self._last_action = None + print(f'[{self.prefix}UMI] Teleoperation has paused') # delta action def get_action(self) -> dict[str, Any]: @@ -124,6 +132,20 @@ class UmiTeleop(UFBaseTeleop): "UmiTeleop is not connected. You need to run `connect()` before `get_action()`." ) + if self._last_action is None: + self._last_action = { + f"{self.prefix}pose.x": self._last_robot_pose[0], + f"{self.prefix}pose.y": self._last_robot_pose[1], + f"{self.prefix}pose.z": self._last_robot_pose[2], + f"{self.prefix}pose.rx": self._last_robot_pose[3], + f"{self.prefix}pose.ry": self._last_robot_pose[4], + f"{self.prefix}pose.rz": self._last_robot_pose[5], + } + if self.config.use_gripper: + self._last_action.update({f"{self.prefix}gripper.pos": self._last_gripper_pos}) + if not self._teleop_enabled: + return self._last_action + if self.tracker is not None: pose_data = self.tracker.get_pose(self.config.vive_tracker_id) if pose_data is None: @@ -133,25 +155,6 @@ class UmiTeleop(UFBaseTeleop): _, 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) @@ -159,43 +162,19 @@ class UmiTeleop(UFBaseTeleop): 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], - } + self._last_action[f"{self.prefix}pose.x"] = robot_target_pose[0] + self._last_action[f"{self.prefix}pose.y"] = robot_target_pose[1] + self._last_action[f"{self.prefix}pose.z"] = robot_target_pose[2] + self._last_action[f"{self.prefix}pose.rx"] = robot_target_pose[3] + self._last_action[f"{self.prefix}pose.ry"] = robot_target_pose[4] + self._last_action[f"{self.prefix}pose.rz"] = robot_target_pose[5] 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}) + self._last_action.update({f"{self.prefix}gripper.pos": gripper_pos}) - return action_dict + return self._last_action def send_feedback(self, feedback: dict[str, float]) -> None: raise NotImplementedError diff --git a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py index bdcdb47..e4ff7e2 100644 --- a/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py +++ b/ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py @@ -16,7 +16,6 @@ from dataclasses import dataclass from typing import Tuple -import numpy as np from lerobot.teleoperators import TeleoperatorConfig @@ -24,12 +23,11 @@ from lerobot.teleoperators import TeleoperatorConfig @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) + tracker_to_robot_eef: Tuple[float, ...] = (0, 0, 0, 0, 0, -90) # [x, y, z, roll(°), pitch(°), yaw(°)] + robot_base_pose: Tuple[float, ...] = (300, 0, 300, 180, -90, 0) # [x, y, z, roll(°), pitch(°), yaw(°)] + + def __post_init__(self): + self.id = 'umi_teleop' if self.id is None else self.id diff --git a/ufactory_lerobot/utils/utils.py b/ufactory_lerobot/utils/utils.py index 2289560..7b6de5b 100644 --- a/ufactory_lerobot/utils/utils.py +++ b/ufactory_lerobot/utils/utils.py @@ -1,4 +1,6 @@ +import logging import importlib +from lerobot.utils.control_utils import is_headless # recursive call, inspired from gello_software: def instantiate_from_dict(cfg, ignore_cameras=False): @@ -16,3 +18,57 @@ def instantiate_from_dict(cfg, ignore_cameras=False): return [instantiate_from_dict(v, ignore_cameras) for v in cfg] else: return cfg + +def init_keyboard_listener(events: dict = None, on_press: callable = None, on_release: callable = None): + """ + Initializes a non-blocking keyboard listener for real-time user interaction. + + This function sets up a listener for specific keys (right arrow, left arrow, escape) to control + the program flow during execution, such as stopping recording or exiting loops. It gracefully + handles headless environments where keyboard listening is not possible. + + Returns: + A tuple containing: + - The `pynput.keyboard.Listener` instance, or `None` if in a headless environment. + - A dictionary of event flags (e.g., `exit_early`) that are set by key presses. + """ + # Allow to exit early while recording an episode or resetting the environment, + # by tapping the right arrow key '->'. This might require a sudo permission + # to allow your terminal to monitor keyboard events. + if events is None: + events = {} + events["exit_early"] = False + events["rerecord_episode"] = False + events["stop_recording"] = False + + if is_headless(): + logging.warning( + "Headless environment detected. On-screen cameras display and keyboard inputs will not be available." + ) + listener = None + return listener, events + + # Only import pynput if not in a headless environment + from pynput import keyboard + + if on_press is None: + def on_press(key): + try: + if key == keyboard.Key.right: + print("Right arrow key pressed. Exiting loop...") + events["exit_early"] = True + elif key == keyboard.Key.left: + print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + events["rerecord_episode"] = True + events["exit_early"] = True + elif key == keyboard.Key.esc: + print("Escape key pressed. Stopping data recording...") + events["stop_recording"] = True + events["exit_early"] = True + except Exception as e: + print(f"Error handling key press: {e}") + + listener = keyboard.Listener(on_press=on_press, on_release=on_release) + listener.start() + + return listener, events