diff --git a/data_gen_dependencies/action/__init__.py b/data_gen_dependencies/action/__init__.py index 74f1133..788cebc 100644 --- a/data_gen_dependencies/action/__init__.py +++ b/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): diff --git a/data_gen_dependencies/action/base.py b/data_gen_dependencies/action/base.py index 2e8612e..d62a349 100644 --- a/data_gen_dependencies/action/base.py +++ b/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 diff --git a/data_gen_dependencies/action/grasp.py b/data_gen_dependencies/action/grasp.py index 74607e6..4857296 100644 --- a/data_gen_dependencies/action/grasp.py +++ b/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]) diff --git a/data_gen_dependencies/action/insert.py b/data_gen_dependencies/action/insert.py index 0a8a78b..6cd80c9 100644 --- a/data_gen_dependencies/action/insert.py +++ b/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) diff --git a/data_gen_dependencies/action/place.py b/data_gen_dependencies/action/place.py index 71dfb2d..0027382 100644 --- a/data_gen_dependencies/action/place.py +++ b/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]) diff --git a/data_gen_dependencies/action/pour.py b/data_gen_dependencies/action/pour.py index 79e6592..a08c6bd 100644 --- a/data_gen_dependencies/action/pour.py +++ b/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 diff --git a/data_gen_dependencies/action/press.py b/data_gen_dependencies/action/press.py new file mode 100644 index 0000000..be65c13 --- /dev/null +++ b/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']) \ No newline at end of file diff --git a/data_gen_dependencies/action/pull.py b/data_gen_dependencies/action/pull.py index 8ada4e6..cb2775d 100644 --- a/data_gen_dependencies/action/pull.py +++ b/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: diff --git a/data_gen_dependencies/action/pull_revolute.py b/data_gen_dependencies/action/pull_revolute.py new file mode 100644 index 0000000..b4cc9ba --- /dev/null +++ b/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 \ No newline at end of file diff --git a/data_gen_dependencies/action/push.py b/data_gen_dependencies/action/push.py index 610795b..75f374f 100644 --- a/data_gen_dependencies/action/push.py +++ b/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']) diff --git a/data_gen_dependencies/action/slide.py b/data_gen_dependencies/action/slide.py index 134c5fa..48d550d 100644 --- a/data_gen_dependencies/action/slide.py +++ b/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 diff --git a/data_gen_dependencies/action/twist.py b/data_gen_dependencies/action/twist.py new file mode 100644 index 0000000..4086aad --- /dev/null +++ b/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 \ No newline at end of file diff --git a/data_gen_dependencies/manip_solver.py b/data_gen_dependencies/manip_solver.py index 79bd279..e952fe2 100644 --- a/data_gen_dependencies/manip_solver.py +++ b/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)) diff --git a/data_gen_dependencies/object.py b/data_gen_dependencies/object.py index 61c8fff..7a35997 100644 --- a/data_gen_dependencies/object.py +++ b/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): diff --git a/data_gen_dependencies/omniagent.py b/data_gen_dependencies/omniagent.py index b6f14fb..08d7f72 100644 --- a/data_gen_dependencies/omniagent.py +++ b/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 diff --git a/task_gen_dependencies/object.py b/task_gen_dependencies/object.py index a58a16b..0231fed 100644 --- a/task_gen_dependencies/object.py +++ b/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]), },