Add torch.compile support, improve IK convergence with adaptive LM damping#67
Merged
Conversation
Split FK into a compilable tensor kernel (_forward_kinematics_tensor) and a non-compiled wrapper. The kernel uses level-by-level BFS traversal with pre-allocated buffers instead of Python loops with dict memoization, enabling torch.compile to trace the full graph without breaks. Key changes: - Precompute static_offsets, parent indices, and BFS level structure at init - Replace None offset checks with pre-multiplied identity matrices - Replace per-frame Python loop with batched level-by-level matmuls - Make axis_and_angle_to_matrix_44 and axis_and_d_to_pris_matrix compile-safe by avoiding .to(tensor) calls and CPU tensor creation - Add test_compile_fk verifying fullgraph=True compilation, correctness, and gradient flow across URDF, SDF, and MJCF robots Closes #66 (phase 1) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…c API) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add jacobian_tensor() on SerialChain as a compilable geometric Jacobian kernel using the forward formula with FK transforms. Inline calc_jacobian logic into SerialChain.jacobian() and reduce jacobian.py to a thin backward-compat stub. Skip unused joint type computation in FK (revolute-only robots no longer compute prismatic transforms), improving large-batch performance. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace boolean indexing with torch.where/torch.clamp/torch.gather in rotation_conversions.py (_sqrt_positive_part, matrix_to_quaternion, quaternion_to_axis_angle, axis_angle_to_quaternion) for torch.compile(fullgraph=True) compatibility. This also improves eager performance 1.5-10x by eliminating intermediate mask allocations. Add fused _ik_step_kernel (delta_pose + DLS) and route IK solver through compiled forward_kinematics_tensor/jacobian_tensor when use_compile=True, giving ~2x end-to-end IK speedup. Guard _check_valid_rotation_matrix with torch.compiler.is_compiling() to skip torch.allclose during tracing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
IKSolution allocated tensors with default dtype (float32), causing RuntimeError when the chain uses float64 and the tensor API code paths assign float64 joint values into float32 solution buffers. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
jacobian_tensor already computes full FK internally. Added ret_eef_pose parameter to return the EEF pose from that same pass, avoiding a second FK call on the same joint angles. Applied in the IK solve loop, _enforce_limits refinement, and the jacobian() wrapper. ~22% IK speedup. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Instead of computing a separate FK after each step just for convergence checking, reuse the dx already computed by _ik_step_kernel at the start of each iteration. One final FK after the loop catches the last step. Same convergence/accuracy, ~34% faster (109.8ms → 81.8ms on 100 problems). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- optimizer_method was ignored; Adam was always created. Now uses the user-specified optimizer class. - BacktrackingLineSearch crashed with UnboundLocalError when all problems were already converged before the search loop ran. Initialize improvement to zeros. - Line search used the slow forward_kinematics() wrapper. Now accepts fk_fn/eef_frame_idx to use the tensor API, falling back to the wrapper for backward compatibility. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
DLS with small regularization is a Gauss-Newton step; taking the full step (lr=1.0) converges in ~6 iterations vs ~30 at lr=0.2, yielding ~2.4x speedup with identical convergence rates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…t FK - Replace boolean indexing in sol.update and no_improve_counter with torch.where for compile compatibility - Fix jacobian() with locations doing FK twice by accepting pre-computed all_transforms in jacobian_tensor - Remove broken optimizer abstraction (q.grad = -dq only works for SGD) - Use _ik_step_kernel in _enforce_limits and final convergence check - Remove dead code: _compute_dq_kernel, compute_dq, apply_mask - Restructure PseudoInverseIKWithSVD to override _ik_step_fn with SVD-based kernel instead of compute_dq method Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…s to IK Inspired by mink's QP-based IK solver: - Adaptive Levenberg-Marquardt damping (lm_damping=0.1 default): scales regularization by error magnitude for stable convergence. Improves 1-retry convergence by ~20pp on typical problems. - Per-coordinate weighting (position_weight, orientation_weight): allows prioritizing position vs orientation accuracy in the DLS objective. - clamp_to_limits option (default False): projected gradient clamping after each step. Disabled by default as hard clamping hurts DLS. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
torch.stack on an empty list fails when a SerialChain has no non-fixed joints (e.g. gripper_link → ee_gripper_link). Use torch.zeros(0, 3) fallback for the empty dof_axes case. Fixes test_extract_serial_chain_from_tree and test_serial_chain_non_root_start. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
torch.compile(fullgraph=True)support for forward kinematics, Jacobian, and rotation conversions via newforward_kinematics_tensorandjacobian_tensorAPIs, addressing Feature Request: Torch compilable/ONNX exportable with latest version #66lm_damping=0.1default), boosting single-retry convergence by ~20pp with negligible speed costtorch.whereposition_weight,orientation_weight) andclamp_to_limitsoptions to IKChanges
torch.compile support (Phases 1-3):
Chain.forward_kinematics_tensor(th)— returns(num_frames, B, 4, 4)tensor, compile-compatibleSerialChain.jacobian_tensor(th)— returns(B, 6, DOF)Jacobian, compile-compatiblematrix_to_quaternion,quaternion_to_axis_angle,axis_angle_to_quaternion) to usetorch.where/torch.gather/torch.clampinstead of boolean indexingforward_kinematicsnow delegates toforward_kinematics_tensorinternallyIK improvements:
mu = lm_damping * ||error||^2scales regularization by error magnitude_ik_step_kernel(delta_pose + DLS in one compilable function)_compute_dq_kernel,compute_dq,apply_mask, unused importsBenchmarks (Kuka IIWA, 500 goals, 10 retries, 30 iters)
Single-retry (hardest case): 67% → 89% convergence on CPU, 65% → 75% on GPU.
Test plan
pytest, CPU + GPU)fullgraph=True