694 lines
30 KiB
Python
694 lines
30 KiB
Python
# -*- coding: utf-8 -*-
|
||
import copy
|
||
import json
|
||
import os
|
||
from sqlite3 import dbapi2
|
||
from this import d
|
||
import time
|
||
|
||
from pyboot.utils.log import Log
|
||
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):
|
||
Log.info('Loading task solution')
|
||
# task_info = json.load(open('%s/task_info.json'%task_dir, 'rb'))
|
||
stages = task_info['stages']
|
||
|
||
objects = {
|
||
'gripper': OmniObject('gripper')
|
||
}
|
||
|
||
for i, obj_info in enumerate(task_info['objects']):
|
||
Log.info(f'Loading object: {obj_info["object_id"]} [Loading task solution: {i+1}/{len(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 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
|
||
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
|
||
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', 'pull_revolute', "twist"]
|
||
|
||
def _load_element(obj, type):
|
||
|
||
if action in ['pick', 'hook']:
|
||
action_mapped = 'grasp'
|
||
elif action in ['pull_revolute']:
|
||
action_mapped = 'pull'
|
||
elif action in ['push_revolute']:
|
||
action_mapped = 'push'
|
||
elif action in ["press_prismatic"]:
|
||
action_mapped = 'press'
|
||
else:
|
||
action_mapped = action
|
||
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] 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
|
||
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]:
|
||
_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', 'press_prismatic']:
|
||
gripper2obj = current_gripper_pose
|
||
|
||
elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute", "twist", "push_revolute"]:
|
||
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
|
||
if stages[0]['action'] == "pull_revolute":
|
||
import ipdb;ipdb.set_trace()
|
||
if stages[0]['action'] == "push_revolute":
|
||
import ipdb;ipdb.set_trace()
|
||
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]
|
||
|
||
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 []
|
||
stage_idx = 0
|
||
for stage in stages:
|
||
stage_idx += 1
|
||
extra_params = stage.get('extra_params', {})
|
||
arm = extra_params.get('arm', 'right')
|
||
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]
|
||
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), objects=objects)
|
||
elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'move',
|
||
'reset', "pull_revolute","push_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,
|
||
objects=objects,
|
||
target_pose=gripper2obj
|
||
)
|
||
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:
|
||
Log.warning(f'{action}: No target_obj_pose can pass isaac curobo IK')
|
||
import ipdb;ipdb.set_trace()
|
||
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', {}),
|
||
objects=objects
|
||
)
|
||
break
|
||
|
||
if substages is None:
|
||
Log.warning(f'{action}:substage is None')
|
||
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
|
||
|
||
|
||
|
||
|
||
|