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

Code before:
        self._dt: float = options.dt
        self._substep_dt: float = options.dt / sim.substeps
        self._init_gravity = getattr(options, "gravity", None)
        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 build(self):
        self._B = self._sim._B

        g_np = np.asarray(
            self._init_gravity if self._init_gravity is not None else self._sim._gravity,
            dtype=np.float32,
        )
        g_np = np.repeat(g_np[None], self._B, axis=0)

        self._gravity = ti.Vector.field(3, dtype=gs.ti_float, shape=self._B)
        self._gravity.from_numpy(g_np)

        if self._B == 1:
            self._gravity.to_numpy = lambda _orig=self._gravity.to_numpy: _orig().squeeze(0)

    def set_gravity(self, gravity, envs_idx=None):
        if self._gravity is None:
            return
        if envs_idx is None:
            self._gravity.copy_from(gravity)
        else:
            self._gravity[envs_idx] = gravity
            
    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)
                    )

Review comment (suggestion):
You cannot do that. Overwriting instance methods is very hacky and should only be used in last resort. Why would this be required in the first place? It is required by `set_gravity`, then just update `set_gravity` accordingly to avoid relying on this mechanism.

Code after:
        self._dt: float = options.dt
        self._substep_dt: float = options.dt / sim.substeps
        self._init_gravity = getattr(options, "gravity", None)
        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 build(self):
        self._B = self._sim._B
        if self._init_gravity is not None:
            g_np = np.asarray(self._init_gravity, dtype=gs.np_float)
            g_np = np.repeat(g_np[None], self._B, axis=0)
            self._gravity = ti.Vector.field(3, dtype=gs.ti_float, shape=self._B)
            self._gravity.from_numpy(g_np)

    @gs.assert_built
    def set_gravity(self, gravity, envs_idx=None):
        if self._gravity is None:
            return
        g = np.asarray(gravity, dtype=gs.np_float)
        if envs_idx is None:
            if g.ndim == 1:
                g = np.repeat(g[None], self._B, axis=0)
            self._gravity.from_numpy(g)
        else:
            self._gravity[envs_idx] = g

    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)