def _convert_transform(pose):
if pose is None:
return tf.Transform3d()
else:
return tf.Transform3d(rot=tf.euler_angles_to_matrix(
torch.tensor(pose[3:]), "ZYX"),
pos=pose[:3])
def _convert_transform(origin):
if origin is None:
return tf.Transform3d()
else:
return tf.Transform3d(rot=torch.tensor(tf2.quaternion_from_euler(
*origin.rpy, "sxyz"),
dtype=torch.float32),
pos=origin.xyz)
def get_transform(self, theta):
dtype = self.joint.axis.dtype
d = self.joint.axis.device
if self.joint.joint_type == 'revolute':
t = tf.Transform3d(rot=tf.axis_angle_to_quaternion(theta * self.joint.axis), dtype=dtype, device=d)
elif self.joint.joint_type == 'prismatic':
t = tf.Transform3d(pos=theta * self.joint.axis, dtype=dtype, device=d)
elif self.joint.joint_type == 'fixed':
t = tf.Transform3d(default_batch_size=theta.shape[0], dtype=dtype, device=d)
else:
raise ValueError("Unsupported joint type %s." % self.joint.joint_type)
return self.joint.offset.compose(t)
def forward_kinematics(self, th, world=tf.Transform3d(), end_only=True):
if world.dtype != self.dtype or world.device != self.device:
world = world.to(dtype=self.dtype, device=self.device, copy=True)
th, N = ensure_2d_tensor(th, self.dtype, self.device)
cnt = 0
link_transforms = {}
trans = tf.Transform3d(matrix=world.get_matrix().repeat(N, 1, 1))
for f in self._serial_frames:
trans = trans.compose(f.get_transform(th[:, cnt].view(N, 1)))
link_transforms[f.link.name] = trans.compose(f.link.offset)
if f.joint.joint_type != "fixed":
cnt += 1
return link_transforms[self._serial_frames[-1].link.name] if end_only else link_transforms
def geoms_to_visuals(geom, base=tf.Transform3d()):
visuals = []
for g in geom:
if g.type == 'capsule':
param = (g.size[0], g.fromto)
elif g.type == 'sphere':
param = g.size[0]
else:
raise ValueError('Invalid geometry type %s.' % g.type)
visuals.append(
frame.Visual(offset=base.compose(
tf.Transform3d(rot=g.quat, pos=g.pos)),
geom_type=g.type,
geom_param=param))
return visuals
def __init__(self,
name=None,
offset=tf.Transform3d(),
joint_type='fixed',
axis=(0.0, 0.0, 1.0),
dtype=torch.float32,
device="cpu"):
self.name = name
self.offset = offset
if joint_type not in self.TYPES:
raise RuntimeError(
"joint specified as {} type not, but we only support {}".
format(joint_type, self.TYPES))
self.joint_type = joint_type
if self.joint_type != 'fixed' and axis is None:
self.axis = torch.tensor([0.0, 0.0, 1.0],
dtype=dtype,
device=device)
else:
if torch.is_tensor(axis):
self.axis = axis.clone().detach().to(dtype=dtype,
device=device)
else:
self.axis = torch.tensor(axis, dtype=dtype, device=device)
# normalize axis to have norm 1 (needed for correct representation scaling with theta)
self.axis = self.axis / self.axis.norm()
def __init__(self,
offset=tf.Transform3d(),
geom_type=None,
geom_param=None):
self.offset = offset
self.geom_type = geom_type
self.geom_param = geom_param
def calc_jacobian(serial_chain, th, tool=transforms.Transform3d()):
"""
Return robot Jacobian J in base frame (N,6,DOF) where dot{x} = J dot{q}
The first 3 rows relate the translational velocities and the
last 3 rows relate the angular velocities.
"""
if not torch.is_tensor(th):
th = torch.tensor(th,
dtype=serial_chain.dtype,
device=serial_chain.device)
if len(th.shape) <= 1:
N = 1
th = th.view(1, -1)
else:
N = th.shape[0]
ndof = th.shape[1]
if tool.dtype != serial_chain.dtype or tool.device != serial_chain.device:
tool = tool.to(device=serial_chain.device,
copy=True,
dtype=serial_chain.dtype)
j_fl = torch.zeros((N, 6, ndof),
dtype=serial_chain.dtype,
device=serial_chain.device)
cur_transform = tool.get_matrix().repeat(N, 1, 1)
cnt = 0
for f in reversed(serial_chain._serial_frames):
if f.joint.joint_type == "revolute":
cnt += 1
d = torch.stack([
-cur_transform[:, 0, 0] * cur_transform[:, 1, 3] +
cur_transform[:, 1, 0] * cur_transform[:, 0, 3],
-cur_transform[:, 0, 1] * cur_transform[:, 1, 3] +
cur_transform[:, 1, 1] * cur_transform[:, 0, 3],
-cur_transform[:, 0, 2] * cur_transform[:, 1, 3] +
cur_transform[:, 1, 2] * cur_transform[:, 0, 3]
]).transpose(0, 1)
delta = cur_transform[:, 2, 0:3]
j_fl[:, :, -cnt] = torch.cat((d, delta), dim=-1)
elif f.joint.joint_type == "prismatic":
cnt += 1
j_fl[:, :3,
-cnt] = f.joint.axis.repeat(N, 1) @ cur_transform[:, :3, :3]
cur_frame_transform = f.get_transform(th[:,
-cnt].view(N,
1)).get_matrix()
cur_transform = cur_frame_transform @ cur_transform
# currently j_fl is Jacobian in flange (end-effector) frame, convert to base/world frame
pose = serial_chain.forward_kinematics(th).get_matrix()
rotation = pose[:, :3, :3]
j_tr = torch.zeros((N, 6, 6),
dtype=serial_chain.dtype,
device=serial_chain.device)
j_tr[:, :3, :3] = rotation
j_tr[:, 3:, 3:] = rotation
j_w = j_tr @ j_fl
return j_w
def add_composite_joint(root_frame, joints, base=tf.Transform3d()):
if len(joints) > 0:
root_frame.children = root_frame.children + (frame.Frame(
link=frame.Link(name=root_frame.link.name + '_child'),
joint=joint_to_joint(joints[0], base)), )
ret, offset = add_composite_joint(root_frame.children[-1], joints[1:])
return ret, root_frame.joint.offset.compose(offset)
else:
return root_frame, root_frame.joint.offset
def _forward_kinematics(root, th_dict, world=tf.Transform3d()):
link_transforms = {}
th, N = ensure_2d_tensor(th_dict.get(root.joint.name, 0.0), world.dtype, world.device)
trans = world.compose(root.get_transform(th.view(N, 1)))
link_transforms[root.link.name] = trans.compose(root.link.offset)
for child in root.children:
link_transforms.update(Chain._forward_kinematics(child, th_dict, trans))
return link_transforms
def forward_kinematics(self, th, world=tf.Transform3d()):
if not isinstance(th, dict):
jn = self.get_joint_parameter_names()
assert len(jn) == len(th)
th_dict = dict((j, th[i]) for i, j in enumerate(jn))
else:
th_dict = th
if world.dtype != self.dtype or world.device != self.device:
world = world.to(dtype=self.dtype, device=self.device, copy=True)
return self._forward_kinematics(self._root, th_dict, world)
def test_euler():
euler_angles = torch.tensor([1, 0, 0.5])
t = tf.Transform3d(rot=euler_angles)
sxyz_matrix = torch.tensor([[0.87758256, -0.47942554, 0., 0., ],
[0.25903472, 0.47415988, -0.84147098, 0.],
[0.40342268, 0.73846026, 0.54030231, 0.],
[0., 0., 0., 1.]])
# from tf.transformations import euler_matrix
# print(euler_matrix(*euler_angles, "rxyz"))
# print(t.get_matrix())
assert torch.allclose(sxyz_matrix, t.get_matrix())
def test_transform_combined():
R = tf.so3_exp_map(torch.randn((1, 3)))
tr = torch.randn((1, 3))
t = tf.Transform3d(rot=R, pos=tr)
N = 10
points = torch.randn((N, 3))
normals = torch.randn((N, 3))
points_out = t.transform_points(points)
normals_out = t.transform_normals(normals)
for i in range(N):
assert torch.allclose(points_out[i], R @ points[i] + tr)
assert torch.allclose(normals_out[i], R @ normals[i])
def test_rotate_axis_angle():
t = tf.Transform3d().rotate_axis_angle(90.0, axis="Z")
points = torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 1.0]]).view(
1, 3, 3
)
normals = torch.tensor(
[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]
).view(1, 3, 3)
points_out = t.transform_points(points)
normals_out = t.transform_normals(normals)
points_out_expected = torch.tensor(
[[0.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [-1.0, 0.0, 1.0]]
).view(1, 3, 3)
normals_out_expected = torch.tensor(
[[0.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]]
).view(1, 3, 3)
assert torch.allclose(points_out, points_out_expected)
assert torch.allclose(normals_out, normals_out_expected)
def test_rotate():
R = tf.so3_exp_map(torch.randn((1, 3)))
t = tf.Transform3d().rotate(R)
points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
1, 3, 3
)
normals = torch.tensor(
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.0, 1.0, 0.0]]
).view(1, 3, 3)
points_out = t.transform_points(points)
normals_out = t.transform_normals(normals)
points_out_expected = torch.bmm(points, R.transpose(-1, -2))
normals_out_expected = torch.bmm(normals, R.transpose(-1, -2))
assert torch.allclose(points_out, points_out_expected)
assert torch.allclose(normals_out, normals_out_expected)
for i in range(3):
assert torch.allclose(points_out[0, i], R @ points[0, i])
assert torch.allclose(normals_out[0, i], R @ normals[0, i])
def test_translations():
t = tf.Translate(1, 2, 3)
points = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.5, 0.5, 0.0]]).view(
1, 3, 3
)
points_out = t.transform_points(points)
points_out_expected = torch.tensor(
[[2.0, 2.0, 3.0], [1.0, 3.0, 3.0], [1.5, 2.5, 3.0]]
).view(1, 3, 3)
assert torch.allclose(points_out, points_out_expected)
N = 20
points = torch.randn((N, N, 3))
translation = torch.randn((N, 3))
transforms = tf.Transform3d(pos=translation)
translated_points = transforms.transform_points(points)
assert torch.allclose(translated_points, translation.repeat(N, 1, 1).transpose(0, 1) + points)
returned_points = transforms.inverse().transform_points(translated_points)
assert torch.allclose(returned_points, points, atol=1e-6)
def _build_chain_recurse(root_frame, lmap, joints):
children = []
for j in joints:
if j.parent == root_frame.link.name:
child_frame = frame.Frame(j.child + "_frame")
link_p = lmap[j.parent]
link_c = lmap[j.child]
t_p = _convert_transform(link_p.pose)
t_c = _convert_transform(link_c.pose)
child_frame.joint = frame.Joint(j.name,
offset=t_p.inverse().compose(t_c),
joint_type=JOINT_TYPE_MAP[j.type],
axis=j.axis.xyz)
child_frame.link = frame.Link(link_c.name,
offset=tf.Transform3d(),
visuals=_convert_visuals(
link_c.visuals))
child_frame.children = _build_chain_recurse(
child_frame, lmap, joints)
children.append(child_frame)
return children
def build_chain_from_sdf(data):
"""
Build a Chain object from SDF data.
Parameters
----------
data : str
SDF string data.
Returns
-------
chain.Chain
Chain object created from SDF.
"""
sdf = SDF.from_xml_string(data)
robot = sdf.model
lmap = robot.link_map
joints = robot.joints
n_joints = len(joints)
has_root = [True for _ in range(len(joints))]
for i in range(n_joints):
for j in range(i + 1, n_joints):
if joints[i].parent == joints[j].child:
has_root[i] = False
elif joints[j].parent == joints[i].child:
has_root[j] = False
for i in range(n_joints):
if has_root[i]:
root_link = lmap[joints[i].parent]
break
root_frame = frame.Frame(root_link.name + "_frame")
root_frame.joint = frame.Joint(offset=_convert_transform(root_link.pose))
root_frame.link = frame.Link(root_link.name, tf.Transform3d(),
_convert_visuals(root_link.visuals))
root_frame.children = _build_chain_recurse(root_frame, lmap, joints)
return chain.Chain(root_frame)
def _convert_visuals(visuals):
vlist = []
for v in visuals:
v_tf = _convert_transform(v.pose)
if isinstance(v.geometry, Mesh):
g_type = "mesh"
g_param = v.geometry.filename
elif isinstance(v.geometry, Cylinder):
g_type = "cylinder"
v_tf = v_tf.compose(
tf.Transform3d(rot=tf.euler_angles_to_matrix(
torch.tensor([0.5 * math.pi, 0, 0]), "ZYX")))
g_param = (v.geometry.radius, v.geometry.length)
elif isinstance(v.geometry, Box):
g_type = "box"
g_param = v.geometry.size
elif isinstance(v.geometry, Sphere):
g_type = "sphere"
g_param = v.geometry.radius
else:
g_type = None
g_param = None
vlist.append(frame.Visual(v_tf, g_type, g_param))
return vlist
def body_to_link(body, base=tf.Transform3d()):
return frame.Link(body.name,
offset=base.compose(
tf.Transform3d(rot=body.quat, pos=body.pos)))
def joint_to_joint(joint, base=tf.Transform3d()):
return frame.Joint(joint.name,
offset=base.compose(tf.Transform3d(pos=joint.pos)),
joint_type=JOINT_TYPE_MAP[joint.type],
axis=joint.axis)
def __init__(self, name=None, offset=tf.Transform3d(), visuals=()):
self.name = name
self.offset = offset
self.visuals = visuals
