finshi press action

data_gen_dependencies/action/__init__.py

@@ -4,7 +4,9 @@ from .insert import InsertStage, HitStage
 from .slide import SlideStage
 from .pour import PourStage
 from .pull import PullStage
-from .push import PushStage, ClickStage
+from .push import PushStage
+from .pull_revolute import PullRevoluteStage
+from .press import PressPrismaticStage
 
 ACTION_STAGE = {
     "grasp": GraspStage,
@@ -19,9 +21,11 @@ ACTION_STAGE = {
     "hit": NotImplemented,
     "pour": PourStage,
     "push": PushStage,
-    'click': ClickStage,
-    'touch': ClickStage,
-    "pull": PullStage
+    'click': PressPrismaticStage,
+    'touch': PressPrismaticStage,
+    'press_prismatic': PressPrismaticStage,
+    "pull": PullStage,
+    "pull_revolute": PullRevoluteStage,
 }
 
 def build_stage(action):

data_gen_dependencies/action/base.py

@@ -27,7 +27,6 @@ def simple_check_completion(goal, objects, last_statement=None, pos_threshold=0.
 
 def solve_target_gripper_pose(stage, objects):
     active_obj_ID, passive_obj_ID, target_pose_canonical, gripper_action, transform_world, motion_type = stage
-    
     anchor_pose = objects[passive_obj_ID].obj_pose
     
     

data_gen_dependencies/action/grasp.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class PickStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         self.grasp_pose_canonical = grasp_pose
         self.use_pre_grasp = False
@@ -104,6 +104,7 @@ class GraspStage(PickStage):
             
             self.sub_stages.append([pre_grasp_pose, None, np.eye(4), motion_type])
         # sub-stage-1   moveTo grasp pose
+
         self.sub_stages.append([grasp_pose, 'close', np.eye(4), motion_type])
 
         

data_gen_dependencies/action/insert.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class InsertStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, extra_params={}, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, extra_params={}, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
 

data_gen_dependencies/action/place.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class PlaceStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, extra_params={}, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, extra_params={}, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         self.pre_transform_up =0.12
         self.place_transform_up = np.array([0, 0, 0.01])

data_gen_dependencies/action/pour.py

@@ -9,7 +9,7 @@ import numpy as np
 
 
 class PourStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, current_pose=None, obj2part=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, current_pose=None, obj2part=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         
         self.pre_pour_height = 0.25

data_gen_dependencies/action/press.py

@@ -0,0 +1,34 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate
+
+class PressPrismaticStage(StageTemplate):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None,active_obj=None, passive_obj=None, **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        self.pre_distance = -0.03
+        
+        self.passive_obj = passive_obj
+        self.active_obj = active_obj
+        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
+        self.joint_position_threshold = self.joint_position_threshold if self.joint_position_threshold is not None else 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.move_distance = -0.005 + abs(self.joint_lower_limit - self.joint_upper_limit) * self.joint_position_threshold
+        self.joint_axis = correspond_joint_info["joint_axis"]
+        self.joint_type = correspond_joint_info["joint_type"]
+        assert self.joint_type == 'prismatic', "joint_type must be prismatic for press action"
+        correspond_joint_info = passive_obj.joints_info[self.correspond_joint_id]
+        self.generate_substage(target_pose, vector_direction)
+
+        
+    def generate_substage(self, target_pose, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        pre_pose = target_pose.copy()
+        pre_pose[:3,3] += vector_direction * self.pre_distance
+        self.sub_stages.append([pre_pose, 'close', np.eye(4), 'AvoidObs'])
+        
+        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'])

data_gen_dependencies/action/pull.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate
 
 
 class PullStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, **kwargs):
+    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)
         self.pull_distance = passive_element['distance']
         assert self.pull_distance > 0
@@ -32,8 +32,6 @@ class PullStage(StageTemplate):
         free_delta_pose[2,3] = -0.03
         self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
 
-        
-
 
     def check_completion(self, objects):
         if self.__len__()==0:

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

data_gen_dependencies/action/push.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, pose_difference
 
 
 class PushStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), vector_direction=None, extra_params=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), vector_direction=None, extra_params=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
         extra_params = {} if extra_params is None else extra_params
