[FEATURE] Add point-specific Jacobian support

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -798,14 +798,17 @@ def _add_equality(self, name, type, objs_name, data, sol_params):
     # ------------------------------------------------------------------------------------
 
     @gs.assert_built
-    def get_jacobian(self, link):
+    def get_jacobian(self, link, local_point=None):
         """
-        Get the Jacobian matrix for a target link.
+        Get the spatial Jacobian for a point on a target link.
 
         Parameters
         ----------
         link : RigidLink
             The target link.
+        local_point : torch.Tensor or None, shape (3,)
+            Coordinates of the point in the link’s *local* frame.
+            If None, the link origin is used (back-compat).
 
         Returns
         -------
@@ -820,31 +823,50 @@ def get_jacobian(self, link):
         if self.n_dofs == 0:
             gs.raise_exception("Entity has zero dofs.")
 
-        self._kernel_get_jacobian(link.idx)
+        if local_point is None:
+            self._kernel_get_jacobian_zero(link.idx)
+        else:
+            if not isinstance(local_point, (list, tuple, np.ndarray)):
+                gs.raise_exception("`local_point` must be length-3 iterable.")
+            if len(local_point) != 3:
+                gs.raise_exception("`local_point` must have 3 elements.")
+            p_local = ti.Vector(local_point, dt=ti.f32)
+            self._kernel_get_jacobian(link.idx, p_local)
 
         jacobian = self._jacobian.to_torch(gs.device).permute(2, 0, 1)
         if self._solver.n_envs == 0:
             jacobian = jacobian.squeeze(0)
 
         return jacobian
 
+    @ti.func
+    def _impl_get_jacobian(self, tgt_link_idx, i_b, p_local):
+        self._func_get_jacobian(
+            tgt_link_idx,
+            i_b,
+            p_local,
+            ti.Vector.one(gs.ti_int, 3),
+            ti.Vector.one(gs.ti_int, 3),
+        )
+
     @ti.kernel
-    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32):
-        ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
+    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.template()):
         for i_b in range(self._solver._B):
-            self._func_get_jacobian(
-                tgt_link_idx,
-                i_b,
-                ti.Vector.one(gs.ti_int, 3),
-                ti.Vector.one(gs.ti_int, 3),
-            )
+            self._impl_get_jacobian(tgt_link_idx, i_b, p_local)
+
+    @ti.kernel
+    def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):
+        zero_vec = ti.Vector.zero(ti.f32, 3)
+        for i_b in range(self._solver._B):
+            self._impl_get_jacobian(tgt_link_idx, i_b, zero_vec)
 
     @ti.func
-    def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
+    def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):
         for i_row, i_d in ti.ndrange(6, self.n_dofs):
             self._jacobian[i_row, i_d, i_b] = 0.0
 
-        tgt_link_pos = self._solver.links_state[tgt_link_idx, i_b].pos
+        tgt_link_state = self._solver.links_state[tgt_link_idx, i_b]
+        tgt_link_pos = tgt_link_state.pos + gu.ti_transform_by_quat(p_local, tgt_link_state.quat)
         i_l = tgt_link_idx
         while i_l > -1:
             I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l
@@ -864,7 +886,7 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
                     I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d
                     i_d_jac = i_d + dof_offset - self._dof_start
                     rotation = gu.ti_transform_by_quat(self._solver.dofs_info[I_d].motion_ang, l_state.quat)
-                    translation = rotation.cross(tgt_link_pos - l_state.pos)
+                    translation = (tgt_link_pos - l_state.pos).cross(rotation)
 
                     self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0]
                     self._jacobian[1, i_d_jac, i_b] = translation[1] * pos_mask[1]
@@ -897,7 +919,7 @@ def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
                         i_d_jac = i_d + dof_offset - self._dof_start
                         I_d = [i_d, i_b] if ti.static(self.solver._options.batch_dofs_info) else i_d
                         rotation = self._solver.dofs_info[I_d].motion_ang
-                        translation = rotation.cross(tgt_link_pos - l_state.pos)
+                        translation = (tgt_link_pos - l_state.pos).cross(rotation)
 
                         self._jacobian[0, i_d_jac, i_b] = translation[0] * pos_mask[0]
                         self._jacobian[1, i_d_jac, i_b] = translation[1] * pos_mask[1]

tests/test_rigid_physics.py

@@ -2299,6 +2299,42 @@ def test_urdf_mimic(show_viewer, tol):
     assert_allclose(gs_qpos[-1], gs_qpos[-2], tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.merge_fixed_links(False)
+@pytest.mark.parametrize("model_name", ["pendulum"])
+def test_pendulum_links_acc(tmp_path, gs_sim, show_viewer, tol):
+    urdf_root = _build_multi_pendulum(n=1)
+    urdf_file = tmp_path / "pendulum.urdf"
+    ET.ElementTree(urdf_root).write(urdf_file)
+
+    scene = gs.Scene(sim=gs_sim, show_viewer=show_viewer, show_FPS=False)
+    ent = scene.add_entity(gs.morphs.URDF(file=str(urdf_file), fixed=True))
+    scene.build()
+
+    angle = 0.7
+    ent.set_qpos(np.array([angle], dtype=np.float32))
+    scene.step()
+
+    link = ent.get_link("PendulumArm_0")
+
+    p_local = np.array([0.05, -0.02, 0.12], dtype=np.float32)
+    J_o = ent.get_jacobian(link).cpu().numpy()
+    J_p = ent.get_jacobian(link, p_local).cpu().numpy()
+
+    c, s = np.cos(angle), np.sin(angle)
+    Rx = np.array([[1, 0, 0], [0, c, -s], [0, s, c]], dtype=np.float32)
+    r_world = Rx @ p_local
+    r_cross = np.array(
+        [[0, -r_world[2], r_world[1]], [r_world[2], 0, -r_world[0]], [-r_world[1], r_world[0], 0]], dtype=np.float32
+    )
+
+    lin_o, ang_o = J_o[:3, 0], J_o[3:, 0]
+    lin_expected = lin_o + r_cross @ ang_o
+
+    assert_allclose(J_p[3:, 0], ang_o, tol=tol)
+    assert_allclose(J_p[:3, 0], lin_expected, tol=tol)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_gravity(show_viewer, tol):