[Kamino] Fix passive universal joint FK constraint (#3720)

Author: rubengrandia

newton/_src/solvers/kamino/_src/solvers/fk/kernels.py

@@ -193,11 +193,15 @@ def _eval_rotation_jacobian_blocks(
 
 @wp.func
 def _eval_passive_universal_jacobian_blocks(
-    X_T: wp.mat33f, q_base: wp.quatf, q_follower: wp.quatf, has_base: wp.bool
+    X_B_T: wp.mat33f,
+    X_F_T: wp.mat33f,
+    q_base: wp.quatf,
+    q_follower: wp.quatf,
+    has_base: wp.bool,
 ) -> tuple[wp.vec4f, wp.vec4f]:
     """Evaluate the base and follower blocks of a passive universal constraint Jacobian."""
-    a_x = X_T[0]
-    a_y = X_T[1]
+    a_x = X_B_T[0]  # x-axis on base
+    a_y = X_F_T[1]  # y-axis on follower
     jac_q_base = wp.vec4f(0.0)
     if has_base:
         a_y_follower = unit_quat_apply(q_follower, a_y)
@@ -970,6 +974,7 @@ def _eval_joint_constraints(
         joints_bid_B: wp.array[wp.int32],
         joints_bid_F: wp.array[wp.int32],
         joints_X_Bj: wp.array[wp.mat33f],
+        joints_X_Fj: wp.array[wp.mat33f],
         joints_B_r_B: wp.array[wp.vec3f],
         joints_F_r_F: wp.array[wp.vec3f],
         bodies_q: wp.array[wp.transformf],
@@ -996,6 +1001,7 @@ def _eval_joint_constraints(
             joints_bid_B: Joint base body id
             joints_bid_F: Joint follower body id
             joints_X_Bj: Joint local frame on base body
+            joints_X_Fj: Joint local frame on follower body
             joints_B_r_B: Joint local position on base body
             joints_F_r_F: Joint local position on follower body
             bodies_q: Body poses
@@ -1059,7 +1065,7 @@ def _eval_joint_constraints(
 
                 # Compute constraint (dot product between x axis on base and y axis on follower)
                 a_x = X_T[0]
-                a_y = X_T[1]
+                a_y = wp.transpose(joints_X_Fj[jt_id_tot])[1]
                 a_x_base = unit_quat_apply(q_base, a_x)
                 a_y_follower = unit_quat_apply(q_follower, a_y)
                 ct = -wp.dot(a_x_base, a_y_follower)
@@ -1168,6 +1174,7 @@ def _eval_joint_constraints_jacobian(
         joints_bid_B: wp.array[wp.int32],
         joints_bid_F: wp.array[wp.int32],
         joints_X_Bj: wp.array[wp.mat33f],
+        joints_X_Fj: wp.array[wp.mat33f],
         joints_B_r_B: wp.array[wp.vec3f],
         joints_F_r_F: wp.array[wp.vec3f],
         bodies_q: wp.array[wp.transformf],
@@ -1191,6 +1198,7 @@ def _eval_joint_constraints_jacobian(
             joints_bid_B: Joint base body id
             joints_bid_F: Joint follower body id
             joints_X_Bj: Joint local frame on base body
+            joints_X_Fj: Joint local frame on follower body
             joints_B_r_B: Joint local position on base body
             joints_F_r_F: Joint local position on follower body
             bodies_q: Body poses
@@ -1272,7 +1280,7 @@ def _eval_joint_constraints_jacobian(
 
                 # Compute constraint Jacobian (cross product between x axis on base and y axis on follower)
                 jac_q_base, jac_q_follower = _eval_passive_universal_jacobian_blocks(
-                    X_T, q_base, q_follower, base_id_tot >= 0
+                    X_T, wp.transpose(joints_X_Fj[jt_id_tot]), q_base, q_follower, base_id_tot >= 0
                 )
 
                 # Write out Jacobian
@@ -1304,6 +1312,7 @@ def _eval_joint_constraints_sparse_jacobian(
         joints_bid_B: wp.array[wp.int32],
         joints_bid_F: wp.array[wp.int32],
         joints_X_Bj: wp.array[wp.mat33f],
+        joints_X_Fj: wp.array[wp.mat33f],
         joints_B_r_B: wp.array[wp.vec3f],
         joints_F_r_F: wp.array[wp.vec3f],
         bodies_q: wp.array[wp.transformf],
@@ -1328,6 +1337,7 @@ def _eval_joint_constraints_sparse_jacobian(
             joints_bid_B: Joint base body id
             joints_bid_F: Joint follower body id
             joints_X_Bj: Joint local frame on base body
+            joints_X_Fj: Joint local frame on follower body
             joints_B_r_B: Joint local position on base body
             joints_F_r_F: Joint local position on follower body
             bodies_q: Body poses
@@ -1413,7 +1423,7 @@ def _eval_joint_constraints_sparse_jacobian(
 
                 # Compute constraint Jacobian (cross product between x axis on base and y axis on follower)
                 jac_q_base, jac_q_follower = _eval_passive_universal_jacobian_blocks(
-                    X_T, q_base, q_follower, base_id >= 0
+                    X_T, wp.transpose(joints_X_Fj[jt_id_tot]), q_base, q_follower, base_id >= 0
                 )
 
                 # Write out Jacobian

newton/_src/solvers/kamino/_src/solvers/fk/solver.py

@@ -1269,6 +1269,7 @@ def _eval_kinematic_constraints(
                 self.joints_bid_B,
                 self.joints_bid_F,
                 self.joints_X_Bj,
+                self.joints_X_Fj,
                 self.joints_B_r_Bj,
                 self.joints_F_r_Fj,
                 bodies_q,
@@ -1344,6 +1345,7 @@ def _eval_kinematic_constraints_jacobian(
                 self.joints_bid_B,
                 self.joints_bid_F,
                 self.joints_X_Bj,
+                self.joints_X_Fj,
                 self.joints_B_r_Bj,
                 self.joints_F_r_Fj,
                 bodies_q,
@@ -1396,6 +1398,7 @@ def _assemble_sparse_jacobian(
                 self.joints_bid_B,
                 self.joints_bid_F,
                 self.joints_X_Bj,
+                self.joints_X_Fj,
                 self.joints_B_r_Bj,
                 self.joints_F_r_Fj,
                 bodies_q,

newton/_src/solvers/kamino/tests/test_solvers_forward_kinematics.py

@@ -35,6 +35,7 @@
     sample_body_poses,
 )
 from newton.tests.utils.basics import build_cartpole
+from newton.tests.utils.testing import build_unary_universal_joint_test
 
 ###
 # Module configs
@@ -130,6 +131,75 @@ def eval_constraints(bodies_q_stepped_np):
         self.assertTrue(success)
 
 
+class PassiveUniversalJointFrameForwardKinematics(unittest.TestCase):
+    def setUp(self):
+        if not test_context.setup_done:
+            setup_tests(clear_cache=False)
+        self.default_device = wp.get_device(test_context.device)
+
+    def tearDown(self):
+        self.default_device = None
+
+    def test_follower_joint_frame(self):
+        """Test kinematic constraints and jacobian for a passive universal joint.
+
+        The follower joint frame is rotated relative to the base joint frame,
+        while the follower body is counter-rotated so both joint frames coincide
+        in world coordinates. The base-frame X axis and follower-frame Y axis
+        are therefore orthogonal, as required by a universal joint.
+        """
+        # Build a single body attached to the world and make both rotational
+        # degrees of freedom passive.
+        builder = build_unary_universal_joint_test(limits=True, ground=False)
+        builder.joint_target_mode[0] = newton.JointTargetMode.NONE
+        builder.joint_target_mode[1] = newton.JointTargetMode.NONE
+        model = ModelKamino.from_newton(builder.finalize(device=self.default_device))
+
+        # Rotate the follower-local joint frame by 90 degrees about Z without
+        # changing the base-local joint frame.
+        q_X_F = wp.quat_from_axis_angle(wp.vec3f(0.0, 0.0, 1.0), 0.5 * wp.pi)
+        X_Fj = model.joints.X_Fj.numpy()
+        X_Fj[0] = np.asarray(wp.quat_to_matrix(q_X_F), dtype=np.float32).reshape(3, 3)
+        model.joints.X_Fj.assign(X_Fj)
+
+        # Counter-rotate and translate the follower body so its joint frame has
+        # the same world-space pose as the base joint frame.
+        q_F = wp.quat_inverse(q_X_F)
+        B_r_Bj = wp.vec3f(model.joints.B_r_Bj.numpy()[0])
+        F_r_Fj = wp.vec3f(model.joints.F_r_Fj.numpy()[0])
+        r_F = B_r_Bj - wp.quat_rotate(q_F, F_r_Fj)
+        bodies_q = wp.array([wp.transformf(r_F, q_F)], dtype=wp.transformf, device=self.default_device)
+
+        # Coincident joint frames satisfy all three anchor constraints and the
+        # universal joint's rotational orthogonality constraint.
+        solver = ForwardKinematicsSolver(model)
+        actuators_q = wp.empty(0, dtype=wp.float32, device=self.default_device)
+        target_transforms = solver.eval_position_control_transformations(actuators_q, None)
+        constraints = solver.eval_kinematic_constraints(bodies_q, target_transforms).numpy()[0]
+        np.testing.assert_allclose(constraints, 0.0, atol=1.0e-6)
+
+        # Validate jacobian with finite differences
+        bodies_q_np = bodies_q.numpy().reshape(-1)
+        jacobian = solver.eval_kinematic_constraints_jacobian(bodies_q, target_transforms).numpy()[0]
+
+        def eval_constraints(bodies_q_stepped_np):
+            bodies_q.assign(bodies_q_stepped_np)
+            stepped_constraints = solver.eval_kinematic_constraints(bodies_q, target_transforms).numpy()[0]
+            bodies_q.assign(bodies_q_np)
+            return stepped_constraints
+
+        self.assertTrue(
+            diff_check(
+                eval_constraints,
+                jacobian,
+                bodies_q_np,
+                epsilon=1.0e-4,
+                tolerance_abs=5.0e-3,
+                tolerance_rel=5.0e-3,
+            )
+        )
+
+
 class SparseJacobianSingleJointCheckForwardKinematics(unittest.TestCase):
     def setUp(self):
         if not test_context.setup_done: