diff --git a/examples/rendering/custom_visual_mesh.py b/examples/rendering/custom_visual_mesh.py new file mode 100644 index 000000000..2bb756936 --- /dev/null +++ b/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() diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index d528c4b78..00470b41b 100644 --- a/genesis/utils/misc.py +++ b/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( diff --git a/genesis/vis/rasterizer_context.py b/genesis/vis/rasterizer_context.py index 9e0f5a7c4..7bf546908 100644 --- a/genesis/vis/rasterizer_context.py +++ b/genesis/vis/rasterizer_context.py @@ -1,3 +1,5 @@ +from dataclasses import dataclass + import numpy as np import torch import trimesh @@ -8,7 +10,7 @@ import genesis.utils.particle as pu from genesis.ext import pyrender from genesis.ext.pyrender.jit_render import JITRenderer -from genesis.utils.misc import tensor_to_array, qd_to_numpy +from genesis.utils.misc import broadcast_array, qd_to_numpy, tensor_to_array class SegmentationColorMap: @@ -74,6 +76,24 @@ def generate_seg_colors(self): self.idxc_to_color = rgb +@dataclass +class _CustomVverts: + """Per-entity state for the custom-vverts render path. + + Encapsulates the three pieces of bookkeeping the rasterizer needs to drive per-frame vertex overrides + set via :meth:`RasterizerContext.set_custom_kinematic_entity_vverts`: + + - ``buffer``: the user-supplied ``(B, n_vverts, 3)`` override array. None means the user has cleared + the override but migration back to the instanced render path hasn't run yet. + - ``active``: True iff the renderer is currently using per-env deformed nodes for this entity. + - ``dirty``: True iff the GL buffer must be re-uploaded on the next :meth:`update_rigid` pass. + """ + + buffer: np.ndarray | None + active: bool = False + dirty: bool = True + + class RasterizerContext: def __init__(self, options): self.show_world_frame = options.show_world_frame @@ -102,6 +122,10 @@ def __init__(self, options): self.link_frame_nodes = dict() self.frustum_nodes = dict() # nodes camera frustums self.rigid_nodes = dict() + # (env_idx, geom.uid) -> per-env pyrender node for vgeoms driven by `set_custom_kinematic_entity_vverts`. + self.custom_vverts_nodes = dict() + # entity.uid -> :class:`_CustomVverts`: buffer + active / dirty bookkeeping for the custom-vverts path. + self._custom_vverts: dict = dict() self.static_nodes = dict() # used across all frames self.dynamic_nodes = dict() # nodes that live within single frame self.external_nodes = dict() # nodes added by external user @@ -175,12 +199,14 @@ def destroy(self): self.link_frame_nodes, self.frustum_nodes, self.rigid_nodes, + self.custom_vverts_nodes, self.static_nodes, self.external_nodes, ): for external_node in node_registry.values(): self.remove_node(external_node) node_registry.clear() + self._custom_vverts.clear() def reset(self): self._t = -1 @@ -207,11 +233,13 @@ def _get_geom_active_envs_idx(self, geom, rendered_envs_idx): return np.intersect1d(geom_active_envs_idx, rendered_envs_idx) return rendered_envs_idx - def add_rigid_node(self, geom, obj, **kwargs): - rigid_node = self.add_node(obj, **kwargs) - self.rigid_nodes[geom.uid] = rigid_node + def add_geom_node(self, geom, obj, **kwargs): + """Add a pyrender node for ``geom``'s mesh and register its segmentation key, returning the node. - # create segemtation id + Callers store the result in whichever per-geom registry they track (``rigid_nodes`` for the + instanced transform path, ``custom_vverts_nodes`` for the per-env custom-vverts path). + """ + node = self.add_node(obj, **kwargs) if self.segmentation_level == "geom": seg_key = (geom.entity.idx, geom.link.idx, geom.idx) elif self.segmentation_level == "link": @@ -220,7 +248,68 @@ def add_rigid_node(self, geom, obj, **kwargs): seg_key = geom.entity.idx else: gs.raise_exception(f"Unsupported segmentation level: {self.segmentation_level}") - self.create_node_seg(seg_key, rigid_node) + self.create_node_seg(seg_key, node) + return node + + def set_custom_kinematic_entity_vverts(self, entity, vverts, envs_idx=None): + """Override the visual vertex positions of a kinematic entity for rendering only, or clear an + existing override by passing ``vverts=None``. + + The rasterizer switches the entity to a per-env render path that uploads ``vverts`` as the entity's + visual mesh positions every time the buffer is updated. The override is independent of physics + simulation - solver state is unaffected. + + Parameters + ---------- + entity : KinematicEntity + The kinematic entity (or any subclass, e.g. ``RigidEntity``) whose visual vertices to override. + vverts : None | float | np.ndarray | torch.Tensor + Vertex positions in world space, broadcastable to ``(B_target, entity.n_vverts, 3)`` where + ``B_target == len(envs_idx)`` (or the scene's environment count when ``envs_idx`` is None). + Scalars, ``(3,)``, and ``(n_vverts, 3)`` are accepted and broadcast. ``None`` clears the + override and restores the standard per-vgeom transform path on the next frame. + envs_idx : None | int | list | np.ndarray, optional + Environment indices to update. None updates every environment. Not supported for + non-parallelized scenes. + """ + # gs.morphs.Plane uses a single-instance render with reflection that the per-env custom-vverts + # path cannot reproduce. Rather than silently degrade the render, refuse the override. + if isinstance(entity._morph, gs.morphs.Plane): + gs.raise_exception("Custom vverts override is not supported for 'gs.morphs.Plane' entities.") + + state = self._custom_vverts.get(entity.uid) + if vverts is None: + if envs_idx is not None: + gs.raise_exception("Cannot clear custom vverts for a subset of environments.") + if state is None: + return + if state.active: + # Migration back will run on the next ``update_rigid``; keep the entry around to signal it. + state.buffer = None + state.dirty = True + else: + # The override was queued but never made it to the render path - just drop it. + del self._custom_vverts[entity.uid] + return + + n_envs = entity._solver.n_envs + B = entity._solver._B + if envs_idx is None: + envs_idx_arr = np.arange(B, dtype=gs.np_int) + else: + if n_envs == 0: + gs.raise_exception("'envs_idx' is not supported for non-parallelized scene.") + envs_idx_arr = np.atleast_1d(envs_idx).astype(gs.np_int, copy=False) + + if state is None: + state = _CustomVverts(buffer=np.zeros((B, entity.n_vverts, 3), dtype=gs.np_float)) + self._custom_vverts[entity.uid] = state + elif state.buffer is None: + state.buffer = np.zeros((B, entity.n_vverts, 3), dtype=gs.np_float) + state.buffer[envs_idx_arr] = broadcast_array( + vverts, gs.np_float, (len(envs_idx_arr), entity.n_vverts, 3), ("envs", "vverts", "xyz") + ) + state.dirty = True def add_static_node(self, entity, obj, i_b, **kwargs): static_node = self.add_node(obj, **kwargs) @@ -444,7 +533,7 @@ def on_rigid(self): geom_T = geom_T[:1] env_shared = True - self.add_rigid_node( + self.rigid_nodes[geom.uid] = self.add_geom_node( geom, pyrender.Mesh.from_trimesh( mesh=mesh, @@ -470,6 +559,80 @@ def update_rigid(self): geoms = entity.geoms geoms_T = solver._geoms_render_T + # Per-env vertex update path (e.g. SMPL skin), driven by ``set_custom_kinematic_entity_vverts``. + # Only takes over in 'visual' mode since the override only writes the visual mesh; collision + # / sdf still need the standard per-vgeom transform update below. + state = self._custom_vverts.get(entity.uid) + needs_custom = state is not None and state.buffer is not None and entity.surface.vis_mode == "visual" + if needs_custom and not state.active: + # Migrate forward: drop the instanced rigid_node, build per-env custom-vverts nodes from + # the static base trimesh. Subsequent frames just refresh their GL position / normal + # buffers. + for geom in entity.vgeoms: + if geom.uid in self.rigid_nodes: + old_node = self.rigid_nodes.pop(geom.uid) + self.remove_node_seg(old_node) + self.remove_node(old_node) + geom_envs_idx = self._get_geom_active_envs_idx(geom, self.rendered_envs_idx) + if len(geom_envs_idx) == 0: + continue + mesh_trimesh = geom.get_trimesh() + for idx in geom_envs_idx: + self.custom_vverts_nodes[(idx, geom.uid)] = self.add_geom_node( + geom, + pyrender.Mesh.from_trimesh( + mesh=mesh_trimesh, + smooth=geom.surface.smooth, + double_sided=geom.surface.double_sided, + ), + ) + state.active = True + elif not needs_custom and state is not None and state.active: + # Migrate back: tear down per-env custom-vverts nodes and recreate the instanced + # rigid_node so the next frame falls into the standard transform path. + for geom in entity.vgeoms: + geom_envs_idx = self._get_geom_active_envs_idx(geom, self.rendered_envs_idx) + for idx in geom_envs_idx: + node = self.custom_vverts_nodes.pop((idx, geom.uid), None) + if node is not None: + self.remove_node_seg(node) + self.remove_node(node) + self.rigid_nodes[geom.uid] = self.add_geom_node( + geom, + pyrender.Mesh.from_trimesh( + mesh=geom.get_trimesh(), + poses=solver._vgeoms_render_T[geom.idx][geom_envs_idx], + smooth=geom.surface.smooth, + double_sided=geom.surface.double_sided, + ), + ) + if state.buffer is None: + del self._custom_vverts[entity.uid] + else: + # vis_mode flipped away from 'visual'; keep the buffer but mark inactive so the next + # transition back into 'visual' re-migrates and re-uploads. + state.active = False + state.dirty = True + if needs_custom and state.dirty: + for geom in entity.vgeoms: + geom_envs_idx = self._get_geom_active_envs_idx(geom, self.rendered_envs_idx) + if len(geom_envs_idx) == 0: + continue + v_start = geom.vvert_start - entity.vvert_start + v_end = geom.vvert_end - entity.vvert_start + for idx in geom_envs_idx: + node = self.custom_vverts_nodes[(idx, geom.uid)] + geom_vverts = state.buffer[idx, v_start:v_end, :] + self.scene.envs_offset[idx] + update_data = self._scene.reorder_vertices(node, geom_vverts) + self.jit.update_buffer(self._scene.get_buffer_id(node, "pos"), update_data) + normal_data = self.jit.update_normal(node, update_data) + if normal_data is not None: + self.jit.update_buffer(self._scene.get_buffer_id(node, "normal"), normal_data) + state.dirty = False + if needs_custom: + continue + + # Standard instanced transform path for geom in geoms: # Skip geoms that weren't added - in heterogeneous simulation, some geoms # may not be rendered in any of the requested environments diff --git a/genesis/vis/visualizer.py b/genesis/vis/visualizer.py index ab1964aaf..55ab0a229 100644 --- a/genesis/vis/visualizer.py +++ b/genesis/vis/visualizer.py @@ -156,6 +156,14 @@ def add_light(self, pos, dir, color, intensity, directional, castshadow, cutoff, else: gs.raise_exception("`add_light` is specific to batch renderer.") + @gs.assert_built + def set_custom_kinematic_entity_vverts(self, entity, vverts, envs_idx=None): + """Override the visual vertex positions of a kinematic entity for rendering only, or clear an + existing override by passing ``vverts=None``. See + :meth:`RasterizerContext.set_custom_kinematic_entity_vverts` for the full signature and broadcasting rules. + """ + self._context.set_custom_kinematic_entity_vverts(entity, vverts, envs_idx=envs_idx) + @gs.assert_built def reset(self): self._t = -1 diff --git a/tests/test_examples.py b/tests/test_examples.py index 0e970395c..3be767c4c 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -17,6 +17,7 @@ "kinematic/**/*.py", "rigid/**/*.py", "render_async/**/*.py", + "rendering/custom_visual_mesh.py", "sap_coupling/**/*.py", "sensors/**/*.py", "tutorials/**/*.py", diff --git a/tests/test_render.py b/tests/test_render.py index b1f2579f3..14694cdcd 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -1987,3 +1987,51 @@ def test_render_offscreen_oversized_resolution(renderer): normal=False, ) assert rgb.shape[:2] == (requested_res[1], requested_res[0]) + + +@pytest.mark.required +@pytest.mark.parametrize("renderer_type", [RENDERER_TYPE.RASTERIZER]) +@pytest.mark.parametrize("n_envs", [0, 3]) +def test_set_vverts(n_envs, renderer, show_viewer): + CAM_RES = (320, 240) + scene = gs.Scene( + renderer=renderer, + show_viewer=show_viewer, + show_FPS=False, + ) + entity = scene.add_entity( + morph=gs.morphs.Mesh( + file="meshes/sphere.obj", + scale=0.2, + pos=(0.0, 0.0, 0.5), + fixed=True, + ), + ) + cam = scene.add_camera( + res=CAM_RES, + pos=(0.0, -1.5, 0.5), + lookat=(0.0, 0.0, 0.5), + ) + scene.build(n_envs=n_envs) + + viz = scene.visualizer + + # Baseline render with the standard transform path. + rgb_baseline = tensor_to_array(cam.render(rgb=True)[0]).astype(np.float32) + + # Override every vertex to a single point well outside the camera frustum. The sphere collapses to a + # degenerate point and largely disappears from the framebuffer - a strong, easy-to-verify signal that the + # override is wired up. + viz.set_custom_kinematic_entity_vverts(entity, (0.0, 0.0, 10.0)) + rgb_deformed = tensor_to_array(cam.render(rgb=True, force_render=True)[0]).astype(np.float32) + assert np.abs(rgb_deformed - rgb_baseline).mean() > 5.0, "deformation should produce a visible pixel diff" + + # Partial batched write with scalar / array-like broadcast. + if n_envs > 0: + viz.set_custom_kinematic_entity_vverts(entity, 0.0, envs_idx=0) + viz.set_custom_kinematic_entity_vverts(entity, 1.0, envs_idx=[0, 2]) + + # ``set_custom_kinematic_entity_vverts(..., None)`` reinstates the instanced rigid_node and reproduces baseline. + viz.set_custom_kinematic_entity_vverts(entity, None) + rgb_restored = tensor_to_array(cam.render(rgb=True, force_render=True)[0]).astype(np.float32) + assert np.abs(rgb_restored - rgb_baseline).mean() < 1.0, "clearing vverts should restore the baseline render"