[FEATURE] Add 'get_weld_constraints' API.

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -5002,6 +5002,41 @@ 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)
+        envs_idx = np.atleast_1d(envs_idx)
+
+        env_col, obj_a_col, obj_b_col = [], [], []
+        WELD = gs.EQUALITY_TYPE.WELD
+
+        for env in 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 == WELD:
+                    env_col.append(env)
+                    obj_a_col.append(int(rec.eq_obj1id))
+                    obj_b_col.append(int(rec.eq_obj2id))
+
+        obj_a = np.asarray(obj_a_col, dtype=np.int32)
+        obj_b = np.asarray(obj_b_col, dtype=np.int32)
+
+        if self.n_envs == 0:
+            return {"obj_a": obj_a, "obj_b": obj_b}
+
+        env_arr = np.asarray(env_col, dtype=np.int32)
+        valid_mask = np.zeros((self.n_envs, obj_a.shape[0]), dtype=bool)
+        for j, e in enumerate(env_arr):
+            valid_mask[e, j] = True
+
+        return {
+            "env": env_arr,
+            "obj_a": obj_a,
+            "obj_b": obj_b,
+            "valid_mask": valid_mask,
+        }
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

tests/test_rigid_physics.py

@@ -2169,6 +2169,29 @@ 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 = np.array([cube1.base_link.idx], dtype=gs.np_int)
+    link_b = np.array([cube2.base_link.idx], dtype=gs.np_int)
+
+    rigid.add_weld_constraint(link_a, link_b)
+    scene.step()
+
+    welds = rigid.get_weld_constraints()
+    row = np.array([welds["env"][0], welds["obj_a"][0], welds["obj_b"][0]], dtype=np.int32)
+    ref = np.array([0, link_a[0], link_b[0]], 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):