Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/solvers/rigid/rigid_solver_decomp.py (Python)

Code before:
        self._kernel_delete_weld_constraint(int(link1_idx), int(link2_idx), envs_idx)

    @ti.kernel
    def _kernel_delete_weld_constraint(
        self,
        link1_idx: ti.i32,
        link2_idx: ti.i32,
        envs_idx: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i_b_ in ti.ndrange(envs_idx.shape[0]):
            i_b = envs_idx[i_b_]
            for i_e in range(self.n_equalities, self.constraint_solver.ti_n_equalities[i_b]):
                if (
                    self.equalities_info[i_e, i_b].eq_type == gs.EQUALITY_TYPE.WELD
                    and self.equalities_info[i_e, i_b].eq_obj1id == link1_idx
                    and self.equalities_info[i_e, i_b].eq_obj2id == link2_idx
                ):
                    if i_e < self.constraint_solver.ti_n_equalities[i_b] - 1:
                        self.equalities_info[i_e, i_b] = self.equalities_info[
                            self.constraint_solver.ti_n_equalities[i_b] - 1, i_b
                        ]
                    self.constraint_solver.ti_n_equalities[i_b] = self.constraint_solver.ti_n_equalities[i_b] - 1

    @gs.assert_built
    def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):
        n_envs = int(self.n_envs) or 1
        n_eqs = self.constraint_solver.ti_n_equalities.to_numpy()
        stride = int(n_eqs.max())
        if stride == 0:
            empty = torch.empty if to_torch else np.empty
            zeros = torch.zeros if to_torch else np.zeros
            obj_a = empty((0,), dtype=np.int32)
            obj_b = empty((0,), dtype=np.int32)
            if n_envs == 1:
                return {"obj_a": obj_a, "obj_b": obj_b}
            return {
                "env": obj_a,
                "obj_a": obj_a,
                "obj_b": obj_b,
                "valid_mask": zeros((n_envs, 0), dtype=bool),
            }
        shape3 = (n_envs, stride, 3)
        if to_torch:
            buf = torch.full(shape3, -1, dtype=gs.tc_int, device=gs.device)
        else:
            buf = np.full(shape3, -1, dtype=np.int32)
        self._kernel_collect_welds(buf)
        if n_envs == 1:
            sub = buf[0]
            sub = sub[sub[:, 1] != -1]

Review comment (bug):
This is still wrong. If you look more closely, there is no such thing as `valid_mask` in `RigidSolver.get_contacts`, only in `RigidEntity.get_contacts`. You should do the same here.

Code after:
                        for j in ti.static(range(4)):
                            rigid_global_info.qpos[q_start + j + 3, i_b] = rot[j]
                    elif joint_type == gs.JOINT_TYPE.FIXED:
                        pass
                    elif joint_type == gs.JOINT_TYPE.SPHERICAL:
                        rot = ti.Vector(
                            [
                                rigid_global_info.qpos[q_start + 0, i_b],
                                rigid_global_info.qpos[q_start + 1, i_b],
                                rigid_global_info.qpos[q_start + 2, i_b],
                                rigid_global_info.qpos[q_start + 3, i_b],
                            ]
                        )
                        ang = (
                            ti.Vector(
                                [
                                    dofs_state.vel[dof_start + 3, i_b],
                                    dofs_state.vel[dof_start + 4, i_b],
                                    dofs_state.vel[dof_start + 5, i_b],
                                ]
                            )
                            * static_rigid_sim_config.substep_dt
                        )
                        qrot = gu.ti_rotvec_to_quat(ang)
                        rot = gu.ti_transform_quat_by_quat(qrot, rot)
                        for j in ti.static(range(4)):
                            rigid_global_info.qpos[q_start + j, i_b] = rot[j]

                    else:
                        for j in range(q_end - q_start):
                            rigid_global_info.qpos[q_start + j, i_b] = (
                                rigid_global_info.qpos[q_start + j, i_b]
                                + dofs_state.vel[dof_start + j, i_b] * static_rigid_sim_config.substep_dt
                            )

    else:
        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_d, i_b in ti.ndrange(n_dofs, _B):
            dofs_state.vel[i_d, i_b] = (
                dofs_state.vel[i_d, i_b] + dofs_state.acc[i_d, i_b] * static_rigid_sim_config.substep_dt
            )

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l, i_b in ti.ndrange(n_links, _B):
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l
            if links_info.n_dofs[I_l] == 0:
                continue

            dof_start = links_info.dof_start[I_l]
            q_start = links_info.q_start[I_l]
            q_end = links_info.q_end[I_l]