wangshuxun/Xarm-DataCollection
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(teleop): 重命名 set_ctrl_status 为 set_teleop_enabled

Browse Source 
启用时支持传入机械臂观测数据以同步初始位姿映射，
重构 Pika/UMI 遥操作内部状态管理。
This commit is contained in:
Vinman 2026-06-12 18:37:53 +08:00
parent 4467e19322
commit d959d82dbe
22 changed files with 507 additions and 369 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
ufactory_lerobot/devices/umi/vive_tracker/transformations.py
View File

@ -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的转换"""

4
ufactory_lerobot/robots/uf_mock_robot/multiple_uf_mock_robot_config.py
View File

@ -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

2
ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot.py
View File

@ -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:

5
ufactory_lerobot/robots/uf_mock_robot/uf_mock_robot_config.py
View File

@ -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

4
ufactory_lerobot/robots/uf_robot/multiple_uf_robot_config.py
View File

@ -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

50
ufactory_lerobot/robots/uf_robot/uf_robot.py
View File

@ -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)

13
ufactory_lerobot/robots/uf_robot/uf_robot_config.py
View File

@ -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

21
ufactory_lerobot/scripts/uf_lerobot_eval.py
View File

@ -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__":

124
ufactory_lerobot/scripts/uf_lerobot_record.py
View File

@ -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}
if is_evt:
from pynput import keyboard
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] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: RESET, <RIGH ARROW>: SAVE')
else:
print('[HELP] <ESC>: EXIT, <LEFT ARROW>: RESET, <RIGH ARROW>: 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] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: RESET, <RIGH ARROW>: SAVE')
else:
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 **********")
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)
else:
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,7 +316,15 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
continue
if events["stop_recording"]:
continue
is_recorded = True
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,
@ -294,33 +343,38 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
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)
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] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: RESET, <RIGH ARROW>: SAVE')
else:
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)
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
is_recorded = False
if is_evt:
print('[HELP] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: RESET, <RIGH ARROW>: SAVE')
else:
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 = 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:

121
ufactory_lerobot/scripts/uf_robot_teleop.py
View File

@ -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)
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] <ESC>: EXIT, <SPACE>: RESET AND START, <LEFT ARROW>: RESET')
else:
print('[HELP] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: RESET')
else:
print('[HELP] <ESC>: EXIT, <SPACE>: PAUSE, <LEFT ARROW>: 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] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: 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] <ESC>: EXIT, <SPACE>: RESET AND START, <LEFT ARROW>: 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] <ESC>: EXIT, <SPACE>: START, <LEFT ARROW>: 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] <ESC>: EXIT, <SPACE>: PAUSE, <LEFT ARROW>: 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():

6
ufactory_lerobot/teleoperators/base_teleop/base_teleop.py
View File

@ -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

8
ufactory_lerobot/teleoperators/gello_teleop/gello_teleop.py
View File

@ -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)

6
ufactory_lerobot/teleoperators/gello_teleop/gello_teleop_config.py
View File

@ -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

220
ufactory_lerobot/teleoperators/pika_teleop/pika_teleop.py
View File

@ -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,167 +97,93 @@ 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
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:
self._ctrl_flag = False
self._need_initial = False
print('停止遥操作')
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]:
if not self.is_connected:
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 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],
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
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

12
ufactory_lerobot/teleoperators/pika_teleop/pika_teleop_config.py
View File

@ -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

3
ufactory_lerobot/teleoperators/space_mouse/space_mouse_config.py
View File

@ -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

3
ufactory_lerobot/teleoperators/uf_mock_teleop/uf_mock_teleop_config.py
View File

@ -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

10
ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop.py
View File

@ -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 = {}

3
ufactory_lerobot/teleoperators/umi_teleop/multiple_umi_teleop_config.py
View File

@ -23,3 +23,6 @@ from .umi_teleop_config import UmiTeleopConfig
@dataclass
class MultipleUmiTeleopConfig(TeleoperatorConfig):
teleops: dict[str, UmiTeleopConfig]
def __post_init__(self):
self.id = 'multiple_umi_teleop' if self.id is None else self.id

129
ufactory_lerobot/teleoperators/umi_teleop/umi_teleop.py
View File

@ -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)
# 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 = [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
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,10 +84,12 @@ 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):
if self.xvlib:
self.xvlib.xv_uninit()
self._is_connected = False
@ -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

12
ufactory_lerobot/teleoperators/umi_teleop/umi_teleop_config.py
View File

@ -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

56
ufactory_lerobot/utils/utils.py
View File

@ -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