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

Code before:
        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]

            links_state.root_COM[i_l, i_b].fill(0.0)
            links_state.mass_sum[i_l, i_b] = 0.0

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            mass = links_info.inertial_mass[I_l] + links_state.mass_shift[i_l, i_b]
            (
                links_state.i_pos[i_l, i_b],
                links_state.i_quat[i_l, i_b],
            ) = gu.ti_transform_pos_quat_by_trans_quat(
                links_info.inertial_pos[I_l] + links_state.i_pos_shift[i_l, i_b],
                links_info.inertial_quat[I_l],
                links_state.pos[i_l, i_b],
                links_state.quat[i_l, i_b],
            )

            i_r = links_info.root_idx[I_l]
            links_state.mass_sum[i_r, i_b] += mass
            links_state.link_COM[i_l, i_b] = links_state.i_pos[i_l, i_b]
            links_state.root_COM[i_r, i_b] += mass * links_state.i_pos[i_l, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            i_r = links_info.root_idx[I_l]
            if i_l == i_r:
                links_state.root_COM[i_l, i_b] = links_state.root_COM[i_l, i_b] / links_state.mass_sum[i_l, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            i_r = links_info.root_idx[I_l]
            links_state.root_COM[i_l, i_b] = links_state.root_COM[i_r, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            i_r = links_info.root_idx[I_l]

Review comment (suggestion):
I don't think storing `link_COM` is useful because it is basically `i_pos`. Having multiple aliases for the same quantity is a bad practice and not memory friendly.

Code after:
    n_links = links_info.root_idx.shape[0]
    if ti.static(static_rigid_sim_config.use_hibernation):
        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]

            links_state.root_COM[i_l, i_b].fill(0.0)
            links_state.mass_sum[i_l, i_b] = 0.0

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            mass = links_info.inertial_mass[I_l] + links_state.mass_shift[i_l, i_b]
            (
                links_state.i_pos[i_l, i_b],
                links_state.i_quat[i_l, i_b],
            ) = gu.ti_transform_pos_quat_by_trans_quat(
                links_info.inertial_pos[I_l] + links_state.i_pos_shift[i_l, i_b],
                links_info.inertial_quat[I_l],
                links_state.pos[i_l, i_b],
                links_state.quat[i_l, i_b],
            )

            i_r = links_info.root_idx[I_l]
            links_state.mass_sum[i_r, i_b] += mass
            links_state.root_COM[i_r, i_b] += mass * links_state.i_pos[i_l, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            i_r = links_info.root_idx[I_l]
            if i_l == i_r:
                links_state.root_COM[i_l, i_b] = links_state.root_COM[i_l, i_b] / links_state.mass_sum[i_l, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

            i_r = links_info.root_idx[I_l]
            links_state.root_COM[i_l, i_b] = links_state.root_COM[i_r, i_b]

        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
        for i_l_ in range(rigid_global_info.n_awake_links[i_b]):
            i_l = rigid_global_info.awake_links[i_l_, i_b]
            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l