PyTorch Kinematics

Highlights

• torch.compile support: New forward_kinematics_tensor and jacobian_tensor APIs are fully compatible with torch.compile(fullgraph=True)
• Improved IK convergence: Adaptive Levenberg-Marquardt damping boosts single-retry convergence by ~20 percentage points
• IK refactoring: Bug fixes, dead code removal, eliminated redundant FK calls

New APIs

• Chain.forward_kinematics_tensor(th) — returns (num_frames, B, 4, 4) tensor, compile-compatible
• SerialChain.jacobian_tensor(th) — returns (B, 6, DOF) Jacobian, compile-compatible
• PseudoInverseIK new parameters: lm_damping, position_weight, orientation_weight, clamp_to_limits

torch.compile support

Forward kinematics, Jacobian, and rotation conversions (matrix_to_quaternion, quaternion_to_axis_angle, axis_angle_to_quaternion) are now compatible with torch.compile(fullgraph=True). Boolean indexing patterns were replaced with torch.where, torch.gather, and torch.clamp.

chain = pk.build_serial_chain_from_urdf(urdf, "end_effector")
fk_compiled = torch.compile(chain.forward_kinematics_tensor, fullgraph=True)
jac_compiled = torch.compile(chain.jacobian_tensor, fullgraph=True)

The IK solver also supports use_compile=True to compile the FK, Jacobian, and IK step kernels.

IK improvements

• Adaptive LM damping (lm_damping=0.1 default): Scales regularization by error magnitude — large errors get more damping (stable), small errors get less (fast convergence)
• Default learning rate changed from 0.2 to 1.0
• Bug fixes: hardcoded Adam optimizer, line search crash on non-converged problems, dtype mismatch with non-default dtypes
• Per-coordinate weighting: position_weight and orientation_weight parameters for task-specific tuning
• Fused IK kernel: delta_pose + DLS solve in a single function, compilable with torch.compile

Performance

All benchmarks: Kuka IIWA 7-DOF, 500 goals, 30 iterations. GPU is NVIDIA RTX 4070.

FK & Jacobian (B=500)

| | v0.8.0 | v0.9.0 eager | v0.9.0 compiled | |---|---|---|---| | CPU FK | 2.07ms | 1.09ms (1.9x) | 0.52ms (4.0x) | | CPU Jacobian | 5.51ms | 1.36ms (4.1x) | 0.59ms (9.3x) | | GPU FK | 1.45ms | 0.64ms (2.3x) | 0.32ms (4.5x) | | GPU Jacobian | 2.76ms | 0.75ms (3.7x) | 0.36ms (7.7x) |

IK — 500 goals, 10 retries

| | v0.8.0 | v0.9.0 eager | v0.9.0 compiled | |---|---|---|---| | CPU | 614ms, 99.6% | 140ms, 99.8% (4.4x) | 92ms, 99.8% (6.7x) | | GPU | 387ms, 98.8% | 76ms, 99.4% (5.1x) | 51ms, 99.4% (7.6x) |

IK — 500 goals, 1 retry (hardest case)

| | v0.8.0 | v0.9.0 eager | v0.9.0 compiled | |---|---|---|---| | CPU | 239ms, 67.4% | 70ms, 89.4% (3.4x) | 32ms, 89.4% (7.5x) | | GPU | 263ms, 71.0% | 51ms, 75.2% (5.2x) | 24ms, 75.2% (11.0x) |

Bug fixes

• Fixed SerialChain crash when all joints are fixed (0 DOF)
• Fixed IKSolution dtype mismatch when chain uses non-default dtype
• Fixed line search crash on non-converged problems
• Eliminated redundant FK calls in IK loop and Jacobian wrapper