[FEATURE] Per-frame visual vertex deformation for KinematicEntity

examples/rendering/custom_visual_mesh.py

@@ -0,0 +1,195 @@
+"""
+Custom Visual Mesh in Genesis
+==============================
+
+Demonstrates how to use ``Visualizer.set_custom_kinematic_entity_vverts`` to drive a kinematic entity's visual vertex
+positions at runtime. Only vertex positions change; mesh topology (faces) stays constant. Ideal for
+SMPL-style body models and other externally-skinned meshes.
+
+Requirements (SMPL demo only)
+-----------------------------
+- ``smplx`` package (``pip install smplx``)
+- SMPL model files (download from https://smpl.is.tue.mpg.de)
+
+Usage
+-----
+    # Wave-deforming box (no external dependencies)
+    python custom_visual_mesh.py -v
+
+    # SMPL body mesh (requires smplx)
+    python custom_visual_mesh.py -v --smpl --model_path /path/to/smpl/models
+
+    # Batched
+    python custom_visual_mesh.py -v -B 4
+"""
+
+import argparse
+import os
+import tempfile
+
+import numpy as np
+import trimesh
+
+import genesis as gs
+from genesis.utils.misc import tensor_to_array
+
+
+def create_box_mesh(size=0.5, subdivisions=3):
+    """Create a subdivided box mesh suitable for vertex deformation."""
+    mesh = trimesh.creation.box(extents=(size, size, size))
+    for _ in range(subdivisions):
+        mesh = mesh.subdivide()
+    return mesh.vertices.astype(gs.np_float), mesh.faces.astype(gs.np_int)
+
+
+def wave_deform(verts, t, amplitude=0.1, freq=3.0):
+    """Apply a sinusoidal wave deformation along the z-axis."""
+    deformed = verts.copy()
+    deformed[:, 2] += amplitude * np.sin(freq * verts[:, 0] + t) * np.cos(freq * verts[:, 1] + t * 0.7)
+    return deformed
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Custom visual mesh in Genesis")
+    parser.add_argument(
+        "-v",
+        "--vis",
+        action="store_true",
+        help="Open interactive viewer",
+    )
+    parser.add_argument(
+        "-B",
+        "--num_envs",
+        type=int,
+        default=0,
+        help="Number of parallel envs (0 = unbatched)",
+    )
+    parser.add_argument(
+        "--smpl",
+        action="store_true",
+        help="Use SMPL body model (requires smplx)",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default=None,
+        help="Path to SMPL model directory (required with --smpl)",
+    )
+    args = parser.parse_args()
+    if args.smpl and args.model_path is None:
+        parser.error("--model_path is required when --smpl is set")
+
+    gs.init(backend=gs.gpu)
+
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(0, -3.0, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+            camera_fov=45,
+            max_FPS=60,
+        ),
+        rigid_options=gs.options.RigidOptions(
+            dt=0.01,
+            gravity=(0, 0, 0),
+        ),
+        show_viewer=args.vis,
+    )
+
+    scene.add_entity(gs.morphs.Plane())
+
+    n_envs = args.num_envs
+    B = max(1, n_envs)
+
+    if args.smpl:
+        import smplx
+        import torch
+
+        smpl = smplx.SMPL(model_path=args.model_path, gender="neutral", batch_size=B)
+
+        # Export T-pose mesh for Genesis to load as a rigid entity.
+        t_pose_verts = tensor_to_array(smpl().vertices.squeeze())
+        faces = smpl.faces.astype(gs.np_int)
+        mesh = trimesh.Trimesh(vertices=t_pose_verts, faces=faces, process=False)
+        tmp = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
+        mesh.export(tmp.name)
+
+        entity = scene.add_entity(
+            morph=gs.morphs.Mesh(
+                file=tmp.name,
+                pos=(0, 0, 0),
+                fixed=True,
+            ),
+            material=gs.materials.Kinematic(),
+        )
+
+        cam = scene.add_camera(
+            res=(640, 480),
+            pos=(0, -3.0, 1.0),
+            lookat=(0, 0, 0.8),
+            fov=45,
+        )
+
+        if n_envs > 0:
+            scene.build(n_envs=n_envs)
+        else:
+            scene.build()
+
+        os.unlink(tmp.name)
+        print(f"Entity n_vverts: {entity.n_vverts}, SMPL verts: {smpl.get_num_verts()}")
+
+        for step in range(500):
+            t = step * 0.03
+            body_pose = torch.zeros(B, 69)
+            body_pose[:, 0] = 0.5 * torch.sin(torch.tensor(t))
+            body_pose[:, 3] = -0.5 * torch.sin(torch.tensor(t))
+            body_pose[:, 9] = 0.3 * torch.clamp(torch.sin(torch.tensor(t)), min=0.0)
+            body_pose[:, 12] = 0.3 * torch.clamp(-torch.sin(torch.tensor(t)), min=0.0)
+
+            verts = smpl(body_pose=body_pose, return_verts=True).vertices  # (B, n_vverts, 3)
+            scene.visualizer.set_custom_kinematic_entity_vverts(entity, verts)
+
+            scene.step()
+            cam.render()
+    else:
+        base_verts, base_faces = create_box_mesh(size=0.5, subdivisions=3)
+        base_verts[:, 2] += 0.5  # lift above ground
+
+        mesh = trimesh.Trimesh(vertices=base_verts, faces=base_faces, process=False)
+        tmp = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
+        mesh.export(tmp.name)
+
+        entity = scene.add_entity(
+            morph=gs.morphs.Mesh(
+                file=tmp.name,
+                pos=(0, 0, 0),
+                fixed=True,
+            ),
+        )
+
+        cam = scene.add_camera(
+            res=(640, 480),
+            pos=(0, -2.5, 1.5),
+            lookat=(0, 0, 0.5),
+            fov=45,
+        )
+
+        if n_envs > 0:
+            scene.build(n_envs=n_envs)
+        else:
+            scene.build()
+
+        os.unlink(tmp.name)
+
+        for step in range(500):
+            t = step * 0.05
+            all_verts = np.stack([wave_deform(base_verts, t + b * 0.5) for b in range(B)], axis=0)
+            scene.visualizer.set_custom_kinematic_entity_vverts(entity, all_verts)
+
+            scene.step()
+            cam.render()
+
+    print("Done!")
+
+
+if __name__ == "__main__":
+    main()

