Source code for gs_nyx_plugin.nyx_renderer

"""Nyx Renderer integration for Genesis simulation."""

# External imports
try:
    import torch
except ImportError as e:
    raise ImportError(
        "'torch' module not available. Please install pytorch manually: https://pytorch.org/get-started/locally/"
    ) from e
from typing import NamedTuple, Optional

# Internal imports
import genesis as gs
import genesis.utils.geom as gu
from genesis.utils.misc import qd_to_torch
from gs_nyx import nyx_py_sdk as nps
from gs_nyx import nyx_py_renderer as npr

# Relative imports
from .nyx_scene_utils import (
    is_mpm_visual_entity,
    is_mpm_recon_entity,
    is_sph_recon_entity,
    particle_recon_budget,
    _particle_solver_radius,
)

# Aliases
NyxRenderer = npr.NyxRenderer
BridgeStartupParams = npr.BridgeStartupParams
BridgeUpdateDesc = npr.BridgeUpdateDesc
BridgeUpdateData = npr.BridgeUpdateData
ERenderMode = npr.ERenderMode
EDebugView = npr.EDebugView


# Excludes the npr re-export aliases above on purpose: they are already
# documented in gs_nyx.nyx_py_renderer and re-stubbing them here would create
# duplicate cross-reference targets in the Sphinx inventory.
__all__ = ["NyxPyRenderer", "entity_to_uuid", "uuid_to_entity"]


class _ReconSlot(NamedTuple):
    """Per-entity bookkeeping for a particle-recon entity in the shared deformable buffer.

    Used for both MPM-recon and SPH-recon entities; the ``solver`` field carries the owning
    solver so the per-frame meshing step can pull particles and the correct radius from the
    right place without the update loop branching on entity type.
    """

    entity: object  # MPMEntity or SPHEntity (kept for surface.recon_backend access)
    solver: object  # MPMSolver or SPHSolver (source of particle_radius + domain bounds)
    deform_idx: int  # Index within the deformable iteration order (FEM, PBD, MPM-visual, MPM-recon, SPH-recon)
    vert_start: int  # Start offset into _deformable_verts_cuda (in vertices)
    tri_start: int  # Start offset into _deformable_indices_cuda (in triangles)
    max_vertices: int  # Pre-allocated vertex budget
    max_triangles: int  # Pre-allocated triangle budget


