solve dependencies problem

data_gen_dependencies/action/slide.py

@@ -0,0 +1,47 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate, pose_difference
+
+
+class SlideStage(StageTemplate):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+
+        self.pre_distance = 0.0
+        self.move_distance = 0.08
+        vector_direction = passive_element['direction']
+        self.generate_substage(target_pose, vector_direction)
+
+
+
+    def generate_substage(self, target_pose, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        # slide
+        move_pose = target_pose.copy()
+        move_pose[:3,3] += vector_direction * self.move_distance
+        self.sub_stages.append([move_pose, None, np.eye(4), 'Simple'])
+
+        # open gripper
+        self.sub_stages.append([None, 'open', np.eye(4), None])
+
+    
+
+    def check_completion(self, objects):
+        if self.__len__()==0:
+            return True
+        sub_stage = self.sub_stages[self.step_id]
+        move_transform, gripper_action = sub_stage[0], sub_stage[1]
+        if gripper_action=='open':
+            self.step_id += 1
+            return True
+        last_pose = self.last_statement['objects'][self.passive_obj_id].obj_pose
+        target_obj_pose = last_pose @ move_transform
+        current_obj_pose = objects[self.passive_obj_id].obj_pose
+        pos_diff, _ = pose_difference(current_obj_pose, target_obj_pose)
+        succ = pos_diff < 0.04
+        print(pos_diff, succ)
+        if succ:
+            self.step_id += 1
+        return succ
+
+