hofee/gen_data_agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

finshi pull revolute action

Browse Source
This commit is contained in:
hofee
2025-09-26 15:25:17 +08:00
parent 330d903047
commit 65a39a4994
10 changed files with 8694 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
data_gen_dependencies/action/__init__.py
View File

@@ -7,6 +7,7 @@ from .pull import PullStage
from .push import PushStage
from .pull_revolute import PullRevoluteStage
from .press import PressPrismaticStage
from .twist import TwistStage
ACTION_STAGE = {
"grasp": GraspStage,
@@ -23,9 +24,14 @@ ACTION_STAGE = {
"push": PushStage,
'click': PressPrismaticStage,
'touch': PressPrismaticStage,
'press_prismatic': PressPrismaticStage,
"pull": PullStage,
# ---- Articulate ----
"pull_revolute": PullRevoluteStage,
"push_revolute": NotImplemented,
"pull_prismatic":PullStage,
"push_prismatic":PushStage,
"twist": TwistStage,
'press_prismatic': PressPrismaticStage,
}
def build_stage(action):

15
data_gen_dependencies/action/press.py
View File

@@ -3,23 +3,23 @@ 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):
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 = passive_obj
self.active_obj = active_obj
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 = passive_obj.joints_info[self.correspond_joint_id]
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 = passive_obj.joints_info[self.correspond_joint_id]
correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
self.generate_substage(target_pose, vector_direction)
@@ -31,4 +31,7 @@ class PressPrismaticStage(StageTemplate):
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'])
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'])

106
data_gen_dependencies/action/pull_revolute.py
View File

