65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
import numpy as np
|
|
|
|
from data_gen_dependencies.action.base import StageTemplate, simple_check_completion
|
|
|
|
|
|
class PlaceStage(StageTemplate):
|
|
def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, extra_params={}, **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])
|
|
self.extra_params = {} if extra_params is None else extra_params
|
|
self.use_pre_place = self.extra_params.get('use_pre_place', True)
|
|
self.generate_substage(target_pose)
|
|
|
|
|
|
|
|
|
|
def generate_substage(self, target_pose):
|
|
target_pose_canonical = target_pose
|
|
gripper_cmd = self.extra_params.get('gripper_state', 'open')
|
|
pre_place_direction = self.extra_params.get('pre_place_direction', 'z')
|
|
#print(gripper_cmd)
|
|
num_against = self.extra_params.get('against', 0)
|
|
if self.use_pre_place:
|
|
# moveTo pre-place position
|
|
transform_up = np.eye(4)
|
|
#print(transform_up)
|
|
if pre_place_direction == "x":
|
|
transform_up[0,3] = self.pre_transform_up
|
|
transform_up[2,3] = 0.02
|
|
elif pre_place_direction == "y":
|
|
transform_up[1,3] = self.pre_transform_up
|
|
else:
|
|
transform_up[2,3] = self.pre_transform_up
|
|
self.sub_stages.append([target_pose_canonical, None, transform_up, 'AvoidObs'])
|
|
|
|
# place
|
|
palce_transform_up = np.eye(4)
|
|
palce_transform_up[:3,3] = self.place_transform_up
|
|
if num_against > 0:
|
|
for index in range(num_against):
|
|
self.sub_stages.append([target_pose_canonical, None, transform_up, 'AvoidObs'])
|
|
self.pick_up_step = num_against + 1
|
|
self.sub_stages.append([target_pose_canonical, gripper_cmd, palce_transform_up, 'Simple'])
|
|
else:
|
|
palce_transform_up = np.eye(4)
|
|
palce_transform_up[:3,3] = self.place_transform_up
|
|
self.sub_stages.append([target_pose_canonical, None, palce_transform_up, 'AvoidObs'])
|
|
|
|
self.sub_stages.append([None, gripper_cmd, np.eye(4), 'Simple'])
|
|
|
|
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==0:
|
|
succ = simple_check_completion(goal_datapack, objects)
|
|
if succ:
|
|
self.step_id += 1
|
|
return succ
|
|
|
|
|
|
|