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

Code before:
        # 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:
        prefix = self.__class__.__name__ + "/"
        for name, val in self.__dict__.items():
            if not self._is_ti_field(val):
                continue
            key = prefix + name
            if key not in arr_dict:
                continue
            arr = arr_dict[key]
            if arr.dtype == object and arr.shape == (1,):
                arr = arr[0]
            val.from_numpy(arr)

    # ------------------------------------------------------------------------------------
    # ----------------------------------- properties -------------------------------------
    # ------------------------------------------------------------------------------------

    @property
    def uid(self):

Review comment (suggestion):
You should rather use `ti_field_to_torch` I think. It would be faster and more consistent with the whole codebase.

Code after:
        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

            key_base = ".".join((self.__class__.__name__, attr_name))
            member_prefix = key_base + "."

            # ---- StructField: gather its members -----------------------------
            member_items = {}
            for saved_key, saved_arr in arr_dict.items():
                if saved_key.startswith(member_prefix):
                    sub_name = saved_key[len(member_prefix) :]
                    member_items[sub_name] = saved_arr

            if member_items:  # we found at least one sub-member
                field.from_numpy(member_items)
                continue

            # ---- Ordinary field ---------------------------------------------
            if key_base not in arr_dict:
                continue  # nothing saved for this attribute