@@ -1,28 +1,36 @@
import os
import numpy as np
from scipy.spatial.transform import Rotation as R
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):
DELTA_DISTANCE = 0.003 # meter
def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects=[], **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
#import ipdb;ipdb.set_trace()
self.objects = objects
self.passive_obj = objects[passive_obj_id]
self.active_obj = objects[active_obj_id]
main_part_name= passive_element["part_id"]
self.obj_base = os.path.dirname(passive_obj_id)
main_part_id = os.path.join(self.obj_base, main_part_name)
if main_part_id not in objects:
raise ValueError(f"Cannot find main part {main_part_name} in objects")
self.root_id = os.path.join(self.obj_base, "root")
self.main_part_obj = objects[main_part_id]
self.root_obj = objects[self.root_id]
self.revolute_radius = np.linalg.norm(self.passive_obj.obj_pose[:3,3] - self.main_part_obj.obj_pose[:3,3])
self.revolute_joint_position = self.main_part_obj.obj_pose[:3,3]
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]
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.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
@@ -36,54 +44,82 @@ class PullRevoluteStage(StageTemplate):
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):
def generate_substage(self, gripper2obj, 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'])
trajectory = self.generate_trajectory(gripper2obj, vector_direction)
self.save_visualized_gripper_trajectory(trajectory, "pull_revolute_traj.txt")
self.sub_stages.append([trajectory, None, np.eye(4), 'Trajectory'])
self.sub_stages.append([None, "open", np.eye(4), 'Simple'])
free_delta_pose = np.eye(4)
free_delta_pose[2,3] = -self.DELTA_DISTANCE
free_delta_pose[2,3] = -0.03
self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
def generate_trajectory(self, init_vector_direction):
def generate_trajectory(self, gripper2passive_obj, 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)
rotation_axis = np.array(self.joint_axis, dtype=float)
rotation_axis /= np.linalg.norm(rotation_axis)
passive_obj_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.passive_obj.obj_pose
gripper_2_root = passive_obj_2_root @ gripper2passive_obj
pos0 = gripper_2_root[:3, 3]
R0 = gripper_2_root[:3, :3]
joint_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.main_part_obj.obj_pose
pivot = joint_2_root[:3, 3]
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
R_delta = R.from_rotvec(rotation_axis * angle).as_matrix()
pos_rot = R_delta @ (pos0 - pivot) + pivot
R_new = R_delta @ R0
waypoint_2_root = np.eye(4)
waypoint_2_root[:3, :3] = R_new
waypoint_2_root[:3, 3] = pos_rot
waypoint_2_obj = np.linalg.inv(passive_obj_2_root) @ waypoint_2_root
trajectory.append(waypoint_2_obj)
return np.asarray(trajectory)
def save_visualized_gripper_trajectory(self, gripper_trajectory, save_path):
#import ipdb; ipdb.set_trace()
N = len(gripper_trajectory)
gripper_pts = np.zeros((1, 60, 3))
for i in range(10):
gripper_pts[0, i] = np.array([0, 0.01, 0.002*i])
gripper_pts[0, i+10] = np.array([0, 0.01, 0.002*i + 0.001])
gripper_pts[0, i+20] = np.array([0, -0.01, 0.002*i])
gripper_pts[0, i+30] = np.array([0, -0.01, 0.002*i + 0.001])
gripper_pts[0, i+40] = np.array([0, -0.01 + 0.002 * i, 0])
for i in range(5):
gripper_pts[0, i+50] = np.array([0, 0, -0.002*i])
gripper_pts[0, i+55] = np.array([0, 0, -0.002*i + 0.001])
gripper_pts -= np.array([0, 0, 0.002*10 + 0.001])
gripper_pts = gripper_pts.repeat(N, axis=0)
gripper_trajectory = gripper_trajectory.reshape(N, 1, 4, 4)
gripper_trajectory = gripper_trajectory.repeat(60, axis=1)
gripper_pts = gripper_pts.reshape(-1, 3)
gripper_trajectory = gripper_trajectory.reshape(-1, 4, 4)
gripper_pts_2_obj = (gripper_trajectory[:, :3, :3] @ gripper_pts[..., None]).squeeze(-1)+ gripper_trajectory[:, :3, 3]
np.savetxt(save_path, gripper_pts_2_obj, fmt="%.6f")
def axis_angle_to_matrix(self, axis, angle):

146
data_gen_dependencies/action/push_revolute.py Normal file
View File

@@ -0,0 +1,146 @@
import os
import numpy as np
from scipy.spatial.transform import Rotation as R
from data_gen_dependencies.action.base import StageTemplate
class PushRevoluteStage(StageTemplate):
DELTA_DISTANCE = 0.003 # meter
def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects=[], **kwargs):
super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
#import ipdb;ipdb.set_trace()
self.objects = objects
self.passive_obj = objects[passive_obj_id]
self.active_obj = objects[active_obj_id]
main_part_name= passive_element["part_id"]
self.obj_base = os.path.dirname(passive_obj_id)
main_part_id = os.path.join(self.obj_base, main_part_name)
if main_part_id not in objects:
raise ValueError(f"Cannot find main part {main_part_name} in objects")
self.root_id = os.path.join(self.obj_base, "root")
self.main_part_obj = objects[main_part_id]
self.root_obj = objects[self.root_id]
self.revolute_radius = np.linalg.norm(self.passive_obj.obj_pose[:3,3] - self.main_part_obj.obj_pose[:3,3])
self.revolute_joint_position = self.main_part_obj.obj_pose[:3,3]
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 = 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.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"
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, gripper2obj, vector_direction):
vector_direction = vector_direction / np.linalg.norm(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}')
trajectory = self.generate_trajectory(gripper2obj, vector_direction)
self.save_visualized_gripper_trajectory(trajectory, "pull_revolute_traj.txt")
self.sub_stages.append([trajectory, None, np.eye(4), 'Trajectory'])
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, gripper2passive_obj, init_vector_direction):
trajectory = []
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, dtype=float)
rotation_axis /= np.linalg.norm(rotation_axis)
passive_obj_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.passive_obj.obj_pose
gripper_2_root = passive_obj_2_root @ gripper2passive_obj
pos0 = gripper_2_root[:3, 3]
R0 = gripper_2_root[:3, :3]
joint_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.main_part_obj.obj_pose
pivot = joint_2_root[:3, 3]
for i in range(1, num_steps + 1):
angle = i * delta_theta
R_delta = R.from_rotvec(rotation_axis * angle).as_matrix()
pos_rot = R_delta @ (pos0 - pivot) + pivot
R_new = R_delta @ R0
waypoint_2_root = np.eye(4)
waypoint_2_root[:3, :3] = R_new
waypoint_2_root[:3, 3] = pos_rot
waypoint_2_obj = np.linalg.inv(passive_obj_2_root) @ waypoint_2_root
trajectory.append(waypoint_2_obj)
return np.asarray(trajectory)
def save_visualized_gripper_trajectory(self, gripper_trajectory, save_path):
#import ipdb; ipdb.set_trace()
N = len(gripper_trajectory)
gripper_pts = np.zeros((1, 60, 3))
for i in range(10):
gripper_pts[0, i] = np.array([0, 0.01, 0.002*i])
gripper_pts[0, i+10] = np.array([0, 0.01, 0.002*i + 0.001])
gripper_pts[0, i+20] = np.array([0, -0.01, 0.002*i])
gripper_pts[0, i+30] = np.array([0, -0.01, 0.002*i + 0.001])
gripper_pts[0, i+40] = np.array([0, -0.01 + 0.002 * i, 0])
for i in range(5):
gripper_pts[0, i+50] = np.array([0, 0, -0.002*i])
gripper_pts[0, i+55] = np.array([0, 0, -0.002*i + 0.001])
gripper_pts -= np.array([0, 0, 0.002*10 + 0.001])
gripper_pts = gripper_pts.repeat(N, axis=0)
gripper_trajectory = gripper_trajectory.reshape(N, 1, 4, 4)
gripper_trajectory = gripper_trajectory.repeat(60, axis=1)
gripper_pts = gripper_pts.reshape(-1, 3)
gripper_trajectory = gripper_trajectory.reshape(-1, 4, 4)
gripper_pts_2_obj = (gripper_trajectory[:, :3, :3] @ gripper_pts[..., None]).squeeze(-1)+ gripper_trajectory[:, :3, 3]
np.savetxt(save_path, gripper_pts_2_obj, fmt="%.6f")
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

35
data_gen_dependencies/action/twist.py
View File

@@ -5,20 +5,24 @@ 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):
def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects = [], **kwargs):
super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
self.passive_obj = passive_obj
self.active_obj = active_obj
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)
if self.joint_position_threshold is None:
self.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]
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"]
if self.joint_lower_limit is None or self.joint_upper_limit is None:
self.twist_degree_range = np.pi
else:
self.twist_degree_range = abs(self.joint_upper_limit - self.joint_lower_limit)
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"
assert self.joint_type == 'continuous', "joint_type must be continuous for twist 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
@@ -34,9 +38,20 @@ class TwistStage(StageTemplate):
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'])
delta_twist_pose = self.axis_angle_to_matrix(np.asarray([0,0,1]), self.twist_degree_range * self.joint_position_threshold * self.joint_direction)
import ipdb;ipdb.set_trace()
self.sub_stages.append([delta_twist_pose, None, np.eye(4), 'local_gripper'])
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: