get_weld_contraints api and unit test

Author: LeonLiu4

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -4966,6 +4966,19 @@ def _kernel_delete_weld_constraint(
                         ]
                     self.constraint_solver.ti_n_equalities[i_b] = self.constraint_solver.ti_n_equalities[i_b] - 1
 
+    def get_weld_constraints(self, envs_idx=None):
+        if envs_idx is None:
+            envs_idx = np.arange(self.n_envs, dtype=np.int32)
+
+        rows = []
+        for env in np.atleast_1d(envs_idx):
+            n_eq = int(self.constraint_solver.ti_n_equalities[env])
+            for i in range(n_eq):
+                rec = self.equalities_info[i, env]
+                if rec.eq_type == gs.EQUALITY_TYPE.WELD:
+                    rows.append((env, int(rec.eq_obj1id), int(rec.eq_obj2id)))
+        return np.asarray(rows, dtype=np.int32)
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

tests/test_rigid_physics.py

@@ -2166,6 +2166,27 @@ 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
+    link1 = np.array([cube1.base_link.idx], dtype=gs.np_int)
+    link2 = np.array([cube2.base_link.idx], dtype=gs.np_int)
+
+    rigid.add_weld_constraint(link1, link2)
+    scene.step()
+
+    welds = rigid.get_weld_constraints()
+    assert_allclose(tuple(welds[0]), (0, link1[0], link2[0]), tol=tol)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_urdf_parsing(show_viewer, tol):