Merge branch 'main' into feature/env-wise-gravity

Author: LeonLiu4

.github/workflows/linux-gpu.yml

@@ -35,6 +35,12 @@ jobs:
 
           mkdir -p "${HOME}/.cache"
 
+          # Prefer idle nodes if any
+          IDLE_NODES=$(sinfo -h -o "%N %t" | awk '$2 == "idle" {print $1}')
+          if [[ -n "$IDLE_NODES" ]]; then
+            NODELIST="--nodelist=$IDLE_NODES"
+          fi
+
           srun \
             --container-image="/mnt/data/images/genesis-v${GENESIS_IMAGE_VER}.sqsh" \
             --container-mounts=\
@@ -44,7 +50,7 @@ jobs:
             --export=\
           HF_TOKEN="${HF_TOKEN}",\
           NVIDIA_DRIVER_CAPABILITIES=all \
-            --partition=hpc-mid --nodes=1 --gpus=1 --time="${TIMEOUT_MINUTES}" \
+            --partition=hpc-mid ${NODELIST} --nodes=1 --time="${TIMEOUT_MINUTES}" \
             --job-name=${SLURM_JOB_NAME} \
             bash -c "
               pip install -e '.[dev,render]' && \
@@ -69,16 +75,16 @@ jobs:
           "${{ github.workspace }}":/root/workspace \
             --no-container-mount-home --container-workdir=/root/workspace \
             --export=${SLURM_ENV_VARS} \
-            --partition=hpc-mid --exclusive --nodes=1 --gpus=1 --time="${TIMEOUT_MINUTES}" \
+            --partition=hpc-mid --exclusive --nodes=1 --time="${TIMEOUT_MINUTES}" \
             --job-name=${SLURM_JOB_NAME} \
             bash -c "
               : # sudo apt install -y tmate && \
               tmate -S /tmp/tmate.sock new-session -d && \
               tmate -S /tmp/tmate.sock wait tmate-ready && \
               tmate -S /tmp/tmate.sock display -p '#{tmate_ssh}'
               pip install -e '.[dev,render]' && \
-              pytest --print -x -m 'benchmarks' --backend gpu ./tests && \
-              cp 'speed_test.txt' '/mnt/data/artifacts/speed_test_${SLURM_JOB_NAME}.txt'
+              pytest --print -x -m 'benchmarks' ./tests && \
+              cat speed_test*.txt > '/mnt/data/artifacts/speed_test_${SLURM_JOB_NAME}.txt'
               : # tmate -S /tmp/tmate.sock wait tmate-exit
             "
 

examples/drone/hover_env.py

@@ -134,7 +134,7 @@ def step(self, actions):
         self.drone.set_propellels_rpm((1 + exec_actions * 0.8) * 14468.429183500699)
         # update target pos
         if self.target is not None:
-            self.target.set_pos(self.commands, zero_velocity=True, envs_idx=list(range(self.num_envs)))
+            self.target.set_pos(self.commands, zero_velocity=True)
         self.scene.step()
 
         # update buffers

genesis/__init__.py

@@ -291,6 +291,7 @@ def _display_greeting(INFO_length):
     wave_width = max(0, min(38, wave_width))
     bar_width = wave_width * 2 + 9
     wave = ("┈┉" * wave_width)[:wave_width]
+    global logger
     logger.info(f"~<╭{'─'*(bar_width)}╮>~")
     logger.info(f"~<│{wave}>~ ~~~~<Genesis>~~~~ ~<{wave}│>~")
     logger.info(f"~<╰{'─'*(bar_width)}╯>~")
@@ -314,9 +315,10 @@ def _custom_excepthook(exctype, value, tb):
     print("".join(traceback.format_exception(exctype, value, tb)))
 
     # Logger the exception right before exit if possible
+    global logger
     try:
         logger.error(f"{exctype.__name__}: {value}")
-    except AttributeError:
+    except (AttributeError, NameError):
         # Logger may not be configured at this point
         pass
 

genesis/engine/entities/hybrid_entity.py