@@ -73,23 +73,5 @@ class PushStage(StageTemplate):
         return succ
 
 
-class ClickStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, **kwargs):
-        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
-        self.pre_distance = -0.03
-        self.move_distance = 0.00
-        self.generate_substage(target_pose, vector_direction)
 
-        
-    def generate_substage(self, target_pose, vector_direction):
-        vector_direction = vector_direction / np.linalg.norm(vector_direction)
-        # moveTo pre-place position
-        pre_pose = target_pose.copy()
-        pre_pose[:3,3] += vector_direction * self.pre_distance
-        self.sub_stages.append([pre_pose, 'close', np.eye(4), 'AvoidObs'])
-        
-        # insert
-        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'])
 

data_gen_dependencies/action/slide.py

@@ -4,7 +4,7 @@ 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):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
         self.pre_distance = 0.0

data_gen_dependencies/action/twist.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate
+
+
+class TwistStage(StageTemplate):
+    DELTA_DISTANCE = 0.01  # 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)
+        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)
+        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 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

data_gen_dependencies/manip_solver.py

@@ -3,6 +3,7 @@ import copy
 import json
 import os
 from sqlite3 import dbapi2
+from this import d
 import time
 
 from pyboot.utils.log import Log
@@ -247,25 +248,30 @@ def load_task_solution(task_info):
             obj_dir = obj_info["data_info_dir"]
             obj = OmniObject.from_obj_dir(obj_dir, obj_info=obj_info)
             objects[obj_id] = obj
-
-            if hasattr(obj, 'part_ids'):
-                for part_id in obj.part_ids:
+            if obj.is_articulated:
+                for part_id in obj.parts_info:
                     id = obj_id + '/%s' % part_id
                     objects[id] = copy.deepcopy(obj)
                     objects[id].name = id
-                    #====== DEBUG START ======
-                    print(f'debug: ')
-                    import ipdb;ipdb.set_trace()
-                    #======  DEBUG END  ======
-                    if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
-                        objects[id].part_joint_limit = obj.part_joint_limits[part_id]
-                if len(obj.part_ids):
+                    objects[id].size = obj.parts_info[part_id]['size']
+                    joint_name = obj.parts_info[part_id]['correspond_joint_id']
+                    objects[id].joint_name = joint_name
+                    if objects[id].joint_name is not None and joint_name in obj.joints_info:
+                        objects[id].joint_type = obj.joints_info[joint_name]['joint_type']
+                        objects[id].joint_axis = obj.joints_info[joint_name]['joint_axis']
+                        objects[id].lower_bound = obj.joints_info[joint_name]['lower_bound']
+                        objects[id].upper_bound = obj.joints_info[joint_name]['upper_bound']
+      
+                if len(obj.parts_info):
                     del objects[obj_id]
+     
     Log.success('Finished Loading task solution')
+
     return stages, objects
 
 
 def parse_stage(stage, objects):
+
     action = stage['action']
     if action in ['reset']:
         return action, 'gripper', 'gripper', None, None, None, None
@@ -280,24 +286,36 @@ 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']
+    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe', 'pull_revolute']
 
     def _load_element(obj, type):
+
         if action in ['pick', 'hook']:
             action_mapped = 'grasp'
+        elif action in ['pull_revolute']:
+            action_mapped = 'pull'
+        elif action in ["press_prismatic"]:
+            action_mapped = 'press'
         else:
             action_mapped = action
-        if action_mapped == 'grasp' and type == 'active':
+        if (action_mapped == 'grasp' or action_mapped == 'pull') and type == 'active':
             return None, None
         elif obj.name == 'gripper':
             element = obj.elements[type][action_mapped]
             return element, 'default'
         primitive = stage[type]['primitive'] if stage[type]['primitive'] is not None else 'default'
         if primitive != 'default' or (action_mapped == 'grasp' and type == 'passive'):
