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

Code before:


class Solver(RBC):
    def __init__(self, scene: "Scene", sim: "Simulator", options):
        self._uid = gs.UID()
        self._sim = sim
        self._scene = scene
        self._dt: float = options.dt
        self._substep_dt: float = options.dt / sim.substeps

        if hasattr(options, "gravity"):
            self._gravity = ti.field(dtype=gs.ti_vec3, shape=())
            self._gravity.from_numpy(np.array(options.gravity, dtype=gs.np_float))
        else:
            self._gravity = None

        self._entities: list[Entity] = gs.List()

        # force fields
        self._ffs = list()

    def _add_force_field(self, force_field):
        self._ffs.append(force_field)

    @staticmethod
    def _is_ti_field(obj) -> bool:
        return callable(getattr(obj, "to_numpy", None)) and callable(getattr(obj, "from_numpy", None))

    def save_ckpt(self, tag: str = "0"):
        for val in self.__dict__.values():
            if self._is_ti_field(val):
                val[tag] = val[None]

    def load_ckpt(self, tag: str = "0"):
        for val in self.__dict__.values():
            if self._is_ti_field(val):
                val[None] = val[tag]

    def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
        arrays, prefix = {}, self.__class__.__name__ + "/"
        for name, val in self.__dict__.items():
            if not self._is_ti_field(val):
                continue
            arr = val.to_numpy()
            if isinstance(arr, np.ndarray):
                arrays[prefix + name] = arr.copy()
            elif isinstance(arr, dict):
                arrays[prefix + name] = np.array([arr], dtype=object)
        return arrays

    def load_ckpt_from_numpy(self, arr_dict: dict[str, np.ndarray]) -> None:

Review comment (refactor):
This is definitely not the right way to check if something is a taichi field. Just do
`isinstance(obj, taichi.Field)` and remove this method entirely. I checked myself and it works perfectly on both `ti.Vector`, `ti.Matrix` and `ti.types.struct` instances.

Code after:
    from genesis.engine.scene import Scene
    from genesis.engine.simulator import Simulator


class Solver(RBC):
    def __init__(self, scene: "Scene", sim: "Simulator", options):
        self._uid = gs.UID()
        self._sim = sim
        self._scene = scene
        self._dt: float = options.dt
        self._substep_dt: float = options.dt / sim.substeps

        if hasattr(options, "gravity"):
            self._gravity = ti.field(dtype=gs.ti_vec3, shape=())
            self._gravity.from_numpy(np.array(options.gravity, dtype=gs.np_float))
        else:
            self._gravity = None

        self._entities: list[Entity] = gs.List()

        # force fields
        self._ffs = list()

    def _add_force_field(self, force_field):
        self._ffs.append(force_field)

    def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
        arrays: dict[str, np.ndarray] = {}

        for attr_name, field in self.__dict__.items():
            if not isinstance(field, ti.Field):
                continue

            key_base = ".".join((self.__class__.__name__, attr_name))
            data = field.to_numpy()

            # StructField → data is a dict: flatten each member
            if isinstance(data, dict):
                for sub_name, sub_arr in data.items():
                    arrays[f"{key_base}.{sub_name}"] = (
                        sub_arr if isinstance(sub_arr, np.ndarray) else np.asarray(sub_arr)
                    )
            else:
                arrays[key_base] = data if isinstance(data, np.ndarray) else np.asarray(data)

        return arrays

    def load_ckpt_from_numpy(self, arr_dict: dict[str, np.ndarray]) -> None:
        for attr_name, field in self.__dict__.items():
            if not isinstance(field, ti.Field):
                continue