37 lines
2.2 KiB
Python
37 lines
2.2 KiB
Python
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, objects=[], **kwargs):
|
|
super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
|
|
self.pre_distance = -0.03
|
|
|
|
self.passive_obj = objects[passive_obj_id]
|
|
self.active_obj = objects[active_obj_id]
|
|
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 = self.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 = self.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'])
|
|
free_delta_pose = np.eye(4)
|
|
free_delta_pose[2,3] = -0.03
|
|
self.sub_stages.append([free_delta_pose, "open", np.eye(4), 'local_gripper']) |