-            if action_mapped not in obj.elements[type]:
+            if action_mapped not in obj.elements[type] and action not in obj.elements[type]:
+                #======= DEBUG START =======
+                print(f'debug: ')
+                import ipdb;ipdb.set_trace()
+                #=======  DEBUG END  =======
                 print('No %s element for %s' % (action_mapped, obj.name))
                 return None, None
-            element = obj.elements[type][action_mapped][primitive]
+            elif action_mapped in obj.elements[type]:
+                element = obj.elements[type][action_mapped][primitive]
+            elif action in obj.elements[type]:
+                element = obj.elements[type][action][primitive]
         else:
             element = []
             for primitive in obj.elements[type][action_mapped]:
@@ -324,11 +342,10 @@ def select_obj(objects, stages, robot):
 
     ''' 初筛抓取pose,得到 grasp_poses_canonical, grasp_poses '''
     grasp_stage_id = None
-
-    if stages[0]['action'] in ['push', 'reset', 'click']:
+    if stages[0]['action'] in ['push', 'reset', 'press_prismatic']:
         gripper2obj = current_gripper_pose
 
-    elif stages[0]['action'] in ['pick', 'grasp', 'hook']:
+    elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute"]:
         action = stages[0]['action']
 
         ''' 筛掉无IK解的grasp pose '''
@@ -386,6 +403,7 @@ def select_obj(objects, stages, robot):
                 direction_canonical = (np.linalg.inv(obj_pose) @ np.array([*direction, 0]))[:3]
                 active_elements = [{'xyz': xyz_canonical, 'direction': direction_canonical}]
 
+
             # import ipdb; ipdb.set_trace()
             t0 = time.time()
             element_ik_score = []
@@ -504,7 +522,6 @@ def split_grasp_stages(stages):
 
 def generate_action_stages(objects, all_stages, robot):
     split_stages = split_grasp_stages(all_stages)
-
     action_stages = []
     for stages in split_stages:
         gripper2obj = select_obj(objects, stages, robot)
@@ -512,30 +529,37 @@ def generate_action_stages(objects, all_stages, robot):
             print('No gripper2obj pose can pass IK')
             gripper2obj = select_obj(objects, stages, robot)
             return []
+        stage_idx = 0
         for stage in stages:
+            stage_idx += 1
             extra_params = stage.get('extra_params', {})
             arm = extra_params.get('arm', 'right')
-            current_gripper_pose = robot.get_ee_pose(id=arm)
             action, active_obj_id, passive_obj_id, active_elements, passive_elements, active_primitive, passive_primitive = parse_stage(
                 stage, objects)
             active_obj = objects[active_obj_id]
             passive_obj = objects[passive_obj_id]
-
-            single_obj = active_obj_id == passive_obj_id
+            Log.info(f'Generating action stage [{stage_idx}/{len(stages)}]: Action({action}), Active Object({active_obj_id}), Passive Object({passive_obj_id})')
 
             substages = None
+            #======= DEBUG START =======
+            print(f'debug: before build_stage {action}')
+            #import ipdb;ipdb.set_trace()
+            #=======  DEBUG END  =======
             if action in ['reset']:
                 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))
-            elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'pull', 'move',
-                            'reset']:  # grasp + 物体自身运动
+                                                gripper2obj, extra_params=stage.get('extra_params', None), active_obj=active_obj, passive_obj=passive_obj)
+            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))]
                 substages = build_stage(action)(
                     active_obj_id=active_obj_id,
                     passive_obj_id=passive_obj_id,
-                    passive_element=passive_element
+                    passive_element=passive_element,
+                    active_obj=active_obj,
+                    passive_obj=passive_obj,
+                    target_pose=gripper2obj
                 )
             else:
                 passive_element = passive_elements[np.random.choice(len(passive_elements))]