genesis/utils/misc.py

@@ -730,6 +730,64 @@ def sanitize_indices(
     return tuple(indices_)
 
 
+def _expand_shape(array, expected_shape):
+    """Reshape ``array`` to match the rank of ``expected_shape`` by inserting singleton dims, without
+    actually broadcasting. Accepts either a ``torch.Tensor`` or a ``numpy.ndarray`` - both expose the
+    ``.shape`` / ``.reshape`` / leading-axis indexing this routine relies on.
+
+    ``-1`` entries in ``expected_shape`` mean "any size"; the routine treats them as wildcards when
+    deciding which input dim aligns with which output dim.
+    """
+    tensor_shape = array.shape
+    tensor_ndim = len(tensor_shape)
+    expected_ndim = len(expected_shape)
+
+    if tensor_ndim == 0:
+        return array[None]
+    if tensor_ndim > expected_ndim:
+        gs.raise_exception(f"Invalid input shape: {tensor_shape}. Expecting at most {expected_ndim}D tensor.")
+    if tensor_ndim == expected_ndim:
+        return array
+    if all(d1 == d2 or d2 == -1 for d1, d2 in zip(tensor_shape, expected_shape[-tensor_ndim:])):
+        return array
+
+    # Search for the most-trailing valid assignment of input dims to output dim positions.
+    for dims_valid in tuple(combinations(range(expected_ndim), tensor_ndim))[::-1]:
+        curr_idx = 0
+        expanded_shape = []
+        for i in range(expected_ndim):
+            if i in dims_valid:
+                dim, size = array.shape[curr_idx], expected_shape[i]
+                if dim == size or dim == 1 or size == -1:
+                    expanded_shape.append(dim)
+                    curr_idx += 1
+                else:
+                    break
+            else:
+                expanded_shape.append(1)
+        else:
+            if curr_idx == tensor_ndim:
+                return array.reshape(expanded_shape)
+    return array
+
+
+def _raise_broadcast_shape_error(input_shape, expected_shape, dim_names, cause):
+    msg_err = f"Invalid input shape: {tuple(input_shape)}."
+    msg_infos: list[str] = []
+    for i, name in enumerate(dim_names):
+        size = expected_shape[i]
+        if size > 0 and i < len(input_shape) and (dim := input_shape[i]) != 1 and dim != size:
+            if name:
+                msg_infos.append(f"Dimension {i} consistent with len({name})={size}")
+            else:
+                msg_infos.append(f"Dimension {i} consistent with required size {size}")
+    if msg_infos:
+        msg_err += f" {' & '.join(msg_infos)}."
+    else:
+        msg_err += f" Expected shape: {tuple(expected_shape)}."
+    gs.raise_exception_from(msg_err, cause)
+
+
 def broadcast_tensor(
     tensor: "np.typing.ArrayLike | None",
     dtype: torch.dtype,
@@ -746,58 +804,38 @@ def broadcast_tensor(
             )
         return torch.empty(expected_shape, dtype=dtype, device=gs.device)
 
-    tensor_ = torch.as_tensor(tensor, dtype=dtype, device=gs.device)
+    tensor_ = _expand_shape(torch.as_tensor(tensor, dtype=dtype, device=gs.device), expected_shape)
+    try:
+        return tensor_.expand(expected_shape)
+    except RuntimeError as e:
+        _raise_broadcast_shape_error(tensor_.shape, expected_shape, dim_names, e)
 
-    tensor_shape = tensor_.shape
-    tensor_ndim = len(tensor_shape)
-    expected_ndim = len(expected_shape)
 
-    # Expand current tensor shape with extra dims of size 1 if necessary before expanding to expected shape
-    if tensor_ndim == 0:
-        tensor_ = tensor_[None]
-    elif tensor_ndim < expected_ndim and not all(
-        [d1 == d2 or d2 == -1 for d1, d2 in zip(tensor_shape, expected_shape[-tensor_ndim:])]
-    ):
-        # Try expanding first dimensions if priority
-        for dims_valid in tuple(combinations(range(expected_ndim), tensor_ndim))[::-1]:
-            curr_idx = 0
-            expanded_shape = []
-            for i in range(expected_ndim):
-                if i in dims_valid:
-                    dim, size = tensor_.shape[curr_idx], expected_shape[i]
-                    if dim == size or dim == 1 or size == -1:
-                        expanded_shape.append(dim)
-                        curr_idx += 1
-                    else:
-                        break
-                else:
-                    expanded_shape.append(1)
-            else:
-                if curr_idx == tensor_ndim:
-                    tensor_ = tensor_.reshape(expanded_shape)
-                    break
-    elif tensor_ndim > expected_ndim:
-        gs.raise_exception(f"Invalid input shape: {tensor_shape}. Expecting at most {expected_ndim}D tensor.")
+def broadcast_array(
+    array: "np.typing.ArrayLike | None",
+    dtype: "np.typing.DTypeLike",
+    expected_shape: tuple[int, ...] | list[int],
+    dim_names: tuple[str, ...] | list[str] | None = None,
+) -> np.ndarray:
+    """Numpy counterpart to :func:`broadcast_tensor`. Matches the same shape-expansion semantics and ``-1``
+    wildcard convention, returning a (read-only) broadcast view via ``np.broadcast_to``.
+    """
+    if dim_names is None:
+        dim_names = ("",) * len(expected_shape)
 
-    try:
-        tensor_ = tensor_.expand(expected_shape)
-    except RuntimeError as e:
-        msg_err = f"Invalid input shape: {tuple(tensor_.shape)}."
-        msg_infos: list[str] = []
-        for i, name in enumerate(dim_names):
-            size = expected_shape[i]
-            if size > 0 and i < tensor_.ndim and (dim := tensor_.shape[i]) != 1 and dim != size:
-                if name:
-                    msg_infos.append(f"Dimension {i} consistent with len({name})={size}")
-                else:
-                    msg_infos.append(f"Dimension {i} consistent with required size {size}")
-        if msg_infos:
-            msg_err += f" {' & '.join(msg_infos)}."
-        else:
-            msg_err += f" Expected shape: {tuple(expected_shape)}."
-        gs.raise_exception_from(msg_err, e)
+    if array is None:
+        if any(size == -1 for size in expected_shape):
+            gs.raise_exception(
+                "Array not pre-allocated and expected shape not fully specified but allocation is not skipped."
+            )
+        return np.empty(expected_shape, dtype=dtype)
 
-    return tensor_
+    array_ = _expand_shape(np.asarray(array, dtype=dtype), expected_shape)
+    resolved_shape = tuple(s if s != -1 else array_.shape[i] for i, s in enumerate(expected_shape))
+    try:
+        return np.broadcast_to(array_, resolved_shape)
+    except ValueError as e:
+        _raise_broadcast_shape_error(array_.shape, expected_shape, dim_names, e)
 
 
 def sanitize_indexed_tensor(