Spaces:
Paused
Paused
| import numpy as np | |
| import sapien | |
| import gymnasium as gym | |
| import torch | |
| from historybench.HistoryBench_env import * | |
| from . import reset_panda | |
| from mani_skill.examples.motionplanning.panda.motionplanner import \ | |
| PandaArmMotionPlanningSolver | |
| from mani_skill.examples.motionplanning.base_motionplanner.utils import ( | |
| compute_grasp_info_by_obb, | |
| get_actor_obb, | |
| ) | |
| from mani_skill.utils.geometry.rotation_conversions import ( | |
| euler_angles_to_matrix, | |
| matrix_to_quaternion, | |
| quaternion_multiply, | |
| ) | |
| from historybench.HistoryBench_env.util import * | |
| import random | |
| from ...logging_utils import logger | |
| # Probability for deliberately triggering a failed hover before pickup. | |
| FAILED_HOVER_PROB = 0.03 | |
| def grasp_and_lift_peg_side(env, planner,obj): | |
| planner.open_gripper() | |
| """Move to the peg tail, close gripper, lift, and keep holding.""" | |
| pose = obj.pose | |
| lift_height=0.2 | |
| grasp_pose_p = pose.p | |
| if isinstance(grasp_pose_p, torch.Tensor): | |
| grasp_pose_p = grasp_pose_p.detach().cpu().numpy() | |
| grasp_pose_p = np.asarray(grasp_pose_p, dtype=np.float32).reshape(-1) | |
| grasp_pose_q = pose.q | |
| if isinstance(grasp_pose_q, torch.Tensor): | |
| grasp_pose_q = grasp_pose_q.detach().cpu().numpy() | |
| grasp_pose_q = np.asarray(grasp_pose_q, dtype=np.float32).reshape(-1) | |
| flip_angles = torch.tensor([[np.pi, 0.0, 0.0]], dtype=torch.float32) | |
| flip_matrix = euler_angles_to_matrix(flip_angles, convention="XYZ") | |
| flip_quat = matrix_to_quaternion(flip_matrix)[0] | |
| grasp_pose_q_tensor = torch.from_numpy(grasp_pose_q).to( | |
| dtype=torch.float32, device=flip_quat.device | |
| ) | |
| rotated_quat = quaternion_multiply( | |
| grasp_pose_q_tensor.unsqueeze(0), flip_quat.unsqueeze(0) | |
| ).squeeze(0) | |
| grasp_pose_q = rotated_quat.detach().cpu().numpy().astype(np.float32) | |
| norm = np.linalg.norm(grasp_pose_q) | |
| if norm > 0: | |
| grasp_pose_q /= norm | |
| lifted_pose_p = grasp_pose_p.copy() | |
| lifted_pose_p[2] = lift_height | |
| planner.move_to_pose_with_screw(sapien.Pose(p=lifted_pose_p, q=grasp_pose_q)) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=grasp_pose_p, q=grasp_pose_q)) | |
| planner.close_gripper() | |
| lifted_pose_p = grasp_pose_p.copy() | |
| lifted_pose_p[2] = lift_height | |
| planner.move_to_pose_with_screw(sapien.Pose(p=lifted_pose_p, q=grasp_pose_q)) | |
| planner.close_gripper() | |
| current_grasp_pose=sapien.Pose(p=grasp_pose_p, q=grasp_pose_q) | |
| env.current_grasp_pose = current_grasp_pose | |
| def return_to_original_pose(env,planner,current_grasp_pose): | |
| grasp_pose_p=current_grasp_pose.p | |
| grasp_pose_q=current_grasp_pose.q | |
| lifted_pose_p = grasp_pose_p.copy() | |
| lifted_pose_p[2] = 0.2 | |
| planner.move_to_pose_with_screw(sapien.Pose(p=lifted_pose_p, q=grasp_pose_q)) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=grasp_pose_p, q=grasp_pose_q)) | |
| planner.open_gripper() | |
| # def insert_peg(env, planner,current_grasp_pose,peg_init_pose,direction,obj): | |
| # """Insert the peg into the box.""" | |
| # grasp_pose = current_grasp_pose | |
| # insert_pose = env.box.pose * peg_init_pose.inv() * grasp_pose | |
| # if obj==-1: | |
| # insert_pose=insert_pose*sapien.Pose(q=[0, 0, 0, 1]) | |
| # if direction==-1: | |
| # insert_pose=insert_pose*sapien.Pose(q=[0, 0, 0, 1]) | |
| # # pre_insert_pose=sapien.Pose(p=[-0.1,0.15,0.2],q=insert_pose.q.tolist()[0]) | |
| # # pre_insert_pose=pre_insert_pose*sapien.Pose([0.05*direction, 0, 0]) | |
| # # planner.move_to_pose_with_screw(pre_insert_pose) | |
| # if obj==-1: | |
| # if direction==-1: | |
| # bias=0.1 | |
| # else: | |
| # bias=0 | |
| # insert_pose_p = np.asarray(insert_pose.p, dtype=np.float32).reshape(-1) | |
| # insert_pose_q = np.asarray(insert_pose.q, dtype=np.float32).reshape(-1) | |
| # pre_pose = sapien.Pose( | |
| # p=[insert_pose_p[0] + 0.2 + bias, insert_pose_p[1], 0.2], | |
| # q=insert_pose_q | |
| # ) | |
| # planner.move_to_pose_with_screw(pre_pose) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([0.2+bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([0.15+bias, 0, 0])) | |
| # for i in range(5): | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([0.05+bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.1+bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([0.15+bias, 0, 0])) | |
| # else: | |
| # pass | |
| # if direction==-1: | |
| # bias=0.1 | |
| # else: | |
| # bias=0 | |
| # dx = -0.2 - bias | |
| # legacy comment | |
| # rel = insert_pose * sapien.Pose([dx, 0, 0]) | |
| # rel_p = np.asarray(rel.p, dtype=np.float32).reshape(-1) | |
| # rel_q = np.asarray(rel.q, dtype=np.float32).reshape(-1) | |
| # legacy comment | |
| # ready = sapien.Pose(p=[rel_p[0], rel_p[1], 0.2], q=rel_q) | |
| # planner.move_to_pose_with_screw(ready) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.2-bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.15-bias, 0, 0])) | |
| # for i in range(5): | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.05-bias, 0, 0])) | |
| # # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.05-bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([0.1-bias, 0, 0])) | |
| # planner.move_to_pose_with_screw(insert_pose * sapien.Pose([-0.15-bias, 0, 0])) | |
| # insert_pose = env.box.pose * insert_obj.pose.inv() * env.agent.tcp.pose | |
| # def insert_peg(env, planner,direction,obj,insert_obj=None): | |
| # """Insert the peg into the box.""" | |
| # if insert_obj is None: | |
| # raise ValueError("insert_obj must be provided to compute the insert pose.") | |
| # def _compute_insert_pose(): | |
| # pose = env.box.pose * insert_obj.pose.inv() * env.agent.tcp.pose | |
| # if obj == -1: | |
| # pose = pose * sapien.Pose(q=[0, 0, 0, 1]) | |
| # if direction == -1: | |
| # pose = pose * sapien.Pose(q=[0, 0, 0, 1]) | |
| # return pose | |
| # def _pose_components(): | |
| # pose = _compute_insert_pose() | |
| # pose_p = np.asarray(pose.p, dtype=np.float32).reshape(-1) | |
| # pose_q = np.asarray(pose.q, dtype=np.float32).reshape(-1) | |
| # return pose, pose_p, pose_q | |
| # def _move_with_offset(offset): | |
| # # Recompute insert pose each time before moving to stay aligned with the box. | |
| # current_pose = _compute_insert_pose() | |
| # planner.move_to_pose_with_screw(current_pose * sapien.Pose(offset)) | |
| # ########################## | |
| # if obj==-1: | |
| # _, insert_pose_p, insert_pose_q = _pose_components() | |
| # pre_pose = sapien.Pose( | |
| # p=[0, 0, 0.2], | |
| # q=insert_pose_q | |
| # ) | |
| # legacy comment | |
| # planner.move_to_pose_with_screw(pre_pose) | |
| # _move_with_offset([0.2 , 0, 0]) | |
| # _move_with_offset([0.15 , 0, 0]) | |
| # for i in range(5): | |
| # _move_with_offset([0.05 , 0, 0]) | |
| # _move_with_offset([-0.1 , 0, 0]) | |
| # _move_with_offset([0.15 , 0, 0]) | |
| # else:#obj=1 | |
| # legacy comment | |
| # _, insert_pose_p, insert_pose_q = _pose_components() | |
| # pre_pose = sapien.Pose( | |
| # p=[0,0, 0.2], | |
| # q=insert_pose_q | |
| # ) | |
| # legacy comment | |
| # planner.move_to_pose_with_screw(pre_pose) | |
| # _move_with_offset([-0.2 , 0, 0]) | |
| # _move_with_offset([-0.15 , 0, 0]) | |
| # for i in range(5): | |
| # _move_with_offset([-0.05 , 0, 0]) | |
| # _move_with_offset([0.1 , 0, 0]) | |
| # _move_with_offset([-0.15 , 0, 0]) | |
| def insert_peg(env, planner,direction,obj,insert_obj=None,cut_retreat=False): | |
| """Insert the peg into the box.""" | |
| if insert_obj is None: | |
| raise ValueError("insert_obj must be provided to compute the insert pose.") | |
| def _compute_insert_pose(): | |
| pose = env.box.pose * insert_obj.pose.inv() * env.agent.tcp.pose | |
| if obj == -1: | |
| pose = pose * sapien.Pose(q=[0, 0, 0, 1]) | |
| if direction == -1: | |
| pose = pose * sapien.Pose(q=[0, 0, 0, 1]) | |
| return pose | |
| def _pose_components(): | |
| pose = _compute_insert_pose() | |
| pose_p = np.asarray(pose.p, dtype=np.float32).reshape(-1) | |
| pose_q = np.asarray(pose.q, dtype=np.float32).reshape(-1) | |
| return pose, pose_p, pose_q | |
| def _move_with_offset(offset): | |
| # Recompute insert pose each time before moving to stay aligned with the box. | |
| current_pose = _compute_insert_pose() | |
| offset_vec = np.asarray(offset, dtype=np.float32).reshape(-1).copy() | |
| if obj == 1 and direction == -1 and offset_vec[0] < 0: | |
| # Tail grasp + left insert: compensate head-tail gap so we still move far enough. | |
| relative_pose = insert_obj.pose.inv() * env.agent.tcp.pose | |
| relative_p = np.asarray(relative_pose.p, dtype=np.float32).reshape(-1) | |
| offset_vec[0] += relative_p[0] | |
| planner.move_to_pose_with_screw(current_pose * sapien.Pose(offset_vec.tolist())) | |
| def _move_with_offset_with_break(offset): | |
| """Move with an interrupt check. Interrupt when elapsed_steps > end_steps + 3; if end_steps is None, execute directly.""" | |
| end_steps = getattr(env, "end_steps", None) | |
| while end_steps is None or int(getattr(env, "elapsed_steps", 0)) <= end_steps + 3: | |
| # legacy comment | |
| current_pose = _compute_insert_pose() | |
| offset_vec = np.asarray(offset, dtype=np.float32) | |
| target_pose = current_pose * sapien.Pose(offset_vec.tolist()) | |
| # legacy comment | |
| pose_for_plan = planner._transform_pose_for_planning(target_pose) | |
| pose_p = np.asarray(pose_for_plan.p, dtype=np.float32).reshape(-1) | |
| pose_q = np.asarray(pose_for_plan.q, dtype=np.float32).reshape(-1) | |
| result = planner.planner.plan_screw( | |
| np.concatenate([pose_p, pose_q]), | |
| planner.robot.get_qpos().cpu().numpy()[0], | |
| time_step=planner.base_env.control_timestep, | |
| use_point_cloud=planner.use_point_cloud, | |
| ) | |
| if result["status"] != "Success": | |
| return False | |
| # legacy comment | |
| n_step = result["position"].shape[0] | |
| for i in range(n_step): | |
| if end_steps is not None and int(getattr(env, "elapsed_steps", 0)) > end_steps + 3: | |
| logger.debug("break early") | |
| return True # legacy comment | |
| qpos = result["position"][i] | |
| if planner.control_mode == "pd_joint_pos_vel": | |
| qvel = result["velocity"][i] | |
| action = np.hstack([qpos, qvel, planner.gripper_state]) | |
| else: | |
| action = np.hstack([qpos, planner.gripper_state]) | |
| planner.env.step(action) | |
| planner.elapsed_steps += 1 | |
| return True # legacy comment | |
| return True # legacy comment | |
| ########################## | |
| if obj==-1: | |
| _, insert_pose_p, insert_pose_q = _pose_components() | |
| _move_with_offset([0.2 , 0, -0.15]) | |
| _move_with_offset([0.2 , 0, 0]) | |
| _move_with_offset([0.15 , 0, 0]) | |
| # for i in range(5): | |
| # _move_with_offset([0.05 , 0, 0]) | |
| if cut_retreat!=True: | |
| _move_with_offset([-0.05 , 0, 0]) | |
| else: | |
| logger.debug(f"cut_retreat mode (obj=-1): elapsed_steps={int(getattr(env, 'elapsed_steps', 0))}, end_steps={env.end_steps}") | |
| _move_with_offset_with_break([-0.05, 0, 0]) | |
| else:#obj=1 | |
| _, insert_pose_p, insert_pose_q = _pose_components() | |
| _move_with_offset([-0.2 , 0, -0.15]) | |
| _move_with_offset([-0.2 , 0, 0]) | |
| _move_with_offset([-0.15 , 0, 0]) | |
| # for i in range(5): | |
| # _move_with_offset([-0.05 , 0, 0]) | |
| if cut_retreat!=True: | |
| _move_with_offset([-0.05 , 0, 0]) | |
| else: | |
| logger.debug(f"cut_retreat mode (obj=1): elapsed_steps={int(getattr(env, 'elapsed_steps', 0))}, end_steps={env.end_steps}") | |
| _move_with_offset_with_break([-0.05, 0, 0]) | |
| def _zero_action_for_space(space): | |
| if isinstance(space, gym.spaces.Box): | |
| return np.zeros(space.shape, dtype=space.dtype) | |
| if isinstance(space, gym.spaces.Dict): | |
| return {k: _zero_action_for_space(subspace) for k, subspace in space.spaces.items()} | |
| sample = space.sample() | |
| if isinstance(sample, np.ndarray): | |
| return np.zeros_like(sample) | |
| if isinstance(sample, dict): | |
| return {k: np.zeros_like(v) for k, v in sample.items()} | |
| raise NotImplementedError("Unsupported action space type for zero action generation") | |
| def _flag_to_bool(flag): | |
| if flag is None: | |
| return False | |
| if isinstance(flag, (bool, np.bool_, np.bool8)): | |
| return bool(flag) | |
| if isinstance(flag, torch.Tensor): | |
| return bool(flag.detach().cpu().bool().any()) | |
| if isinstance(flag, np.ndarray): | |
| return bool(flag.any()) | |
| return bool(flag) | |
| def solve_liftup_Xdistance(env,planner,distance): | |
| original_pose = env.agent.tcp.pose | |
| lift_pose_p=original_pose.p.tolist()[0] | |
| lift_pose_q=original_pose.q.tolist()[0] | |
| lift_pose_p[2]+=distance | |
| planner.move_to_pose_with_screw(sapien.Pose(p=lift_pose_p,q=lift_pose_q)) | |
| # def solve_push_to_target(env, planner, obj=None, target=None): | |
| # planner.open_gripper() | |
| # FINGER_LENGTH = 0.025 | |
| # env = env.unwrapped | |
| # legacy comment | |
| # obj_pos = obj.pose.sp.p if hasattr(obj.pose.sp.p, '__iter__') else np.array(obj.pose.sp.p) | |
| # target_pos = target.pose.sp.p if hasattr(target.pose.sp.p, '__iter__') else np.array(target.pose.sp.p) | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # push_direction_3d = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # x_axis = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| # z_axis = np.array([0, 0, -1]) | |
| # legacy comment | |
| # y_axis = np.cross(z_axis, x_axis) | |
| # y_norm = np.linalg.norm(y_axis) | |
| # if y_norm < 1e-6: | |
| # raise ValueError("Push direction invalid; cannot construct gripper frame.") | |
| # y_axis = y_axis / y_norm | |
| # legacy comment | |
| # rotation_matrix = np.column_stack([x_axis, y_axis, z_axis]) | |
| # legacy comment | |
| # rotation_matrix_torch = torch.from_numpy(rotation_matrix).float().unsqueeze(0) | |
| # push_quat = matrix_to_quaternion(rotation_matrix_torch)[0] | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # # -------------------------------------------------------------------------- # | |
| # if target_pos[0] < obj_pos[0] : | |
| # z_rotation_angles = torch.deg2rad( | |
| # torch.tensor([0.0, 0.0, 180], dtype=torch.float32) | |
| # ) | |
| # z_rotation_matrix = euler_angles_to_matrix(z_rotation_angles, convention="XYZ") | |
| # z_rotation_quat = matrix_to_quaternion(z_rotation_matrix.unsqueeze(0))[0] | |
| # legacy comment | |
| # push_quat = quaternion_multiply( | |
| # push_quat.unsqueeze(0), z_rotation_quat.unsqueeze(0) | |
| # )[0] | |
| # legacy comment | |
| # push_pose = sapien.Pose(p=obj_pos, q=push_quat.detach().cpu().numpy()) | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # start_pos = obj_pos - push_direction_3d * offset_distance | |
| # legacy comment | |
| # reach_pose_q = push_pose.q.tolist() if hasattr(push_pose.q, 'tolist') else list(push_pose.q) | |
| # legacy comment | |
| # #planner.move_to_pose_with_screw(sapien.Pose(p=[0,0,0.1], q=reach_pose_q)) | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=start_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| # planner.close_gripper() | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # end_pos = target_pos.copy() | |
| # legacy comment | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=end_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| # planner.open_gripper() | |
| def solve_push_to_target(env, planner, obj=None, target=None): | |
| planner.open_gripper() | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # legacy comment | |
| obj_pos = obj.pose.sp.p if hasattr(obj.pose.sp.p, '__iter__') else np.array(obj.pose.sp.p) | |
| target_pos = target.pose.sp.p if hasattr(target.pose.sp.p, '__iter__') else np.array(target.pose.sp.p) | |
| # legacy comment | |
| push_direction_xy = target_pos[:2] - obj_pos[:2] # legacy comment | |
| push_direction_xy = push_direction_xy / np.linalg.norm(push_direction_xy) # legacy comment | |
| # legacy comment | |
| push_direction_3d = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| x_axis = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| z_axis = np.array([0, 0, -1]) | |
| # legacy comment | |
| y_axis = np.cross(z_axis, x_axis) | |
| y_norm = np.linalg.norm(y_axis) | |
| if y_norm < 1e-6: | |
| raise ValueError("Push direction invalid; cannot construct gripper frame.") | |
| y_axis = y_axis / y_norm | |
| # legacy comment | |
| rotation_matrix = np.column_stack([x_axis, y_axis, z_axis]) | |
| # legacy comment | |
| rotation_matrix_torch = torch.from_numpy(rotation_matrix).float().unsqueeze(0) | |
| push_quat = matrix_to_quaternion(rotation_matrix_torch)[0] | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| if target_pos[0] < obj_pos[0] : | |
| z_rotation_angles = torch.deg2rad( | |
| torch.tensor([0.0, 0.0, 180], dtype=torch.float32) | |
| ) | |
| z_rotation_matrix = euler_angles_to_matrix(z_rotation_angles, convention="XYZ") | |
| z_rotation_quat = matrix_to_quaternion(z_rotation_matrix.unsqueeze(0))[0] | |
| # legacy comment | |
| push_quat = quaternion_multiply( | |
| push_quat.unsqueeze(0), z_rotation_quat.unsqueeze(0) | |
| )[0] | |
| # legacy comment | |
| push_pose = sapien.Pose(p=obj_pos, q=push_quat.detach().cpu().numpy()) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| offset_distance = 0.05 # legacy comment | |
| start_pos = obj_pos - push_direction_3d * offset_distance | |
| start_pos[2] = push_pose.p[2] # legacy comment | |
| reach_pose_q = push_pose.q.tolist() if hasattr(push_pose.q, 'tolist') else list(push_pose.q) | |
| # legacy comment | |
| #planner.move_to_pose_with_screw(sapien.Pose(p=[0,0,0.1], q=reach_pose_q)) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=start_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| planner.close_gripper() | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| end_pos = target_pos.copy() - push_direction_3d * env.cube_half_size | |
| end_pos[2] = start_pos[2] # legacy comment | |
| planner.move_to_pose_with_screw(sapien.Pose(p=end_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| planner.open_gripper() | |
| # def solve_push_to_target_with_peg(env, planner, obj=None, target=None,direction=None,obj_flag=None): | |
| # FINGER_LENGTH = 0.025 | |
| # env = env.unwrapped | |
| # legacy comment | |
| # obj_pos = obj.pose.sp.p if hasattr(obj.pose.sp.p, '__iter__') else np.array(obj.pose.sp.p) | |
| # target_pos = target.pose.sp.p if hasattr(target.pose.sp.p, '__iter__') else np.array(target.pose.sp.p) | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # push_direction_3d = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # x_axis = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # legacy comment | |
| # z_axis = np.array([0, 0, -1]) | |
| # legacy comment | |
| # y_axis = np.cross(z_axis, x_axis) | |
| # y_norm = np.linalg.norm(y_axis) | |
| # if y_norm < 1e-6: | |
| # raise ValueError("Push direction invalid; cannot construct gripper frame.") | |
| # y_axis = y_axis / y_norm | |
| # legacy comment | |
| # rotation_matrix = np.column_stack([x_axis, y_axis, z_axis]) | |
| # legacy comment | |
| # rotation_matrix_torch = torch.from_numpy(rotation_matrix).float().unsqueeze(0) | |
| # base_quat = matrix_to_quaternion(rotation_matrix_torch)[0] | |
| # legacy comment | |
| # z_rotation_angles = torch.deg2rad(torch.tensor([0.0, 0.0, 90.0*direction*obj_flag], dtype=torch.float32)) | |
| # z_rotation_matrix = euler_angles_to_matrix(z_rotation_angles, convention="XYZ") | |
| # z_rotation_quat = matrix_to_quaternion(z_rotation_matrix.unsqueeze(0))[0] | |
| # legacy comment | |
| # push_quat = quaternion_multiply(base_quat.unsqueeze(0), z_rotation_quat.unsqueeze(0))[0].cpu().numpy() | |
| # legacy comment | |
| # push_pose = sapien.Pose(p=obj_pos, q=push_quat) | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # start_pos = obj_pos - push_direction_3d * offset_distance | |
| # legacy comment | |
| # reach_pose_q = push_pose.q.tolist() if hasattr(push_pose.q, 'tolist') else list(push_pose.q) | |
| # legacy comment | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=[0, 0, 0.1], q=reach_pose_q)) | |
| # start_pos=start_pos.tolist() | |
| # start_pos[1]-=0.1*direction | |
| # legacy comment | |
| # start_ready_pos[2]=0.1 | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=start_ready_pos, q=reach_pose_q)) | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=start_pos, q=reach_pose_q)) | |
| # legacy comment | |
| # planner.close_gripper() | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # end_pos = target_pos.copy() | |
| # legacy comment | |
| # end_pos[1]-=0.1*direction | |
| # planner.move_to_pose_with_screw(sapien.Pose(p=end_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| # planner.open_gripper() | |
| def solve_push_to_target_with_peg(env, planner, obj=None, target=None, direction=None, obj_flag=None): | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| obj_pos = obj.pose.sp.p if hasattr(obj.pose.sp.p, '__iter__') else np.array(obj.pose.sp.p) | |
| target_pos = target.pose.sp.p if hasattr(target.pose.sp.p, '__iter__') else np.array(target.pose.sp.p) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| push_direction_xy = target_pos[:2] - obj_pos[:2] | |
| push_direction_xy = push_direction_xy / np.linalg.norm(push_direction_xy) | |
| # legacy comment | |
| push_direction_3d = np.array([push_direction_xy[0], push_direction_xy[1], 0]) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| x_axis = np.array([push_direction_xy[0], push_direction_xy[1], 0]) # legacy comment | |
| z_axis = np.array([0, 0, -1]) # legacy comment | |
| y_axis = np.cross(z_axis, x_axis) # legacy comment | |
| # legacy comment | |
| y_norm = np.linalg.norm(y_axis) | |
| if y_norm < 1e-6: | |
| raise ValueError("Invalid push direction; failed to construct gripper frame.") | |
| y_axis = y_axis / y_norm | |
| # legacy comment | |
| rotation_matrix = np.column_stack([x_axis, y_axis, z_axis]) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| rotation_matrix_torch = torch.from_numpy(rotation_matrix).float().unsqueeze(0) | |
| base_quat = matrix_to_quaternion(rotation_matrix_torch)[0] | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| z_rotation_angles = torch.deg2rad( | |
| torch.tensor([0.0, 0.0, 90.0 * direction * obj_flag], dtype=torch.float32) | |
| ) | |
| z_rotation_matrix = euler_angles_to_matrix(z_rotation_angles, convention="XYZ") | |
| z_rotation_quat = matrix_to_quaternion(z_rotation_matrix.unsqueeze(0))[0] | |
| # legacy comment | |
| push_quat = quaternion_multiply( | |
| base_quat.unsqueeze(0), z_rotation_quat.unsqueeze(0) | |
| )[0].cpu().numpy() | |
| # legacy comment | |
| push_pose = sapien.Pose(p=obj_pos, q=push_quat) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| offset_distance = 0.1 # legacy comment | |
| start_pos = obj_pos - push_direction_3d * offset_distance | |
| start_pos[2] = obj_pos[2] # legacy comment | |
| # legacy comment | |
| lateral_unit = np.array([-push_direction_xy[1], push_direction_xy[0], 0], dtype=np.float32) | |
| lateral_norm = np.linalg.norm(lateral_unit[:2]) | |
| # legacy comment | |
| if lateral_norm < 1e-6: | |
| lateral_unit = np.array([0.0, 1.0, 0.0], dtype=np.float32) | |
| else: | |
| lateral_unit /= lateral_norm | |
| # legacy comment | |
| lateral_distance = 0.1 * direction | |
| start_pos = start_pos - lateral_unit * lateral_distance | |
| # legacy comment | |
| reach_pose_q = push_pose.q.tolist() if hasattr(push_pose.q, "tolist") else list(push_pose.q) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| start_ready_pos = start_pos.copy() | |
| # legacy comment | |
| planner.move_to_pose_with_screw(sapien.Pose(p=start_ready_pos.tolist(), q=reach_pose_q)) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=start_pos.tolist(), q=reach_pose_q)) | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| planner.close_gripper() | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| # -------------------------------------------------------------------------- # | |
| # legacy comment | |
| end_pos = target_pos - push_direction_3d * 0.03 # legacy comment | |
| end_pos[2] = start_pos[2] # legacy comment | |
| end_pos = end_pos - lateral_unit * lateral_distance | |
| planner.move_to_pose_with_screw(sapien.Pose(p=end_pos.tolist(), q=reach_pose_q)) | |
| # legacy comment | |
| planner.open_gripper() | |
| def move_to_avoid(env, planner): | |
| original_pose = env.agent.tcp.pose | |
| lift_pose_p=original_pose.p.tolist()[0] | |
| lift_pose_q=original_pose.q.tolist()[0] | |
| lift_pose_p=[-0.1,0,0.1] | |
| planner.move_to_pose_with_screw(sapien.Pose(p=lift_pose_p,q=lift_pose_q)) | |
| return None | |
| def solve_pickup_fail(env, planner, obj=None,z_offset=None,xy_offset=None,obj_type="cube",mode=None): | |
| """Hover directly above grasp pose with slight random +z, close then reopen gripper.""" | |
| if obj is None: | |
| return None | |
| env = getattr(env, "unwrapped", env) | |
| planner.open_gripper() | |
| # Build the same grasp pose as the normal pickup, but stop above it with random z lift <= 0.1. | |
| obb = get_actor_obb(obj) | |
| approaching = np.array([0, 0, -1]) | |
| target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| grasp_info = compute_grasp_info_by_obb( | |
| obb, | |
| approaching=approaching, | |
| target_closing=target_closing, | |
| depth=0.025, | |
| ) | |
| closing = grasp_info["closing"] | |
| grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| if obj_type=="bin": | |
| grasp_pose = grasp_pose * sapien.Pose([0, 0, -0.01]) | |
| fail_pose_p = np.asarray(grasp_pose.p, dtype=np.float32).reshape(-1).tolist() | |
| fail_pose_q = np.asarray(grasp_pose.q, dtype=np.float32).reshape(-1).tolist() | |
| normalized_mode = mode.lower() if isinstance(mode, str) else mode | |
| # Randomly pick whether to perturb the failed hover in XY or Z. | |
| #mode = random.choice(["xy", "z"]) | |
| #mode="xy" | |
| if normalized_mode == "xy": | |
| env.fail="xy" | |
| xy_offset = np.asarray(xy_offset, dtype=np.float32).reshape(-1) | |
| if xy_offset.size == 1: | |
| xy_offset = np.repeat(xy_offset, 2) | |
| # allow +/- adjustments or no change per axis, but ensure at least one axis moves | |
| signs = [random.choice((-1, 0, 1)) for _ in range(2)] | |
| while signs[0] == 0 and signs[1] == 0: | |
| signs = [random.choice((-1, 0, 1)) for _ in range(2)] | |
| signed_offset = xy_offset * np.array(signs, dtype=np.float32) | |
| fail_pose_p[0] += float(signed_offset[0]) | |
| fail_pose_p[1] += float(signed_offset[1]) | |
| elif normalized_mode == "z": | |
| env.fail="z" | |
| z_shift = z_offset | |
| fail_pose_p[2] += z_shift | |
| else: | |
| raise ValueError(f"Invalid fail mode: {mode}") | |
| ready_pose_p = fail_pose_p.copy() | |
| ready_pose_p[2] = 0.15 | |
| if obj_type=="bin": | |
| ready_pose_p[2] = 0.2 | |
| ready_pose = sapien.Pose(p=ready_pose_p, q=fail_pose_q) | |
| planner.move_to_pose_with_screw(ready_pose) | |
| fail_pose = sapien.Pose(p=fail_pose_p, q=fail_pose_q) | |
| planner.move_to_pose_with_screw(fail_pose) | |
| planner.close_gripper() | |
| planner.move_to_pose_with_screw(ready_pose) | |
| planner.open_gripper() | |
| env.use_fail_planner=True | |
| return None | |
| def solve_pickup(env, planner, obj=None,fail_grasp=False,mode=None): | |
| # 10% chance to perform a deliberate failed hover before the normal pickup. | |
| planner.open_gripper() | |
| if(env.use_demonstrationwrapper==False): | |
| if fail_grasp==True: | |
| solve_pickup_fail(env, planner, obj,z_offset=env.cube_half_size*2,xy_offset=env.cube_half_size*2,obj_type="cube",mode=mode) | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # retrieves the object oriented bounding box (trimesh box object) | |
| obb = get_actor_obb(obj) | |
| approaching = np.array([0, 0, -1]) | |
| # get transformation matrix of the tcp pose, is default batched and on torch | |
| target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # we can build a simple grasp pose using this information for Panda | |
| grasp_info = compute_grasp_info_by_obb( | |
| obb, | |
| approaching=approaching, | |
| target_closing=target_closing, | |
| depth=FINGER_LENGTH, | |
| ) | |
| closing, center = grasp_info["closing"], grasp_info["center"] | |
| grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| # -------------------------------------------------------------------------- # | |
| # Reach | |
| # -------------------------------------------------------------------------- # | |
| reach_pose_p = grasp_pose.p.tolist() if hasattr(grasp_pose.p, 'tolist') else list(grasp_pose.p) | |
| reach_pose_q = grasp_pose.q.tolist() if hasattr(grasp_pose.q, 'tolist') else list(grasp_pose.q) | |
| reach_pose_p[2]=0.15 | |
| planner.move_to_pose_with_screw(sapien.Pose(p=reach_pose_p,q=reach_pose_q)) | |
| planner.open_gripper() | |
| # -------------------------------------------------------------------------- # | |
| # Grasp | |
| # -------------------------------------------------------------------------- # | |
| planner.move_to_pose_with_screw(grasp_pose) | |
| planner.close_gripper() | |
| # -------------------------------------------------------------------------- # | |
| # Move to goal pose | |
| # -------------------------------------------------------------------------- # | |
| goal_pose_P=obj.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.15 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose.q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| def solve_pickup_bin(env, planner, obj=None, fail_grasp=False, mode=None): | |
| planner.open_gripper() | |
| # 10% chance to perform a deliberate failed hover before the normal pickup. | |
| if(env.use_demonstrationwrapper==False): | |
| if fail_grasp==True: | |
| solve_pickup_fail(env, planner, obj,z_offset=0.035,xy_offset=0.035,obj_type="bin", mode=mode) | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # retrieves the object oriented bounding box (trimesh box object) | |
| obb = get_actor_obb(obj) | |
| approaching = np.array([0, 0, -1]) | |
| # get transformation matrix of the tcp pose, is default batched and on torch | |
| target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # we can build a simple grasp pose using this information for Panda | |
| grasp_info = compute_grasp_info_by_obb( | |
| obb, | |
| approaching=approaching, | |
| target_closing=target_closing, | |
| depth=FINGER_LENGTH, | |
| ) | |
| closing, center = grasp_info["closing"], grasp_info["center"] | |
| grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| # -------------------------------------------------------------------------- # | |
| # Reach | |
| # -------------------------------------------------------------------------- # | |
| reach_pose = grasp_pose * sapien.Pose([0, 0, -0.15]) | |
| reach_pose_p=reach_pose.p.tolist() | |
| reach_pose_p[2]=0.2 | |
| reach_pose_q=reach_pose.q.tolist() | |
| reach_pose_fix=sapien.Pose(reach_pose_p,reach_pose_q) | |
| planner.move_to_pose_with_screw(reach_pose_fix) | |
| planner.open_gripper() | |
| # -------------------------------------------------------------------------- # | |
| # Grasp | |
| # -------------------------------------------------------------------------- # | |
| grasp_pose_up=grasp_pose * sapien.Pose([0, 0, -0.01]) | |
| planner.move_to_pose_with_screw(grasp_pose_up) | |
| planner.close_gripper() | |
| # grasp_pose_up=grasp_pose * sapien.Pose([0, 0.1,0])#test | |
| # planner.move_to_pose_with_screw(grasp_pose_up)#test | |
| # -------------------------------------------------------------------------- # | |
| # Move to goal pose | |
| # -------------------------------------------------------------------------- # | |
| goal_pose_P=obj.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.2 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose.q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| # def solve_putdown_bin(env, planner, obj=None): | |
| # planner.open_gripper() | |
| # FINGER_LENGTH = 0.025 | |
| # env = env.unwrapped | |
| # planner.close_gripper() | |
| # goal_pose_p=env.agent.tcp.pose.p.tolist()[0] | |
| # goal_pose_p[2]=0.15 | |
| # goal_pose_p[0]+=0.1 | |
| # goal_pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| # goal_pose = sapien.Pose(goal_pose_p,goal_pose_q) | |
| # res = planner.move_to_pose_with_screw(goal_pose) | |
| # goal_pose_p[2]=0 | |
| # goal_pose = sapien.Pose(goal_pose_p,goal_pose_q) | |
| # res = planner.move_to_pose_with_screw(goal_pose) | |
| # planner.open_gripper() | |
| # planner.close() | |
| # return res | |
| # def solve_putdown_whenhold(env, planner, obj=None): | |
| # FINGER_LENGTH = 0.025 | |
| # env = env.unwrapped | |
| # # retrieves the object oriented bounding box (trimesh box object) | |
| # obb = get_actor_obb(obj) | |
| # approaching = np.array([0, 0, -1]) | |
| # # get transformation matrix of the tcp pose, is default batched and on torch | |
| # target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # # we can build a simple grasp pose using this information for Panda | |
| # grasp_info = compute_grasp_info_by_obb( | |
| # obb, | |
| # approaching=approaching, | |
| # target_closing=target_closing, | |
| # depth=FINGER_LENGTH, | |
| # ) | |
| # closing, center = grasp_info["closing"], grasp_info["center"] | |
| # grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| # goal_pose_P=obj.pose.p.tolist()[0] | |
| # goal_pose_P[2]=0 | |
| # goal_pose = sapien.Pose(goal_pose_P, grasp_pose.q) | |
| # res = planner.move_to_pose_with_screw(goal_pose) | |
| # planner.open_gripper() | |
| # planner.close() | |
| # return res | |
| def solve_putonto_whenhold(env, planner,target=None): | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # # retrieves the object oriented bounding box (trimesh box object) | |
| # obb = get_actor_obb(obj) | |
| # approaching = np.array([0, 0, -1]) | |
| # # get transformation matrix of the tcp pose, is default batched and on torch | |
| # target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # # we can build a simple grasp pose using this information for Panda | |
| # grasp_info = compute_grasp_info_by_obb( | |
| # obb, | |
| # approaching=approaching, | |
| # target_closing=target_closing, | |
| # depth=FINGER_LENGTH, | |
| # ) | |
| # closing, center = grasp_info["closing"], grasp_info["center"] | |
| # grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| grasp_pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| goal_pose_P_prepare=target.pose.p.tolist()[0] | |
| goal_pose_P_prepare[2]=0.15 | |
| goal_pose = sapien.Pose(goal_pose_P_prepare, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.open_gripper() | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.15 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| def solve_swingonto_whenhold(env, planner,target=None,height=0.05): | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # # retrieves the object oriented bounding box (trimesh box object) | |
| # obb = get_actor_obb(obj) | |
| # approaching = np.array([0, 0, -1]) | |
| # # get transformation matrix of the tcp pose, is default batched and on torch | |
| # target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # # we can build a simple grasp pose using this information for Panda | |
| # grasp_info = compute_grasp_info_by_obb( | |
| # obb, | |
| # approaching=approaching, | |
| # target_closing=target_closing, | |
| # depth=FINGER_LENGTH, | |
| # ) | |
| # closing, center = grasp_info["closing"], grasp_info["center"] | |
| # grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| grasp_pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose_P[2]=height | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| def solve_swingonto_withDirection(env, planner, target=None, radius=0.1, direction="counterclockwise"): | |
| """Planar arc motion at z=0.07 from current TCP to target. | |
| direction: "counterclockwise" means the left side of t_start->t_end (positive cross product); "clockwise" means the right side.""" | |
| if target is None: | |
| raise ValueError("target must be provided for swing onto motion.") | |
| start_pos = env.agent.tcp.pose.p.reshape(-1, 3)[0] | |
| end_pos = target.pose.p.reshape(-1, 3)[0] | |
| start_xy = np.asarray(start_pos[:2], dtype=np.float32) | |
| end_xy = np.asarray(end_pos[:2], dtype=np.float32) | |
| chord_vec = end_xy - start_xy | |
| chord_len = np.linalg.norm(chord_vec) | |
| current_qpos = env.agent.tcp.pose.q.reshape(-1, 4)[0].tolist() | |
| # Initial joint pose used to dry-run each segment in order, then concatenate into one trajectory for execution | |
| init_qpos_tensor = planner.robot.get_qpos() | |
| qpos_device = init_qpos_tensor.device if hasattr(init_qpos_tensor, "device") else None | |
| qpos_dtype = init_qpos_tensor.dtype if hasattr(init_qpos_tensor, "dtype") else torch.float32 | |
| init_qpos = ( | |
| init_qpos_tensor.detach().cpu().numpy() if hasattr(init_qpos_tensor, "detach") else np.asarray(init_qpos_tensor) | |
| ).reshape(-1) | |
| plan_start_qpos = init_qpos.copy() | |
| waypoints = [] | |
| if chord_len < 1e-6: | |
| goal_p = end_pos.tolist() | |
| goal_p[2] = 0.07 | |
| waypoints.append(sapien.Pose(goal_p, current_qpos)) | |
| else: | |
| # radius only controls lateral offset; forward span comes from the chord length. | |
| lateral_offset = float(max(radius, 1e-4)) | |
| dir_unit = chord_vec / chord_len | |
| perp = np.array([-dir_unit[1], dir_unit[0]]) | |
| direction_l = str(direction).lower() | |
| if direction_l == "counterclockwise": | |
| sign = -1.0 | |
| elif direction_l == "clockwise": | |
| sign = 1.0 | |
| else: | |
| # compatible with legacy left/right notation | |
| sign = 1.0 if direction_l == "left" else -1.0 | |
| control_xy = (start_xy + end_xy) / 2.0 + sign * lateral_offset * perp | |
| # Use a quadratic Bezier curve (start -> control -> end) to create a smoother bend than a circular arc. | |
| num_steps = 45 | |
| for t in np.linspace(0.0, 1.0, num_steps): | |
| one_minus_t = 1.0 - t | |
| waypoint_xy = ( | |
| (one_minus_t ** 2) * start_xy | |
| + 2 * one_minus_t * t * control_xy | |
| + (t ** 2) * end_xy | |
| ) | |
| goal_p = [float(waypoint_xy[0]), float(waypoint_xy[1]), 0.07] | |
| waypoints.append(sapien.Pose(goal_p, current_qpos)) | |
| if waypoints: | |
| # Hold the final pose a bit longer by repeating it 5 times. | |
| last_p = np.asarray(waypoints[-1].p, dtype=np.float32).reshape(-1).tolist() | |
| last_q = np.asarray(waypoints[-1].q, dtype=np.float32).reshape(-1).tolist() | |
| for _ in range(5): | |
| waypoints.append(sapien.Pose(last_p, last_q)) | |
| logger.debug(" get waypoint") | |
| # Directly connect each waypoint IK solution into a discrete path without extra interpolation/planning | |
| positions = [] | |
| last_res = None | |
| # legacy comment | |
| plan_start_qpos_full = planner.planner.pad_qpos(plan_start_qpos.copy()) | |
| for idx, wp in enumerate(waypoints): | |
| pose_for_plan = planner._transform_pose_for_planning(wp) | |
| pose_p = np.asarray(pose_for_plan.p, dtype=np.float32).reshape(-1) | |
| pose_q = np.asarray(pose_for_plan.q, dtype=np.float32).reshape(-1) | |
| goal_world = np.concatenate([pose_p, pose_q]) | |
| goal_base = planner.planner.transform_goal_to_wrt_base(goal_world) | |
| ik_status, ik_solutions = planner.planner.IK( | |
| goal_base, | |
| plan_start_qpos_full.copy(), | |
| ) | |
| if ik_status != "Success" or len(ik_solutions) == 0: | |
| logger.debug(f"IK failed at waypoint {idx}: {ik_status}") | |
| continue | |
| # Take the first IK solution and append directly; do not call plan_qpos_to_qpos for interpolation | |
| qpos_sol = ik_solutions[0] | |
| padded_qpos = plan_start_qpos_full.copy() | |
| padded_qpos[: qpos_sol.shape[0]] = qpos_sol | |
| positions.append(padded_qpos) | |
| # Update full qpos as the next segment start to avoid base/joint drift between segments | |
| plan_start_qpos_full = padded_qpos | |
| if len(positions) == 0: | |
| logger.debug("No IK solutions found for waypoints, aborting follow_path.") | |
| return None | |
| traj_res = { | |
| "status": "Success", | |
| "position": np.stack(positions, axis=0), | |
| } | |
| if planner.control_mode == "pd_joint_pos_vel": | |
| traj_res["velocity"] = np.zeros_like(traj_res["position"]) | |
| last_res = planner.follow_path(traj_res) | |
| return last_res | |
| # def solve_swingonto_withDirection(env, planner, target=None, radius=0.1, direction="counterclockwise"): | |
| # """Planar arc motion at z=0.07 from current TCP to target. | |
| # direction: "counterclockwise" means the left side of t_start->t_end (positive cross product); "clockwise" means the right side.""" | |
| # if target is None: | |
| # raise ValueError("target must be provided for swing onto motion.") | |
| # start_pos = env.agent.tcp.pose.p.reshape(-1, 3)[0] | |
| # end_pos = target.pose.p.reshape(-1, 3)[0] | |
| # start_xy = np.asarray(start_pos[:2], dtype=np.float32) | |
| # end_xy = np.asarray(end_pos[:2], dtype=np.float32) | |
| # chord_vec = end_xy - start_xy | |
| # chord_len = np.linalg.norm(chord_vec) | |
| # current_qpos = env.agent.tcp.pose.q.reshape(-1, 4)[0].tolist() | |
| # if chord_len < 1e-6: | |
| # goal_p = end_pos.tolist() | |
| # goal_p[2] = 0.07 | |
| # try: | |
| # return planner.move_to_pose_with_screw(sapien.Pose(goal_p, current_qpos)) | |
| # except Exception as e: | |
| # print(f"move_to_pose_with_screw failed: {e}") | |
| # return None | |
| # # Keep radius feasible for chord length | |
| # radius = float(max(radius, chord_len / 2.0 + 1e-4)) | |
| # half_chord = chord_len / 2.0 | |
| # sagitta = math.sqrt(max(radius ** 2 - half_chord ** 2, 0.0)) | |
| # dir_unit = chord_vec / chord_len | |
| # perp = np.array([-dir_unit[1], dir_unit[0]]) | |
| # direction_l = str(direction).lower() | |
| # if direction_l == "counterclockwise": | |
| # sign = 1.0 | |
| # elif direction_l == "clockwise": | |
| # sign = -1.0 | |
| # else: | |
| # # compatible with legacy left/right notation | |
| # sign = 1.0 if direction_l == "left" else -1.0 | |
| # center_xy = (start_xy + end_xy) / 2.0 + sign * sagitta * perp | |
| # start_angle = math.atan2(start_xy[1] - center_xy[1], start_xy[0] - center_xy[0]) | |
| # end_angle = math.atan2(end_xy[1] - center_xy[1], end_xy[0] - center_xy[0]) | |
| # angle_diff = end_angle - start_angle | |
| # if sign > 0 and angle_diff <= 0: | |
| # angle_diff += 2 * math.pi | |
| # if sign < 0 and angle_diff >= 0: | |
| # angle_diff -= 2 * math.pi | |
| # num_steps = 10 | |
| # res = None | |
| # for ang in np.linspace(start_angle, start_angle + angle_diff, num_steps): | |
| # waypoint_xy = center_xy + radius * np.array([math.cos(ang), math.sin(ang)]) | |
| # goal_p = [float(waypoint_xy[0]), float(waypoint_xy[1]), 0.07] | |
| # try: | |
| # res = planner.move_to_pose_with_screw(sapien.Pose(goal_p, current_qpos)) | |
| # except Exception as e: | |
| # print(f"move_to_pose_with_screw failed: {e}") | |
| # res = None | |
| # return res | |
| def solve_swingonto(env, planner,target=None,record_swing_qpos=False): | |
| env = env.unwrapped | |
| # if horizontal==True: | |
| # # use current gripper pose and rotate 90° around Z axis | |
| # current_qpos = torch.tensor([-7.3356e-08, 1.0000e+00, -2.0862e-07, -1.8728e-09]) | |
| # z_rot = torch.tensor([[0.0, 0.0, np.pi / 2]], dtype=torch.float32, device=current_qpos.device) | |
| # z_rot = matrix_to_quaternion(euler_angles_to_matrix(z_rot, convention="XYZ"))[0] | |
| # current_qpos = quaternion_multiply(current_qpos.unsqueeze(0), z_rot.unsqueeze(0))[0] | |
| # current_qpos = current_qpos.tolist() | |
| # else: | |
| # current_qpos = env.agent.tcp.pose.q.reshape(-1, 4)[0] | |
| current_qpos = env.agent.tcp.pose.q.reshape(-1, 4)[0] | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.07 | |
| goal_pose = sapien.Pose(goal_pose_P, current_qpos) | |
| for i in range(2): | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| try: | |
| planner.close_gripper() | |
| except: | |
| AttributeError | |
| if record_swing_qpos==True: | |
| env.swing_qpos=env.agent.robot.qpos | |
| return res | |
| def solve_strong_reset(env, planner,timestep=30,gripper=None,action=None): | |
| try: | |
| planner.open_gripper() | |
| except: | |
| AttributeError | |
| if action==None: | |
| action=reset_panda.get_reset_panda_param("action",gripper=gripper) | |
| for i in range(timestep): | |
| env.step(action) | |
| logger.debug("strong reset!!") | |
| env.unwrapped.reset_in_proecess=True | |
| env.unwrapped.after_demo=True | |
| env.unwrapped.reset_in_proecess=False | |
| def solve_reset(env, planner): | |
| pose_p=[0,0,0.2] | |
| pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| planner.move_to_pose_with_screw(sapien.Pose(p=pose_p,q=pose_q)) | |
| planner.open_gripper() | |
| def solve_putdown_whenhold(env, planner,release_z=0.07): | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # # retrieves the object oriented bounding box (trimesh box object) | |
| # obb = get_actor_obb(obj) | |
| # approaching = np.array([0, 0, -1]) | |
| # # get transformation matrix of the tcp pose, is default batched and on torch | |
| # target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # # we can build a simple grasp pose using this information for Panda | |
| # grasp_info = compute_grasp_info_by_obb( | |
| # obb, | |
| # approaching=approaching, | |
| # target_closing=target_closing, | |
| # depth=FINGER_LENGTH, | |
| # ) | |
| # closing, center = grasp_info["closing"], grasp_info["center"] | |
| # grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| grasp_pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| goal_pose_P=env.agent.tcp.pose.p.tolist()[0] | |
| goal_pose_P[2]=release_z | |
| #goal_pose_P[0]+=0.1#test | |
| goal_pose = sapien.Pose(goal_pose_P,grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.open_gripper() | |
| goal_pose_P=env.agent.tcp.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.15 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| def solve_putonto_whenhold_binspecial(env, planner,target=None): | |
| FINGER_LENGTH = 0.025 | |
| env = env.unwrapped | |
| # # retrieves the object oriented bounding box (trimesh box object) | |
| # obb = get_actor_obb(obj) | |
| # approaching = np.array([0, 0, -1]) | |
| # # get transformation matrix of the tcp pose, is default batched and on torch | |
| # target_closing = env.agent.tcp.pose.to_transformation_matrix()[0, :3, 1].cpu().numpy() | |
| # # we can build a simple grasp pose using this information for Panda | |
| # grasp_info = compute_grasp_info_by_obb( | |
| # obb, | |
| # approaching=approaching, | |
| # target_closing=target_closing, | |
| # depth=FINGER_LENGTH, | |
| # ) | |
| # closing, center = grasp_info["closing"], grasp_info["center"] | |
| #grasp_pose = env.agent.build_grasp_pose(approaching, closing, obj.pose.sp.p) | |
| grasp_pose_q=env.agent.tcp.pose.q.tolist()[0] | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose_P[2]=0.2 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.open_gripper() | |
| goal_pose_P=target.pose.p.tolist()[0] | |
| goal_pose_P[0]=goal_pose_P[0]-0.1 | |
| goal_pose_P[2]=0.2 | |
| goal_pose = sapien.Pose(goal_pose_P, grasp_pose_q) | |
| res = planner.move_to_pose_with_screw(goal_pose) | |
| planner.close() | |
| return res | |
| def solve_hold_obj(env, planner, static_steps,close=False): | |
| start_step = int(getattr(env, "elapsed_steps", 0)) | |
| target_step = start_step + static_steps | |
| while int(getattr(env, "elapsed_steps", 0)) < target_step: | |
| if close: | |
| try: | |
| planner.close_gripper() | |
| except: | |
| AttributeError | |
| else: | |
| try: | |
| planner.open_gripper() | |
| except: | |
| AttributeError | |
| current_step = int(getattr(env, "elapsed_steps", 0)) | |
| #print(f"<plannner>:Holding object:{current_step}/{target_step})") | |
| return None | |
| def solve_hold_obj_absTimestep(env,planner,absTimestep): | |
| while int(getattr(env, "elapsed_steps", 0)) < absTimestep: | |
| planner.close_gripper() | |
| return None | |
| def solve_button(env, planner,obj,steps_press=None,interval=20,without_hold=False): | |
| # if steps_press: | |
| # while env.elapsed_steps<steps_press-interval: | |
| # current_step = int(getattr(env, "elapsed_steps", 0)) | |
| # print(f"Waiting to press button:{current_step}/{steps_press-interval})") | |
| # planner.close_gripper() | |
| FINGER_LENGTH = 0.025 | |
| env=env.unwrapped | |
| position=obj.pose.p.tolist()[0] | |
| ready_position=position.copy() | |
| ready_position[2]=0.15 | |
| angles = torch.deg2rad(torch.tensor([180.0, 0.0, 0.0], dtype=torch.float32)) # (3,) | |
| rotate = matrix_to_quaternion( | |
| euler_angles_to_matrix(angles, convention="XYZ") | |
| ) | |
| if without_hold==False: | |
| planner.move_to_pose_with_screw(sapien.Pose(p=ready_position,q=rotate)) | |
| planner.close_gripper() | |
| steps=env.elapsed_steps.item() | |
| logger.debug(f"press button at step {steps}") | |
| planner.move_to_pose_with_screw(sapien.Pose(p=position,q=rotate)) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=ready_position, q=rotate)) | |
| def solve_button_ready(env, planner,obj): | |
| FINGER_LENGTH = 0.025 | |
| env=env.unwrapped | |
| position=obj.pose.p.tolist()[0] | |
| ready_position=position.copy() | |
| ready_position[2]=0.15 | |
| angles = torch.deg2rad(torch.tensor([180.0, 0.0, 0.0], dtype=torch.float32)) # (3,) | |
| rotate = matrix_to_quaternion( | |
| euler_angles_to_matrix(angles, convention="XYZ") | |
| ) | |
| planner.move_to_pose_with_screw(sapien.Pose(p=ready_position,q=rotate)) | |
| planner.close_gripper() | |