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

Code before:
        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] = {}
        cls_prefix = f"{self.__class__.__name__}."

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

            key_base = cls_prefix + attr_name

            # ---- fast path: torch → numpy ------------------------------------
            try:
                arrays[key_base] = ti_field_to_torch(field).cpu().numpy()
                continue
            except (AttributeError, ti.TaichiRuntimeError):
                pass  # fall back

            # ---- generic path ------------------------------------------------
            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:
        cls_prefix = f"{self.__class__.__name__}."

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

Review comment (suggestion):
You can remove the fallback if using `to_torch` or `to_numpy` directly.

Ignore my comment about using `ti_field_to_torch`. That was a bad idea. Sorry about that...

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

            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