@@ -575,7 +599,7 @@ def generate_action_stages(objects, all_stages, robot):
                     ik_jacobian_score = ik_info["jacobian_score"][ik_success]
 
                     if len(target_gripper_poses_pass_ik) == 0:
-                        print(action, ': No target_obj_pose can pass isaac curobo IK')
+                        Log.warning(f'{action}: No target_obj_pose can pass isaac curobo IK')
                         continue
 
                     target_joint_positions = []
@@ -617,12 +641,14 @@ def generate_action_stages(objects, all_stages, robot):
                         obj2part=obj2part,
                         vector_direction=passive_element['direction'],
                         passive_element=passive_element,
-                        extra_params=stage.get('extra_params', {})
+                        extra_params=stage.get('extra_params', {}),
+                        active_obj=active_obj,
+                        passive_obj=passive_obj
                     )
                     break
 
             if substages is None:
-                print(action, ': No target_obj_pose can pass IK')
+                Log.warning(f'{action}:substage is None')
                 return []
             action_stages.append((action, substages))
 

data_gen_dependencies/object.py

@@ -49,9 +49,9 @@ class OmniObject:
                         "direction": np.array([-1, 0, 0.3]),
                     },
                 ],
-                'click': {
+                'press': {
                     "part": "finger edge",
-                    "task": "click",
+                    "task": "press",
                     "xyz": np.array([0, 0, 0]),
                     "direction": np.array([0, 0, 1]),
                 },
@@ -80,7 +80,6 @@ class OmniObject:
         if "interaction" in obj_info:
             obj_info = obj_info
             interaction_info = obj_info["interaction"]
-            part_joint_limits_info = obj_info.get("part_joint_limits", None)
 
         else:
 
@@ -93,7 +92,6 @@ class OmniObject:
             obj_info = json.load(open(obj_info_file))
             interaction_data = json.load(open(interaction_label_file))
             interaction_info = interaction_data['interaction']
-            part_joint_limits_info = interaction_data.get("part_joint_limits", None)
 
         obj = cls(
             name=obj_info['object_id'],
@@ -104,10 +102,9 @@ class OmniObject:
         if os.path.exists(mesh_file):
             obj_info['mesh_file'] = mesh_file
 
-        obj.part_joint_limits = part_joint_limits_info
 
         ''' Load interaction labels '''
-        obj.part_ids = []
+        
         for type in ['active', 'passive']:
             if type not in interaction_info:
                 continue
@@ -197,15 +194,14 @@ class OmniObject:
                         action_info[grasp_part] = grasp_data
 
                     interaction_info[type][action] = action_info
-                else:
-                    for primitive in action_info:
-                        for primitive_info in action_info[primitive]:
-                            if 'part_id' in primitive_info:
-                                obj.part_ids.append(primitive_info['part_id'])
+        obj.is_articulated = obj_info.get("is_articulated", False)
+        if obj.is_articulated:
+            obj.parts_info = obj_info["parts_info"]
+            obj.joints_info = obj_info["joints_info"]
         obj.elements = interaction_info
         obj.info = obj_info
         # import ipdb;ipdb.set_trace()
-        obj.is_articulated = True if part_joint_limits_info is not None else False
+        
         return obj
 
     def set_element(self, element):

data_gen_dependencies/omniagent.py

@@ -133,17 +133,16 @@ class Agent(BaseAgent):
                     objects['fix_pose'].obj_pose[:3, :3] = rotation_matrix
                 continue
                 # TODO(unify part_name and obj_name)
-            if '/' in obj_id:
+            obj = objects[obj_id]
+            if obj.is_articulated:
                 obj_name = obj_id.split('/')[0]
-                part_name = obj_id.split('/')[1]
-
                 object_joint_state = self.robot.client.get_object_joint('/World/Objects/%s' % obj_name)
                 for joint_name, joint_position, joint_velocity in zip(object_joint_state.joint_names,
                                                                       object_joint_state.joint_positions,
                                                                       object_joint_state.joint_velocities):
-                    if joint_name[-1] == part_name[-1]:
-                        objects[obj_id].joint_position = joint_position
-                        objects[obj_id].joint_velocity = joint_velocity
+                    if joint_name == obj.joint_name:
+                        obj.joint_position = joint_position
+                        obj.joint_velocity = joint_velocity
 
             objects[obj_id].obj_pose = self.robot.get_prim_world_pose('/World/Objects/%s' % obj_id)
             if hasattr(objects[obj_id], 'info') and 'simple_place' in objects[obj_id].info and objects[obj_id].info[
@@ -239,14 +238,15 @@ class Agent(BaseAgent):
                                      camera_prim_list=camera_list, fps=fps,
                                      render_semantic=render_semantic)  # TODO 录制判断
 
-            stage_id = -1
+            
             success = False
             substages = None
             for _stages in split_stages:
+                stage_id = -1
                 extra_params = _stages[0].get('extra_params', {})
                 active_id, passive_id = _stages[0]['active']['object_id'], _stages[0]['passive']['object_id']
                 arm = extra_params.get('arm', 'right')
-                Log.info(f'Generating action stages [Executing task: {stage_id+1}/{len(split_stages)}]')
+                Log.info(f'Generating action stages')
                 action_stages = generate_action_stages(objects, _stages, self.robot)
                 if not len(action_stages):
                     success = False
@@ -255,12 +255,14 @@ class Agent(BaseAgent):
 
                 # Execution
                 success = True
-
+                last_gripper_state = None
                 for action, substages in action_stages:
                     stage_id += 1
                     Log.info(f'Start action stage: {action} [Executing task: {stage_id+1}/{len(action_stages)}]')
+
                     active_id, passive_id = _stages[stage_id]['active']['object_id'], _stages[stage_id]['passive'][
                         'object_id']
+                    
                     if action in ['reset']:
                         init_pose = self.robot.reset_pose[arm]
                         curr_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
@@ -277,7 +279,6 @@ class Agent(BaseAgent):
                         self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
                                                           self.attached_obj_id is not None, arm=arm,
                                                           target_pose=target_gripper_pose)
-
                         if target_gripper_pose is not None:
                             self.robot.move_pose(target_gripper_pose, motion_type, arm=arm, block=True)
 
@@ -286,13 +287,15 @@ class Agent(BaseAgent):
                         self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
                                                           self.attached_obj_id is not None, set_gripper_open,
                                                           set_gripper_close, arm=arm, target_pose=target_gripper_pose)
