import os import numpy as np from scipy.spatial.transform import Rotation as R from data_gen_dependencies.action.base import StageTemplate class PullRevoluteStage(StageTemplate): DELTA_DISTANCE = 0.003 # meter def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects=[], **kwargs): super().__init__(active_obj_id, passive_obj_id, active_element, passive_element) #import ipdb;ipdb.set_trace() self.objects = objects self.passive_obj = objects[passive_obj_id] self.active_obj = objects[active_obj_id] main_part_name= passive_element["part_id"] self.obj_base = os.path.dirname(passive_obj_id) main_part_id = os.path.join(self.obj_base, main_part_name) if main_part_id not in objects: raise ValueError(f"Cannot find main part {main_part_name} in objects") self.root_id = os.path.join(self.obj_base, "root") self.main_part_obj = objects[main_part_id] self.root_obj = objects[self.root_id] self.revolute_radius = np.linalg.norm(self.passive_obj.obj_pose[:3,3] - self.main_part_obj.obj_pose[:3,3]) self.revolute_joint_position = self.main_part_obj.obj_pose[:3,3] self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7) self.correspond_joint_id = passive_element.get('correspond_joint_id', None) correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id] self.joint_lower_limit = correspond_joint_info["lower_bound"] self.joint_upper_limit = correspond_joint_info["upper_bound"] self.joint_axis = correspond_joint_info["joint_axis"] self.joint_type = correspond_joint_info["joint_type"] assert self.joint_type == 'revolute', "joint_type must be revolute for pull_revolute action" if self.joint_position_threshold is None: self.joint_position_threshold = 0.7 assert self.joint_position_threshold >= 0 and self.joint_position_threshold <= 1 self.joint_direction = passive_element.get('joint_direction', 1) assert self.joint_direction in [-1, 1] self.joint_velocity_threshold = passive_element.get('joint_velocity_threshold', 999) vector_direction = passive_element['direction'] self.extra_params = {} if extra_params is None else extra_params self.generate_substage(target_pose, vector_direction) def generate_substage(self, gripper2obj, vector_direction): vector_direction = vector_direction / np.linalg.norm(vector_direction) total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit) arc_length = abs(self.revolute_radius * total_angle) step = int(arc_length / self.DELTA_DISTANCE) from pyboot.utils.log import Log Log.debug(f'PullRevoluteStage: total_angle={total_angle}, arc_length={arc_length}, step={step}') trajectory = self.generate_trajectory(gripper2obj, vector_direction) self.save_visualized_gripper_trajectory(trajectory, "pull_revolute_traj.txt") self.sub_stages.append([trajectory, None, np.eye(4), 'Trajectory']) self.sub_stages.append([None, "open", np.eye(4), 'Simple']) free_delta_pose = np.eye(4) free_delta_pose[2,3] = -0.03 self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper']) def generate_trajectory(self, gripper2passive_obj, init_vector_direction): trajectory = [] total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit) total_angle *= self.joint_direction arc_length = abs(self.revolute_radius * total_angle) num_steps = int(np.ceil(arc_length / self.DELTA_DISTANCE)) if num_steps < 1: num_steps = 1 delta_theta = total_angle / num_steps rotation_axis = np.array(self.joint_axis, dtype=float) rotation_axis /= np.linalg.norm(rotation_axis) passive_obj_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.passive_obj.obj_pose gripper_2_root = passive_obj_2_root @ gripper2passive_obj pos0 = gripper_2_root[:3, 3] R0 = gripper_2_root[:3, :3] joint_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.main_part_obj.obj_pose pivot = joint_2_root[:3, 3] for i in range(1, num_steps + 1): angle = i * delta_theta R_delta = R.from_rotvec(rotation_axis * angle).as_matrix() pos_rot = R_delta @ (pos0 - pivot) + pivot R_new = R_delta @ R0 waypoint_2_root = np.eye(4) waypoint_2_root[:3, :3] = R_new waypoint_2_root[:3, 3] = pos_rot waypoint_2_obj = np.linalg.inv(passive_obj_2_root) @ waypoint_2_root trajectory.append(waypoint_2_obj) return np.asarray(trajectory) def save_visualized_gripper_trajectory(self, gripper_trajectory, save_path): #import ipdb; ipdb.set_trace() N = len(gripper_trajectory) gripper_pts = np.zeros((1, 60, 3)) for i in range(10): gripper_pts[0, i] = np.array([0, 0.01, 0.002*i]) gripper_pts[0, i+10] = np.array([0, 0.01, 0.002*i + 0.001]) gripper_pts[0, i+20] = np.array([0, -0.01, 0.002*i]) gripper_pts[0, i+30] = np.array([0, -0.01, 0.002*i + 0.001]) gripper_pts[0, i+40] = np.array([0, -0.01 + 0.002 * i, 0]) for i in range(5): gripper_pts[0, i+50] = np.array([0, 0, -0.002*i]) gripper_pts[0, i+55] = np.array([0, 0, -0.002*i + 0.001]) gripper_pts -= np.array([0, 0, 0.002*10 + 0.001]) gripper_pts = gripper_pts.repeat(N, axis=0) gripper_trajectory = gripper_trajectory.reshape(N, 1, 4, 4) gripper_trajectory = gripper_trajectory.repeat(60, axis=1) gripper_pts = gripper_pts.reshape(-1, 3) gripper_trajectory = gripper_trajectory.reshape(-1, 4, 4) gripper_pts_2_obj = (gripper_trajectory[:, :3, :3] @ gripper_pts[..., None]).squeeze(-1)+ gripper_trajectory[:, :3, 3] np.savetxt(save_path, gripper_pts_2_obj, fmt="%.6f") def axis_angle_to_matrix(self, axis, angle): axis = axis / np.linalg.norm(axis) K = np.array([[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]) R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K) T = np.eye(4) T[:3, :3] = R return T def check_completion(self, objects): if self.__len__() == 0: return True succ = False if self.step_id >= 0: succ = True if succ: self.step_id += 1 return succ