Merge branch 'main' into bvh

Author: Libero0809

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/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
@@ -34,6 +36,53 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
     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
+                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/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:

genesis/utils/misc.py

@@ -179,8 +179,8 @@ def get_device(backend: gs_backend):
             gs.raise_exception("torch cuda not available")
 
         device_idx = torch.cuda.current_device()
-        device = torch.device(f"cuda:{device_idx}")
-        device_property = torch.cuda.get_device_properties(device_idx)
+        device = torch.device("cuda", device_idx)
+        device_property = torch.cuda.get_device_properties(device)
         device_name = device_property.name
         total_mem = device_property.total_memory / 1024**3
 
@@ -190,14 +190,14 @@ def get_device(backend: gs_backend):
 
         # on mac, cpu and gpu are in the same device
         _, device_name, total_mem, _ = get_device(gs_backend.cpu)
-        device = torch.device("mps:0")
+        device = torch.device("mps", 0)
 
     elif backend == gs_backend.vulkan:
         if torch.cuda.is_available():
             device, device_name, total_mem, _ = get_device(gs_backend.cuda)
         elif torch.xpu.is_available():  # pytorch 2.5+ supports Intel XPU device
             device_idx = torch.xpu.current_device()
-            device = torch.device(f"xpu:{device_idx}")
+            device = torch.device("xpu", device_idx)
             device_property = torch.xpu.get_device_properties(device_idx)
             device_name = device_property.name
             total_mem = device_property.total_memory / 1024**3

genesis/utils/terrain.py

@@ -150,10 +150,12 @@ def parse_terrain(morph: Terrain, surface):
                         i * subterrain_rows : (i + 1) * subterrain_rows, j * subterrain_cols : (j + 1) * subterrain_cols
                     ] = subterrain_height_field
 
+        need_uvs = getattr(surface, "diffuse_texture", None) is not None
         tmesh, sdf_tmesh = convert_heightfield_to_watertight_trimesh(
             heightfield,
             horizontal_scale=morph.horizontal_scale,
             vertical_scale=morph.vertical_scale,
+            uv_scale=morph.uv_scale if need_uvs else None,
         )
 
         terrain_dir = os.path.join(get_assets_dir(), f"terrain/{morph.name}")
@@ -169,7 +171,7 @@ def parse_terrain(morph: Terrain, surface):
             }
             pickle.dump(info, f)
 
-    vmesh = gs.Mesh.from_trimesh(mesh=tmesh, surface=surface)
+    vmesh = gs.Mesh.from_trimesh(mesh=tmesh, surface=surface, metadata={})
     mesh = gs.Mesh.from_trimesh(
         mesh=tmesh,
         surface=gs.surfaces.Collision(),
@@ -211,7 +213,9 @@ def fractal_terrain(terrain, levels=8, scale=1.0):
     return terrain
 
 
-def convert_heightfield_to_watertight_trimesh(height_field_raw, horizontal_scale, vertical_scale, slope_threshold=None):
+def convert_heightfield_to_watertight_trimesh(
+    height_field_raw, horizontal_scale, vertical_scale, slope_threshold=None, uv_scale=None
+):
     """
     Adapted from Issac Gym's `convert_heightfield_to_trimesh` function.
     Convert a heightfield array to a triangle mesh represented by vertices and triangles.
@@ -263,8 +267,8 @@ def convert_heightfield_to_watertight_trimesh(height_field_raw, horizontal_scale
 
     # create triangle mesh vertices and triangles from the heightfield grid
     vertices_top = np.zeros((num_rows * num_cols, 3), dtype=np.float32)
-    vertices_top[:, 0] = xx.flatten()
-    vertices_top[:, 1] = yy.flatten()
+    vertices_top[:, 0] = xx.flat
+    vertices_top[:, 1] = yy.flat
     vertices_top[:, 2] = hf.flatten() * vertical_scale
     triangles_top = -np.ones((2 * (num_rows - 1) * (num_cols - 1), 3), dtype=np.uint32)
     for i in range(num_rows - 1):
@@ -285,8 +289,8 @@ def convert_heightfield_to_watertight_trimesh(height_field_raw, horizontal_scale
     z_min = np.min(vertices_top[:, 2]) - 1.0
 
     vertices_bottom = np.zeros((num_rows * num_cols, 3), dtype=np.float32)
-    vertices_bottom[:, 0] = xx.flatten()
-    vertices_bottom[:, 1] = yy.flatten()
+    vertices_bottom[:, 0] = xx.flat
+    vertices_bottom[:, 1] = yy.flat
     vertices_bottom[:, 2] = z_min
     triangles_bottom = -np.ones((2 * (num_rows - 1) * (num_cols - 1), 3), dtype=np.uint32)
     for i in range(num_rows - 1):
@@ -347,15 +351,45 @@ def convert_heightfield_to_watertight_trimesh(height_field_raw, horizontal_scale
         axis=0,
     )
 
-    # This a uniformly-distributed full mesh, which gives faster sdf generation
-    sdf_mesh = trimesh.Trimesh(vertices, triangles, process=False)
+    if uv_scale is not None:
+        uv_top = np.zeros((num_rows * num_cols, 2), dtype=np.float32)
+        uv_top[:, 0] = (xx.flat - xx.min()) / (xx.max() - xx.min()) * uv_scale
+        uv_top[:, 1] = (yy.flat - yy.min()) / (yy.max() - yy.min()) * uv_scale
+
+        uvs = np.concatenate([uv_top, uv_top], axis=0)
+        visual = trimesh.visual.TextureVisuals(uv=uvs)
+    else:
+        uvs = None
+        visual = None
+
+    sdf_mesh = trimesh.Trimesh(vertices, triangles, process=False, visual=visual)
 
     # This is the mesh used for non-sdf purposes.
     # It's losslessly simplified from the full mesh, to save memory cost for storing verts and faces.
-    mesh = trimesh.Trimesh(
-        *fast_simplification.simplify(sdf_mesh.vertices, sdf_mesh.faces, target_count=0, lossless=True)
+
+    v_simp, f_simp = fast_simplification.simplify(
+        sdf_mesh.vertices,
+        sdf_mesh.faces,
+        target_count=0,
+        lossless=True,
     )
 
+    if uvs is not None:
+        idx_map = np.empty(len(v_simp), dtype=np.int64)
+        for i, v in enumerate(v_simp):
+            dists = np.square(vertices - v).sum(axis=1)
+            idx_map[i] = np.argmin(dists)
+
+        uv_simp = uvs[idx_map]
+
+        mesh = trimesh.Trimesh(
+            v_simp,
+            f_simp,
+            visual=trimesh.visual.TextureVisuals(uv=uv_simp),
+        )
+    else:
+        mesh = trimesh.Trimesh(v_simp, f_simp)
+
     return mesh, sdf_mesh