@@ -431,8 +431,8 @@ def _kernel_update_soft_part_mpm(self, f: ti.i32):
                 acc = vel_d / dt_for_rigid_acc
                 frc_vel = mass_real * acc
                 frc_ang = (x_pos - link.COM).cross(frc_vel)
-                self._solver_rigid.links_state[link_idx, i_b].cfrc_ext_vel += frc_vel
-                self._solver_rigid.links_state[link_idx, i_b].cfrc_ext_ang += frc_ang
+                self._solver_rigid.links_state[link_idx, i_b].cfrc_applied_vel += frc_vel
+                self._solver_rigid.links_state[link_idx, i_b].cfrc_applied_ang += frc_ang
 
                 # rigid-to-soft coupling # NOTE: this may lead to unstable feedback loop
                 self._solver_soft.particles[f_, i_global, i_b].vel += vel_d * self.material.soft_dv_coef

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -2748,7 +2748,7 @@ def get_contacts(self, with_entity=None, exclude_self_contact=False):
         if self._solver.n_envs == 0:
             contacts_info = {key: value[valid_mask] for key, value in contacts_info.items()}
         else:
-            contacts_info = {key: value[:, valid_mask] for key, value in contacts_info.items()}
+            contacts_info["valid_mask"] = valid_mask
 
         contacts_info["force_a"] = -contacts_info["force"]
         contacts_info["force_b"] = +contacts_info["force"]

genesis/engine/scene.py

@@ -2,6 +2,9 @@
 
 import numpy as np
 import torch
+import pickle
+import time
+import taichi as ti
 
 import genesis as gs
 import genesis.utils.geom as gu
@@ -1057,6 +1060,69 @@ def _backward(self):
         self._backward_ready = False
         self._forward_ready = False
 
+    def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
+        """
+        Collect every Taichi field in the **scene and its active solvers** and
+        return them as a flat ``{key: ndarray}`` dictionary.
+
+        Returns
+        -------
+        dict[str, np.ndarray]
+            Mapping ``"Class.attr[.member]" → array`` with raw field data.
+        """
+        arrays: dict[str, np.ndarray] = {}
+
+        for name, field in self.__dict__.items():
+            if isinstance(field, ti.Field):
+                arrays[".".join((self.__class__.__name__, name))] = field.to_numpy()
+
+        for solver in self.active_solvers:
+            arrays.update(solver.dump_ckpt_to_numpy())
+
+        return arrays
+
+    def save_checkpoint(self, path: str | os.PathLike) -> None:
+        """
+        Pickle the full physics state to *one* file.
+
+        Parameters
+        ----------
+        path : str | os.PathLike
+            Destination filename.
+        """
+        state = {
+            "timestamp": time.time(),
+            "step_index": self.t,
+            "arrays": self.dump_ckpt_to_numpy(),
+        }
+        with open(path, "wb") as f:
+            pickle.dump(state, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+    def load_checkpoint(self, path: str | os.PathLike) -> None:
+        """
+        Restore a file produced by :py:meth:`save_checkpoint`.
+
+        Parameters
+        ----------
+        path : str | os.PathLike
+            Path to the checkpoint pickle.
+        """
+        with open(path, "rb") as f:
+            state = pickle.load(f)
+
+        arrays = state["arrays"]
+
+        for name, field in self.__dict__.items():
+            if isinstance(field, ti.Field):
+                key = ".".join((self.__class__.__name__, name))
+                if key in arrays:
+                    field.from_numpy(arrays[key])
+
+        for solver in self.active_solvers:
+            solver.load_ckpt_from_numpy(arrays)
+
+        self._t = state.get("step_index", self._t)
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/solvers/base_solver.py

@@ -1,6 +1,8 @@
 from typing import TYPE_CHECKING
 import numpy as np
 import taichi as ti
+import torch
+from genesis.utils.misc import ti_field_to_torch
 
 import genesis as gs
 from genesis.engine.entities.base_entity import Entity
@@ -40,6 +42,53 @@ def _finalize_batch(self, B: int):
             g = np.tile(g[None, :], (B, 1))
         self._gravity.from_numpy(g[..., None])
 
+    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
+
+            arr = arr_dict[key_base]
+            field.from_numpy(arr)
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/solvers/rigid/collider_decomp.py

@@ -2088,19 +2088,20 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True):
         # Allocate output buffer
         if to_torch:
             iout = torch.full((out_size, 4), -1, dtype=gs.tc_int, device=gs.device)
