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._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

Review comment (suggestion):
If think you should rather completely ignore gravity if `self._init_gravity` is None:

```python
if self._init_gravity is not None:
    gravity = np.asarray(self._init_gravity, dtype=np.float32)
    gravity = 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(gravity)
```

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))
            data = field.to_numpy()