finshi pull revolute action

data_gen_dependencies/manip_solver.py

@@ -286,7 +286,7 @@ def parse_stage(stage, objects):
     active_obj = objects[active_obj_id]
     passive_obj = objects[passive_obj_id]
 
-    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe', 'pull_revolute']
+    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe', 'pull_revolute', "twist"]
 
     def _load_element(obj, type):
 
@@ -345,7 +345,7 @@ def select_obj(objects, stages, robot):
     if stages[0]['action'] in ['push', 'reset', 'press_prismatic']:
         gripper2obj = current_gripper_pose
 
-    elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute"]:
+    elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute", "twist"]:
         action = stages[0]['action']
 
         ''' 筛掉无IK解的grasp pose '''
@@ -375,7 +375,6 @@ def select_obj(objects, stages, robot):
         if len(grasp_poses) == 0:
             print(action, 'No grasp_gripper_pose can pass Isaac IK')
             return []
-
     ''' 基于有IK解的grasp pose分数,选择最优的passive primitive element,同时选出最优的一个grasp pose'''
     if grasp_stage_id is not None:
         next_stage_id = grasp_stage_id + 1
@@ -404,7 +403,7 @@ def select_obj(objects, stages, robot):
                 active_elements = [{'xyz': xyz_canonical, 'direction': direction_canonical}]
 
 
-            # import ipdb; ipdb.set_trace()
+            #import ipdb; ipdb.set_trace()
             t0 = time.time()
             element_ik_score = []
             grasp_pose_ik_score = []
@@ -549,7 +548,7 @@ def generate_action_stages(objects, all_stages, robot):
                 substages = True
             elif action in ['pick', 'grasp', 'hook']:
                 substages = build_stage(action)(active_obj_id, passive_obj_id, active_elements, passive_elements,
-                                                gripper2obj, extra_params=stage.get('extra_params', None), active_obj=active_obj, passive_obj=passive_obj)
+                                                gripper2obj, extra_params=stage.get('extra_params', None), objects=objects)
             elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'move',
                             'reset', "pull_revolute", "twist"]:  # grasp + 物体自身运动
                 passive_element = passive_elements[np.random.choice(len(passive_elements))]
@@ -557,8 +556,7 @@ def generate_action_stages(objects, all_stages, robot):
                     active_obj_id=active_obj_id,
                     passive_obj_id=passive_obj_id,
                     passive_element=passive_element,
-                    active_obj=active_obj,
-                    passive_obj=passive_obj,
+                    objects=objects,
                     target_pose=gripper2obj
                 )
             else:
@@ -642,8 +640,7 @@ def generate_action_stages(objects, all_stages, robot):
                         vector_direction=passive_element['direction'],
                         passive_element=passive_element,
                         extra_params=stage.get('extra_params', {}),
-                        active_obj=active_obj,
-                        passive_obj=passive_obj
+                        objects=objects
                     )
                     break