-            fout = torch.empty((out_size, 10), dtype=gs.tc_float, device=gs.device)
+            fout = torch.zeros((out_size, 10), dtype=gs.tc_float, device=gs.device)
         else:
             iout = np.full((out_size, 4), -1, dtype=gs.np_int)
-            fout = np.empty((out_size, 10), dtype=gs.np_float)
+            fout = np.zeros((out_size, 10), dtype=gs.np_float)
 
         # Copy contact data
-        self._kernel_get_contacts(as_tensor, iout, fout)
+        if n_contacts_max > 0:
+            self._kernel_get_contacts(as_tensor, iout, fout)
 
-        # Return structured view (no copy)
+        # Build structured view (no copy)
         if as_tensor:
             if self._solver.n_envs > 0:
-                iout = iout.reshape((n_contacts_max, n_envs, -1))
-                fout = fout.reshape((n_contacts_max, n_envs, -1))
+                iout = iout.reshape((n_envs, n_contacts_max, 4))
+                fout = fout.reshape((n_envs, n_contacts_max, 10))
             iout_chunks = (iout[..., 0], iout[..., 1], iout[..., 2], iout[..., 3])
             fout_chunks = (fout[..., 0], fout[..., 1:4], fout[..., 4:7], fout[..., 7:])
             values = (*iout_chunks, *fout_chunks)
@@ -2135,10 +2136,7 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True):
                     values = (*iout_chunks, *fout_chunks)
 
         contacts_info = dict(
-            zip(
-                ("link_a", "link_b", "geom_a", "geom_b", "penetration", "position", "normal", "force"),
-                (value.swapaxes(0, 1) for value in values) if as_tensor and self._solver.n_envs > 0 else values,
-            )
+            zip(("link_a", "link_b", "geom_a", "geom_b", "penetration", "position", "normal", "force"), values)
         )
 
         # Cache contact information before returning

genesis/options/morphs.py

@@ -893,6 +893,8 @@ class Terrain(Morph):
         The size of each cell in the subterrain in meters. Defaults to 0.25.
     vertical_scale : float, optional
         The height of each step in the subterrain in meters. Defaults to 0.005.
+    uv_scale : float, optional
+        The scale of the UV mapping for the terrain. Defaults to 1.0.
     subterrain_types : str or 2D list of str, optional
         The types of subterrains to generate. If a string, it will be repeated for all subterrains. If a 2D list, it should have the same shape as `n_subterrains`.
     height_field : array-like, optional
@@ -911,6 +913,7 @@ class Terrain(Morph):
     subterrain_size: Tuple[float, float] = (12.0, 12.0)  # meter
     horizontal_scale: float = 0.25  # meter size of each cell in the subterrain
     vertical_scale: float = 0.005  # meter height of each step in the subterrain
+    uv_scale: float = 1.0
     subterrain_types: Any = [
         ["flat_terrain", "random_uniform_terrain", "stepping_stones_terrain"],
         ["pyramid_sloped_terrain", "discrete_obstacles_terrain", "wave_terrain"],

genesis/utils/geom.py

@@ -734,11 +734,7 @@ def quat_to_xyz(quat, rpy=False, degrees=False):
             sinp = 2 * (qw * qy - qz * qx)
         else:
             sinp = 2 * (qx * qz + qw * qy)
-        pitch = torch.where(
-            torch.abs(sinp) >= 1,
-            torch.sign(sinp) * torch.tensor(torch.pi / 2),
-            torch.asin(sinp),
-        )
+        pitch = torch.where(torch.abs(sinp) >= 1, torch.sign(sinp) * (torch.pi / 2), torch.asin(sinp))
 
         # Yaw (z-axis rotation)
         if rpy: