solve dependencies problem

data_gen_dependencies/manip_solver.py

@@ -0,0 +1,657 @@
+# -*- coding: utf-8 -*-
+import copy
+import json
+import os
+import time
+
+import trimesh
+
+curPath = os.path.abspath(os.path.dirname(__file__))
+rootPath = os.path.split(curPath)[0]
+import numpy as np
+
+from data_gen_dependencies.transforms import calculate_rotation_matrix, rotate_around_axis
+from data_gen_dependencies.object import OmniObject, transform_coordinates_3d
+from data_gen_dependencies.fix_rotation import rotate_along_axis
+
+from data_gen_dependencies.action import build_stage
+
+from scipy.interpolate import interp1d
+from scipy.spatial.transform import Rotation as R
+
+
+def draw_axis(robot, pose):
+    if pose.shape[0] == 4 and pose.shape[1] == 4:
+        pose = np.concatenate([pose[:3, 3], R.from_matrix(pose[:3, :3]).as_euler('xyz')])
+
+    def generate_orthogonal_lines(point, angles, scale):
+        assert len(point) == 3 and len(angles) == 3
+
+        # Define rotation matrices for each axis
+        Rx = np.array([[1, 0, 0],
+                       [0, np.cos(angles[0]), -np.sin(angles[0])],
+                       [0, np.sin(angles[0]), np.cos(angles[0])]])
+
+        Ry = np.array([[np.cos(angles[1]), 0, np.sin(angles[1])],
+                       [0, 1, 0],
+                       [-np.sin(angles[1]), 0, np.cos(angles[1])]])
+
+        Rz = np.array([[np.cos(angles[2]), -np.sin(angles[2]), 0],
+                       [np.sin(angles[2]), np.cos(angles[2]), 0],
+                       [0, 0, 1]])
+
+        # Combined rotation matrix
+        R = Rz @ Ry @ Rx
+
+        # Define the unit vectors for the axes
+        unit_vectors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+
+        # Rotate the unit vectors
+        rotated_vectors = R @ unit_vectors.T
+
+        # Draw the lines
+        lines = []
+        for vec in rotated_vectors.T:
+            start_point = point
+            end_point = point + vec * scale
+            lines.append(np.concatenate([start_point, end_point]))
+        return np.stack(lines)
+
+    draw_colors = [
+        [0, 0, 1],
+        [0, 1, 0],
+        [1, 0, 0],
+    ]
+    line_sizes = [1, 1, 1]
+    xyz_axis = generate_orthogonal_lines(point=pose[:3], angles=pose[3:], scale=0.1)
+    start_point = xyz_axis[:, :3]
+    end_point = xyz_axis[:, 3:]
+    robot.client.DrawLine(start_point, end_point, draw_colors, line_sizes)
+
+
+def load_gripper_width_interp_func(file):
+    data = json.load(open(file, 'r'))
+
+    digits = []
+    widths = []
+    depths = []
+    for frame in data:
+        digits.append(frame["angel"])
+        widths.append(frame["width"])
+        depths.append(frame["depth"])
+    digits = np.array(digits)
+    widths = np.array(widths)
+    depths = np.array(depths)
+
+    func_width2depth = interp1d(widths, depths, kind="linear", fill_value="extrapolate")
+    return func_width2depth
+
+
+func_width2depth = load_gripper_width_interp_func(os.path.dirname(__file__) + "/gripper_width_120s_fixed.json")
+
+
+def format_object(obj, distance, type='active'):
+    if obj is None:
+        return None
+    xyz, direction = obj.xyz, obj.direction
+
+    direction = direction / np.linalg.norm(direction) * distance
+    type = type.lower()
+    if type == 'active':
+        xyz_start = xyz
+        xyz_end = xyz_start + direction
+    elif type == 'passive' or type == 'plane':
+        xyz_end = xyz
+        xyz_start = xyz_end - direction
+
+    part2obj = np.eye(4)
+    part2obj[:3, 3] = xyz_start
+    obj.obj2part = np.linalg.inv(part2obj)
+
+    obj_info = {
+        'pose': obj.obj_pose,
+        'length': obj.obj_length,
+        'xyz_start': xyz_start,
+        'xyz_end': xyz_end,
+        'obj2part': obj.obj2part
+    }
+    return obj_info
+
+
+def obj2world(obj_info):
+    obj_pose = obj_info['pose']
+    obj_length = obj_info['length']
+    obj2part = obj_info['obj2part']
+    xyz_start = obj_info['xyz_start']
+    xyz_end = obj_info['xyz_end']
+
+    arrow_in_obj = np.array([xyz_start, xyz_end]).transpose(1, 0)
+    arrow_in_world = transform_coordinates_3d(arrow_in_obj, obj_pose).transpose(1, 0)
+
+    xyz_start_world, xyz_end_world = arrow_in_world
+    direction_world = xyz_end_world - xyz_start_world
+    direction_world = direction_world / np.linalg.norm(direction_world)
+
+    obj_info_world = {
+        'pose': obj_pose,
+        'length': obj_length,
+        'obj2part': obj2part,
+        'xyz_start': xyz_start_world,
+        'xyz_end': xyz_end_world,
+        'direction': direction_world,
+
+    }
+    return obj_info_world
+
+
+def get_aligned_fix_pose(active_obj, passive_obj, distance=0.01, N=1):
+    try:
+        active_object = format_object(active_obj, type='active', distance=distance)
+        passive_object = format_object(passive_obj, type='passive', distance=distance)
+    except:
+        print('error')
+    active_obj_world = obj2world(active_object)
+    current_obj_pose = active_obj_world['pose']
+    if passive_object is None:
+        return current_obj_pose[np.newaxis, ...]
+
+    passive_obj_world = obj2world(passive_object)
+    passive_obj_world['direction'] = passive_obj.direction
+
+    R = calculate_rotation_matrix(active_obj_world['direction'], passive_obj_world['direction'])
+    T = passive_obj_world['xyz_end'] - R @ active_obj_world['xyz_start']
+    transform_matrix = np.eye(4)
+    transform_matrix[:3, :3] = R
+    transform_matrix[:3, 3] = T
+    target_obj_pose = transform_matrix @ current_obj_pose
+    target_obj_pose[:3, 3] = passive_obj.obj_pose[:3, 3]
+
+    poses = []
+    for angle in [i * 360 / N for i in range(N)]:
+        pose_rotated = rotate_around_axis(
+            target_obj_pose,
+            passive_obj_world['xyz_start'],
+            passive_obj_world['direction'],
+            angle)
+        poses.append(pose_rotated)
+    return np.stack(poses)
+
+
+def get_aligned_pose(active_obj, passive_obj, distance=0.01, N=1):
+    try:
+        active_object = format_object(active_obj, type='active', distance=distance)
+        passive_object = format_object(passive_obj, type='passive', distance=distance)
+    except:
+        print('error')
+
+    active_obj_world = obj2world(active_object)
+    current_obj_pose = active_obj_world['pose']
+    if passive_object is None:
+        return current_obj_pose[np.newaxis, ...]
+
+    passive_obj_world = obj2world(passive_object)
+
+    R = calculate_rotation_matrix(active_obj_world['direction'], passive_obj_world['direction'])
+    T = passive_obj_world['xyz_end'] - R @ active_obj_world['xyz_start']
+    transform_matrix = np.eye(4)
+    transform_matrix[:3, :3] = R
+    transform_matrix[:3, 3] = T
+    target_obj_pose = transform_matrix @ current_obj_pose
+
+    poses = []
+    for angle in [i * 360 / N for i in range(N)]:
+        pose_rotated = rotate_around_axis(
+            target_obj_pose,
+            passive_obj_world['xyz_start'],
+            passive_obj_world['direction'],
+            angle)
+        poses.append(pose_rotated)
+    return np.stack(poses)
+
+
+def load_task_solution(task_info):
+    # task_info = json.load(open('%s/task_info.json'%task_dir, 'rb'))
+    stages = task_info['stages']
+
+    objects = {
+        'gripper': OmniObject('gripper')
+    }
+
+    for obj_info in task_info['objects']:
+        obj_id = obj_info['object_id']
+        if obj_id == 'fix_pose':
+            obj = OmniObject('fix_pose')
+            if 'position' not in obj_info or 'direction' not in obj_info:
+                print(f"Error: Missing position/direction in object {obj_id}")
+                continue
+            obj.set_pose(np.array(obj_info['position']), np.array([0.001, 0.001, 0.001]))
+            obj.elements = {
+                'active': {},
+                'passive': {
+                    'place': {
+                        'default': [
+                            {
+                                'xyz': np.array([0, 0, 0]),
+                                'direction': np.array(obj_info['direction'])
+                            }
+                        ]
+                    }
+                }
+            }
+            objects[obj_id] = obj
+        else:
+            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'):
+                if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
+                    obj_parts_joint_limits = obj.part_joint_limits
+                for part_id in obj.part_ids:
+                    id = obj_id + '/%s' % part_id
+                    objects[id] = copy.deepcopy(obj)
+                    objects[id].name = id
+                    objects[id].part_joint_limit = obj_parts_joint_limits[part_id]
+                if len(obj.part_ids):
+                    del objects[obj_id]
+    return stages, objects
+
+
+def parse_stage(stage, objects):
+    action = stage['action']
+    if action in ['reset']:
+        return action, 'gripper', 'gripper', None, None, None, None
+    active_obj_id = stage['active']['object_id']
+    if 'part_id' in stage['active']:
+        active_obj_id += '/%s' % stage['active']['part_id']
+
+    passive_obj_id = stage['passive']['object_id']
+    if 'part_id' in stage['passive']:
+        passive_obj_id += '/%s' % stage['passive']['part_id']
+
+    active_obj = objects[active_obj_id]
+    passive_obj = objects[passive_obj_id]
+
+    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe']
+
+    def _load_element(obj, type):
+        if action in ['pick', 'hook']:
+            action_mapped = 'grasp'
+        else:
+            action_mapped = action
+        if action_mapped == 'grasp' 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]:
+                print('No %s element for %s' % (action_mapped, obj.name))
+                return None, None
+            element = obj.elements[type][action_mapped][primitive]
+        else:
+            element = []
+            for primitive in obj.elements[type][action_mapped]:
+                _element = obj.elements[type][action_mapped][primitive]
+                if isinstance(_element, list):
+                    element += _element
+                else:
+                    element.append(_element)
+        return element, primitive
+
+    passive_element, passive_primitive = _load_element(passive_obj, 'passive')
+    if not single_obj:
+        active_element, active_primitive = _load_element(active_obj, 'active')
+    else:
+        active_element, active_primitive = passive_element, passive_primitive
+    return action, active_obj_id, passive_obj_id, active_element, passive_element, active_primitive, passive_primitive
+
+
+def select_obj(objects, stages, robot):
+    gripper2obj = None
+    extra_params = stages[0].get('extra_params', {})
+    arm = extra_params.get('arm', 'right')
+    current_gripper_pose = robot.get_ee_pose(id=arm)
+
+    ''' 初筛抓取pose,得到 grasp_poses_canonical, grasp_poses '''
+    grasp_stage_id = None
+
+    if stages[0]['action'] in ['push', 'reset', 'click']:
+        gripper2obj = current_gripper_pose
+
+    elif stages[0]['action'] in ['pick', 'grasp', 'hook']:
+        action = stages[0]['action']
+
+        ''' 筛掉无IK解的grasp pose '''
+        grasp_stage_id = 0
+        grasp_stage = parse_stage(stages[0], objects)
+        _, _, passive_obj_id, _, passive_element, _, _ = grasp_stage
+        grasp_obj_id = passive_obj_id
+        grasp_poses_canonical = passive_element['grasp_pose'].copy()
+        grasp_widths = passive_element['width']
+
+        grasp_poses_canonical[:, :3, :3] = grasp_poses_canonical[:, :3, :3] @ robot.robot_gripper_2_grasp_gripper[
+            np.newaxis, ...]
+
+        grasp_poses_canonical_flip = []
+        for _i in range(grasp_poses_canonical.shape[0]):
+            grasp_poses_canonical_flip.append(rotate_along_axis(grasp_poses_canonical[_i], 180, 'z', use_local=True))
+        grasp_poses_canonical_flip = np.stack(grasp_poses_canonical_flip)
+        grasp_poses_canonical = np.concatenate([grasp_poses_canonical, grasp_poses_canonical_flip], axis=0)
+        grasp_widths = np.concatenate([grasp_widths, grasp_widths], axis=0)
+        grasp_poses = objects[passive_obj_id].obj_pose[np.newaxis, ...] @ grasp_poses_canonical
+
+        # filter with IK-checking
+        ik_success, _ = robot.solve_ik(grasp_poses, ee_type='gripper', arm=arm, type='Simple')
+        grasp_poses_canonical, grasp_poses = grasp_poses_canonical[ik_success], grasp_poses[ik_success]
+        grasp_widths = grasp_widths[ik_success]
+        print('%s, %s, Filtered grasp pose with isaac-sim IK: %d/%d' % (action, passive_obj_id, grasp_poses.shape[0],
+                                                                        ik_success.shape[0]))
+
+        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
+        if next_stage_id < len(stages):
+            action, active_obj_id, passive_obj_id, active_elements, passive_elements, active_primitive, passive_primitive = parse_stage(
+                stages[next_stage_id], objects)
+
+            single_obj = active_obj_id == passive_obj_id
+
+            active_obj = objects[active_obj_id]
+            passive_obj = objects[passive_obj_id]
+            passive_element = passive_elements[np.random.choice(len(passive_elements))]
+
+            if action == 'place':  # TODO A2D暂时这样搞，Franka要取消
+                # import ipdb; ipdb.set_trace()
+                obj_pose = active_obj.obj_pose
+                mesh = trimesh.load(active_obj.info['mesh_file'], force="mesh")
+                mesh.apply_scale(0.001)
+                mesh.apply_transform(obj_pose)
+                pts, _ = trimesh.sample.sample_surface(mesh, 200)  # 表面采样
+                xyz = np.array([np.mean(pts[:, 0]), np.mean(pts[:, 1]), np.percentile(pts[:, 2], 1)])
+
+                direction = np.array([0, 0, -1])
+                xyz_canonical = (np.linalg.inv(obj_pose) @ np.array([*xyz, 1]))[:3]
+                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 = []
+            grasp_pose_ik_score = []
+            for active_element in active_elements:
+                # interaction between two rigid objects
+                obj_pose = active_obj.obj_pose
+
+                N_align = 12
+                if not single_obj:
+                    active_obj.xyz, active_obj.direction = active_element['xyz'], active_element['direction']
+                    if passive_obj_id == 'fix_pose':
+                        passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
+                        target_obj_poses = get_aligned_fix_pose(active_obj, passive_obj, N=N_align)
+                    else:
+                        passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
+                        target_obj_poses = get_aligned_pose(active_obj, passive_obj, N=N_align)
+                else:  # 物体自身移动
+                    transform = np.eye(4)
+                    transform[:3, 3] = active_element['xyz']
+                    target_obj_poses = (obj_pose @ transform)[np.newaxis, ...]
+                    N_align = 1
+
+                N_obj_pose = target_obj_poses.shape[0]
+                N_grasp_pose = grasp_poses_canonical.shape[0]
+                target_gripper_poses = (
+                            target_obj_poses[:, np.newaxis, ...] @ grasp_poses_canonical[np.newaxis, ...]).reshape(-1,
+                                                                                                                   4, 4)
+
+                ik_success, _ = robot.solve_ik(target_gripper_poses, ee_type='gripper', type='Simple', arm=arm)
+                element_ik_score.append(np.max(ik_success.reshape(N_obj_pose, N_grasp_pose).sum(axis=1)))
+
+                grasp_pose_ik = ik_success.reshape(N_obj_pose, N_grasp_pose)
+                grasp_pose_ik_score.append(np.sum(grasp_pose_ik, axis=0))
+
+            # print("ik solve time: ", time.time() - t0)
+            best_element_id = np.argmax(element_ik_score)
+            best_active_element = active_elements[best_element_id]
+
+            if not single_obj:
+                active_obj.elements['active'][action] = {active_primitive: best_active_element}
+
+            grasp_ik_score = grasp_pose_ik_score[best_element_id]
+
+            # import ipdb;ipdb.set_trace()
+            _mask = grasp_ik_score >= np.median(grasp_ik_score) / 2
+            best_grasp_poses = grasp_poses[_mask]
+
+            # downsample grasp pose
+            downsample_num = 100
+            grasps_num = best_grasp_poses.shape[0]
+            selected_index = np.random.choice(grasps_num, size=min(downsample_num, grasps_num), replace=False)
+            best_grasp_poses = best_grasp_poses[selected_index]
+            # if best_grasp_poses.shape[0] > downsample_num:
+            #     best_grasp_poses = best_grasp_poses[:downsample_num]
+            if best_grasp_poses.shape[0] > 1:
+                joint_names = robot.joint_names[arm]
+
+                ik_success, ik_info = robot.solve_ik(best_grasp_poses, ee_type='gripper', type='Simple', arm=arm)
+                best_grasp_poses = best_grasp_poses[ik_success]
+                ik_joint_positions = ik_info["joint_positions"][ik_success]
+                ik_joint_names = ik_info["joint_names"][ik_success]
+                ik_jacobian_score = ik_info["jacobian_score"][ik_success]
+
+                if len(best_grasp_poses) == 0:
+                    print(action, 'No best_grasp_poses can pass curobo IK')
+                    return []
+                target_joint_positions = []
+                for ik_joint_position, ik_joint_name in zip(ik_joint_positions, ik_joint_names):
+                    temp_target_joint_positions = []
+                    for joint_name in joint_names:
+                        temp_target_joint_positions.append(ik_joint_position[list(ik_joint_name).index(joint_name)])
+                    target_joint_positions.append(np.array(temp_target_joint_positions))
+                target_joint_positions = np.array(target_joint_positions)
+
+                # choose target pose based on the ik jacobian score and ik joint positions
+                cur_joint_states = robot.client.get_joint_positions().states
+                cur_joint_positions = []
+                for key in cur_joint_states:
+                    if key.name in joint_names:
+                        cur_joint_positions.append(key.position)
+                cur_joint_positions = np.array(cur_joint_positions)
+                joint_pos_dist = np.linalg.norm(target_joint_positions - cur_joint_positions[np.newaxis, :], axis=1)
+                # joint_pos_dist = (joint_pos_dist - np.min(joint_pos_dist))/(np.max(joint_pos_dist) - np.min(joint_pos_dist))
+                # ik_jacobian_score = (ik_jacobian_score - np.min(ik_jacobian_score))/(np.max(ik_jacobian_score) - np.min(ik_jacobian_score))
+                cost = joint_pos_dist  # - ik_jacobian_score
+                idx_sorted = np.argsort(cost)
+                best_grasp_pose = best_grasp_poses[idx_sorted][0]
+
+            else:
+                best_grasp_pose = best_grasp_poses[0]
+
+            best_grasp_pose_canonical = np.linalg.inv(objects[grasp_obj_id].obj_pose) @ best_grasp_pose
+            gripper2obj = best_grasp_pose_canonical
+        else:
+            gripper2obj = grasp_poses_canonical[0]
+    return gripper2obj
+
+
+def split_grasp_stages(stages):
+    split_stages = []
+    i = 0
+    while i < len(stages):
+        if stages[i]['action'] in ['pick', 'grasp', 'hook']:
+            if (i + 1) < len(stages) and stages[i + 1]['action'] not in ['pick', 'grasp', 'hook']:
+                split_stages.append([stages[i], stages[i + 1]])
+                i += 2
+            else:
+                split_stages.append([stages[i]])
+                i += 1
+        else:
+            split_stages.append([stages[i]])
+            i += 1
+    return split_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)
+        if gripper2obj is None or len(gripper2obj) == 0:
+            print('No gripper2obj pose can pass IK')
+            gripper2obj = select_obj(objects, stages, robot)
+            return []
+        for stage in stages:
+            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
+
+            substages = None
+            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 + 物体自身运动
+                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
+                )
+            else:
+                passive_element = passive_elements[np.random.choice(len(passive_elements))]
+                # active_element = active_elements[np.random.choice(len(active_elements))] if isinstance(active_elements, list) else active_elements
+                if not isinstance(active_elements, list):
+                    active_elements = [active_elements]
+
+                for active_element in active_elements:
+                    joint_names = robot.joint_names[arm]
+
+                    # interaction between two rigid objects
+                    obj_pose = active_obj.obj_pose
+                    anchor_pose = passive_obj.obj_pose
+                    current_obj_pose_canonical = np.linalg.inv(anchor_pose) @ obj_pose
+                    active_obj.xyz, active_obj.direction = active_element['xyz'], active_element['direction']
+                    passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
+                    if active_obj.name == 'gripper':
+                        gripper2obj = np.eye(4)
+
+                    if passive_obj_id == 'fix_pose':
+                        # import ipdb;ipdb.set_trace()
+                        target_obj_poses = get_aligned_fix_pose(active_obj, passive_obj, N=18)
+                    else:
+                        target_obj_poses = get_aligned_pose(active_obj, passive_obj, N=18)
+                    target_gripper_poses = target_obj_poses @ gripper2obj[np.newaxis, ...]
+                    # import ipdb;ipdb.set_trace()
+
+                    # downsample target_gripper_poses
+                    downsample_num = 100
+                    if target_gripper_poses.shape[0] > downsample_num:
+                        target_gripper_poses = target_gripper_poses[:downsample_num]
+                    ik_success, ik_info = robot.solve_ik(target_gripper_poses, ee_type='gripper', type='AvoidObs',
+                                                         arm=arm)
+                    target_gripper_poses_pass_ik = target_gripper_poses[ik_success]
+                    ik_joint_positions = ik_info["joint_positions"][ik_success]
+                    ik_joint_names = ik_info["joint_names"][ik_success]
+                    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')
+                        continue
+
+                    target_joint_positions = []
+                    for ik_joint_position, ik_joint_name in zip(ik_joint_positions, ik_joint_names):
+                        temp_target_joint_positions = []
+                        for joint_name in joint_names:
+                            temp_target_joint_positions.append(ik_joint_position[list(ik_joint_name).index(joint_name)])
+                        target_joint_positions.append(np.array(temp_target_joint_positions))
+                    target_joint_positions = np.array(target_joint_positions)
+
+                    # choose target pose based on the jacobian score and ik joint positions difference
+                    cur_joint_states = robot.client.get_joint_positions().states
+                    cur_joint_positions = []
+                    for key in cur_joint_states:
+                        if key.name in joint_names:
+                            cur_joint_positions.append(key.position)
+                    cur_joint_positions = np.array(cur_joint_positions)
+                    joint_pos_dist = np.linalg.norm(target_joint_positions - cur_joint_positions[np.newaxis, :], axis=1)
+                    # joint_pos_dist = (joint_pos_dist - np.min(joint_pos_dist))/(np.max(joint_pos_dist) - np.min(joint_pos_dist))
+                    # ik_jacobian_score = (ik_jacobian_score - np.min(ik_jacobian_score))/(np.max(ik_jacobian_score) - np.min(ik_jacobian_score))
+                    cost = joint_pos_dist  # - ik_jacobian_score
+                    idx_sorted = np.argsort(cost)
+
+                    best_target_gripper_pose = target_gripper_poses_pass_ik[idx_sorted][0]
+
+                    best_target_obj_pose = best_target_gripper_pose @ np.linalg.inv(gripper2obj)
+                    target_obj_pose_canonical = np.linalg.inv(
+                        anchor_pose) @ best_target_obj_pose  # TODO 暂时只取一个可行解，后面要结合grasp pose做整条trajectory的joint solve
+
+                    part2obj = np.eye(4)
+                    part2obj[:3, 3] = active_obj.xyz
+                    obj2part = np.linalg.inv(part2obj)
+
+                    substages = build_stage(action)(
+                        active_obj_id=active_obj_id,
+                        passive_obj_id=passive_obj_id,
+                        target_pose=target_obj_pose_canonical,
+                        current_pose=current_obj_pose_canonical,
+                        obj2part=obj2part,
+                        vector_direction=passive_element['direction'],
+                        passive_element=passive_element,
+                        extra_params=stage.get('extra_params', {})
+                    )
+                    break
+
+            if substages is None:
+                print(action, ': No target_obj_pose can pass IK')
+                return []
+            action_stages.append((action, substages))
+
+    return action_stages
+
+
+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
+
+    if motion_type == 'Trajectory':
+        assert len(target_pose_canonical.shape) == 3, 'The target_pose should be a list of poses'
+        target_pose = anchor_pose[np.newaxis, ...] @ target_pose_canonical
+        target_pose = transform_world[np.newaxis, ...] @ target_pose
+    else:
+        target_pose = anchor_pose @ target_pose_canonical
+        target_pose = transform_world @ target_pose
+    assert 'gripper' in objects, 'The gripper should be the first one in the object list'
+    current_gripper_pose = objects['gripper'].obj_pose
+
+    if active_obj_ID == 'gripper':
+        target_gripper_pose = target_pose
+    else:
+        current_obj_pose = objects[active_obj_ID].obj_pose
+        gripper2obj = np.linalg.inv(current_obj_pose) @ current_gripper_pose
+        if len(target_pose.shape) == 3:
+            gripper2obj = gripper2obj[np.newaxis, ...]
+
+        target_obj_pose = target_pose
+        target_gripper_pose = target_obj_pose @ gripper2obj
+
+    return target_gripper_pose
+
+
+
+
+