finshi press action

data_gen_dependencies/action/pull_revolute.py

@@ -0,0 +1,110 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate
+
+
+class PullRevoluteStage(StageTemplate):
+    DELTA_DISTANCE = 0.005  # meter
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, active_obj=None, passive_obj=None, **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        #======= DEBUG START =======
+        print(f'debug: PullRevoluteStage init')
+        import ipdb;ipdb.set_trace()
+        #=======  DEBUG END  =======
+        self.passive_obj = passive_obj
+        self.active_obj = active_obj
+        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 = 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"
+        self.revolute_radius = passive_element.get('revolute_radius', None)
+        assert self.revolute_radius is not None, "revolute_radius is required 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, target_pose, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        #trajectory = self.test_generate_trajectory(vector_direction)
+        #trajectory = self.test_generate_trajectory(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}')
+        for i in range(10):
+            free_delta_pose = np.eye(4)
+            free_delta_pose[2,3] = -self.DELTA_DISTANCE
+
+            self.sub_stages.append([target_pose, None, np.eye(4), 'AvoidObs'])
+            self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
+        self.sub_stages.append([None, "open", np.eye(4), 'Simple'])
+
+        free_delta_pose = np.eye(4)
+        free_delta_pose[2,3] = -self.DELTA_DISTANCE
+        self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
+
+
+    def generate_trajectory(self, init_vector_direction):
+        trajectory = []
+        #======= DEBUG START =======
+        print(f'debug: ')
+        import ipdb;ipdb.set_trace()
+        #=======  DEBUG END  =======
+        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)
+        current_pose = start_pose.copy()
+        trajectory.append(current_pose)
+
+        for i in range(1, num_steps + 1):
+            angle = i * delta_theta
+            R = self.axis_angle_to_matrix(rotation_axis, angle)
+            move_pose = start_pose @ R
+            trajectory.append(move_pose)
+        return trajectory
+
+    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