# ========================== Utility Functions ==========================
[docs] def entity_to_uuid(pair_list, entity): """Look up the Nyx UUID assigned to a Genesis entity. Parameters ---------- pair_list : list of tuple ``(entity, uuid)`` pairs produced by :class:`NyxSceneExporter` during scene export. The list is ordered by export order and uses object identity for the entity field. entity : object The Genesis entity to look up. Compared by identity (``is``), not equality, because two distinct entities can compare equal under Genesis's ``__eq__``. Returns ------- uuid or None The Nyx UUID registered for that entity, or ``None`` if the entity was not exported (e.g. its morph had ``visualization=False``). """ for e, u in pair_list: if e is entity: # identity comparison (important) return u return None
[docs] def uuid_to_entity(pair_list, uuid): """Reverse lookup: return the Genesis entity that owns a Nyx UUID. Used to translate the ``targetUUID`` returned by Nyx's pixel-pick raycast back into the Genesis-side entity object. Parameters ---------- pair_list : list of tuple ``(entity, uuid)`` pairs from :class:`NyxSceneExporter`. uuid : object UUID value to match (Nyx's :class:`gs_nyx.nyx_py_sdk.uuid` type). Returns ------- object or None The Genesis entity, or ``None`` if the UUID is not present in the pair list. """ for e, u in pair_list: if u == uuid: return e return None
# ========================== Nyx Renderer ==========================
[docs] class NyxPyRenderer: """Python-side wrapper around the native Nyx renderer. Owns the CUDA tensors that bridge Genesis simulation state (rigid transforms, deformable vertex positions, camera poses) to the Nyx GPU renderer, and drives the per-frame upload + render. End users typically interact with the renderer indirectly via :class:`NyxCameraSensor`; this class is exposed for advanced use cases that need to drive the renderer manually. Coordinate convention --------------------- Genesis is Z-up, Nyx is Y-up. Static scene data is converted on the CPU at export time by :class:`NyxSceneExporter`; per-frame transforms are uploaded as raw Genesis-space (Z-up, WXYZ) values and converted on the GPU. Lifecycle --------- 1. Construct with the scene, exported scene-description path, and max viewport size. 2. Call :meth:`build` once after the Genesis scene has been built, passing the entity-to-UUID mapping from :class:`NyxSceneExporter`. This starts the renderer, loads the scene, and allocates GPU buffers sized to the current geometry counts. 3. Each frame: :meth:`update_scene` for each environment, then :meth:`render` per camera. :meth:`update` once per frame drives the window event loop and animation clock. 4. :meth:`destroy` releases the native renderer and GPU buffers. """
[docs] def __init__( self, scene, scene_file_path, max_width, max_height, n_envs=1, open_window=False, camera_defs=None, render_mode=ERenderMode.FastPathTracer, debug_view=EDebugView.Meshlet, ): """Construct the renderer wrapper. Does **not** start the native renderer. The native renderer is started inside :meth:`build`; this constructor only captures parameters and inspects the Genesis solvers so the geometry-count properties can be queried before build. Parameters ---------- scene : genesis.Scene The built Genesis scene. Its solvers (rigid / FEM / PBD / MPM) are queried for geometry counts and per-step state. scene_file_path : str Path to the Nyx scene-description JSON produced by :class:`NyxSceneExporter`. Loaded once during :meth:`build`. max_width, max_height : int Maximum viewport size in pixels. Nyx pre-allocates framebuffers up to this size; cameras with smaller resolutions render into a sub-region. Set this to the largest ``res`` across all cameras that will be registered. n_envs : int, optional Number of Genesis parallel environments (i.e. ``scene._B``). Stored for reference; the renderer itself processes one env at a time, with :meth:`update_scene` selecting the env via its ``env_index`` argument. open_window : bool, optional If ``True``, open a native GUI window for live preview. Default ``False`` (offscreen). camera_defs : list of dict or None, optional Per-camera definition dicts (built by :class:`NyxCameraSensor.build`). When provided, the renderer uses the dict-driven path and skips Genesis ``Camera`` objects. Pass ``None`` when no Nyx sensors are registered to fall back to the Genesis visualizer's cameras. render_mode : gs_nyx.nyx_py_renderer.ERenderMode, optional Rendering algorithm. Defaults to :attr:`~gs_nyx.nyx_py_renderer.ERenderMode.FastPathTracer`. debug_view : gs_nyx.nyx_py_renderer.EDebugView, optional Debug visualization channel (e.g. meshlet IDs, normals). Only takes effect when ``render_mode`` is the debug mode. """ # Keep track of the parameters. self._scene = scene self._visualizer = self._scene.visualizer self._scene_file_path = scene_file_path self._max_width = max_width self._max_height = max_height self._n_envs = n_envs self._open_window = open_window self._render_mode = render_mode self._debug_view = debug_view # Create nyx instance self._renderer = NyxRenderer() # Get the supported solvers self._rigid_solver = self._scene.rigid_solver self._fem_solver = self._scene.fem_solver self._pbd_solver = self._scene.pbd_solver self._mpm_solver = self._scene.mpm_solver self._sph_solver = self._scene.sph_solver # Camera tensor data self._num_cameras = 0 self._cam_pos_tensor_cuda = None self._cam_rot_tensor_cuda = None # Pre-allocated single-env CUDA tensors for rigid geometry (shape: [n_vgeoms, 3/4]) self._geom_pos_tensor_cuda: torch.Tensor self._geom_rot_tensor_cuda: torch.Tensor # Shared deformable GPU buffers for all the supported deformable solvers self._deformable_verts_cuda: torch.Tensor self._deformable_uvs_cuda: torch.Tensor self._deformable_indices_cuda: torch.Tensor # Per-section views into the shared vertex buffer self._deformable_fem_verts_view: torch.Tensor self._deformable_pbd_verts_view: torch.Tensor self._deformable_mpm_visual_verts_view: torch.Tensor # MPM-recon and SPH-recon entities: per-entity slot info; populated in _allocate_cuda_buffers. # Kept as separate lists so the per-section vert/tri budgets and the reference-table # ordering remain easy to follow, but a single _ReconSlot type drives both via slot.solver. self._mpm_recon_slots: list[_ReconSlot] = [] self._sph_recon_slots: list[_ReconSlot] = [] # Cached MPM-visual entity list (in solver order) for per-step updates self._mpm_visual_entities: list = [] # MPM-visual section offsets (set when allocating) self._mpm_visual_vert_start = 0 self._mpm_visual_tri_start = 0 # Camera definitions from sensor (if using NyxCameraSensor) # Used only at build time for scene export. Not mutated at runtime. self._camera_defs = camera_defs
# ------------------------------------------------------------------------- # Solver geometry counts (per single Genesis env) # ------------------------------------------------------------------------- @property def _num_rigid_geoms(self) -> int: return self._rigid_solver.n_vgeoms @property def _num_fem_verts(self) -> int: return self._fem_solver.n_vertices if self._fem_solver.is_active else 0 @property def _num_fem_tris(self) -> int: return self._fem_solver.n_surfaces if self._fem_solver.is_active else 0 @property def _num_fem_entities(self) -> int: return len(self._fem_solver.entities) if self._fem_solver.is_active else 0 @property def _num_pbd_verts(self) -> int: return self._pbd_solver.n_vverts if self._pbd_solver.is_active else 0 @property def _num_pbd_tris(self) -> int: return self._pbd_solver.n_vfaces if self._pbd_solver.is_active else 0 @property def _num_pbd_entities(self) -> int: return len(self._pbd_solver.entities) if self._pbd_solver.is_active else 0 def _mpm_entities(self, predicate) -> list: if not self._mpm_solver.is_active: return [] return [e for e in self._mpm_solver.entities if predicate(e)] @property def _mpm_visual_entity_list(self) -> list: return self._mpm_entities(is_mpm_visual_entity) @property def _mpm_recon_entity_list(self) -> list: return self._mpm_entities(is_mpm_recon_entity) @property def _num_mpm_visual_verts(self) -> int: return sum(e.n_vverts for e in self._mpm_visual_entity_list) @property def _num_mpm_visual_tris(self) -> int: return sum(e.n_vfaces for e in self._mpm_visual_entity_list) @property def _num_mpm_visual_entities(self) -> int: return len(self._mpm_visual_entity_list) @property def _num_mpm_recon_entities(self) -> int: return len(self._mpm_recon_entity_list) # SPH (recon-only — SPH ships a fluid material exclusively, no skinning path). def _sph_entities(self, predicate) -> list: if not self._sph_solver.is_active: return [] return [e for e in self._sph_solver.entities if predicate(e)] @property def _sph_recon_entity_list(self) -> list: return self._sph_entities(is_sph_recon_entity) @property def _num_sph_recon_entities(self) -> int: return len(self._sph_recon_entity_list) @property def _num_deform_verts(self) -> int: return ( self._num_fem_verts + self._num_pbd_verts + self._num_mpm_visual_verts + self._mpm_recon_verts_budget + self._sph_recon_verts_budget ) @property def _num_deform_tris(self) -> int: return ( self._num_fem_tris + self._num_pbd_tris + self._num_mpm_visual_tris + self._mpm_recon_tris_budget + self._sph_recon_tris_budget ) @property def _num_deform_objs(self) -> int: return ( self._num_fem_entities + self._num_pbd_entities + self._num_mpm_visual_entities + self._num_mpm_recon_entities + self._num_sph_recon_entities ) # MPM-recon worst-case budgets — computed once and cached so the property is stable @property def _mpm_recon_verts_budget(self) -> int: return sum( particle_recon_budget(self._mpm_solver, e)[0] for e in self._mpm_recon_entity_list ) @property def _mpm_recon_tris_budget(self) -> int: return sum( particle_recon_budget(self._mpm_solver, e)[1] for e in self._mpm_recon_entity_list ) @property def _sph_recon_verts_budget(self) -> int: return sum( particle_recon_budget(self._sph_solver, e)[0] for e in self._sph_recon_entity_list ) @property def _sph_recon_tris_budget(self) -> int: return sum( particle_recon_budget(self._sph_solver, e)[1] for e in self._sph_recon_entity_list ) # -------------------------------------------------------------------------
[docs] def build(self, pair_list): """Start the native renderer and allocate all GPU buffers. Must be called exactly once after the Genesis scene has been built and the scene-description JSON has been exported by :class:`NyxSceneExporter`. Inspects the active Genesis solvers to determine vertex / triangle / camera counts, then: 1. Starts the Nyx native renderer with the configured viewport size, render mode, and debug view. 2. Loads the scene description from ``scene_file_path``. 3. Allocates Nyx scene buffers sized for current geometry. 4. Allocates and pins this wrapper's interop CUDA tensors. 5. Uploads the static deformable data (UVs, triangle indices) for FEM / PBD / MPM-visual entities. MPM-recon meshes are streamed per frame. Parameters ---------- pair_list : list of tuple ``(entity, uuid)`` pairs from :attr:`NyxSceneExporter._entity_uuid_pairs`. Used to populate the rigid- and deformable-geometry reference tables so per-frame transforms can be matched to the right scene instances. Notes ----- If any MPM-recon entities are present this also primes the configured reconstruction backend (openvdb or splashsurf) with a one-particle call so missing dependencies fail at build time rather than on the first rendered frame. """ self._cameras = gs.List( [camera for camera in self._visualizer._cameras if not camera.debug] ) # Cache the MPM-visual entity list for per-step updates (solver order) self._mpm_visual_entities = self._mpm_visual_entity_list # Log per-entity geometry counts so the user can sanity-check the allocated buffers self._log_mpm_entity_counts() # If any particle-recon entities (MPM or SPH) are in play, smoke-test the reconstruction # backend so a missing ParticleMesher / pysplashsurf surfaces here instead of every frame. if self._num_mpm_recon_entities > 0 or self._num_sph_recon_entities > 0: self._check_recon_backend_available() # Startup the renderer startup_params = BridgeStartupParams() startup_params.cudaDevice = torch.cuda.current_device() startup_params.renderMode = self._render_mode startup_params.debugView = self._debug_view startup_params.maxViewportWidth = self._max_width startup_params.maxViewportHeight = self._max_height startup_params.openWindow = self._open_window self._renderer.startup(startup_params) # Load the scene from file (exported by NyxSceneExporter) self._renderer.load_scene_from_file(self._scene_file_path) # Allocate scene buffers based on geometry counts self._allocate_scene_buffers(pair_list) # Initialize update data structure with buffer pointers self._init_update_data()
def _log_mpm_entity_counts(self): """Print per-particle-entity vertex/triangle counts (visual: fixed; recon: worst-case budget).""" if self._mpm_solver.is_active: for i, entity in enumerate(self._mpm_visual_entities): name = getattr(entity, "name", None) or f"entity[{i}]" print( f"[NyxRenderer] MPM-visual '{name}': " f"n_vverts={entity.n_vverts}, n_vfaces={entity.n_vfaces}" ) for i, entity in enumerate(self._mpm_recon_entity_list): name = getattr(entity, "name", None) or f"entity[{i}]" max_v, max_t = particle_recon_budget(self._mpm_solver, entity) backend = getattr(entity.surface, "recon_backend", "openvdb") print( f"[NyxRenderer] MPM-recon '{name}': budget verts={max_v}, " f"budget triangles={max_t}, backend='{backend}'" ) if self._sph_solver.is_active: for i, entity in enumerate(self._sph_recon_entity_list): name = getattr(entity, "name", None) or f"entity[{i}]" max_v, max_t = particle_recon_budget(self._sph_solver, entity) backend = getattr(entity.surface, "recon_backend", "openvdb") print( f"[NyxRenderer] SPH-recon '{name}': budget verts={max_v}, " f"budget triangles={max_t}, backend='{backend}'" ) def _check_recon_backend_available(self): """Smoke-test the reconstruction backend for each particle-recon entity (MPM or SPH). `particles_to_mesh()` raises if the configured backend isn't usable on the current platform (openvdb is Linux+py3.9 only; splashsurf needs pysplashsurf installed). Run a one-particle call here so the user gets a clean failure at build time, not on the first render. """ from genesis.utils.particle import particles_to_mesh import numpy as np probe_positions = np.zeros((1, 3), dtype=np.float32) recon_targets = [ ("MPM", self._mpm_solver, self._mpm_recon_entity_list), ("SPH", self._sph_solver, self._sph_recon_entity_list), ] for solver_name, solver, entity_list in recon_targets: if not entity_list: continue radius = _particle_solver_radius(solver) for entity in entity_list: backend = getattr(entity.surface, "recon_backend", "openvdb") try: particles_to_mesh(probe_positions, radius, backend) except Exception as e: gs.raise_exception_from( f"{solver_name} recon entity '{getattr(entity, 'name', '<unnamed>')}' requested " f"backend '{backend}', but it isn't available on this platform. openvdb requires " "Linux + Python 3.9; splashsurf requires 'pysplashsurf' to be installed. Use " "vis_mode='visual' if the material supports it, or skip this scene on this platform.", e, ) def _allocate_scene_buffers(self, pair_list): """Allocate Nyx scene buffers based on geometry counts.""" # Create the update description update_desc = BridgeUpdateDesc() # Determine number of cameras (from sensor definitions or Genesis visualizer) if self._camera_defs is not None: self._num_cameras = len(self._camera_defs) else: self._num_cameras = len(self._cameras) update_desc.numCameras = self._num_cameras # Rigid geometry update_desc.numRigidGeom = self._num_rigid_geoms # Deformable counts (includes MPM budgets) update_desc.numDeformVertices = self._num_deform_verts update_desc.numDeformTriangles = self._num_deform_tris update_desc.numDeformObjects = self._num_deform_objs # Build the reference table mapping geometry to UUIDs for Nyx self._build_reference_table(update_desc, pair_list) # Allocate the scene buffers in Nyx based on the update description self._renderer.allocate_scene_buffers(update_desc) # Allocate internal CUDA buffers for the interop between Genesis and Nyx self._allocate_cuda_buffers() # Upload static deformable data (UVs + indices) to the internal CUDA buffers self._upload_static_deformable_data() def _allocate_cuda_buffers(self): """Allocate all Python-side CUDA tensors for camera, rigid, and deformable geometry.""" # Camera tensors if self._num_cameras > 0: self._cam_pos_tensor_cuda = torch.zeros( (self._num_cameras, 3), dtype=torch.float32, device="cuda" ) self._cam_rot_tensor_cuda = torch.zeros( (self._num_cameras, 4), dtype=torch.float32, device="cuda" ) # Rigid geometry tensors self._geom_pos_tensor_cuda = torch.zeros( (self._num_rigid_geoms, 3), dtype=torch.float32, device="cuda" ) self._geom_rot_tensor_cuda = torch.zeros( (self._num_rigid_geoms, 4), dtype=torch.float32, device="cuda" ) # Deformable geometry tensors if self._num_deform_verts == 0: return # Section offsets (in vertices / triangles): FEM | PBD | MPM-visual | MPM-recon | SPH-recon. # The SPH-recon section starts where MPM-recon ends; the slot-build loop below carries the # running offset forward through both recon sections, so SPH start offsets aren't named here. fem_verts = self._num_fem_verts fem_tris = self._num_fem_tris pbd_verts = self._num_pbd_verts pbd_tris = self._num_pbd_tris mpm_v_verts = self._num_mpm_visual_verts mpm_v_tris = self._num_mpm_visual_tris pbd_vert_start = fem_verts pbd_tri_start = fem_tris mpm_visual_vert_start = pbd_vert_start + pbd_verts mpm_visual_tri_start = pbd_tri_start + pbd_tris mpm_recon_vert_start = mpm_visual_vert_start + mpm_v_verts mpm_recon_tri_start = mpm_visual_tri_start + mpm_v_tris self._mpm_visual_vert_start = mpm_visual_vert_start self._mpm_visual_tri_start = mpm_visual_tri_start # Nyx expects flat, contiguous arrays: float3 positions, float2 UVs, int32 indices self._deformable_verts_cuda = torch.zeros( (self._num_deform_verts, 3), dtype=torch.float32, device="cuda" ) self._deformable_uvs_cuda = torch.zeros( (self._num_deform_verts, 2), dtype=torch.float32, device="cuda" ) self._deformable_indices_cuda = torch.zeros( self._num_deform_tris * 3, dtype=torch.int32, device="cuda" ) # Calculate the starting index for PBD vertices in the shared vertex buffer. pbd_vert_start = self._num_fem_verts # Create per-solver views into the shared vertex buffer if self._fem_solver.is_active and fem_verts > 0: self._deformable_fem_verts_view = self._deformable_verts_cuda[ :pbd_vert_start ] if self._pbd_solver.is_active and pbd_verts > 0: self._deformable_pbd_verts_view = self._deformable_verts_cuda[ pbd_vert_start:mpm_visual_vert_start ] if mpm_v_verts > 0: self._deformable_mpm_visual_verts_view = self._deformable_verts_cuda[ mpm_visual_vert_start:mpm_recon_vert_start ] # Bookkeeping for particle-recon entities (MPM then SPH). deform_idx must match the order # references are added in _build_reference_table: # FEM (n_fem) -> PBD (n_pbd) -> MPM-visual (n_mpm_v) -> MPM-recon -> SPH-recon. self._mpm_recon_slots = [] self._sph_recon_slots = [] mpm_recon_deform_base = ( self._num_fem_entities + self._num_pbd_entities + self._num_mpm_visual_entities ) cur_vert_start = mpm_recon_vert_start cur_tri_start = mpm_recon_tri_start for i, entity in enumerate(self._mpm_recon_entity_list): max_v, max_t = particle_recon_budget(self._mpm_solver, entity) self._mpm_recon_slots.append( _ReconSlot( entity=entity, solver=self._mpm_solver, deform_idx=mpm_recon_deform_base + i, vert_start=cur_vert_start, tri_start=cur_tri_start, max_vertices=max_v, max_triangles=max_t, ) ) cur_vert_start += max_v cur_tri_start += max_t # SPH-recon section starts where MPM-recon ends (matches sph_recon_vert_start above; we # carry cur_* across instead of resetting to keep the running offset honest if budgets # are clamped per entity later). sph_recon_deform_base = mpm_recon_deform_base + self._num_mpm_recon_entities for i, entity in enumerate(self._sph_recon_entity_list): max_v, max_t = particle_recon_budget(self._sph_solver, entity) self._sph_recon_slots.append( _ReconSlot( entity=entity, solver=self._sph_solver, deform_idx=sph_recon_deform_base + i, vert_start=cur_vert_start, tri_start=cur_tri_start, max_vertices=max_v, max_triangles=max_t, ) ) cur_vert_start += max_v cur_tri_start += max_t def _build_reference_table(self, update_desc, pair_list): """Build the reference table mapping geometry to UUIDs for Nyx. Deformable order MUST match _allocate_cuda_buffers's section layout: FEM, PBD, MPM-visual, MPM-recon. The deform_idx fed to set_deform_entity_active_triangles() is the index in this iteration. """ # Build the reference table for all geometry (rigid + all deformables) total_refs = update_desc.numRigidGeom + update_desc.numDeformObjects update_desc.reference_resize(total_refs) # Rigid geometry references for refIdx in range(update_desc.numRigidGeom): currentRef = update_desc.get_reference(refIdx) current_vgeom = self._rigid_solver.vgeoms[refIdx] current_entity = current_vgeom.entity currentRef.targetUUID = entity_to_uuid(pair_list, current_entity) if isinstance(current_entity.morph, gs.morphs.URDF): currentRef.nodeName = current_vgeom.link.name elif isinstance(current_entity.morph, gs.morphs.MJCF): currentRef.nodeName = f"{current_vgeom.link.name}_{current_vgeom.link.vgeoms.index(current_vgeom)}" else: currentRef.nodeName = "" update_desc.set_reference(refIdx, currentRef) # Deformable references in order: FEM -> PBD -> MPM-visual -> MPM-recon -> SPH-recon deform_idx = update_desc.numRigidGeom deformable_entity_lists = [] if self._fem_solver.is_active: deformable_entity_lists.append(self._fem_solver.entities) if self._pbd_solver.is_active: deformable_entity_lists.append(self._pbd_solver.entities) deformable_entity_lists.append(self._mpm_visual_entity_list) deformable_entity_lists.append(self._mpm_recon_entity_list) deformable_entity_lists.append(self._sph_recon_entity_list) for entity_list in deformable_entity_lists: for entity in entity_list: currentRef = update_desc.get_reference(deform_idx) currentRef.targetUUID = entity_to_uuid(pair_list, entity) currentRef.nodeName = "" update_desc.set_reference(deform_idx, currentRef) deform_idx += 1 def _upload_static_deformable_data(self): """Upload static deformable data (UVs + indices) for FEM, PBD, and MPM-visual entities. MPM-recon UVs and indices are streamed per frame, so we leave their buffer slots zeroed. """ if self._num_deform_verts == 0: return pbd_vert_start = self._num_fem_verts pbd_tris_start = self._num_fem_tris mpm_v_vert_start = self._mpm_visual_vert_start mpm_v_tri_start = self._mpm_visual_tri_start # FEM if self._fem_solver.is_active and self._num_fem_verts > 0: _, fem_ti, fem_uv = self._fem_solver.get_state_render( self._fem_solver._sim.cur_substep_local ) fem_uvs_src = qd_to_torch(fem_uv, transpose=False, copy=True).to( torch.float32 ) fem_idx_src = qd_to_torch(fem_ti, transpose=False, copy=True).reshape(-1) self._deformable_uvs_cuda[:pbd_vert_start].copy_(fem_uvs_src) self._deformable_indices_cuda[: pbd_tris_start * 3].copy_(fem_idx_src) # PBD (indices need remapping into the shared vertex space) if self._pbd_solver.is_active and self._num_pbd_verts > 0: _, pbd_vu, pbd_fi = self._pbd_solver.get_state_render() pbd_uvs_src = qd_to_torch(pbd_vu, transpose=False, copy=True).to( torch.float32 ) pbd_idx_src = qd_to_torch(pbd_fi, transpose=False, copy=True).reshape(-1) pbd_uv_dst = self._deformable_uvs_cuda[pbd_vert_start:mpm_v_vert_start] pbd_idx_dst = self._deformable_indices_cuda[ pbd_tris_start * 3 : mpm_v_tri_start * 3 ] pbd_uv_dst.copy_(pbd_uvs_src) pbd_idx_dst.copy_(pbd_idx_src) pbd_idx_dst.add_(pbd_vert_start) # MPM-visual: per-entity static face indices + (zero) UVs, remapped into shared vertex space cur_vert = mpm_v_vert_start cur_tri = mpm_v_tri_start for entity in self._mpm_visual_entities: n_v = entity.n_vverts n_t = entity.n_vfaces if n_v == 0 or n_t == 0: continue # Genesis exposes the static face indices via _vfaces (numpy, local to the entity) faces = torch.as_tensor( entity._vfaces, dtype=torch.int32, device="cuda" ).reshape(-1) self._deformable_indices_cuda[cur_tri * 3 : (cur_tri + n_t) * 3].copy_( faces + cur_vert ) cur_vert += n_v cur_tri += n_t def _init_update_data(self): """Initialize the BridgeUpdateData structure with all buffer pointers.""" self._update_data = BridgeUpdateData() # Note that it's safe to use data_ptr() only at init time because we allocate # the buffers ourselves and their address won't change. It's not safe to # do a similar thing with buffers from Genesis. # Camera data self._update_data.numCameras = self._num_cameras self._update_data.camPosPtr = ( self._cam_pos_tensor_cuda.data_ptr() if self._cam_pos_tensor_cuda is not None else 0 ) self._update_data.camRotPtr = ( self._cam_rot_tensor_cuda.data_ptr() if self._cam_rot_tensor_cuda is not None else 0 ) # Rigid geometry self._update_data.numRigidGeom = self._num_rigid_geoms self._update_data.rigidGeomPosPtr = self._geom_pos_tensor_cuda.data_ptr() self._update_data.rigidGeomRotPtr = self._geom_rot_tensor_cuda.data_ptr() # Combined deformable geometry for all supported solvers if self._num_deform_verts > 0: self._update_data.numDeformVertices = self._num_deform_verts self._update_data.numDeformTriangles = self._num_deform_tris self._update_data.numDeformObjects = self._num_deform_objs self._update_data.deformVertexPosPtr = ( self._deformable_verts_cuda.data_ptr() ) self._update_data.deformVertexUVsPtr = self._deformable_uvs_cuda.data_ptr() self._update_data.deformTriangleIndicesPtr = ( self._deformable_indices_cuda.data_ptr() ) def _update_geometry_tensors(self, env_index): """Update geometry buffers each frame: rigid transforms + deformable vertex positions.""" # Get the current frame rigid geometry data from Genesis geom_pos_tensor = qd_to_torch( self._rigid_solver.vgeoms_state.pos, transpose=True, copy=False, ) geom_rot_tensor = qd_to_torch( self._rigid_solver.vgeoms_state.quat, transpose=True, copy=False, ) # Copy env n data into pre-allocated single-env CUDA tensors # Handles CPU to GPU transfer and float64 to float32 conversion self._geom_pos_tensor_cuda.copy_(geom_pos_tensor[env_index].to(torch.float32)) self._geom_rot_tensor_cuda.copy_(geom_rot_tensor[env_index].to(torch.float32)) # Copy deformable vertex positions (env n) to the pre-allocated views if self._fem_solver.is_active and self._num_fem_verts > 0: fem_sv, _, _ = self._fem_solver.get_state_render( self._fem_solver._sim.cur_substep_local ) fem_src = qd_to_torch(fem_sv, transpose=True, copy=False) self._deformable_fem_verts_view.copy_(fem_src[env_index].to(torch.float32)) if self._pbd_solver.is_active and self._num_pbd_verts > 0: pbd_vp, _, _ = self._pbd_solver.get_state_render() pbd_src = qd_to_torch(pbd_vp, transpose=True, copy=False) self._deformable_pbd_verts_view.copy_(pbd_src[env_index].to(torch.float32)) # MPM-visual: skinned vertices on GPU (mpm_solver.vverts_render.pos). Topology fixed. if self._num_mpm_visual_verts > 0: self._mpm_solver.update_render_fields() mpm_vverts = qd_to_torch( self._mpm_solver.vverts_render.pos, transpose=True, copy=False ) # Shape: (n_envs, total_vverts, 3) — slice the per-env data for the visual section self._deformable_mpm_visual_verts_view.copy_( mpm_vverts[env_index].to(torch.float32) ) # Particle-recon: CPU reconstruction per entity per env, then upload + per-entity active # count. MPM-recon and SPH-recon share the same code path; the slot carries the owning # solver so the meshing call resolves the correct particle radius. for slot in self._mpm_recon_slots: self._update_recon_slot(slot, env_index) for slot in self._sph_recon_slots: self._update_recon_slot(slot, env_index) def _update_recon_slot(self, slot: _ReconSlot, env_index: int): """Reconstruct a particle-fluid surface and stream it into the shared deformable buffer.""" from genesis.utils.particle import particles_to_mesh import numpy as np # Get particle positions for this env (numpy) poss = slot.entity.get_particles_pos(envs_idx=env_index) if hasattr(poss, "cpu"): poss = poss.cpu().numpy() poss = np.ascontiguousarray(poss.reshape(-1, 3), dtype=np.float32) backend = getattr(slot.entity.surface, "recon_backend", "openvdb") radius = _particle_solver_radius(slot.solver) mesh = particles_to_mesh(poss, radius, backend) n_v = int(mesh.vertices.shape[0]) n_t = int(mesh.faces.shape[0]) # Clamp to the pre-allocated budget; warn (once) if exceeded so the user can resize. if n_v > slot.max_vertices or n_t > slot.max_triangles: gs.logger.warning( f"Recon reconstruction exceeded budget ({n_v}v/{n_t}t > " f"{slot.max_vertices}v/{slot.max_triangles}t); clamping. Coarsen recon voxel scale " "or tighten the solver domain." ) n_v = min(n_v, slot.max_vertices) n_t = min(n_t, slot.max_triangles) if n_v > 0: verts = torch.as_tensor( mesh.vertices[:n_v], dtype=torch.float32, device="cuda" ) self._deformable_verts_cuda[slot.vert_start : slot.vert_start + n_v].copy_( verts ) if n_t > 0: # Remap face indices into the shared vertex space (faces are local to the recon mesh) faces = torch.as_tensor( mesh.faces[:n_t], dtype=torch.int32, device="cuda" ).reshape(-1) faces = faces + slot.vert_start self._deformable_indices_cuda[ slot.tri_start * 3 : (slot.tri_start + n_t) * 3 ].copy_(faces) # Push the live triangle count to Nyx so the shader knows where to stop reading indices # and starts zeroing the tail to keep the BLAS clean. self._renderer.set_deform_entity_active_triangles(slot.deform_idx, n_t)
[docs] def update_camera_tensor(self, cam_index, pos, lookat, up): """Write one camera's pos + quaternion directly into the CUDA tensors. Called by NyxCameraSensor per-env before update_scene(). Parameters ---------- cam_index : int Camera index in the shared tensor. pos, lookat, up : array-like (3,) Camera pose in Z-up coordinates. """ cam_pos_t = torch.as_tensor(pos, dtype=torch.float32) cam_lookat_t = torch.as_tensor(lookat, dtype=torch.float32) cam_up_t = torch.as_tensor(up, dtype=torch.float32) transform_mat = gu.pos_lookat_up_to_T(cam_pos_t, cam_lookat_t, cam_up_t) cam_quat_zup_wxyz = gu.R_to_quat(transform_mat[..., :3, :3]) self._cam_pos_tensor_cuda[cam_index].copy_(cam_pos_t) self._cam_rot_tensor_cuda[cam_index].copy_(cam_quat_zup_wxyz)
def _update_camera_tensors(self, env_index): """Update camera tensors from Genesis Camera objects (Z-up, WXYZ format). Only used when no Nyx sensor is registered (camera_defs is None). When Nyx sensors are active, they call update_camera_tensor() directly. """ if self._cam_pos_tensor_cuda is None: return # Nyx sensors manage their own cameras via update_camera_tensor() if self._camera_defs is not None: return # Genesis Camera objects — select data for the target env for cam_idx, camera in enumerate(self._cameras): cam_pos = camera.get_pos(envs_idx=env_index) cam_quat = camera.get_quat(envs_idx=env_index) # WXYZ format self._cam_pos_tensor_cuda[cam_idx].copy_(cam_pos.to(torch.float32)) self._cam_rot_tensor_cuda[cam_idx].copy_(cam_quat.to(torch.float32))
[docs] def update_scene(self, env_index): """Sync one Genesis environment's state to the Nyx GPU buffers. Per-frame entry point. Pulls the current rigid-body transforms, deformable vertex positions, and (when no Nyx sensor is registered) Genesis camera poses for the selected env, copies them into the pre-allocated CUDA tensors, hands their pointers to Nyx via :class:`~gs_nyx.nyx_py_renderer.BridgeUpdateData`, and selects the environment map slot. MPM-recon entities also get their per-frame surface reconstruction run here: particles are pulled to the CPU, meshed via :func:`genesis.utils.particle.particles_to_mesh`, and the resulting vertices + remapped indices are streamed into the shared deformable buffer along with a per-entity active-triangle count. Parameters ---------- env_index : int Which parallel Genesis environment to render. Must be in ``[0, n_envs)``. Notes ----- Z-up (Genesis) to Y-up (Nyx) conversion for per-frame transforms is deferred to the Nyx GPU shaders; the tensors uploaded here remain in Genesis-native coordinates with WXYZ quaternions. """ # Update camera tensors with current Genesis state (raw Z-up data) self._update_camera_tensors(env_index) # Update all geometry buffer pointers self._update_geometry_tensors(env_index) # Sync the updated tensors to Nyx (GPU shader handles Z-up to Y-up conversion) self._renderer.update_scene_buffers(self._update_data) # Set the env map we need to set self._renderer.set_env_map(env_index)
[docs] def update(self): """Drive the native renderer's per-frame tick (window + animation clock). Call once per simulation step, before the per-env render loop. When :attr:`open_window` is ``True`` this pumps window events; otherwise it advances the renderer's internal animation timing. """ self._renderer.update()
[docs] def render(self, camera_index=0): """Render one camera and return the resulting RGB image as a CUDA tensor. Synchronizes the CUDA stream first so any pending Genesis kernel work on the input tensors has completed before the renderer reads them; without this the renderer can race with the simulation and sample stale geometry. Parameters ---------- camera_index : int, optional Index into the registered camera list (matches the order in which :class:`NyxCameraSensor` instances were built, or the Genesis visualizer's non-debug camera order when no sensors are registered). Default 0. Returns ------- torch.Tensor ``uint8`` image of shape ``(H, W, 3)`` on the CUDA device, where ``(H, W)`` matches the camera's configured resolution. The buffer is owned by the renderer — copy it if you need to keep the data past the next :meth:`render` call. """ # Ensure all prior CUDA work is finished (necessary if all the input data is on the GPU) torch.cuda.synchronize() color_tensor = self._renderer.render_frame(camera_index) return color_tensor
[docs] def pick_pixel( self, pair_list, camera_index: int, x: int, y: int ) -> Optional[ tuple[gs.engine.entities.base_entity.Entity, str, tuple[float, float, float]] ]: """Cast a ray through one pixel and return the entity and hit point. Wraps the native :meth:`gs_nyx.nyx_py_renderer.NyxRenderer.pick_pixel` and translates Nyx's UUID + Y-up hit location back to Genesis world space. Parameters ---------- pair_list : list of tuple ``(entity, uuid)`` pairs from :class:`NyxSceneExporter`. Used to convert the picked UUID back into a Genesis entity. camera_index : int Camera to cast through. Matches the index used by :meth:`render`. x, y : int Pixel coordinates with the origin at the top-left of the camera's framebuffer. Returns ------- tuple or None ``None`` if the ray missed all geometry. Otherwise a 3-tuple ``(entity, node_name, position)``: - ``entity`` — the Genesis entity that was hit - ``node_name`` — link / subgeometry name (URDF: link name; MJCF: ``"<link>_<vgeom_idx>"``; otherwise empty) - ``position`` — world-space hit point as ``(x, y, z)`` in Genesis Z-up coordinates """ # Try to intersect with the scene. ray_result = nps.float3(0.0, 0.0, 0.0) scene_ref = npr.SceneReference() trace_result = self._renderer.pick_pixel( camera_index, nps.uint2(x, y), ray_result, scene_ref ) # If we actually did hit something, evaluate the intersection data if trace_result: # Convert the return uuid to a Genesis entity and return the entity, node name, and hit location entity = uuid_to_entity(pair_list, scene_ref.targetUUID) # Nyx is Y-up, Genesis is Z-up return ( entity, scene_ref.nodeName, nps.float3_y_up_a_to_z_up(ray_result), ) # If we didn't hit anything, return None return None
[docs] def unload_scene(self): """Release the loaded scene without shutting down the renderer. The native renderer remains alive; another scene can be loaded by calling :meth:`build` again with a different scene-description path. Most users should call :meth:`destroy` instead. """ self._renderer.unload_scene()
[docs] def destroy(self): """Shut down the native renderer and release the interop buffers. After this call the wrapper is no longer usable: calling any of the per-frame methods will fail. Safe to call from a sensor's ``destroy()`` hook. """ del self._update_data self._update_data = None self._renderer.shutdown() del self._renderer self._renderer = None