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:

15
data_gen_dependencies/manip_solver.py
View File

@@ -286,7 +286,7 @@ 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', 'pull_revolute']
single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe', 'pull_revolute', "twist"]
def _load_element(obj, type):
@@ -345,7 +345,7 @@ def select_obj(objects, stages, robot):
if stages[0]['action'] in ['push', 'reset', 'press_prismatic']:
gripper2obj = current_gripper_pose
elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute"]:
elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute", "twist"]:
action = stages[0]['action']
''' 筛掉无IK解的grasp pose '''
@@ -375,7 +375,6 @@ def select_obj(objects, stages, robot):
if len(grasp_poses) == 0:
print(action, 'No grasp_gripper_pose can pass Isaac IK')
return []
''' 基于有IK解的grasp pose分数,选择最优的passive primitive element,同时选出最优的一个grasp pose'''
if grasp_stage_id is not None:
next_stage_id = grasp_stage_id + 1
@@ -404,7 +403,7 @@ def select_obj(objects, stages, robot):
active_elements = [{'xyz': xyz_canonical, 'direction': direction_canonical}]
# import ipdb; ipdb.set_trace()
#import ipdb; ipdb.set_trace()
t0 = time.time()
element_ik_score = []
grasp_pose_ik_score = []
@@ -549,7 +548,7 @@ def generate_action_stages(objects, all_stages, robot):
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), active_obj=active_obj, passive_obj=passive_obj)
gripper2obj, extra_params=stage.get('extra_params', None), objects=objects)
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))]
@@ -557,8 +556,7 @@ def generate_action_stages(objects, all_stages, robot):
active_obj_id=active_obj_id,
passive_obj_id=passive_obj_id,
passive_element=passive_element,
active_obj=active_obj,
passive_obj=passive_obj,
objects=objects,
target_pose=gripper2obj
)
else:
@@ -642,8 +640,7 @@ def generate_action_stages(objects, all_stages, robot):
vector_direction=passive_element['direction'],
passive_element=passive_element,
extra_params=stage.get('extra_params', {}),
active_obj=active_obj,
passive_obj=passive_obj
objects=objects
)
break

4
data_gen_dependencies/object.py
View File

@@ -67,9 +67,9 @@ class OmniObject:
"xyz": np.array([0, 0, 0]),
"direction": np.array([0, 0, 0.08]),
},
'rotate': {
'twist': {
"part": "finger edge",
"task": "rotate",
"task": "twist",
"xyz": np.array([0, 0, 0]),
"direction": np.array([0, 0, 0.08]),
}

46
data_gen_dependencies/omniagent.py
View File

@@ -82,26 +82,27 @@ class Agent(BaseAgent):
target_lookat_point = np.mean(np.stack(target_lookat_point), axis=0)
self.robot.client.SetTargetPoint(target_lookat_point.tolist())
Log.info('Setting material [Generating layout: 3/5]')
#Log.info('Setting material [Generating layout: 3/5]')
Log.warning('Pass setting material [Generating layout: 3/5]')
''' Set material '''
material_infos = []
if "object_with_material" in task_info:
for key in task_info['object_with_material']:
material_infos += task_info['object_with_material'][key]
if len(material_infos):
self.robot.client.SetMaterial(material_infos)
time.sleep(0.3)
Log.info('Setting light [Generating layout: 4/5]')
# material_infos = []
# if "object_with_material" in task_info:
# for key in task_info['object_with_material']:
# material_infos += task_info['object_with_material'][key]
# if len(material_infos):
# self.robot.client.SetMaterial(material_infos)
# time.sleep(0.3)
#Log.info('Setting light [Generating layout: 4/5]')
Log.warning('Pass setting light [Generating layout: 4/5]')
''' Set light '''
light_infos = []
if "lights" in task_info:
for key in task_info['lights']:
light_infos += task_info['lights'][key]
if len(light_infos):
self.robot.client.SetLight(light_infos)
time.sleep(0.3)
# light_infos = []
# if "lights" in task_info:
# for key in task_info['lights']:
# light_infos += task_info['lights'][key]
# if len(light_infos):
# self.robot.client.SetLight(light_infos)
# time.sleep(0.3)
Log.info('Setting camera [Generating layout: 5/5]')
''' Set camera'''
@@ -228,7 +229,7 @@ class Agent(BaseAgent):
#print('Reset pose:', self.robot.reset_pose)
task_info = json.load(open(task_file, 'rb'))
import ipdb; ipdb.set_trace()
stages, objects = load_task_solution(task_info)
objects = self.update_objects(objects)
split_stages = split_grasp_stages(stages)
@@ -257,6 +258,7 @@ class Agent(BaseAgent):
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)}]')
@@ -270,11 +272,15 @@ class Agent(BaseAgent):
interp_pose[:3, 3] = curr_pose[:3, 3] + (init_pose[:3, 3] - curr_pose[:3, 3]) * 0.25
success = self.robot.move_pose(self.robot.reset_pose[arm], type='AvoidObs', arm=arm, block=True)
continue
substage_id = 0
substage_len = len(substages)
while len(substages):
objects = self.update_objects(objects, arm=arm)
substage_id += 1
target_gripper_pose, motion_type, gripper_action, arm = substages.get_action(objects)
target_gripper_pos = None if target_gripper_pose is None else target_gripper_pose[:3, 3]
Log.info(f'-- ({action}) Sub-stage {substage_id}/{substage_len}: target_pos={target_gripper_pos}, motion_type={motion_type}, gripper_action={gripper_action}, arm={arm}')
arm = extra_params.get('arm', 'right')
self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
self.attached_obj_id is not None, arm=arm,