137 lines
5.0 KiB
Python
137 lines
5.0 KiB
Python
import numpy as np
|
|
|
|
from data_gen_dependencies.action.base import StageTemplate, simple_check_completion, pose_difference
|
|
|
|
|
|
class PickStage(StageTemplate):
|
|
def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=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
|
|
|
|
self.extra_params = {} if extra_params is None else extra_params
|
|
self.pick_up_step = 999
|
|
self.generate_substage(grasp_pose)
|
|
|
|
|
|
def generate_substage(self, grasp_pose):
|
|
pick_up_distance = self.extra_params.get('pick_up_distance', 0.12)
|
|
pick_up_type = self.extra_params.get('pick_up_type', 'Simple')
|
|
pick_up_direction = self.extra_params.get('pick_up_direction', 'z')
|
|
num_against = self.extra_params.get('against', 0)
|
|
if self.use_pre_grasp:
|
|
pre_grasp_distance = self.extra_params.get('pre_grasp_distance', 0.08)
|
|
# sub-stage-0 moveTo pregrasp pose
|
|
pre_pose = np.array([
|
|
[1., 0, 0, 0],
|
|
[0, 1., 0, 0],
|
|
[0, 0, 1., -pre_grasp_distance],
|
|
[0, 0, 0, 1]])
|
|
pre_grasp_pose = grasp_pose @ pre_pose
|
|
self.sub_stages.append([pre_grasp_pose, None, np.eye(4), 'AvoidObs'])
|
|
# sub-stage-1 moveTo grasp pose
|
|
self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'Simple'])
|
|
self.pick_up_step = 2
|
|
elif num_against > 0:
|
|
for index in range(num_against):
|
|
self.sub_stages.append([grasp_pose, 'open', np.eye(4), 'AvoidObs'])
|
|
self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'AvoidObs'])
|
|
self.pick_up_step = num_against + 1
|
|
gripper_action = None
|
|
motion_type = pick_up_type
|
|
transform_up = np.eye(4)
|
|
if pick_up_direction == "x":
|
|
transform_up[0,3] = pick_up_distance
|
|
transform_up[2,3] = 0.02
|
|
elif pick_up_direction == "y":
|
|
transform_up[1,3] = pick_up_distance
|
|
else:
|
|
transform_up[2,3] = pick_up_distance
|
|
self.sub_stages.append([grasp_pose, gripper_action, transform_up, motion_type])
|
|
else:
|
|
# sub-stage-0 moveTo grasp pose
|
|
self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'AvoidObs'])
|
|
self.pick_up_step = 1
|
|
|
|
# pick-up
|
|
gripper_action = None
|
|
motion_type = pick_up_type
|
|
transform_up = np.eye(4)
|
|
if pick_up_direction == "x":
|
|
transform_up[0,3] = pick_up_distance
|
|
transform_up[2,3] = 0.02
|
|
elif pick_up_direction == "y":
|
|
transform_up[1,3] = pick_up_distance
|
|
else:
|
|
transform_up[2,3] = pick_up_distance
|
|
self.sub_stages.append([grasp_pose, gripper_action, transform_up, motion_type])
|
|
|
|
|
|
|
|
def check_completion(self, objects):
|
|
if self.__len__()==0:
|
|
return True
|
|
|
|
goal_datapack = [self.active_obj_id, self.passive_obj_id] + self.sub_stages[self.step_id]
|
|
|
|
succ = True
|
|
if self.step_id < self.pick_up_step:
|
|
succ = simple_check_completion(goal_datapack, objects, is_grasped=True)
|
|
|
|
if succ:
|
|
self.step_id += 1
|
|
return succ
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class GraspStage(PickStage):
|
|
def generate_substage(self, grasp_pose):
|
|
motion_type = self.extra_params.get('move_type', 'AvoidObs')
|
|
|
|
if self.extra_params.get('use_pre_grasp', False):
|
|
pre_pose = np.array([
|
|
[1., 0, 0, 0],
|
|
[0, 1., 0, 0],
|
|
[0, 0, 1., -0.08],
|
|
[0, 0, 0, 1]])
|
|
pre_grasp_pose = grasp_pose @ pre_pose
|
|
gripper_action = None
|
|
|
|
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])
|
|
|
|
|
|
|
|
|
|
|
|
def check_completion(self, objects):
|
|
if self.__len__()==0:
|
|
return True
|
|
|
|
goal_datapack = [self.active_obj_id, self.passive_obj_id] + self.sub_stages[self.step_id]
|
|
|
|
pos_threshold = self.extra_params.get('position_threshold', 0.06)
|
|
succ = simple_check_completion(goal_datapack, objects, pos_threshold=pos_threshold, is_grasped=True)
|
|
|
|
if self.step_id >= 2:
|
|
gripper_pose = objects['gripper'].obj_pose
|
|
grasp_pose = objects[self.passive_obj_id].obj_pose @ self.grasp_pose_canonical
|
|
pos_diff, _ = pose_difference(gripper_pose, grasp_pose)
|
|
grasp_succ = pos_diff < pos_threshold
|
|
succ = succ and grasp_succ
|
|
|
|
if succ:
|
|
self.step_id += 1
|
|
return succ
|
|
|
|
|
|
class HookStage(PickStage):
|
|
def generate_substage(self, grasp_pose, *args, **kwargs):
|
|
self.sub_stages.append([grasp_pose, None, np.eye(4), 'AvoidObs'])
|