Fix Jacobian computation in FeatherstoneIntegrator

CHANGELOG.md

@@ -53,6 +53,7 @@
 - Fix an aliasing issue with zero-copy array initialization from numpy introduced in 1.3.0
 - Fix a bug in `wp.sim.collide.triangle_closest_point_barycentric()` where the returned barycentric coordinates may be incorrect when the closest point lies on an edge.
 - Unexposed `wp.rand*()`, `wp.sample*()`, and `wp.poisson()` from Python's runtime.
+- Fix bug in `FeatherstoneIntegrator` where `eval_rigid_jacobian` could give incorrect results or reach an infinite loop when the body and joint indices were not in the same order. Added `Model.joint_ancestor` to fix the indexing from a joint to its parent joint in the articulation.
 
 ## [1.3.1] - 2024-07-27
 

warp/sim/integrator_featherstone.py

@@ -1015,7 +1015,7 @@ def eval_rigid_tau(
 def eval_rigid_jacobian(
     articulation_start: wp.array(dtype=int),
     articulation_J_start: wp.array(dtype=int),
-    joint_parent: wp.array(dtype=int),
+    joint_ancestor: wp.array(dtype=int),
     joint_qd_start: wp.array(dtype=int),
     joint_S_s: wp.array(dtype=wp.spatial_vector),
     # outputs
@@ -1051,7 +1051,7 @@ def eval_rigid_jacobian(
                 for k in range(6):
                     J[J_offset + dense_index(articulation_dof_count, row_start + k, col)] = S[k]
 
-            j = joint_parent[j]
+            j = joint_ancestor[j]
 
 
 @wp.func
@@ -1769,7 +1769,7 @@ def simulate(self, model: Model, state_in: State, state_out: State, dt: float, c
                             inputs=[
                                 model.articulation_start,
                                 self.articulation_J_start,
-                                model.joint_parent,
+                                model.joint_ancestor,
                                 model.joint_qd_start,
                                 state_aug.joint_S_s,
                             ],

warp/sim/model.py

@@ -559,6 +559,7 @@ class Model:
         joint_type (array): Joint type, shape [joint_count], int
         joint_parent (array): Joint parent body indices, shape [joint_count], int
         joint_child (array): Joint child body indices, shape [joint_count], int
+        joint_ancestor (array): Maps from joint index to the index of the joint that has the current joint parent body as child (-1 if no such joint ancestor exists), shape [joint_count], int
         joint_X_p (array): Joint transform in parent frame, shape [joint_count, 7], float
         joint_X_c (array): Joint mass frame in child frame, shape [joint_count, 7], float
         joint_axis (array): Joint axis in child frame, shape [joint_axis_count, 3], float
@@ -729,6 +730,7 @@ def __init__(self, device=None):
         self.joint_type = None
         self.joint_parent = None
         self.joint_child = None
+        self.joint_ancestor = None
         self.joint_X_p = None
         self.joint_X_c = None
         self.joint_axis = None
@@ -4520,6 +4522,14 @@ def finalize(self, device=None, requires_grad=False) -> Model:
             m.joint_q = wp.array(self.joint_q, dtype=wp.float32, requires_grad=requires_grad)
             m.joint_qd = wp.array(self.joint_qd, dtype=wp.float32, requires_grad=requires_grad)
             m.joint_name = self.joint_name
+            # compute joint ancestors
+            child_to_joint = {}
+            for i, child in enumerate(self.joint_child):
+                child_to_joint[child] = i
+            parent_joint = []
+            for parent in self.joint_parent:
+                parent_joint.append(child_to_joint.get(parent, -1))
+            m.joint_ancestor = wp.array(parent_joint, dtype=wp.int32)
 
             # dynamics properties
             m.joint_armature = wp.array(self.joint_armature, dtype=wp.float32, requires_grad=requires_grad)