[FEATURE] Add 'get_weld_constraints' API.

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1834,6 +1834,32 @@ def set_pos(self, pos, envs_idx=None, *, relative=False, zero_velocity=True, uns
         if zero_velocity:
             self.zero_all_dofs_velocity(envs_idx, unsafe=unsafe)
 
+    @gs.assert_built
+    def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):
+        welds = self.scene.sim.rigid_solver.get_weld_constraints(as_tensor=as_tensor, to_torch=to_torch)
+        obj_a = welds["obj_a"]
+        obj_b = welds["obj_b"]
+        if as_tensor:
+            mask = (obj_a == self.idx) | (obj_b == self.idx)
+            if self._solver.n_envs > 0:
+                welds["valid_mask"] = mask
+            for k in ("obj_a", "obj_b"):
+                welds[k] = welds[k][mask]
+            if "env" in welds:
+                welds["env"] = welds["env"][mask]
+        else:
+            obj_a_new, obj_b_new, env_new = [], [], []
+            for ea, eb, env in zip(obj_a, obj_b, welds.get("env", [0] * len(obj_a))):
+                keep = (ea == self.idx) | (eb == self.idx)
+                obj_a_new.append(ea[keep])
+                obj_b_new.append(eb[keep])
+                if "env" in welds:
+                    env_new.append(env[keep])
+            welds["obj_a"], welds["obj_b"] = obj_a_new, obj_b_new
+            if "env" in welds:
+                welds["env"] = env_new
+        return welds
+
     @gs.assert_built
     def set_quat(self, quat, envs_idx=None, *, relative=False, zero_velocity=True, unsafe=False):
         """

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -2296,6 +2296,54 @@ def update_verts_for_geom(self, i_g):
             self.fixed_verts_state,
         )
 
+    @gs.assert_built
+    def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):
+        n_welds = tuple(self.constraint_solver.ti_n_equalities.to_numpy())
+        n_envs = len(n_welds)
+        n_welds_max = max(n_welds) if n_welds else 0
+        out_size = n_welds_max * n_envs
+
+        if to_torch:
+            buf = torch.full((out_size, 3), -1, dtype=gs.tc_int, device=gs.device)
+        else:
+            buf = np.full((out_size, 3), -1, dtype=np.int32)
+
+        if n_welds_max > 0:
+            self._kernel_collect_welds(buf)
+
+        if to_torch:
+            buf_view = buf.view(n_envs, n_welds_max, 3)
+        else:
+            buf_view = buf.reshape(n_envs, n_welds_max, 3)
+        env_idx = buf_view[..., 0]
+        obj_a = buf_view[..., 1]
+        obj_b = buf_view[..., 2]
+
+        if as_tensor:
+            return {"env": env_idx, "obj_a": obj_a, "obj_b": obj_b}
+        result_a = []
+        result_b = []
+        for e, count in enumerate(n_welds):
+            result_a.append(obj_a[e, :count].copy())
+            result_b.append(obj_b[e, :count].copy())
+        if n_envs == 1:
+            return {"obj_a": result_a[0], "obj_b": result_b[0]}
+        return {"obj_a": result_a, "obj_b": result_b}
+
+    @ti.kernel
+    def _kernel_collect_welds(self, buf: ti.types.ndarray()):
+        for env in range(self.n_envs):
+            base = env * self.n_equalities_candidate
+            out = 0
+            n_eq = self.constraint_solver.ti_n_equalities[env]
+            for j in range(n_eq):
+                rec = self.equalities_info[j, env]
+                if rec.eq_type == gs.EQUALITY_TYPE.WELD and out < self.n_equalities_candidate:
+                    buf[base + out, 0] = env
+                    buf[base + out, 1] = rec.eq_obj1id
+                    buf[base + out, 2] = rec.eq_obj2id
+                    out += 1
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

tests/test_rigid_physics.py

@@ -2352,6 +2352,62 @@ def test_drone_advanced(show_viewer):
     assert abs(quat_1[2] - quat_2[2]) < tol
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_get_weld_constraints_basic(show_viewer, tol):
+    scene = gs.Scene(
+        sim_options=gs.options.SimOptions(gravity=(0.0, 0.0, 0.0)),
+        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)
+
+    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)
+
+    scene.sim.rigid_solver.add_weld_constraint(link_a, link_b)
+    scene.step()
+
+    welds = scene.sim.rigid_solver.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)
+
+    cube1.set_dofs_velocity(
+        velocity=np.array([1.0, 0.0, 0.0], dtype=np.float32),
+        dofs_idx_local=np.array([0, 1, 2], dtype=np.int32),
+        envs_idx=None,
+    )
+    for _ in range(200):
+        scene.step()
+    assert_allclose(cube1.get_vel(), cube2.get_vel(), tol=tol)
+
+
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_get_weld_constraints_entity(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)
+
+    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)
+
+    scene.sim.rigid_solver.add_weld_constraint(link_a, link_b)
+    scene.step()
+
+    welds_single = cube1.get_weld_constraints(as_tensor=True, to_torch=False)
+    assert_allclose(
+        [(welds_single["obj_a"][0]), (welds_single["obj_b"][0])],
+        [link_a.item(), link_b.item()],
+        tol=tol,
+    )
+
+
 @pytest.mark.parametrize(
     "n_envs, batched, backend",
     [