diff --git a/dex_retargeting/__init__.py b/dex_retargeting/__init__.py index f6b7e26..cd1ee63 100644 --- a/dex_retargeting/__init__.py +++ b/dex_retargeting/__init__.py @@ -1 +1 @@ -__version__ = "0.4.3" +__version__ = "0.4.4" diff --git a/dex_retargeting/constants.py b/dex_retargeting/constants.py index 9983a75..4b7147a 100644 --- a/dex_retargeting/constants.py +++ b/dex_retargeting/constants.py @@ -2,6 +2,24 @@ from pathlib import Path from typing import Optional +import numpy as np + +OPERATOR2MANO_RIGHT = np.array( + [ + [0, 0, -1], + [-1, 0, 0], + [0, 1, 0], + ] +) + +OPERATOR2MANO_LEFT = np.array( + [ + [0, 0, -1], + [1, 0, 0], + [0, -1, 0], + ] +) + class RobotName(enum.Enum): allegro = enum.auto() @@ -59,3 +77,9 @@ def get_default_config_path( else: config_name = f"{robot_name_str}_{hand_type_str}.yml" return config_path / config_name + + +OPERATOR2MANO = { + HandType.right: OPERATOR2MANO_RIGHT, + HandType.left: OPERATOR2MANO_LEFT, +} diff --git a/dex_retargeting/robot_wrapper.py b/dex_retargeting/robot_wrapper.py index 4635219..335fb43 100644 --- a/dex_retargeting/robot_wrapper.py +++ b/dex_retargeting/robot_wrapper.py @@ -60,7 +60,19 @@ def get_joint_index(self, name: str): def get_link_index(self, name: str): if name not in self.link_names: raise ValueError(f"{name} is not a link name. Valid link names: \n{self.link_names}") - return self.model.getFrameId(name) + return self.model.getFrameId(name, pin.BODY) + + def get_joint_parent_child_frames(self, joint_name: str): + joint_id = self.model.getFrameId(joint_name) + parent_id = self.model.frames[joint_id].parent + child_id = -1 + for idx, frame in enumerate(self.model.frames): + if frame.previousFrame == joint_id: + child_id = idx + if child_id == -1: + raise ValueError(f"Can not find child link of {joint_name}") + + return parent_id, child_id # -------------------------------------------------------------------------- # # Kinematics function diff --git a/dex_retargeting/seq_retarget.py b/dex_retargeting/seq_retarget.py index c243295..cbc947d 100644 --- a/dex_retargeting/seq_retarget.py +++ b/dex_retargeting/seq_retarget.py @@ -4,6 +4,7 @@ import numpy as np from pytransform3d import rotations +from dex_retargeting.constants import OPERATOR2MANO, HandType from dex_retargeting.optimizer import Optimizer from dex_retargeting.optimizer_utils import LPFilter @@ -39,40 +40,37 @@ def __init__( # Warm started self.is_warm_started = False - # TODO: hack here - self.scene = None - - def warm_start(self, wrist_pos: np.ndarray, wrist_orientation: np.ndarray, global_rot: np.array): + def warm_start( + self, + wrist_pos: np.ndarray, + wrist_quat: np.ndarray, + hand_type: HandType = HandType.right, + is_mano_convention: bool = False, + ): """ Initialize the wrist joint pose using analytical computation instead of retargeting optimization. This function is specifically for position retargeting with the flying robot hand, i.e. has 6D free joint You are not expected to use this function for vector retargeting, e.g. when you are working on teleoperation + Args: wrist_pos: position of the hand wrist, typically from human hand pose - wrist_orientation: orientation of the hand orientation, typically from human hand pose in MANO convention - global_rot: - + wrist_quat: quaternion of the hand wrist, the same convention as the operator frame definition if not is_mano_convention + hand_type: hand type, used to determine the operator2mano matrix + is_mano_convention: whether the wrist_quat is in mano convention """ # This function can only be used when the first joints of robot are free joints - if len(wrist_pos) != 3: - raise ValueError(f"Wrist pos:{wrist_pos} is not a 3-dim vector.") - if len(wrist_orientation) != 3: - raise ValueError(f"Wrist orientation:{wrist_orientation} is not a 3-dim vector.") - if np.linalg.norm(wrist_orientation) < 1e-3: - mat = np.eye(3) - else: - mat = rotations.matrix_from_compact_axis_angle(wrist_orientation) + if len(wrist_pos) != 3: + raise ValueError(f"Wrist pos: {wrist_pos} is not a 3-dim vector.") + if len(wrist_quat) != 4: + raise ValueError(f"Wrist quat: {wrist_quat} is not a 4-dim vector.") + operator2mano = OPERATOR2MANO[hand_type] if is_mano_convention else np.eye(3) robot = self.optimizer.robot - operator2mano = np.array([[0, 0, -1], [-1, 0, 0], [0, 1, 0]]) - mat = global_rot.T @ mat @ operator2mano target_wrist_pose = np.eye(4) - target_wrist_pose[:3, :3] = mat + target_wrist_pose[:3, :3] = rotations.matrix_from_quaternion(wrist_quat) @ operator2mano.T target_wrist_pose[:3, 3] = wrist_pos - wrist_link_name = self.optimizer.wrist_link_name - wrist_link_id = self.optimizer.robot.get_link_index(wrist_link_name) name_list = [ "dummy_x_translation_joint", "dummy_y_translation_joint", @@ -81,6 +79,9 @@ def warm_start(self, wrist_pos: np.ndarray, wrist_orientation: np.ndarray, globa "dummy_y_rotation_joint", "dummy_z_rotation_joint", ] + wrist_link_id = robot.get_joint_parent_child_frames(name_list[5])[1] + + # Set the dummy joints angles to zero old_qpos = robot.q0 new_qpos = old_qpos.copy() for num, joint_name in enumerate(self.optimizer.target_joint_names): @@ -128,6 +129,13 @@ def set_qpos(self, robot_qpos: np.ndarray): target_qpos = robot_qpos[self.optimizer.idx_pin2target] self.last_qpos = target_qpos + def get_qpos(self, fixed_qpos: np.ndarray | None = None): + robot_qpos = np.zeros(self.optimizer.robot.dof) + robot_qpos[self.optimizer.idx_pin2target] = self.last_qpos + if fixed_qpos is not None: + robot_qpos[self.optimizer.idx_pin2fixed] = fixed_qpos + return robot_qpos + def verbose(self): min_value = self.optimizer.opt.last_optimum_value() print(f"Retargeting {self.num_retargeting} times takes: {self.accumulated_time}s") diff --git a/example/position_retargeting/hand_robot_viewer.py b/example/position_retargeting/hand_robot_viewer.py index 3ca0930..4690f5d 100644 --- a/example/position_retargeting/hand_robot_viewer.py +++ b/example/position_retargeting/hand_robot_viewer.py @@ -2,38 +2,22 @@ from pathlib import Path from typing import Dict, List -import numpy as np import cv2 -from tqdm import trange +import numpy as np import sapien -import transforms3d.quaternions +from hand_viewer import HandDatasetSAPIENViewer +from pytransform3d import rotations +from tqdm import trange from dex_retargeting import yourdfpy as urdf -from dex_retargeting.constants import RobotName, HandType, get_default_config_path, RetargetingType +from dex_retargeting.constants import ( + HandType, + RetargetingType, + RobotName, + get_default_config_path, +) from dex_retargeting.retargeting_config import RetargetingConfig from dex_retargeting.seq_retarget import SeqRetargeting -from hand_viewer import HandDatasetSAPIENViewer - -ROBOT2MANO = np.array( - [ - [0, 0, -1], - [-1, 0, 0], - [0, 1, 0], - ] -) -ROBOT2MANO_POSE = sapien.Pose(q=transforms3d.quaternions.mat2quat(ROBOT2MANO)) - - -def prepare_position_retargeting(joint_pos: np.array, link_hand_indices: np.ndarray): - link_pos = joint_pos[link_hand_indices] - return link_pos - - -def prepare_vector_retargeting(joint_pos: np.array, link_hand_indices_pairs: np.ndarray): - joint_pos = joint_pos @ ROBOT2MANO - origin_link_pos = joint_pos[link_hand_indices_pairs[0]] - task_link_pos = joint_pos[link_hand_indices_pairs[1]] - return task_link_pos - origin_link_pos class RobotHandDatasetSAPIENViewer(HandDatasetSAPIENViewer): @@ -45,6 +29,7 @@ def __init__(self, robot_names: List[RobotName], hand_type: HandType, headless=F self.robot_file_names: List[str] = [] self.retargetings: List[SeqRetargeting] = [] self.retarget2sapien: List[np.ndarray] = [] + self.hand_type = hand_type # Load optimizer and filter loader = self.scene.create_urdf_loader() @@ -126,7 +111,22 @@ def render_dexycb_data(self, data: Dict, fps=5, y_offset=0.8): robot_names = "_".join(robot_names) video_path = Path(__file__).parent.resolve() / f"data/{robot_names}_video.mp4" writer = cv2.VideoWriter( - str(video_path), cv2.VideoWriter_fourcc(*"mp4v"), 30.0, (self.camera.get_width(), self.camera.get_height()) + str(video_path), + cv2.VideoWriter_fourcc(*"mp4v"), + 30.0, + (self.camera.get_width(), self.camera.get_height()), + ) + + # Warm start + hand_pose_start = hand_pose[start_frame] + wrist_quat = rotations.quaternion_from_compact_axis_angle(hand_pose_start[0, 0:3]) + vertex, joint = self._compute_hand_geometry(hand_pose_start) + for robot, retargeting, retarget2sapien in zip(self.robots, self.retargetings, self.retarget2sapien): + retargeting.warm_start( + joint[0, :], + wrist_quat, + hand_type=self.hand_type, + is_mano_convention=True, ) # Loop rendering