-                        self.robot.set_gripper_action(gripper_action, arm=arm)
-                        if gripper_action == 'open':
-                            time.sleep(1)
+                        if last_gripper_state != gripper_action and gripper_action is not None:
+                            self.robot.set_gripper_action(gripper_action, arm=arm)
+                            last_gripper_state = gripper_action
+                            if gripper_action == 'open':
+                                time.sleep(1)
 
-                        self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                          self.attached_obj_id is not None, arm=arm,
-                                                          target_pose=target_gripper_pose)
+                            self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                            self.attached_obj_id is not None, arm=arm,
+                                                            target_pose=target_gripper_pose)
 
                         # check sub-stage completion
                         objects['gripper'].obj_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
@@ -305,7 +308,7 @@ class Agent(BaseAgent):
                         if success == False:
                             # import ipdb;ipdb.set_trace()
                             self.attached_obj_id = None
-                            Log.error('Failed at sub-stage %d' % substages.step_id)
+                            Log.warning('Failed at sub-stage %d' % substages.step_id)
                             break
 
                         # attach grasped object to gripper           # TODO avoid articulated objects

task_gen_dependencies/object.py

@@ -48,9 +48,9 @@ class OmniObject:
                         "direction": np.array([-1,0,0.3]),
                     },
                 ],
-                'click': {
+                'press': {
                     "part": "finger edge",
-                    "task": "click",
+                    "task": "press",
                     "xyz": np.array([0,0,0]),
                     "direction": np.array([0,0,1]),
                 },