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

Code before:

if TYPE_CHECKING:
    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
        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, B: int):
        self._B = B
        g_np = np.asarray(self._sim._gravity)
        g_np = np.repeat(g_np[None], B, axis=0)
        self._gravity = ti.Vector.field(3, dtype=gs.ti_float, shape=B)
        self._gravity.from_numpy(g_np)

    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():

Review comment (suggestion):
```python
if self._init_gravity is not None:
    gravity = np.asarray(self._init_gravity, dtype=gs.np_float)
    [...]
```

Code after:

if TYPE_CHECKING:
    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
        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))