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

Code before:
                self.entities_state[i, b].hibernated = False
                self.awake_entities[i, b] = i

            ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
            for b in range(self._B):
                self.n_awake_entities[b] = self.n_entities

    def _init_equality_fields(self):
        struct_equality_info = ti.types.struct(
            equality_type=gs.ti_int,
            eq_obj1id=gs.ti_int,
            eq_obj2id=gs.ti_int,
            eq_data=gs.ti_vec11,
            eq_type=gs.ti_int,
            entity_idx=gs.ti_int,
            sol_params=gs.ti_vec7,
        )
        self.equality_info = struct_equality_info.field(
            shape=self.n_equalities_, needs_grad=False, layout=ti.Layout.SOA
        )
        if self.n_equalities > 0:
            equalities = self.equalities

            equalities_sol_params = np.array([equality.sol_params for equality in equalities], dtype=gs.np_float)
            equalities_sol_params = _sanitize_sol_params(
                equalities_sol_params, self._substep_dt, self._sol_contact_resolve_time
            )

            self._kernel_init_equality_fields(
                equalities_type=np.array([equality.type for equality in equalities], dtype=gs.np_int),
                equalities_eq_obj1id=np.array([equality.eq_obj1id for equality in equalities], dtype=gs.np_int),
                equalities_eq_obj2id=np.array([equality.eq_obj2id for equality in equalities], dtype=gs.np_int),
                equalities_eq_data=np.array([equality.eq_data for equality in equalities], dtype=gs.np_float),
                equalities_eq_type=np.array([equality.type for equality in equalities], dtype=gs.np_int),
                equalities_sol_params=equalities_sol_params,
            )
            if self._use_contact_island:
                gs.logger.warn("contact island is not supported for equality constraints yet")

    @ti.kernel
    def _kernel_init_equality_fields(
        self,
        equalities_type: ti.types.ndarray(),
        equalities_eq_obj1id: ti.types.ndarray(),
        equalities_eq_obj2id: ti.types.ndarray(),
        equalities_eq_data: ti.types.ndarray(),
        equalities_eq_type: ti.types.ndarray(),
        equalities_sol_params: ti.types.ndarray(),
    ):
        ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
        for i in range(self.n_equalities):

Review comment (security):
You should not specify `self._sol_contact_resolve_time` because equality constraints are NOT contact. Just leave this field empty: `_sanitize_sol_params(equalities_sol_params, self._substep_dt)`

Code after:
                for i_d in range(entities_dof_start[i], entities_dof_end[i]):
                    self.dofs_info[i_d].dof_start = entities_dof_start[i]

        if ti.static(self._use_hibernation):
            ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
            for i, b in ti.ndrange(self.n_entities, self._B):
                self.entities_state[i, b].hibernated = False
                self.awake_entities[i, b] = i

            ti.loop_config(serialize=self._para_level < gs.PARA_LEVEL.PARTIAL)
            for b in range(self._B):
                self.n_awake_entities[b] = self.n_entities

    def _init_equality_fields(self):
        struct_equality_info = ti.types.struct(
            equality_type=gs.ti_int,
            eq_obj1id=gs.ti_int,
            eq_obj2id=gs.ti_int,
            eq_data=gs.ti_vec11,
            eq_type=gs.ti_int,
            entity_idx=gs.ti_int,
            sol_params=gs.ti_vec7,
        )
        self.equality_info = struct_equality_info.field(
            shape=self.n_equalities_, needs_grad=False, layout=ti.Layout.SOA
        )
        if self.n_equalities > 0:
            equalities = self.equalities

            equalities_sol_params = np.array([equality.sol_params for equality in equalities], dtype=gs.np_float)
            equalities_sol_params = _sanitize_sol_params(
                equalities_sol_params, self._substep_dt, self._sol_constraint_resolve_time
            )

            self._kernel_init_equality_fields(
                equalities_type=np.array([equality.type for equality in equalities], dtype=gs.np_int),
                equalities_eq_obj1id=np.array([equality.eq_obj1id for equality in equalities], dtype=gs.np_int),
                equalities_eq_obj2id=np.array([equality.eq_obj2id for equality in equalities], dtype=gs.np_int),
                equalities_eq_data=np.array([equality.eq_data for equality in equalities], dtype=gs.np_float),
                equalities_eq_type=np.array([equality.type for equality in equalities], dtype=gs.np_int),
                equalities_sol_params=equalities_sol_params,
            )
            if self._use_contact_island:
                gs.logger.warn("contact island is not supported for equality constraints yet")

    @ti.kernel
    def _kernel_init_equality_fields(
        self,
        equalities_type: ti.types.ndarray(),
        equalities_eq_obj1id: ti.types.ndarray(),
        equalities_eq_obj2id: ti.types.ndarray(),