Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/couplers/sap_coupler.py (Python)

Code before:
                self.fem_state_v.impulse[i_b, i_v].norm_sqr() / self.fem_solver.elements_v_info[i_v].mass
            )

    @ti.func
    def add_rigid_norms(self, rigid_global_info: array_class.RigidGlobalInfo):
        for i_b, i_d in ti.ndrange(self._B, self.rigid_solver.n_dofs):
            if not self.batch_active[i_b]:
                continue
            self.sap_state[i_b].gradient_norm += (
                self.rigid_state_dof.gradient[i_b, i_d] ** 2 / rigid_global_info.mass_mat[i_d, i_d, i_b]
            )
            self.sap_state[i_b].momentum_norm += (
                self.rigid_state_dof.v[i_b, i_d] ** 2 * rigid_global_info.mass_mat[i_d, i_d, i_b]
            )
            self.sap_state[i_b].impulse_norm += (
                self.rigid_state_dof.impulse[i_b, i_d] ** 2 / rigid_global_info.mass_mat[i_d, i_d, i_b]
            )

    @ti.func
    def update_batch_active(self):
        for i_b in range(self._B):
            if not self.batch_active[i_b]:
                continue
            self._sap_convergence_atol + self._sap_convergence_rtol * ti.max(
                self.sap_state[i_b].momentum_norm, self.sap_state[i_b].impulse_norm
            )

    @ti.kernel
    def compute_regularization(
        self,
        dofs_state: array_class.DofsState,
        entities_info: array_class.EntitiesInfo,
        rigid_global_info: array_class.RigidGlobalInfo,
    ):
        for contact in ti.static(self.contact_handlers):
            contact.compute_regularization(entities_info=entities_info, rigid_global_info=rigid_global_info)
        if ti.static(self.rigid_solver.is_active and self.rigid_solver.n_equalities > 0):
            self.equality_constraint_handler.compute_regularization(dofs_state=dofs_state)

    @ti.kernel
    def _init_sap_solve(self, i_step: ti.i32, dofs_state: array_class.DofsState):
        self._init_v(i_step, dofs_state=dofs_state)
        self.batch_active.fill(True)

    @ti.func
    def _init_v(self, i_step: ti.i32, dofs_state: array_class.DofsState):
        if ti.static(self.fem_solver.is_active):
            self._init_v_fem(i_step)
        if ti.static(self.rigid_solver.is_active):
            self._init_v_rigid(i_step, dofs_state=dofs_state)


Review comment (bug):
WTF?? The original code is garbage. This bug was introduced in this PR:

https://github.com/Genesis-Embodied-AI/Genesis/pull/1565/changes

Code after:
                self.fem_state_v.impulse[i_b, i_v].norm_sqr() / self.fem_solver.elements_v_info[i_v].mass
            )

    @ti.func
    def add_rigid_norms(self, rigid_global_info: array_class.RigidGlobalInfo):
        for i_b, i_d in ti.ndrange(self._B, self.rigid_solver.n_dofs):
            if not self.batch_active[i_b]:
                continue
            self.sap_state[i_b].gradient_norm += (
                self.rigid_state_dof.gradient[i_b, i_d] ** 2 / rigid_global_info.mass_mat[i_d, i_d, i_b]
            )
            self.sap_state[i_b].momentum_norm += (
                self.rigid_state_dof.v[i_b, i_d] ** 2 * rigid_global_info.mass_mat[i_d, i_d, i_b]
            )
            self.sap_state[i_b].impulse_norm += (
                self.rigid_state_dof.impulse[i_b, i_d] ** 2 / rigid_global_info.mass_mat[i_d, i_d, i_b]
            )

    @ti.func
    def update_batch_active(self):
        for i_b in range(self._B):
            if not self.batch_active[i_b]:
                continue
            norm_thr = self._sap_convergence_atol + self._sap_convergence_rtol * ti.max(
                self.sap_state[i_b].momentum_norm, self.sap_state[i_b].impulse_norm
            )
            self.batch_active[i_b] = self.sap_state[i_b].gradient_norm >= norm_thr

    @ti.kernel
    def compute_regularization(
        self,
        dofs_state: array_class.DofsState,
        entities_info: array_class.EntitiesInfo,
        rigid_global_info: array_class.RigidGlobalInfo,
    ):
        for contact in ti.static(self.contact_handlers):
            contact.compute_regularization(entities_info=entities_info, rigid_global_info=rigid_global_info)
        if ti.static(self.rigid_solver.is_active and self.rigid_solver.n_equalities > 0):
            self.equality_constraint_handler.compute_regularization(dofs_state=dofs_state)

    @ti.kernel
    def _init_sap_solve(self, i_step: ti.i32, dofs_state: array_class.DofsState):
        self._init_v(i_step, dofs_state=dofs_state)
        self.batch_active.fill(True)

    @ti.func
    def _init_v(self, i_step: ti.i32, dofs_state: array_class.DofsState):
        if ti.static(self.fem_solver.is_active):
            self._init_v_fem(i_step)
        if ti.static(self.rigid_solver.is_active):
            self._init_v_rigid(i_step, dofs_state=dofs_state)