File size: 4,951 Bytes
f96995c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import os
import copy
import time
import numpy as np
import transforms3d
from pathlib import Path
import open3d as o3d
import warnings
warnings.filterwarnings("always", category=RuntimeWarning)
import sapien.core as sapien
from nclaw.utils import get_root
root: Path = get_root(__file__)
def np2o3d(pcd, color=None):
# pcd: (n, 3)
# color: (n, 3)
pcd_o3d = o3d.geometry.PointCloud()
pcd_o3d.points = o3d.utility.Vector3dVector(pcd)
if color is not None:
assert pcd.shape[0] == color.shape[0]
assert color.max() <= 1
assert color.min() >= 0
pcd_o3d.colors = o3d.utility.Vector3dVector(color)
return pcd_o3d
class KinHelper():
def __init__(self, robot_name, headless=True):
# load robot
if "xarm7" in robot_name:
urdf_path = str(root / "assets/robots/xarm7/xarm7.urdf")
self.eef_name = 'link7'
else:
raise RuntimeError('robot name not supported')
self.robot_name = robot_name
# load sapien robot
self.engine = sapien.Engine()
self.scene = self.engine.create_scene()
loader = self.scene.create_urdf_loader()
self.sapien_robot = loader.load(urdf_path)
self.robot_model = self.sapien_robot.create_pinocchio_model()
self.sapien_eef_idx = -1
for link_idx, link in enumerate(self.sapien_robot.get_links()):
if link.name == self.eef_name:
self.sapien_eef_idx = link_idx
break
# load meshes and offsets from urdf_robot
self.meshes = {}
self.scales = {}
self.offsets = {}
self.pcd_dict = {}
self.tool_meshes = {}
def compute_fk_sapien_links(self, qpos, link_idx):
fk = self.robot_model.compute_forward_kinematics(qpos)
link_pose_ls = []
for i in link_idx:
link_pose_ls.append(self.robot_model.get_link_pose(i).to_transformation_matrix())
return link_pose_ls
def compute_ik_sapien(self, initial_qpos, cartesian, verbose=False):
"""
Compute IK using sapien
initial_qpos: (N, ) numpy array
cartesian: (6, ) numpy array, x,y,z in meters, r,p,y in radians
"""
tf_mat = np.eye(4)
tf_mat[:3, :3] = transforms3d.euler.euler2mat(ai=cartesian[3], aj=cartesian[4], ak=cartesian[5], axes='sxyz')
tf_mat[:3, 3] = cartesian[0:3]
pose = sapien.Pose.from_transformation_matrix(tf_mat)
if 'xarm7' in self.robot_name:
active_qmask = np.array([True, True, True, True, True, True, True])
qpos = self.robot_model.compute_inverse_kinematics(
link_index=self.sapien_eef_idx,
pose=pose,
initial_qpos=initial_qpos,
active_qmask=active_qmask,
)
if verbose:
print('ik qpos:', qpos)
# verify ik
fk_pose = self.compute_fk_sapien_links(qpos[0], [self.sapien_eef_idx])[0]
if verbose:
print('target pose for IK:', tf_mat)
print('fk pose for IK:', fk_pose)
pose_diff = np.linalg.norm(fk_pose[:3, 3] - tf_mat[:3, 3])
rot_diff = np.linalg.norm(fk_pose[:3, :3] - tf_mat[:3, :3])
if pose_diff > 0.01 or rot_diff > 0.01:
print('ik pose diff:', pose_diff)
print('ik rot diff:', rot_diff)
warnings.warn('ik pose diff or rot diff too large. Return initial qpos.', RuntimeWarning, stacklevel=2, )
return initial_qpos
return qpos[0]
def test_fk():
robot_name = 'xarm7'
init_qpos = np.array([0, 0, 0, 0, 0, 0, 0])
end_qpos = np.array([0, 0, 0, 0, 0, 0, 0])
kin_helper = KinHelper(robot_name=robot_name, headless=False)
START_ARM_POSE = [0, 0, 0, 0, 0, 0, 0]
for i in range(100):
curr_qpos = init_qpos + (end_qpos - init_qpos) * i / 100
fk = kin_helper.compute_fk_sapien_links(curr_qpos, [kin_helper.sapien_eef_idx])[0]
fk_euler = transforms3d.euler.mat2euler(fk[:3, :3], axes='sxyz')
if i == 0:
init_ik_qpos = np.array(START_ARM_POSE)
ik_qpos = kin_helper.compute_ik_sapien(init_ik_qpos, np.array(list(fk[:3, 3]) + list(fk_euler)).astype(np.float32))
re_fk_pos_mat = kin_helper.compute_fk_sapien_links(ik_qpos, [kin_helper.sapien_eef_idx])[0]
re_fk_euler = transforms3d.euler.mat2euler(re_fk_pos_mat[:3, :3], axes='sxyz')
re_fk_pos = re_fk_pos_mat[:3, 3]
print('re_fk_pos diff:', np.linalg.norm(re_fk_pos - fk[:3, 3]))
print('re_fk_euler diff:', np.linalg.norm(np.array(re_fk_euler) - np.array(fk_euler)))
init_ik_qpos = ik_qpos.copy()
qpos_diff = np.linalg.norm(ik_qpos[:6] - curr_qpos[:6])
if qpos_diff > 0.01:
warnings.warn('qpos diff too large', RuntimeWarning, stacklevel=2, )
time.sleep(0.1)
if __name__ == "__main__":
test_fk()
|