Fix translated floating-base IK Jacobians (#4088)

Author: ooctipus

changelog/3866.fixed.md

@@ -0,0 +1 @@
+Fix distance-dependent drift in analytic IK Jacobians for FREE and DISTANCE joints.

newton/_src/sim/articulation.py

@@ -934,25 +934,23 @@ def eval_ik(
 @wp.func
 def write_free_distance_motion_subspace(
     X_pa_world: wp.transform,
-    x_child_com_world: wp.vec3,
+    pivot_world: wp.vec3,
     qd_start: int,
     # outputs
     joint_S_s: wp.array[wp.spatial_vector],
 ):
     """Write the 6 motion-subspace columns for a FREE/DISTANCE joint.
 
-    Used by both the Featherstone inverse-dynamics path (``jcalc_motion``) and
-    the IK/Jacobian path (``jcalc_motion_subspace``) so they agree on the exact
-    column layout. Linear DOFs act at the child body's COM; angular DOFs are
-    world-aligned axes expressed through ``X_pa_world``.
+    Linear DOFs follow the parent-anchor axes. Angular columns describe
+    rotations about ``pivot_world`` as world-origin spatial twists.
 
     Args:
         X_pa_world: Parent-anchor world transform (``X_wp * joint_X_p``) used
             to rotate the joint's parent-anchor axes into the world frame.
             This is *not* the classical Featherstone ``X_sc`` (spatial-to-
             child); Newton's FREE/DISTANCE joint coordinates live in the
             parent-anchor basis.
-        x_child_com_world: World-space position of the child body's COM.
+        pivot_world: World-space point about which angular columns rotate.
         qd_start: Starting velocity-DOF index for this joint.
         joint_S_s: Output spatial-vector subspace array; six slots starting at
             ``qd_start`` are overwritten.
@@ -964,9 +962,9 @@ def write_free_distance_motion_subspace(
     joint_S_s[qd_start + 0] = wp.spatial_vector(axis_world_x, wp.vec3())
     joint_S_s[qd_start + 1] = wp.spatial_vector(axis_world_y, wp.vec3())
     joint_S_s[qd_start + 2] = wp.spatial_vector(axis_world_z, wp.vec3())
-    joint_S_s[qd_start + 3] = wp.spatial_vector(-wp.cross(axis_world_x, x_child_com_world), axis_world_x)
-    joint_S_s[qd_start + 4] = wp.spatial_vector(-wp.cross(axis_world_y, x_child_com_world), axis_world_y)
-    joint_S_s[qd_start + 5] = wp.spatial_vector(-wp.cross(axis_world_z, x_child_com_world), axis_world_z)
+    joint_S_s[qd_start + 3] = wp.spatial_vector(-wp.cross(axis_world_x, pivot_world), axis_world_x)
+    joint_S_s[qd_start + 4] = wp.spatial_vector(-wp.cross(axis_world_y, pivot_world), axis_world_y)
+    joint_S_s[qd_start + 5] = wp.spatial_vector(-wp.cross(axis_world_z, pivot_world), axis_world_z)
 
 
 @wp.func

newton/_src/sim/ik/ik_lbfgs_optimizer.py

@@ -854,14 +854,12 @@ def _compute_motion_subspace(
             inputs=[
                 self.model.joint_type,
                 self.model.joint_parent,
-                self.model.joint_child,
                 self.model.joint_q_start,
                 self.model.joint_qd_start,
                 joint_q_in,
                 self.model.joint_axis,
                 self.model.joint_dof_dim,
                 body_q,
-                self.model.body_com,
                 self.model.joint_X_p,
             ],
             outputs=[
@@ -1348,7 +1346,10 @@ def _select_best_step_template(
         _select_best_step_tiled = wp.kernel(enable_backward=False, module="unique")(_select_best_step_template)
 
         # late-import jcalc_* helpers to avoid circular import error
-        from ...sim.articulation import jcalc_motion_subspace  # noqa: PLC0415
+        from ...sim.articulation import (  # noqa: PLC0415
+            jcalc_motion_subspace,
+            write_free_distance_motion_subspace,
+        )
         from ...solvers.featherstone.kernels import (  # noqa: PLC0415
             jcalc_integrate,
             jcalc_transform,
@@ -1427,14 +1428,12 @@ def _integrate_dq_dof(
         def _compute_motion_subspace_2d(
             joint_type: wp.array[wp.int32],  # (n_joints)
             joint_parent: wp.array[wp.int32],  # (n_joints)
-            joint_child: wp.array[wp.int32],  # (n_joints)
             joint_q_start: wp.array[wp.int32],  # (n_joints + 1)
             joint_qd_start: wp.array[wp.int32],  # (n_joints + 1)
             joint_q: wp.array2d[wp.float32],  # (n_batch, n_coords)
             joint_axis: wp.array[wp.vec3],  # (n_joint_dof_count)
             joint_dof_dim: wp.array2d[wp.int32],  # (n_joints, 2)
             body_q: wp.array2d[wp.transform],  # (n_batch, n_bodies)
-            body_com: wp.array[wp.vec3],  # (n_bodies)
             joint_X_p: wp.array[wp.transform],  # (n_joints)
             # outputs
             joint_S_s: wp.array2d[wp.spatial_vector],  # (n_batch, n_joint_dof_count)
@@ -1443,7 +1442,6 @@ def _compute_motion_subspace_2d(
 
             type = joint_type[joint_idx]
             parent = joint_parent[joint_idx]
-            child = joint_child[joint_idx]
             q_start = joint_q_start[joint_idx]
             qd_start = joint_qd_start[joint_idx]
 
@@ -1459,16 +1457,10 @@ def _compute_motion_subspace_2d(
             S_s_out = joint_S_s[row]
 
             if type == JointType.FREE or type == JointType.DISTANCE:
-                jcalc_motion_subspace(
-                    type,
-                    joint_axis,
-                    joint_q_1d,
-                    lin_axis_count,
-                    ang_axis_count,
+                # IK integration applies angular increments about the fixed parent anchor.
+                write_free_distance_motion_subspace(
                     X_wpj,
-                    body_q[row, child],
-                    body_com[child],
-                    q_start,
+                    wp.transform_get_translation(X_wpj),
                     qd_start,
                     S_s_out,
                 )

newton/_src/sim/ik/ik_lm_optimizer.py

@@ -568,14 +568,12 @@ def _compute_motion_subspace(
             inputs=[
                 self.model.joint_type,
                 self.model.joint_parent,
-                self.model.joint_child,
                 self.model.joint_q_start,
                 self.model.joint_qd_start,
                 joint_q_in,
                 self.model.joint_axis,
                 self.model.joint_dof_dim,
                 body_q,
-                self.model.body_com,
                 self.model.joint_X_p,
             ],
             outputs=[
@@ -790,7 +788,10 @@ def _template(
         _lm_solve_tiled = wp.kernel(enable_backward=False, module="unique")(_template)
 
         # late-import jcalc_* helpers to avoid circular import error
-        from ...sim.articulation import jcalc_motion_subspace  # noqa: PLC0415
+        from ...sim.articulation import (  # noqa: PLC0415
+            jcalc_motion_subspace,
+            write_free_distance_motion_subspace,
+        )
         from ...solvers.featherstone.kernels import (  # noqa: PLC0415
             jcalc_integrate,
             jcalc_transform,
@@ -869,14 +870,12 @@ def _integrate_dq_dof(
         def _compute_motion_subspace_2d(
             joint_type: wp.array[wp.int32],  # (n_joints)
             joint_parent: wp.array[wp.int32],  # (n_joints)
-            joint_child: wp.array[wp.int32],  # (n_joints)
             joint_q_start: wp.array[wp.int32],  # (n_joints + 1)
             joint_qd_start: wp.array[wp.int32],  # (n_joints + 1)
             joint_q: wp.array2d[wp.float32],  # (n_batch, n_coords)
             joint_axis: wp.array[wp.vec3],  # (n_joint_dof_count)
             joint_dof_dim: wp.array2d[wp.int32],  # (n_joints, 2)
             body_q: wp.array2d[wp.transform],  # (n_batch, n_bodies)
-            body_com: wp.array[wp.vec3],  # (n_bodies)
             joint_X_p: wp.array[wp.transform],  # (n_joints)
             # outputs
             joint_S_s: wp.array2d[wp.spatial_vector],  # (n_batch, n_joint_dof_count)
@@ -885,7 +884,6 @@ def _compute_motion_subspace_2d(
 
             type = joint_type[joint_idx]
             parent = joint_parent[joint_idx]
-            child = joint_child[joint_idx]
             q_start = joint_q_start[joint_idx]
             qd_start = joint_qd_start[joint_idx]
 
@@ -901,16 +899,10 @@ def _compute_motion_subspace_2d(
             S_s_out = joint_S_s[row]
 
             if type == JointType.FREE or type == JointType.DISTANCE:
-                jcalc_motion_subspace(
-                    type,
-                    joint_axis,
-                    joint_q_1d,
-                    lin_axis_count,
-                    ang_axis_count,
+                # IK integration applies angular increments about the fixed parent anchor.
+                write_free_distance_motion_subspace(
                     X_wpj,
-                    body_q[row, child],
-                    body_com[child],
-                    q_start,
+                    wp.transform_get_translation(X_wpj),
                     qd_start,
                     S_s_out,
                 )

newton/tests/test_ik.py

@@ -163,6 +163,25 @@ def _build_descendant_free_distance(device, joint_type) -> tuple[newton.Model, i
     return builder.finalize(device=device, requires_grad=True), child
 
 
+def _build_root_free_distance(device, joint_type) -> tuple[newton.Model, int]:
+    builder = newton.ModelBuilder()
+    body = builder.add_link(mass=1.0, inertia=wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0))
+    builder.body_com[body] = wp.vec3(0.2, -0.1, 0.3)
+    joint = _add_free_distance_joint(
+        builder=builder,
+        joint_type=joint_type,
+        parent=-1,
+        child=body,
+        parent_xform=wp.transform(
+            wp.vec3(0.3, -0.4, 0.6),
+            wp.quat_from_axis_angle(wp.normalize(wp.vec3(0.2, 1.0, -0.3)), 0.7),
+        ),
+        child_xform=wp.transform_identity(),
+    )
+    builder.add_articulation([joint])
+    return builder.finalize(device=device, requires_grad=True), body
+
+
 # ----------------------------------------------------------------------------
 # helpers - D6
 # ----------------------------------------------------------------------------
@@ -711,6 +730,54 @@ def test_d6_jacobian_compare(test, device):
     _jacobian_compare(test, device, _d6_objective_builder)
 
 
+def test_free_distance_translated_jacobian_compare(test, device, joint_type, optimizer):
+    """Match analytic and autodiff Jacobians for a translated floating body."""
+    with wp.ScopedDevice(device):
+        translations = (0.0, 1.0, 5.0, 20.0)
+        model, body = _build_root_free_distance(device, joint_type)
+        joint_q = np.zeros((len(translations), model.joint_coord_count), dtype=np.float32)
+        joint_q[:, 0] = translations
+        joint_q[:, 6] = 1.0
+        target_positions = wp.zeros(len(translations), dtype=wp.vec3, device=device)
+
+        solver_auto = ik.IKSolver(
+            model,
+            len(translations),
+            [ik.IKObjectivePosition(body, wp.vec3(), target_positions)],
+            optimizer=optimizer,
+            jacobian_mode=ik.IKJacobianType.AUTODIFF,
+        )
+        solver_ana = ik.IKSolver(
+            model,
+            len(translations),
+            [ik.IKObjectivePosition(body, wp.vec3(), target_positions)],
+            optimizer=optimizer,
+            jacobian_mode=ik.IKJacobianType.ANALYTIC,
+        )
+        q_auto = wp.array(joint_q, device=device, requires_grad=True)
+        q_ana = wp.array(joint_q, device=device)
+
+        solver_auto._impl._compute_residuals(q_auto)
+        solver_ana._impl._compute_residuals(q_ana)
+        if optimizer == ik.IKOptimizer.LM:
+            jacobian_auto = solver_auto._impl._jacobian_at(solver_auto._impl._ctx_solver(q_auto)).numpy()
+            jacobian_ana = solver_ana._impl._jacobian_at(solver_ana._impl._ctx_solver(q_ana)).numpy()
+            for problem_idx, translation in enumerate(translations):
+                with test.subTest(translation=translation):
+                    assert_np_equal(jacobian_auto[problem_idx], jacobian_ana[problem_idx], tol=1e-4)
+        else:
+            gradient_auto = wp.zeros((len(translations), model.joint_dof_count), dtype=wp.float32, device=device)
+            gradient_ana = wp.zeros_like(gradient_auto)
+            solver_auto._impl._gradient_at(solver_auto._impl._ctx_solver(q_auto), gradient_auto)
+            solver_ana._impl._gradient_at(solver_ana._impl._ctx_solver(q_ana), gradient_ana)
+            assert_np_equal(gradient_auto.numpy(), gradient_ana.numpy(), tol=1e-4)
+
+        motion_subspace = solver_ana._impl.joint_S_s.numpy()
+        for problem_idx, translation in enumerate(translations[1:], start=1):
+            with test.subTest(motion_subspace_translation=translation):
+                assert_np_equal(motion_subspace[0], motion_subspace[problem_idx], tol=1e-6)
+
+
 # ----------------------------------------------------------------------------
 # 3.  Test-class registration per device
 # ----------------------------------------------------------------------------
@@ -761,6 +828,16 @@ class TestIKModes(unittest.TestCase):
 add_function_test(TestIKModes, "test_rotation_jacobian_compare", test_rotation_jacobian_compare, cuda_devices)
 add_function_test(TestIKModes, "test_joint_limit_jacobian_compare", test_joint_limit_jacobian_compare, devices)
 add_function_test(TestIKModes, "test_d6_jacobian_compare", test_d6_jacobian_compare, cuda_devices)
+for optimizer, optimizer_name in ((ik.IKOptimizer.LM, "lm"), (ik.IKOptimizer.LBFGS, "lbfgs")):
+    for joint_type in (newton.JointType.FREE, newton.JointType.DISTANCE):
+        add_function_test(
+            TestIKModes,
+            f"test_translated_{_joint_type_name(joint_type)}_{optimizer_name}_jacobian_compare",
+            test_free_distance_translated_jacobian_compare,
+            devices,
+            joint_type=joint_type,
+            optimizer=optimizer,
+        )
 
 
 if __name__ == "__main__":