[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,48 @@ 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:
+            p_local = torch.as_tensor(local_point, dtype=torch.float64, device="cpu").numpy()
+            if p_local.shape != (3,):
+                gs.raise_exception("Must be a vector of length 3")
+            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_vec):
+        self._func_get_jacobian(
+            tgt_link_idx,
+            i_b,
+            p_vec,
+            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.types.ndarray()):
+        p_vec = ti.Vector([p_local[0], p_local[1], p_local[2]], dt=gs.ti_float)
         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_vec)
+
+    @ti.kernel
+    def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):
+        for i_b in range(self._solver._B):
+            self._impl_get_jacobian(tgt_link_idx, i_b, ti.Vector.zero(gs.ti_float, 3))
 
     @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
@@ -1320,7 +1340,7 @@ def _kernel_inverse_kinematics(
                         # update jacobian for ee link
                         i_l_ee = links_idx[i_ee]
                         self._func_get_jacobian(
-                            i_l_ee, i_b, pos_mask, rot_mask
+                            i_l_ee, i_b, ti.Vector.zero(gs.ti_float, 3), pos_mask, rot_mask
                         )  # NOTE: we still compute jacobian for all dofs as we haven't found a clean way to implement this
 
                         # copy to multi-link jacobian (only for the effective n_dofs instead of self.n_dofs)

tests/test_rigid_physics.py

@@ -2169,6 +2169,37 @@ def test_terrain_size(show_viewer, tol):
     assert_allclose((height_ref * 2.0), height_test, tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_get_weld_constraints_basic(show_viewer, tol):
+    scene = gs.Scene(show_viewer=show_viewer)
+
+    cube1 = scene.add_entity(gs.morphs.Box(size=(0.05,) * 3, pos=(0.0, 0.0, 0.05)))
+    cube2 = scene.add_entity(gs.morphs.Box(size=(0.05,) * 3, pos=(0.2, 0.0, 0.05)))
+
+    scene.build(n_envs=1)
+
+    rigid = scene.sim.rigid_solver
+
+    link_a = torch.tensor([cube1.base_link.idx], dtype=gs.tc_int, device=gs.device)
+    link_b = torch.tensor([cube2.base_link.idx], dtype=gs.tc_int, device=gs.device)
+
+    rigid.add_weld_constraint(link_a, link_b)
+    scene.step()
+
+    welds = rigid.get_weld_constraints(as_tensor=True, to_torch=False)
+
+    env_id = 0 if "env" not in welds else int(welds["env"][0])
+
+    row = np.array(
+        [env_id, int(welds["obj_a"][0]), int(welds["obj_b"][0])],
+        dtype=np.int32,
+    )
+    ref = np.array([0, link_a.item(), link_b.item()], dtype=np.int32)
+
+    assert_allclose(row, ref, tol=tol)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_urdf_parsing(show_viewer, tol):