address review comments

Author: Peter

src/pytorch_kinematics/transforms/__init__.py

@@ -4,10 +4,10 @@
     axis_and_angle_to_matrix,
     axis_angle_to_quaternion,
     euler_angles_to_matrix,
+    matrix_to_axis_angle,
     matrix_to_euler_angles,
     matrix_to_quaternion,
     matrix_to_rotation_6d,
-    matrix_to_axis_angle,
     quaternion_apply,
     quaternion_invert,
     quaternion_multiply,
@@ -18,8 +18,10 @@
     random_rotations,
     rotation_6d_to_matrix,
     standardize_quaternion,
+    tensor_axis_and_angle_to_matrix,
+    tensor_axis_and_d_to_pris_matrix,
+    wxyz_to_xyzw,
     xyzw_to_wxyz,
-    wxyz_to_xyzw
 )
 from .so3 import (
     so3_exp_map,

src/pytorch_kinematics/transforms/rotation_conversions.py

@@ -452,6 +452,7 @@ def quaternion_apply(quaternion, point):
 
 def tensor_axis_and_d_to_pris_matrix(axis, d):
     """
+    Creates a 4x4 matrix that represents a translation along an axis of a distance d
     Works with any number of batch dimensions.
 
     Args:
@@ -471,6 +472,7 @@ def tensor_axis_and_d_to_pris_matrix(axis, d):
 
 def tensor_axis_and_angle_to_matrix(axis, theta):
     """
+    Creates a 4x4 matrix that represents a rotation around an axis by an angle theta.
     Works with any number of batch dimensions.
 
     Args:

tests/gen_fk_perf.py

@@ -11,12 +11,11 @@ def main():
     np.set_printoptions(precision=3, suppress=True, linewidth=220)
     torch.set_printoptions(precision=3, sci_mode=False, linewidth=220)
 
-    chains = [
-        pk.build_chain_from_mjcf(open('val.xml').read()),
-        pk.build_serial_chain_from_mjcf(open('val.xml').read(), end_link_name='left_tool'),
-        pk.build_serial_chain_from_urdf(open('kuka_iiwa.urdf').read(), end_link_name='lbr_iiwa_link_7'),
-    ]
-    names = ['val', 'val_serial', 'kuka_iiwa']
+    chains = {
+        'val':        pk.build_chain_from_mjcf(open('val.xml').read()),
+        'val_serial': pk.build_serial_chain_from_mjcf(open('val.xml').read(), end_link_name='left_tool'),
+        'kuka_iiwa':  pk.build_serial_chain_from_urdf(open('kuka_iiwa.urdf').read(), end_link_name='lbr_iiwa_link_7'),
+    }
 
     devices = ['cpu', 'cuda']
     dtypes = [torch.float32, torch.float64]
@@ -27,27 +26,19 @@ def main():
     headers = ['method', 'chain', 'device', 'dtype', 'batch_size', 'time']
     data = []
 
-    def _fk_cpp(th):
+    def _fk(th):
         return chain.forward_kinematics(th)
 
-    @torch.compile(backend='eager')
-    def _fk_torch_compile(th):
-        return chain.forward_kinematics_py(th)
-
-    method_names = ['fk_cpp', 'fk_torch_compile']
-    methods = [_fk_cpp, _fk_torch_compile]
-
-    for chain, name in zip(chains, names):
+    for name, chain in chains.items():
         for device in devices:
             for dtype in dtypes:
                 for batch_size in batch_sizes:
-                    for method_name, method in zip(method_names, methods):
-                        chain = chain.to(dtype=dtype, device=device)
-                        th = torch.zeros(batch_size, chain.n_joints).to(dtype=dtype, device=device)
+                    chain = chain.to(dtype=dtype, device=device)
+                    th = torch.zeros(batch_size, chain.n_joints).to(dtype=dtype, device=device)
 
-                        dt = timeit.timeit(lambda: method(th), number=number)
-                        data.append([name, device, dtype, batch_size, dt / number])
-                        print(f"{method_name} {name=} {device=} {dtype=} {batch_size=} {dt / number:.4f}")
+                    dt = timeit.timeit(lambda: _fk(th), number=number)
+                    data.append([name, device, dtype, batch_size, dt / number])
+                    print(f"{name=} {device=} {dtype=} {batch_size=} {dt / number:.4f}")
 
     # pickle the data for visualization in jupyter notebook
     import pickle

tests/test_kinematics.py

@@ -27,16 +27,9 @@ def test_fk_mjcf():
     chain = chain.to(dtype=torch.float64)
     print(chain)
     print(chain.get_joint_parameter_names())
-    th = {
-        'hip_1':   1.0,
-        'ankle_1': 1,
-        'hip_2':   0.0,
-        'ankle_2': 0.0,
-        'hip_3':   0.0,
-        'ankle_3': 0.0,
-        'hip_4':   0.0,
-        'ankle_4': 0.0,
-    }
+
+    th = {joint: 0.0 for joint in chain.get_joint_parameter_names()}
+    th.update({'hip_1': 1.0, 'ankle_1': 1})
     ret = chain.forward_kinematics(th)
     tg = ret['aux_1']
     pos, rot = quat_pos_from_transform3d(tg)
@@ -173,12 +166,7 @@ def test_fk_simple_arm():
     assert torch.allclose(pos, torch.tensor([1.05, 0.55, 0.5], dtype=torch.float64))
 
     N = 100
-    ret = chain.forward_kinematics({
-        'arm_shoulder_pan_joint': torch.rand(N),
-        'arm_elbow_pan_joint':    torch.rand(N),
-        'arm_wrist_lift_joint':   torch.rand(N),
-        'arm_wrist_roll_joint':   torch.rand(N),
-    })
+    ret = chain.forward_kinematics({k: torch.rand(N) for k in chain.get_joint_parameter_names()})
     tg = ret['arm_wrist_roll']
     assert list(tg.get_matrix().shape) == [N, 4, 4]