diff --git a/docs/source/api/lab_newton/isaaclab_newton.assets.rst b/docs/source/api/lab_newton/isaaclab_newton.assets.rst index f044f9b089fb..160f6fcfe32b 100644 --- a/docs/source/api/lab_newton/isaaclab_newton.assets.rst +++ b/docs/source/api/lab_newton/isaaclab_newton.assets.rst @@ -15,6 +15,9 @@ RigidObjectCollectionData DeformableObject DeformableObjectData + MPMObject + MPMObjectCfg + MPMObjectData .. currentmodule:: isaaclab_newton.assets @@ -81,3 +84,22 @@ Deformable Object :inherited-members: :show-inheritance: :exclude-members: __init__ + +MPM Object +---------- + +.. autoclass:: MPMObject + :members: + :inherited-members: + :show-inheritance: + +.. autoclass:: MPMObjectCfg + :members: + :show-inheritance: + :exclude-members: __init__ + +.. autoclass:: MPMObjectData + :members: + :inherited-members: + :show-inheritance: + :exclude-members: __init__ diff --git a/docs/source/api/lab_newton/isaaclab_newton.physics.rst b/docs/source/api/lab_newton/isaaclab_newton.physics.rst index b452f3ad538b..8570e55f8806 100644 --- a/docs/source/api/lab_newton/isaaclab_newton.physics.rst +++ b/docs/source/api/lab_newton/isaaclab_newton.physics.rst @@ -14,6 +14,7 @@ XPBDSolverCfg FeatherstoneSolverCfg KaminoSolverCfg + MPMSolverCfg NewtonCollisionPipelineCfg HydroelasticSDFCfg NewtonShapeCfg @@ -21,6 +22,7 @@ NewtonXPBDManager NewtonFeatherstoneManager NewtonKaminoManager + NewtonMPMManager .. currentmodule:: isaaclab_newton.physics @@ -65,6 +67,11 @@ Physics Configuration :show-inheritance: :exclude-members: __init__ +.. autoclass:: MPMSolverCfg + :members: + :show-inheritance: + :exclude-members: __init__ + .. autoclass:: NewtonCollisionPipelineCfg :members: :show-inheritance: @@ -102,3 +109,8 @@ Solver Managers :members: :inherited-members: :show-inheritance: + +.. autoclass:: NewtonMPMManager + :members: + :inherited-members: + :show-inheritance: diff --git a/docs/source/api/lab_newton/isaaclab_newton.sim.spawners.rst b/docs/source/api/lab_newton/isaaclab_newton.sim.spawners.rst index d60b953ce5c8..f42cef3d8167 100644 --- a/docs/source/api/lab_newton/isaaclab_newton.sim.spawners.rst +++ b/docs/source/api/lab_newton/isaaclab_newton.sim.spawners.rst @@ -31,3 +31,41 @@ Newton provides the backend-specific deformable material cfgs. Deformable materi :members: :show-inheritance: :exclude-members: __init__, func + +.. automodule:: isaaclab_newton.sim.spawners.mpm + + .. rubric:: Classes + + .. autosummary:: + + MPMParticleSpawnerCfg + MPMGridCfg + MPMPointsCfg + MPMParticleMaterialCfg + +MPM Particles +------------- + +Declarative particle generation for :class:`~isaaclab_newton.assets.MPMObject`. +The spawner creates a placeholder ``Xform`` prim; the particles themselves are +emitted into the Newton model builder during replication. + +.. autoclass:: MPMParticleSpawnerCfg + :members: + :show-inheritance: + :exclude-members: __init__, func + +.. autoclass:: MPMGridCfg + :members: + :show-inheritance: + :exclude-members: __init__, func + +.. autoclass:: MPMPointsCfg + :members: + :show-inheritance: + :exclude-members: __init__, func + +.. autoclass:: MPMParticleMaterialCfg + :members: + :show-inheritance: + :exclude-members: __init__ diff --git a/docs/source/overview/core-concepts/physical-backends/newton/newton-manager-abstraction.rst b/docs/source/overview/core-concepts/physical-backends/newton/newton-manager-abstraction.rst index 39909caa27c3..381a98f08a91 100644 --- a/docs/source/overview/core-concepts/physical-backends/newton/newton-manager-abstraction.rst +++ b/docs/source/overview/core-concepts/physical-backends/newton/newton-manager-abstraction.rst @@ -94,6 +94,22 @@ solver actually needs it: notification before delegating to the base manager. * ``start_simulation()`` or ``instantiate_builder_from_stage()``: customize model building or post-finalize setup. +* ``_register_builder_attributes(builder)``: register solver-specific Newton + custom attributes (particle, shape, body) on the builder before particles or + finalize run. The active manager class invokes this hook from + ``create_builder()``, ``start_simulation()``, and + ``instantiate_builder_from_stage()``. + :class:`~isaaclab_newton.physics.NewtonMPMManager` is the in-tree example — + it registers ``mpm:young_modulus`` and the rest of the implicit MPM + particle attributes. +* ``_prepare_builder_for_finalize(builder)``: normalize imported or replicated + builder data right before ``ModelBuilder.finalize()``. + :class:`~isaaclab_newton.physics.NewtonMPMManager` uses this to clear mass and + inertia on kinematic bodies so implicit MPM treats them as massless colliders. +* ``_supports_cuda_graph_capture()``: return ``False`` to opt the solver out of + CUDA graph capture and fall back to eager execution. Defaults to ``True``; + :class:`~isaaclab_newton.physics.NewtonMPMManager` returns ``True`` only for a + fixed grid, since sparse/dense MPM grids reallocate as particles move. * ``_solver_specific_clear()``: release any class-level state owned by the solver manager. diff --git a/scripts/demos/mpm/newton_mpm_granular.py b/scripts/demos/mpm/newton_mpm_granular.py new file mode 100644 index 000000000000..fdb7b8208256 --- /dev/null +++ b/scripts/demos/mpm/newton_mpm_granular.py @@ -0,0 +1,224 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Isaac Lab port of Newton's granular implicit-MPM example. + +A block of granular material is dropped onto static box colliders. The demo +shows the intended Isaac Lab MPM path: configure +:class:`~isaaclab_newton.physics.MPMSolverCfg`, add an +:class:`~isaaclab_newton.assets.mpm_object.MPMObject` to an +:class:`~isaaclab.scene.InteractiveScene`, and let +:class:`~isaaclab_newton.physics.NewtonMPMManager` own solver creation and +stepping. + +.. code-block:: bash + + ./isaaclab.sh -p scripts/demos/mpm/newton_mpm_granular.py --visualizer newton +""" + +from __future__ import annotations + +import argparse +import math + +from isaaclab.app import add_launcher_args, launch_simulation + +parser = argparse.ArgumentParser(description="Newton implicit MPM granular demo.") +parser.add_argument("--max-steps", type=int, default=-1, help="Stop after this many frames; negative runs forever.") +parser.add_argument("--collider", default="cube", choices=["cube", "wedge", "concave", "none"], help="Collider set.") +parser.add_argument("--voxel-size", type=float, default=0.1, help="MPM grid voxel size [m].") +parser.add_argument("--substeps", type=int, default=1, help="Solver substeps per frame.") +add_launcher_args(parser) +parser.set_defaults(visualizer=["newton"]) +args_cli = parser.parse_args() + + +FPS = 100.0 +GRAVITY = (0.0, 0.0, -10.0) +VOXEL_SIZE = float(args_cli.voxel_size) + +# Solver settings that differ from the MPMSolverCfg defaults. A fixed grid has a +# static topology, so the MPM step can be captured as a CUDA graph +# (``NewtonMPMManager._supports_cuda_graph_capture``). ``grid_padding`` sizes the +# frozen grid to contain the block as it falls and spreads over the colliders, +# and ``max_active_cell_count`` bounds the per-step active set so the +# graph-captured allocations stay static. +GRID_TYPE = "fixed" +GRID_PADDING = 48 +MAX_ACTIVE_CELL_COUNT = 1 << 18 + +# Granular block, emitted as a jittered particle grid. +EMIT_LO = (-1.0, -1.0, 2.0) +EMIT_HI = (1.0, 1.0, 3.5) +PARTICLES_PER_CELL = 3.0 +PARTICLE_JITTER = VOXEL_SIZE / PARTICLES_PER_CELL +NEWTON_VISUAL_UPDATE_FREQUENCY = 1 +KIT_PARTICLE_VISUAL_UPDATE_FREQUENCY = 4 + +PARTICLE_COLOR = (0.7, 0.6, 0.4) + +IDENTITY_ROT = (0.0, 0.0, 0.0, 1.0) +Y_ROT_45_DEG = (0.0, math.sin(math.pi / 8.0), 0.0, math.cos(math.pi / 8.0)) +Y_ROT_NEG_45_DEG = (0.0, -math.sin(math.pi / 8.0), 0.0, math.cos(math.pi / 8.0)) + + +def create_visualizer_cfgs(): + """Create demo-specific visualizer configs for the requested backends.""" + if "newton" not in args_cli.visualizer: + return [] + + from isaaclab_visualizers.newton import NewtonVisualizerCfg + + return [ + NewtonVisualizerCfg( + show_particles=True, + particle_color=PARTICLE_COLOR, + update_frequency=NEWTON_VISUAL_UPDATE_FREQUENCY, + ) + ] + + +def create_sim_cfg(): + """Create the Isaac Lab simulation config that the MPM manager drives.""" + from isaaclab_newton.physics import MPMSolverCfg, NewtonCfg + + import isaaclab.sim as sim_utils + + return sim_utils.SimulationCfg( + dt=1.0 / FPS, + device=args_cli.device, + gravity=GRAVITY, + visualizer_cfgs=create_visualizer_cfgs(), + physics=NewtonCfg( + solver_cfg=MPMSolverCfg( + voxel_size=VOXEL_SIZE, + grid_type=GRID_TYPE, + grid_padding=GRID_PADDING, + max_active_cell_count=MAX_ACTIVE_CELL_COUNT, + project_outside_colliders=True, + ), + num_substeps=args_cli.substeps, + ), + ) + + +def preview_material(color): + """Return a preview-surface material for Kit runs; Kit-less runs spawn no USD materials.""" + if "kit" not in args_cli.visualizer: + return None + + import isaaclab.sim as sim_utils + + return sim_utils.PreviewSurfaceCfg(diffuse_color=color) + + +def create_scene_cfg(): + """Create an Isaac Lab scene config using declarative assets.""" + from isaaclab_newton.assets.mpm_object import MPMObjectCfg + from isaaclab_newton.sim.spawners.mpm import MPMGridCfg + + import isaaclab.sim as sim_utils + from isaaclab.assets import AssetBaseCfg + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.utils.configclass import configclass + + def collider_cfg(prim_path: str, center, half_extents, orientation, friction: float = 0.1) -> AssetBaseCfg: + return AssetBaseCfg( + prim_path=prim_path, + spawn=sim_utils.CuboidCfg( + size=(2.0 * half_extents[0], 2.0 * half_extents[1], 2.0 * half_extents[2]), + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + physics_material=sim_utils.NewtonMaterialPropertiesCfg( + static_friction=friction, + dynamic_friction=friction, + ), + visual_material=preview_material((0.45, 0.45, 0.45)), + ), + init_state=AssetBaseCfg.InitialStateCfg(pos=center, rot=orientation), + ) + + @configclass + class GranularSceneCfg(InteractiveSceneCfg): + """Scene containing static colliders and one Newton MPM object.""" + + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.GroundPlaneCfg(size=(6.0, 6.0), color=(0.30, 0.30, 0.30)), + ) + + dome_light = AssetBaseCfg( + prim_path="/World/DomeLight", + spawn=sim_utils.DomeLightCfg(intensity=2500.0, color=(0.78, 0.78, 0.78)), + ) + + media = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/GranularMedia", + spawn=MPMGridCfg( + lower=EMIT_LO, + upper=EMIT_HI, + voxel_size=VOXEL_SIZE, + particles_per_cell=PARTICLES_PER_CELL, + jitter=PARTICLE_JITTER, + visual_color=PARTICLE_COLOR, + visual_update_frequency=KIT_PARTICLE_VISUAL_UPDATE_FREQUENCY, + ), + ) + + if args_cli.collider == "cube": + cube = collider_cfg("/World/Colliders/Cube", (0.75, 0.0, 0.8), (0.5, 2.0, 0.8), IDENTITY_ROT) + elif args_cli.collider == "wedge": + wedge = collider_cfg("/World/Colliders/Wedge", (0.1, 0.0, 0.5), (0.5, 2.0, 0.5), Y_ROT_45_DEG) + elif args_cli.collider == "concave": + left_ramp = collider_cfg("/World/Colliders/LeftRamp", (-0.7, 0.0, 0.8), (1.0, 2.0, 0.25), Y_ROT_45_DEG) + right_ramp = collider_cfg("/World/Colliders/RightRamp", (0.7, 0.0, 0.8), (1.0, 2.0, 0.25), Y_ROT_NEG_45_DEG) + + return GranularSceneCfg(num_envs=1, env_spacing=0.0) + + +def particle_count(scene) -> int: + """Return the number of MPM particles in the scene.""" + media = scene["media"] + return media.num_instances * media.particles_per_object + + +def keep_running(sim, count: int) -> bool: + """Return whether the demo loop should continue this frame.""" + if args_cli.max_steps >= 0 and count >= args_cli.max_steps: + return False + return sim.is_headless_or_exist_active_visualizer() + + +def run_simulator(sim, scene) -> None: + """Run the MPM simulation using Isaac Lab's scene and physics-manager loop.""" + sim_dt = sim.get_physics_dt() + count = 0 + while keep_running(sim, count): + sim.step(render=False) + scene.update(sim_dt) + if sim.is_rendering: + sim.render() + count += 1 + + +def main() -> None: + """Launch and run the Isaac Lab MPM granular demo.""" + sim_cfg = create_sim_cfg() + with launch_simulation(sim_cfg, args_cli): + import isaaclab.sim as sim_utils + from isaaclab.scene import InteractiveScene + + sim = sim_utils.SimulationContext(sim_cfg) + sim.set_camera_view(eye=(4.0, -6.0, 4.0), target=(0.0, 0.0, 1.0)) + scene = InteractiveScene(create_scene_cfg()) + sim.reset() + print( + f"[INFO]: Isaac Lab Newton granular MPM demo ready. Spawned {particle_count(scene)} particles.", + flush=True, + ) + run_simulator(sim, scene) + + +if __name__ == "__main__": + main() diff --git a/scripts/demos/mpm/particle_pour.py b/scripts/demos/mpm/particle_pour.py new file mode 100644 index 000000000000..b50f4eae05a1 --- /dev/null +++ b/scripts/demos/mpm/particle_pour.py @@ -0,0 +1,536 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton implicit MPM particle-pour demo. + +This demo shows the Isaac Lab MPM scene path: + +* configure :class:`~isaaclab_newton.physics.MPMSolverCfg`; +* add fluid as an :class:`~isaaclab_newton.assets.mpm_object.MPMObject`; +* use standard Isaac Lab USD and mesh assets as MPM colliders; +* drive the pouring container through the standard rigid-object API. + +.. code-block:: bash + + ./isaaclab.sh -p scripts/demos/mpm/particle_pour.py --device cuda:0 --visualizer newton +""" + +from __future__ import annotations + +import argparse +import math +from collections.abc import Callable +from dataclasses import MISSING + +import numpy as np +import torch + +from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +DEFAULT_VOXEL_SIZE = 0.003 + +parser = argparse.ArgumentParser(description="Newton implicit MPM particle-pour demo.") +parser.add_argument("--max-steps", type=int, default=3000, help="Stop after this many frames; negative runs forever.") +parser.add_argument( + "--voxel-size", + type=float, + default=DEFAULT_VOXEL_SIZE, + help=f"MPM grid voxel size in meters. Defaults to {DEFAULT_VOXEL_SIZE:g}.", +) +parser.add_argument( + "--container-usd", + type=str, + default=f"{ISAAC_NUCLEUS_DIR}/Props/Teapot/utah_teapot.usdc", + help="USD asset used as the pouring container (kinematic mesh collider).", +) +add_launcher_args(parser) +parser.set_defaults(visualizer=["newton"]) +args_cli = parser.parse_args() + + +FPS = 200 +VOXEL_SIZE = float(args_cli.voxel_size) +NEWTON_VISUAL_UPDATE_FREQUENCY = 1 +KIT_PARTICLE_VISUAL_UPDATE_FREQUENCY = 4 + +GRID_TYPE = "fixed" +GRID_PADDING = 64 +MAX_ACTIVE_CELL_COUNT = 1 << 17 +# With a fixed grid the solver loop is captured in a CUDA graph, so the rheology +# solve always runs exactly ``max_iterations`` (the convergence tolerance cannot +# trigger an early exit inside the graph). 100 iterations measures ~1.6x faster +# than the 250 default with an end state identical to within noise. +MPM_MAX_ITERATIONS = 100 + +PARTICLES_PER_CELL = 2.0 +PARTICLE_DENSITY = 1000.0 + +PIPE_EMITTER_CENTER_XY = (0.0, 0.0) +PIPE_EMITTER_RADIUS = 0.035 +PIPE_EMITTER_Z_RANGE = (0.018, 0.064) +PIPE_EMITTER_LO = ( + PIPE_EMITTER_CENTER_XY[0] - PIPE_EMITTER_RADIUS, + PIPE_EMITTER_CENTER_XY[1] - PIPE_EMITTER_RADIUS, + PIPE_EMITTER_Z_RANGE[0], +) +PIPE_EMITTER_HI = ( + PIPE_EMITTER_CENTER_XY[0] + PIPE_EMITTER_RADIUS, + PIPE_EMITTER_CENTER_XY[1] + PIPE_EMITTER_RADIUS, + PIPE_EMITTER_Z_RANGE[1], +) + +COLLIDER_MARGIN = 0.0003 +CONTAINER_MARGIN = 0.001 +CONTAINER_FRICTION = 0.0 +BOWL_MARGIN = 0.0025 +BOWL_FRICTION = 0.05 +TABLE_FRICTION = 0.5 + +HOLD_TIME = 0.55 +TILT_TIME = 2.0 +POUR_ANGLE = math.radians(65.0) +CONTAINER_LIFT_HEIGHT = 0.24 +CONTAINER_LIFT_TIME = 3.0 + +TABLE_TOP_Z = 0.255 +TABLE_HALF_EXTENTS = (0.255, 0.165, 0.009) +BOWL_BASE_POS = (0.066, 0.0, TABLE_TOP_Z + 0.006) +BOWL_HEIGHT = 0.039 +BOWL_RIM_Z = BOWL_BASE_POS[2] + BOWL_HEIGHT +CONTAINER_BASE_POS = (-0.105, 0.0, BOWL_RIM_Z + 0.115) + +BOWL_INNER_BOTTOM_RADIUS = 0.0135 +BOWL_INNER_TOP_RADIUS = 0.057 +BOWL_WALL_THICKNESS = 0.0075 +BOWL_BOTTOM_THICKNESS = 0.0075 +BOWL_COLOR = (1.0, 1.0, 1.0) +TABLE_COLOR = (0.48, 0.38, 0.26) +PARTICLE_COLOR = (0.12, 0.35, 0.78) + +CAMERA_EYE = (0.0, -0.36, 0.46) +CAMERA_TARGET = (-0.01, 0.0, 0.38) + + +def create_visualizer_cfgs(): + """Create demo-specific visualizer configs for requested backends.""" + if "newton" not in args_cli.visualizer: + return [] + + from isaaclab_visualizers.newton import NewtonVisualizerCfg + + return [ + NewtonVisualizerCfg( + show_particles=True, + particle_color=PARTICLE_COLOR, + update_frequency=NEWTON_VISUAL_UPDATE_FREQUENCY, + ) + ] + + +def quat_y(angle_rad: float) -> tuple[float, float, float, float]: + """Return an XYZW quaternion for a rotation about +Y.""" + half = 0.5 * angle_rad + return (0.0, math.sin(half), 0.0, math.cos(half)) + + +def smoothstep(value: float) -> float: + """Smoothly remap a 0..1 value for the scripted tilt.""" + t = max(0.0, min(1.0, value)) + return t * t * (3.0 - 2.0 * t) + + +def container_pose_at_time(sim_time: float): + """Return teapot ``(position, orientation, twist)`` for the scripted pour.""" + raw = (sim_time - HOLD_TIME) / TILT_TIME + clamped = max(0.0, min(1.0, raw)) + alpha = smoothstep(clamped) + alpha_dot = (6.0 * clamped * (1.0 - clamped)) / TILT_TIME if 0.0 < raw < 1.0 else 0.0 + + lift_raw = (sim_time - HOLD_TIME - TILT_TIME) / CONTAINER_LIFT_TIME + lift_alpha = max(0.0, min(1.0, lift_raw)) + lift_speed = CONTAINER_LIFT_HEIGHT / CONTAINER_LIFT_TIME if 0.0 < lift_raw < 1.0 else 0.0 + + angle = POUR_ANGLE * alpha + angular_speed = POUR_ANGLE * alpha_dot + pos = ( + CONTAINER_BASE_POS[0], + CONTAINER_BASE_POS[1], + CONTAINER_BASE_POS[2] + CONTAINER_LIFT_HEIGHT * lift_alpha, + ) + # Newton spatial vectors are (linear, angular). + twist = (0.0, 0.0, lift_speed, 0.0, angular_speed, 0.0) + return pos, quat_y(angle), twist + + +def launch_omniverse_asset_resolver(): + """Start Kit for Newton-only runs that need to resolve remote USD layers. + + The default teapot container is served from Nucleus over HTTPS, which plain + ``pxr`` cannot resolve; Kit ships the asset resolver that can. Kit-visualizer + runs boot Kit anyway, so this only applies to Newton-only runs. + """ + if "kit" in args_cli.visualizer: + return None + + from isaaclab.utils import has_kit + + if has_kit(): + return None + + from isaaclab.app import AppLauncher + + return AppLauncher(args_cli) + + +def create_demo_bowl_mesh(num_segments: int = 96): + """Build one local-space open bowl mesh used by the catch-bowl collider.""" + theta = np.linspace(0.0, 2.0 * math.pi, num_segments, endpoint=False) + cos_t = np.cos(theta) + sin_t = np.sin(theta) + outer_bottom_radius = BOWL_INNER_BOTTOM_RADIUS + BOWL_WALL_THICKNESS + outer_top_radius = BOWL_INNER_TOP_RADIUS + BOWL_WALL_THICKNESS + + def ring(radius: float, z: float): + return np.column_stack([radius * cos_t, radius * sin_t, np.full(num_segments, z)]) + + vertices = np.vstack( + [ + ring(BOWL_INNER_BOTTOM_RADIUS, BOWL_BOTTOM_THICKNESS), + ring(BOWL_INNER_TOP_RADIUS, BOWL_HEIGHT), + ring(outer_top_radius, BOWL_HEIGHT), + ring(outer_bottom_radius, 0.0), + np.array([[0.0, 0.0, BOWL_BOTTOM_THICKNESS], [0.0, 0.0, 0.0]], dtype=np.float32), + ] + ).astype(np.float32) + + inner_center_id = 4 * num_segments + outer_center_id = inner_center_id + 1 + indices: list[int] = [] + for i in range(num_segments): + j = (i + 1) % num_segments + ib_i, ib_j = i, j + it_i, it_j = i + num_segments, j + num_segments + ot_i, ot_j = i + 2 * num_segments, j + 2 * num_segments + ob_i, ob_j = i + 3 * num_segments, j + 3 * num_segments + indices.extend([ib_i, it_i, ib_j, ib_j, it_i, it_j]) + indices.extend([ob_i, ob_j, ot_i, ot_i, ob_j, ot_j]) + indices.extend([it_i, ot_i, it_j, it_j, ot_i, ot_j]) + indices.extend([inner_center_id, ib_i, ib_j, outer_center_id, ob_j, ob_i]) + + return vertices, np.asarray(indices, dtype=np.int32).reshape((-1, 3)) + + +def spawn_demo_bowl_mesh( + prim_path: str, + cfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn the demo bowl as a standard Isaac Lab mesh asset.""" + from isaaclab.sim import schemas + from isaaclab.sim.utils import bind_physics_material, bind_visual_material, create_prim, get_current_stage + + stage = get_current_stage() + vertices = np.asarray(cfg.vertices, dtype=np.float32) + faces = np.asarray(cfg.faces, dtype=np.int32) + + create_prim(prim_path, prim_type="Xform", translation=translation, orientation=orientation, stage=stage) + geom_prim_path = f"{prim_path}/geometry" + mesh_prim_path = f"{geom_prim_path}/mesh" + create_prim(geom_prim_path, prim_type="Xform", stage=stage) + create_prim( + mesh_prim_path, + prim_type="Mesh", + attributes={ + "points": vertices, + "faceVertexIndices": faces.reshape(-1), + "faceVertexCounts": np.full(faces.shape[0], 3, dtype=np.int32), + "subdivisionScheme": "bilinear", + }, + stage=stage, + ) + + if cfg.collision_props is not None: + schemas.define_collision_properties(mesh_prim_path, cfg.collision_props, stage=stage) + if cfg.mesh_collision_props is not None: + schemas.define_mesh_collision_properties(mesh_prim_path, cfg.mesh_collision_props, stage=stage) + + if cfg.visual_material is not None: + material_path = cfg.visual_material_path + if not material_path.startswith("/"): + material_path = f"{geom_prim_path}/{material_path}" + cfg.visual_material.func(material_path, cfg.visual_material) + bind_visual_material(mesh_prim_path, material_path, stage=stage) + + if cfg.physics_material is not None: + material_path = cfg.physics_material_path + if not material_path.startswith("/"): + material_path = f"{geom_prim_path}/{material_path}" + cfg.physics_material.func(material_path, cfg.physics_material) + bind_physics_material(mesh_prim_path, material_path, stage=stage) + + return stage.GetPrimAtPath(prim_path) + + +def create_fluid_particles(): + """Return local-space MPM particle points seeded inside the teapot.""" + center_xy = np.array(PIPE_EMITTER_CENTER_XY, dtype=np.float32) + particle_lo = np.asarray(PIPE_EMITTER_LO, dtype=np.float32) + particle_hi = np.asarray(PIPE_EMITTER_HI, dtype=np.float32) + resolution = np.maximum(np.ceil(PARTICLES_PER_CELL * (particle_hi - particle_lo) / VOXEL_SIZE), 1).astype(int) + cell_size = (particle_hi - particle_lo) / resolution + cell_volume = float(np.prod(cell_size)) + radius = float(np.max(cell_size) * 0.45) + mass = float(cell_volume * PARTICLE_DENSITY) + + px = np.arange(int(resolution[0]) + 1) * cell_size[0] + py = np.arange(int(resolution[1]) + 1) * cell_size[1] + pz = np.arange(int(resolution[2]) + 1) * cell_size[2] + points = np.stack(np.meshgrid(px, py, pz, indexing="ij")).reshape(3, -1).T + + rng = np.random.default_rng(7) + points += (rng.random(points.shape) - 0.5) * (0.10 * np.max(cell_size)) + points += particle_lo + + normalized_xy = (points[:, :2] - center_xy) / PIPE_EMITTER_RADIUS + points = points[np.sum(normalized_xy * normalized_xy, axis=1) < 1.0] + if points.shape[0] == 0: + raise RuntimeError("Particle initialization produced no particles; reduce --voxel-size.") + + return points.astype(np.float32, copy=False), radius, mass + + +def create_sim_cfg(): + """Create the Isaac Lab simulation config using the MPM manager.""" + from isaaclab_newton.physics import MPMSolverCfg, NewtonCfg + + import isaaclab.sim as sim_utils + + return sim_utils.SimulationCfg( + dt=1.0 / FPS, + device=args_cli.device, + gravity=(0.0, 0.0, -9.81), + visualizer_cfgs=create_visualizer_cfgs(), + physics=NewtonCfg( + solver_cfg=MPMSolverCfg( + voxel_size=VOXEL_SIZE, + grid_type=GRID_TYPE, + grid_padding=GRID_PADDING, + max_active_cell_count=MAX_ACTIVE_CELL_COUNT, + max_iterations=MPM_MAX_ITERATIONS, + air_drag=0.2, + collider_velocity_mode="backward", + project_outside_colliders=True, + ), + use_cuda_graph=True, + simplify_meshes=False, + ), + ) + + +def preview_material(color): + """Return a preview-surface material for Kit runs; Kit-less runs spawn no USD materials.""" + if "kit" not in args_cli.visualizer: + return None + + import isaaclab.sim as sim_utils + + return sim_utils.PreviewSurfaceCfg(diffuse_color=color) + + +def create_scene_cfg(): + """Create the particle-pour scene using declarative Isaac Lab assets.""" + from isaaclab_newton.assets.mpm_object import MPMObjectCfg + from isaaclab_newton.sim.spawners.mpm import MPMParticleMaterialCfg, MPMPointsCfg + + import isaaclab.sim as sim_utils + from isaaclab.assets import AssetBaseCfg, RigidObjectCfg + from isaaclab.scene import InteractiveSceneCfg + from isaaclab.sim.utils import clone + from isaaclab.utils.configclass import configclass + + container_pos, container_rot, _ = container_pose_at_time(0.0) + fluid_points, particle_radius, particle_mass = create_fluid_particles() + bowl_vertices, bowl_faces = create_demo_bowl_mesh() + + @configclass + class DemoBowlMeshCfg(sim_utils.MeshCfg): + """Demo-local arbitrary mesh asset config for the catch bowl.""" + + func: Callable | str = clone(spawn_demo_bowl_mesh) + vertices: list[list[float]] = MISSING + faces: list[list[int]] = MISSING + mesh_collision_props: sim_utils.NewtonMeshCollisionPropertiesCfg | None = None + + @configclass + class PourSceneCfg(InteractiveSceneCfg): + """Scene containing MPM colliders and one MPM fluid object.""" + + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + spawn=sim_utils.CuboidCfg( + size=(2.0 * TABLE_HALF_EXTENTS[0], 2.0 * TABLE_HALF_EXTENTS[1], 2.0 * TABLE_HALF_EXTENTS[2]), + collision_props=sim_utils.NewtonCollisionPropertiesCfg( + collision_enabled=True, + contact_margin=COLLIDER_MARGIN, + ), + physics_material=sim_utils.NewtonMaterialPropertiesCfg( + static_friction=TABLE_FRICTION, + dynamic_friction=TABLE_FRICTION, + ), + physics_material_path="physicsMaterial", + visual_material=preview_material(TABLE_COLOR), + visual_material_path="visualMaterial", + ), + init_state=AssetBaseCfg.InitialStateCfg( + pos=(0.0, 0.0, TABLE_TOP_Z - TABLE_HALF_EXTENTS[2]), + ), + ) + + catch_bowl = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/CatchBowl", + spawn=DemoBowlMeshCfg( + vertices=bowl_vertices.tolist(), + faces=bowl_faces.tolist(), + collision_props=sim_utils.NewtonCollisionPropertiesCfg( + collision_enabled=True, + contact_margin=BOWL_MARGIN, + ), + mesh_collision_props=sim_utils.NewtonMeshCollisionPropertiesCfg(mesh_approximation_name="none"), + physics_material=sim_utils.NewtonMaterialPropertiesCfg( + static_friction=BOWL_FRICTION, + dynamic_friction=BOWL_FRICTION, + ), + physics_material_path="physicsMaterial", + visual_material=preview_material(BOWL_COLOR), + visual_material_path="visualMaterial", + ), + init_state=AssetBaseCfg.InitialStateCfg(pos=BOWL_BASE_POS), + ) + + # Utah Teapot: free for any use; credit as the (Modified) Utah Teapot + # (Univ. of Utah); provided "as is", no warranty. + container = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/PourContainer", + spawn=sim_utils.UsdFileCfg( + usd_path=args_cli.container_usd, + rigid_props=sim_utils.NewtonRigidBodyPropertiesCfg( + rigid_body_enabled=True, + kinematic_enabled=True, + disable_gravity=True, + ), + collision_props=sim_utils.NewtonCollisionPropertiesCfg( + collision_enabled=True, + contact_margin=CONTAINER_MARGIN, + ), + physics_material=sim_utils.NewtonMaterialPropertiesCfg( + static_friction=CONTAINER_FRICTION, + dynamic_friction=CONTAINER_FRICTION, + ), + physics_material_path="physicsMaterial", + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=container_pos, rot=container_rot), + ) + + fluid = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Fluid", + spawn=MPMPointsCfg( + positions=fluid_points.tolist(), + mass=particle_mass, + radius=particle_radius, + material=MPMParticleMaterialCfg( + viscosity=0.1, + friction=0.0, + damping=0.02, + yield_pressure=1.0e15, + tensile_yield_ratio=5.0, + ), + visual_color=PARTICLE_COLOR, + visual_update_frequency=KIT_PARTICLE_VISUAL_UPDATE_FREQUENCY, + ), + init_state=MPMObjectCfg.InitialStateCfg(pos=container_pos), + ) + + ground = AssetBaseCfg( + prim_path="/World/Ground", + spawn=sim_utils.GroundPlaneCfg(size=(1.0, 1.0), color=(0.30, 0.30, 0.30)), + ) + + dome_light = AssetBaseCfg( + prim_path="/World/DomeLight", + spawn=sim_utils.DomeLightCfg(intensity=2500.0, color=(0.78, 0.78, 0.78)), + ) + + return PourSceneCfg(num_envs=1, env_spacing=0.0) + + +def particle_count(scene) -> int: + """Return the number of MPM particles in the scene.""" + fluid = scene["fluid"] + return fluid.num_instances * fluid.particles_per_object + + +def keep_running(sim, count: int) -> bool: + """Return whether the demo loop should continue this frame.""" + if args_cli.max_steps >= 0 and count >= args_cli.max_steps: + return False + return sim.is_headless_or_exist_active_visualizer() + + +def write_container_state(container, sim_time: float) -> None: + """Write the scripted container pose and velocity through the rigid-object API.""" + pos, quat, qd = container_pose_at_time(sim_time) + pose = torch.tensor([tuple(pos) + tuple(quat)], dtype=torch.float32, device=container.device) + velocity = torch.tensor([qd], dtype=torch.float32, device=container.device) + container.write_root_link_pose_to_sim_index(root_pose=pose) + container.write_root_link_velocity_to_sim_index(root_velocity=velocity) + + +def run_simulator(sim, scene) -> None: + """Run the scripted particle-pour MPM loop.""" + sim_dt = sim.get_physics_dt() + container = scene["container"] + count = 0 + while keep_running(sim, count): + write_container_state(container, count / FPS) + scene.write_data_to_sim() + sim.step(render=False) + scene.update(sim_dt) + if sim.is_rendering: + sim.render() + count += 1 + + +def main() -> None: + """Set up and run the Isaac Lab Newton MPM particle-pour demo.""" + app_launcher = launch_omniverse_asset_resolver() + try: + sim_cfg = create_sim_cfg() + with launch_simulation(sim_cfg, args_cli): + import isaaclab.sim as sim_utils + from isaaclab.scene import InteractiveScene + + sim = sim_utils.SimulationContext(sim_cfg) + scene = InteractiveScene(create_scene_cfg()) + sim.reset() + sim.set_camera_view(eye=CAMERA_EYE, target=CAMERA_TARGET) + + print( + "[INFO]: Isaac Lab Newton particle-pour MPM demo ready." + f" Spawned {particle_count(scene)} MPM particles;" + f" voxel size {VOXEL_SIZE:.4g} m;" + f" the teapot will tilt after {HOLD_TIME:.2f}s.", + flush=True, + ) + run_simulator(sim, scene) + finally: + if app_launcher is not None: + app_launcher.app.close() + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab_newton/changelog.d/max-newton-mpm-manager.minor.rst b/source/isaaclab_newton/changelog.d/max-newton-mpm-manager.minor.rst new file mode 100644 index 000000000000..4dbff686ca56 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/max-newton-mpm-manager.minor.rst @@ -0,0 +1,41 @@ +Added +^^^^^ + +* Added :class:`~isaaclab_newton.physics.MPMSolverCfg` and + :class:`~isaaclab_newton.physics.NewtonMPMManager` for Newton implicit MPM + simulations. +* Added :class:`~isaaclab_newton.assets.MPMObject` (with + :class:`~isaaclab_newton.assets.MPMObjectCfg` and + :class:`~isaaclab_newton.assets.MPMObjectData`) exposing Newton MPM particles + through the deformable-object interface, together with the declarative + particle spawner configs :class:`~isaaclab_newton.sim.MPMGridCfg`, + :class:`~isaaclab_newton.sim.MPMPointsCfg`, and + :class:`~isaaclab_newton.sim.MPMParticleMaterialCfg`. +* Added the :class:`~isaaclab_newton.physics.NewtonManager` subclass hooks + ``_register_builder_attributes`` (register a solver's Newton custom builder + attributes), ``_prepare_builder_for_finalize`` (normalize imported builder + data before finalization), and ``_supports_cuda_graph_capture`` (opt a solver + out of CUDA graph capture). +* Added :attr:`~isaaclab_newton.physics.MPMSolverCfg.project_outside_colliders` + (default ``False``): when set, + :class:`~isaaclab_newton.physics.NewtonMPMManager` runs + ``SolverImplicitMPM.project_outside`` after each substep to push particles out + of collider interiors. +* Added :attr:`~isaaclab_newton.physics.NewtonCfg.simplify_meshes` to control + whether Newton replication approximates mesh colliders with convex hulls. + Disable it for thin or hollow MPM colliders that need exact triangle meshes. +* Added ``visual_update_frequency`` to MPM particle spawner configs so Kit USD + point-cloud visualization can be throttled independently from physics. + +Changed +^^^^^^^ + +* :meth:`~isaaclab_newton.physics.NewtonManager.sync_particles_to_usd` now also + writes registered ``UsdGeom.Points`` prims (used for MPM particle clouds) in + addition to the existing Fabric mesh-points sync for deformable visuals. +* :meth:`~isaaclab_newton.physics.NewtonManager.create_builder` and the model + build path now invoke the active manager's solver-specific builder hooks so + MPM custom attributes (``mpm:young_modulus``, ...) are registered on the + builder before particles are added or the model is finalized. +* CUDA graph capture is skipped when the active solver reports it is + unsupported, so sparse/dense-grid MPM falls back to eager execution. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi index 5476609e4eac..bc824be54e5e 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/assets/__init__.pyi @@ -8,6 +8,9 @@ __all__ = [ "ArticulationData", "DeformableObject", "DeformableObjectData", + "MPMObject", + "MPMObjectCfg", + "MPMObjectData", "RigidObject", "RigidObjectData", "RigidObjectCollection", @@ -18,3 +21,4 @@ from .articulation import Articulation, ArticulationData from .rigid_object import RigidObject, RigidObjectData from .rigid_object_collection import RigidObjectCollection, RigidObjectCollectionData from .deformable_object import DeformableObject, DeformableObjectData +from .mpm_object import MPMObject, MPMObjectCfg, MPMObjectData diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.py new file mode 100644 index 000000000000..e14e0f6d52c5 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.pyi new file mode 100644 index 000000000000..7b029922024b --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/__init__.pyi @@ -0,0 +1,14 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "MPMObject", + "MPMObjectCfg", + "MPMObjectData", +] + +from .mpm_object import MPMObject +from .mpm_object_cfg import MPMObjectCfg +from .mpm_object_data import MPMObjectData diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/kernels.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/kernels.py new file mode 100644 index 000000000000..d8c9036412fd --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/kernels.py @@ -0,0 +1,107 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Warp kernels for Newton MPM object gather/scatter operations.""" + +import warp as wp + +vec6f = wp.types.vector(length=6, dtype=wp.float32) + + +@wp.kernel +def gather_particles_vec3f( + src: wp.array(dtype=wp.vec3f), + offsets: wp.array(dtype=wp.int32), + dst: wp.array2d(dtype=wp.vec3f), +): + i, j = wp.tid() + dst[i, j] = src[offsets[i] + j] + + +@wp.kernel +def scatter_particles_vec3f_index( + src: wp.array2d(dtype=wp.vec3f), + env_ids: wp.array(dtype=wp.int32), + offsets: wp.array(dtype=wp.int32), + full_data: bool, + dst: wp.array(dtype=wp.vec3f), +): + i, j = wp.tid() + env_id = env_ids[i] + if full_data: + dst[offsets[env_id] + j] = src[env_id, j] + else: + dst[offsets[env_id] + j] = src[i, j] + + +@wp.kernel +def scatter_particles_vec3f_mask( + src: wp.array2d(dtype=wp.vec3f), + env_mask: wp.array(dtype=wp.bool), + offsets: wp.array(dtype=wp.int32), + dst: wp.array(dtype=wp.vec3f), +): + i, j = wp.tid() + if env_mask[i]: + dst[offsets[i] + j] = src[i, j] + + +@wp.kernel +def scatter_particles_state_vec6f_index( + src: wp.array2d(dtype=vec6f), + env_ids: wp.array(dtype=wp.int32), + offsets: wp.array(dtype=wp.int32), + full_data: bool, + particle_q: wp.array(dtype=wp.vec3f), + particle_qd: wp.array(dtype=wp.vec3f), +): + i, j = wp.tid() + env_id = env_ids[i] + src_id = env_id if full_data else i + state = src[src_id, j] + flat_idx = offsets[env_id] + j + particle_q[flat_idx] = wp.vec3f(state[0], state[1], state[2]) + particle_qd[flat_idx] = wp.vec3f(state[3], state[4], state[5]) + + +@wp.kernel +def scatter_particles_state_vec6f_mask( + src: wp.array2d(dtype=vec6f), + env_mask: wp.array(dtype=wp.bool), + offsets: wp.array(dtype=wp.int32), + particle_q: wp.array(dtype=wp.vec3f), + particle_qd: wp.array(dtype=wp.vec3f), +): + i, j = wp.tid() + if env_mask[i]: + state = src[i, j] + flat_idx = offsets[i] + j + particle_q[flat_idx] = wp.vec3f(state[0], state[1], state[2]) + particle_qd[flat_idx] = wp.vec3f(state[3], state[4], state[5]) + + +@wp.kernel +def compute_particle_state_w( + particle_pos: wp.array2d(dtype=wp.vec3f), + particle_vel: wp.array2d(dtype=wp.vec3f), + particle_state: wp.array2d(dtype=vec6f), +): + i, j = wp.tid() + p = particle_pos[i, j] + v = particle_vel[i, j] + particle_state[i, j] = vec6f(p[0], p[1], p[2], v[0], v[1], v[2]) + + +@wp.kernel +def compute_mean_vec3f_over_particles( + data: wp.array2d(dtype=wp.vec3f), + num_particles: int, + result: wp.array(dtype=wp.vec3f), +): + i = wp.tid() + acc = wp.vec3f(0.0, 0.0, 0.0) + for j in range(num_particles): + acc = acc + data[i, j] + result[i] = acc / float(num_particles) diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py new file mode 100644 index 000000000000..c5007f69cfc1 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py @@ -0,0 +1,444 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import numpy as np +import torch +import warp as wp + +from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject +from isaaclab.physics import PhysicsEvent +from isaaclab.utils.warp import ProxyArray + +from isaaclab_newton.cloner import queue_newton_physics_replication +from isaaclab_newton.physics import NewtonManager as SimulationManager +from isaaclab_newton.sim.spawners.mpm import create_mpm_particle_visualization, emit_mpm_particles + +from .kernels import ( + compute_particle_state_w, + gather_particles_vec3f, + scatter_particles_state_vec6f_index, + scatter_particles_state_vec6f_mask, + scatter_particles_vec3f_index, + scatter_particles_vec3f_mask, + vec6f, +) +from .mpm_object_data import MPMObjectData + +if TYPE_CHECKING: + from .mpm_object_cfg import MPMObjectCfg + +logger = logging.getLogger(__name__) + + +@dataclass +class MPMObjectRegistryEntry: + """Particle object registration consumed by Newton builder replication.""" + + cfg: MPMObjectCfg + particle_offsets: list[int] = field(default_factory=list) + particles_per_object: int = 0 + + +def add_mpm_entry_to_builder( + builder, + entry: MPMObjectRegistryEntry, + env_idx: int, + env_position: list[float], + env_rotation: list[float] | tuple[float, float, float, float], +) -> None: + """Emit one registered MPM object into one Newton builder world.""" + if env_idx == 0: + entry.particle_offsets.clear() + entry.particles_per_object = 0 + + before_count = builder.particle_count + position, orientation = _compose_env_asset_pose(entry.cfg, env_position, env_rotation) + emit_mpm_particles(builder, entry.cfg.spawn, position=position, orientation=orientation) + delta = builder.particle_count - before_count + + entry.particle_offsets.append(before_count) + if env_idx == 0: + entry.particles_per_object = delta + elif entry.particles_per_object != delta: + raise RuntimeError( + f"MPM object '{entry.cfg.prim_path}' produced {delta} particles in env {env_idx}, " + f"but env 0 produced {entry.particles_per_object}." + ) + + +def add_registered_mpm_objects_to_builder( + builder, + world_idx: int, + env_position: list[float], + env_rotation: list[float] | tuple[float, float, float, float], +) -> None: + """Emit all registered MPM objects into one Newton builder world.""" + for entry in SimulationManager._mpm_object_registry: + add_mpm_entry_to_builder(builder, entry, world_idx, env_position, env_rotation) + + +class MPMObject(BaseDeformableObject): + """Newton MPM particle object asset. + + The object is presented through Isaac Lab's deformable-object interface so it + can participate in existing scene reset/update/state workflows while exposing + particle-specific aliases on :attr:`data`. + """ + + cfg: MPMObjectCfg + __backend_name__: str = "newton" + + _DTYPE_TO_TORCH_TRAILING_DIMS = {**BaseDeformableObject._DTYPE_TO_TORCH_TRAILING_DIMS, vec6f: (6,)} + + def __init__(self, cfg: MPMObjectCfg): + super().__init__(cfg) + queue_newton_physics_replication(cfg) + self._registry_entry = MPMObjectRegistryEntry(self.cfg) + SimulationManager._mpm_object_registry.append(self._registry_entry) + self._physics_ready_handle = None + + @property + def data(self) -> MPMObjectData: + return self._data + + @property + def num_instances(self) -> int: + return self._num_instances + + @property + def num_bodies(self) -> int: + return 1 + + @property + def max_sim_vertices_per_body(self) -> int: + return self._particles_per_object + + @property + def particles_per_object(self) -> int: + """Number of particles generated for each environment instance.""" + return self._particles_per_object + + def reset(self, env_ids: Sequence[int] | None = None, env_mask: wp.array | None = None) -> None: + """Reset selected particle instances to their default particle state.""" + if env_mask is not None: + self.write_nodal_state_to_sim_mask(self.data.default_nodal_state_w.warp, env_mask=env_mask) + else: + self.write_nodal_state_to_sim_index(self.data.default_nodal_state_w.warp, env_ids=env_ids, full_data=True) + + def write_data_to_sim(self): + """No-op; MPM particle writes are applied immediately by write methods.""" + + def update(self, dt: float): + self._data.update(dt) + + def write_nodal_state_to_sim_index( + self, + nodal_state: torch.Tensor | wp.array | ProxyArray, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + self._scatter_to_sim_index( + nodal_state, + env_ids, + full_data, + vec6f, + scatter_particles_state_vec6f_index, + ("particle_q", "particle_qd"), + "nodal_state", + ) + self._invalidate_caches(pos=True, vel=True) + + def write_nodal_pos_to_sim_index( + self, + nodal_pos: torch.Tensor | wp.array | ProxyArray, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + self._scatter_to_sim_index( + nodal_pos, env_ids, full_data, wp.vec3f, scatter_particles_vec3f_index, ("particle_q",), "nodal_pos" + ) + self._invalidate_caches(pos=True) + + def write_nodal_velocity_to_sim_index( + self, + nodal_vel: torch.Tensor | wp.array | ProxyArray, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + self._scatter_to_sim_index( + nodal_vel, env_ids, full_data, wp.vec3f, scatter_particles_vec3f_index, ("particle_qd",), "nodal_vel" + ) + self._invalidate_caches(vel=True) + + def write_nodal_kinematic_target_to_sim_index( + self, + targets: torch.Tensor | wp.array | ProxyArray, + env_ids: Sequence[int] | torch.Tensor | wp.array | None = None, + full_data: bool = False, + ) -> None: + raise NotImplementedError("MPMObject does not support deformable kinematic targets.") + + def write_nodal_state_to_sim_mask( + self, + nodal_state: torch.Tensor | wp.array | ProxyArray, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + self._scatter_to_sim_mask( + nodal_state, + env_mask, + vec6f, + scatter_particles_state_vec6f_mask, + ("particle_q", "particle_qd"), + "nodal_state", + ) + self._invalidate_caches(pos=True, vel=True) + + def write_nodal_pos_to_sim_mask( + self, + nodal_pos: torch.Tensor | wp.array | ProxyArray, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + self._scatter_to_sim_mask( + nodal_pos, env_mask, wp.vec3f, scatter_particles_vec3f_mask, ("particle_q",), "nodal_pos" + ) + self._invalidate_caches(pos=True) + + def write_nodal_velocity_to_sim_mask( + self, + nodal_vel: torch.Tensor | wp.array | ProxyArray, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + self._scatter_to_sim_mask( + nodal_vel, env_mask, wp.vec3f, scatter_particles_vec3f_mask, ("particle_qd",), "nodal_vel" + ) + self._invalidate_caches(vel=True) + + def write_nodal_kinematic_target_to_sim_mask( + self, + targets: torch.Tensor | wp.array | ProxyArray, + env_mask: wp.array | torch.Tensor | None = None, + ) -> None: + raise NotImplementedError("MPMObject does not support deformable kinematic targets.") + + write_particle_state_to_sim_index = write_nodal_state_to_sim_index + write_particle_pos_to_sim_index = write_nodal_pos_to_sim_index + write_particle_velocity_to_sim_index = write_nodal_velocity_to_sim_index + write_particle_state_to_sim_mask = write_nodal_state_to_sim_mask + write_particle_pos_to_sim_mask = write_nodal_pos_to_sim_mask + write_particle_velocity_to_sim_mask = write_nodal_velocity_to_sim_mask + + def _scatter_to_sim_index(self, data, env_ids, full_data: bool, dtype, kernel, targets, name: str) -> None: + """Scatter per-environment particle data into the Newton state arrays in ``targets``.""" + env_ids = self._resolve_env_ids(env_ids) + num_rows = self.num_instances if full_data else env_ids.shape[0] + data = self._as_warp(data, dtype, (num_rows, self._particles_per_object), name) + for state in self._iter_particle_states(): + wp.launch( + kernel, + dim=(env_ids.shape[0], self._particles_per_object), + inputs=[data, env_ids, self._particle_offsets, full_data], + outputs=[getattr(state, target) for target in targets], + device=self.device, + ) + + def _scatter_to_sim_mask(self, data, env_mask, dtype, kernel, targets, name: str) -> None: + """Scatter masked per-environment particle data into the Newton state arrays in ``targets``.""" + env_mask = self._resolve_mask(env_mask) + data = self._as_warp(data, dtype, (env_mask.shape[0], self._particles_per_object), name) + for state in self._iter_particle_states(): + wp.launch( + kernel, + dim=(env_mask.shape[0], self._particles_per_object), + inputs=[data, env_mask, self._particle_offsets], + outputs=[getattr(state, target) for target in targets], + device=self.device, + ) + + def _as_warp(self, data, dtype, shape: tuple[int, int], name: str) -> wp.array: + """Validate user data and return it as a Warp array of ``dtype``.""" + if isinstance(data, ProxyArray): + data = data.warp + self.assert_shape_and_dtype(data, shape, dtype, name) + if isinstance(data, torch.Tensor): + data = wp.from_torch(data.contiguous(), dtype=dtype) + return data + + def _initialize_impl(self): + entry = self._registry_entry + self._num_instances = len(entry.particle_offsets) + self._particles_per_object = entry.particles_per_object + self._recorded_particle_offsets = entry.particle_offsets + + if self._num_instances == 0 or self._particles_per_object == 0: + raise RuntimeError( + f"No MPM particle instances found for '{self.cfg.prim_path}'. " + "Ensure Newton replication processed the MPM object registry." + ) + + logger.info( + "Newton MPM object initialized at '%s': %d instances x %d particles.", + self.cfg.prim_path, + self._num_instances, + self._particles_per_object, + ) + + self._particle_offsets = wp.array(self._recorded_particle_offsets, dtype=wp.int32, device=self.device) + self._data = MPMObjectData( + particle_offsets=self._particle_offsets, + particles_per_object=self._particles_per_object, + num_instances=self._num_instances, + device=self.device, + ) + self._create_buffers() + self.update(0.0) + + self._physics_ready_handle = SimulationManager.register_callback( + lambda _: self._data._create_simulation_bindings(), + PhysicsEvent.PHYSICS_READY, + name=f"mpm_object_rebind_{self.cfg.prim_path}", + ) + + def _create_buffers(self): + self._ALL_INDICES = wp.array(np.arange(self._num_instances, dtype=np.int32), device=self.device) + self._ALL_ENV_MASK = wp.ones((self._num_instances,), dtype=wp.bool, device=self.device) + + state = SimulationManager.get_state_0() + if state is None or state.particle_q is None or state.particle_qd is None: + raise RuntimeError("Cannot initialize MPMObject buffers before Newton particle state exists.") + + default_pos = wp.zeros((self._num_instances, self._particles_per_object), dtype=wp.vec3f, device=self.device) + default_vel = wp.zeros((self._num_instances, self._particles_per_object), dtype=wp.vec3f, device=self.device) + default_state = wp.zeros((self._num_instances, self._particles_per_object), dtype=vec6f, device=self.device) + wp.launch( + gather_particles_vec3f, + dim=(self._num_instances, self._particles_per_object), + inputs=[state.particle_q, self._particle_offsets], + outputs=[default_pos], + device=self.device, + ) + wp.launch( + gather_particles_vec3f, + dim=(self._num_instances, self._particles_per_object), + inputs=[state.particle_qd, self._particle_offsets], + outputs=[default_vel], + device=self.device, + ) + wp.launch( + compute_particle_state_w, + dim=(self._num_instances, self._particles_per_object), + inputs=[default_pos, default_vel], + outputs=[default_state], + device=self.device, + ) + self._data.default_nodal_state_w = ProxyArray(default_state) + self._data.default_particle_state_w = self._data.default_nodal_state_w + self._create_kit_points() + + def _create_kit_points(self) -> None: + """Create Kit-visible ``UsdGeom.Points`` prims for the particles when the Kit visualizer is active.""" + from isaaclab.sim import SimulationContext # noqa: PLC0415 + + sim = SimulationContext.instance() + if sim is None or "kit" not in sim.resolve_visualizer_types() or not self.cfg.spawn.visible: + return + + first_offset = self._recorded_particle_offsets[0] + radii = ( + SimulationManager.get_model() + .particle_radius[first_offset : first_offset + self._particles_per_object] + .numpy() + ) + base_path = _create_kit_visualization_path(self.cfg.prim_path) + prim_paths = create_mpm_particle_visualization( + prim_path=base_path, + positions=self.data.particle_pos_w.warp.numpy(), + widths=2.0 * radii, + color=self.cfg.spawn.visual_color, + ) + for env_idx, prim_path in enumerate(prim_paths): + SimulationManager.register_particle_visual_prim( + prim_path, + particle_offset=self._recorded_particle_offsets[env_idx], + particle_count=self._particles_per_object, + sync_frequency=self.cfg.spawn.visual_update_frequency, + ) + logger.info("Kit MPM particle visualization initialized at: %s", base_path) + + def _resolve_env_ids(self, env_ids): + if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): + return self._ALL_INDICES + if isinstance(env_ids, torch.Tensor): + return wp.from_torch(env_ids.to(device=self.device, dtype=torch.int32), dtype=wp.int32) + if isinstance(env_ids, Sequence): + return wp.array(list(env_ids), dtype=wp.int32, device=self.device) + return env_ids + + def _resolve_mask(self, mask: wp.array | torch.Tensor | None) -> wp.array: + if mask is None: + return self._ALL_ENV_MASK + if isinstance(mask, torch.Tensor): + if mask.dtype != torch.bool: + mask = mask.to(torch.bool) + return wp.from_torch(mask.to(device=self.device).contiguous(), dtype=wp.bool) + return mask + + def _iter_particle_states(self): + """Yield the Newton states whose particle arrays must receive writes.""" + state_0 = SimulationManager.get_state_0() + state_1 = SimulationManager.get_state_1() + yield state_0 + if state_1 is not None and state_1 is not state_0: + yield state_1 + + def _invalidate_caches(self, pos: bool = False, vel: bool = False) -> None: + """Invalidate gathered data buffers after a particle write and flag the render sync.""" + if pos: + self._data._particle_pos_w.timestamp = -1.0 + self._data._root_pos_w.timestamp = -1.0 + if vel: + self._data._particle_vel_w.timestamp = -1.0 + self._data._root_vel_w.timestamp = -1.0 + self._data._particle_state_w.timestamp = -1.0 + SimulationManager._mark_particles_dirty() + + def _set_debug_vis_impl(self, debug_vis: bool): + raise NotImplementedError("Debug visualization is not implemented for MPMObject.") + + def _debug_vis_callback(self, event): + raise NotImplementedError("Debug visualization is not implemented for MPMObject.") + + def _clear_callbacks(self) -> None: + super()._clear_callbacks() + if self._physics_ready_handle is not None: + self._physics_ready_handle.deregister() + self._physics_ready_handle = None + registry = SimulationManager._mpm_object_registry + if self._registry_entry in registry: + registry.remove(self._registry_entry) + + +def _compose_env_asset_pose( + cfg: MPMObjectCfg, + env_position: list[float], + env_rotation: list[float] | tuple[float, float, float, float], +) -> tuple[tuple[float, float, float], tuple[float, float, float, float]]: + """Compose the environment transform with the asset's initial pose (both ``xyzw``).""" + env_pos = wp.vec3(*env_position) + env_rot = wp.quat(*env_rotation) + pos = env_pos + wp.quat_rotate(env_rot, wp.vec3(*cfg.init_state.pos)) + rot = env_rot * wp.quat(*cfg.init_state.rot) + return (float(pos[0]), float(pos[1]), float(pos[2])), (float(rot[0]), float(rot[1]), float(rot[2]), float(rot[3])) + + +def _create_kit_visualization_path(prim_path: str) -> str: + sanitized = "".join(char if char.isalnum() else "_" for char in prim_path.strip("/")) + return f"/World/Visuals/MPMParticles/{sanitized or 'Object'}" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_cfg.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_cfg.py new file mode 100644 index 000000000000..e1af8957ff5d --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_cfg.py @@ -0,0 +1,27 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING + +from isaaclab.assets.deformable_object.deformable_object_cfg import DeformableObjectCfg +from isaaclab.utils.configclass import configclass + +from isaaclab_newton.sim.spawners.mpm import MPMParticleSpawnerCfg + +if TYPE_CHECKING: + from .mpm_object import MPMObject + + +@configclass +class MPMObjectCfg(DeformableObjectCfg): + """Configuration parameters for a Newton MPM particle object.""" + + class_type: type[MPMObject] | str = "{DIR}.mpm_object:MPMObject" + + spawn: MPMParticleSpawnerCfg = MISSING + """Particle generation configuration for this MPM object.""" diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_data.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_data.py new file mode 100644 index 000000000000..2b39f5dd7faa --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object_data.py @@ -0,0 +1,147 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import warp as wp + +from isaaclab.assets.deformable_object.base_deformable_object_data import BaseDeformableObjectData +from isaaclab.utils.buffers import TimestampedBufferWarp as TimestampedBuffer +from isaaclab.utils.warp import ProxyArray + +from isaaclab_newton.physics import NewtonManager as SimulationManager + +from .kernels import compute_mean_vec3f_over_particles, compute_particle_state_w, gather_particles_vec3f, vec6f + + +class MPMObjectData(BaseDeformableObjectData): + """Data container for a Newton MPM particle object.""" + + __backend_name__: str = "newton" + + def __init__(self, particle_offsets: wp.array, particles_per_object: int, num_instances: int, device: str): + super().__init__(device) + self._particle_offsets = particle_offsets + self._particles_per_object = particles_per_object + self._num_instances = num_instances + + self._particle_pos_w = TimestampedBuffer((num_instances, particles_per_object), device, wp.vec3f) + self._particle_vel_w = TimestampedBuffer((num_instances, particles_per_object), device, wp.vec3f) + self._particle_state_w = TimestampedBuffer((num_instances, particles_per_object), device, vec6f) + self._root_pos_w = TimestampedBuffer((num_instances,), device, wp.vec3f) + self._root_vel_w = TimestampedBuffer((num_instances,), device, wp.vec3f) + self._particle_pos_w_ta = ProxyArray(self._particle_pos_w.data) + self._particle_vel_w_ta = ProxyArray(self._particle_vel_w.data) + self._particle_state_w_ta = ProxyArray(self._particle_state_w.data) + self._root_pos_w_ta = ProxyArray(self._root_pos_w.data) + self._root_vel_w_ta = ProxyArray(self._root_vel_w.data) + + self.default_nodal_state_w: ProxyArray | None = None + self.default_particle_state_w: ProxyArray | None = None + self.nodal_kinematic_target: ProxyArray | None = None + + self._create_simulation_bindings() + + def _create_simulation_bindings(self) -> None: + """Validate current Newton particle arrays and invalidate gathered buffers.""" + self._get_current_particle_state() + self._particle_pos_w.timestamp = -1.0 + self._particle_vel_w.timestamp = -1.0 + self._particle_state_w.timestamp = -1.0 + self._root_pos_w.timestamp = -1.0 + self._root_vel_w.timestamp = -1.0 + + def _get_current_particle_state(self): + state = SimulationManager.get_state_0() + if state is None or state.particle_q is None or state.particle_qd is None: + raise RuntimeError( + "Failed to access Newton MPM particle state. Ensure the Newton model has been finalized and contains " + "particle position and velocity arrays." + ) + return state + + @property + def particle_pos_w(self) -> ProxyArray: + """Particle positions in simulation world frame [m].""" + if self._particle_pos_w.timestamp < self._sim_timestamp: + state = self._get_current_particle_state() + wp.launch( + gather_particles_vec3f, + dim=(self._num_instances, self._particles_per_object), + inputs=[state.particle_q, self._particle_offsets], + outputs=[self._particle_pos_w.data], + device=self.device, + ) + self._particle_pos_w.timestamp = self._sim_timestamp + return self._particle_pos_w_ta + + @property + def particle_vel_w(self) -> ProxyArray: + """Particle velocities in simulation world frame [m/s].""" + if self._particle_vel_w.timestamp < self._sim_timestamp: + state = self._get_current_particle_state() + wp.launch( + gather_particles_vec3f, + dim=(self._num_instances, self._particles_per_object), + inputs=[state.particle_qd, self._particle_offsets], + outputs=[self._particle_vel_w.data], + device=self.device, + ) + self._particle_vel_w.timestamp = self._sim_timestamp + return self._particle_vel_w_ta + + @property + def particle_state_w(self) -> ProxyArray: + """Particle state ``[pos, vel]`` in simulation world frame [m, m/s].""" + if self._particle_state_w.timestamp < self._sim_timestamp: + wp.launch( + compute_particle_state_w, + dim=(self._num_instances, self._particles_per_object), + inputs=[self.particle_pos_w.warp, self.particle_vel_w.warp], + outputs=[self._particle_state_w.data], + device=self.device, + ) + self._particle_state_w.timestamp = self._sim_timestamp + return self._particle_state_w_ta + + @property + def nodal_pos_w(self) -> ProxyArray: + return self.particle_pos_w + + @property + def nodal_vel_w(self) -> ProxyArray: + return self.particle_vel_w + + @property + def nodal_state_w(self) -> ProxyArray: + return self.particle_state_w + + @property + def root_pos_w(self) -> ProxyArray: + """Mean particle position per instance in simulation world frame [m].""" + if self._root_pos_w.timestamp < self._sim_timestamp: + wp.launch( + compute_mean_vec3f_over_particles, + dim=(self._num_instances,), + inputs=[self.particle_pos_w.warp, self._particles_per_object], + outputs=[self._root_pos_w.data], + device=self.device, + ) + self._root_pos_w.timestamp = self._sim_timestamp + return self._root_pos_w_ta + + @property + def root_vel_w(self) -> ProxyArray: + """Mean particle velocity per instance in simulation world frame [m/s].""" + if self._root_vel_w.timestamp < self._sim_timestamp: + wp.launch( + compute_mean_vec3f_over_particles, + dim=(self._num_instances,), + inputs=[self.particle_vel_w.warp, self._particles_per_object], + outputs=[self._root_vel_w.data], + device=self.device, + ) + self._root_vel_w.timestamp = self._sim_timestamp + return self._root_vel_w_ta diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index 87184a383079..224e7cd8058d 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -16,6 +16,8 @@ from pxr import Usd from isaaclab.cloner.replicate_session import REPLICATION_QUEUE +from isaaclab.physics import PhysicsManager +from isaaclab.sim import SimulationContext from isaaclab_newton.physics import NewtonManager @@ -55,8 +57,10 @@ def _build_newton_builder_from_mapping( quaternions[:, 3] = 1.0 schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] + sim_ctx = SimulationContext.instance() + manager_cls = sim_ctx.physics_manager if sim_ctx is not None else NewtonManager - builder = NewtonManager.create_builder(up_axis=up_axis) + builder = manager_cls.create_builder(up_axis=up_axis) stage_info = builder.add_usd( stage, ignore_paths=["/World/envs", *sources], @@ -87,7 +91,7 @@ def _build_newton_builder_from_mapping( protos: dict[str, ModelBuilder] = {} for src_path in sources: - p = NewtonManager.create_builder(up_axis=up_axis) + p = manager_cls.create_builder(up_axis=up_axis) solvers.SolverMuJoCo.register_custom_attributes(p) p.add_usd( stage, @@ -147,6 +151,12 @@ def _build_newton_builder_from_mapping( for proto_shape_idx in proto_shape_indices: local_site_map[label][col].append(offset + proto_shape_idx) + # Emit registered MPM particle objects into this Newton world. + if NewtonManager._mpm_object_registry: + from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder + + add_registered_mpm_objects_to_builder(builder, col, positions[col].tolist(), quaternions[col].tolist()) + # Run per-world builder hooks (e.g. deformable body registration). if hasattr(NewtonManager, "_per_world_builder_hooks"): for hook in NewtonManager._per_world_builder_hooks: @@ -300,7 +310,7 @@ def __init__( *, device: str = "cpu", up_axis: str = "Z", - simplify_meshes: bool = True, + simplify_meshes: bool | None = None, commit_to_manager: bool = True, ): """Initialize the context. @@ -309,13 +319,19 @@ def __init__( stage: USD stage containing source assets. device: Device used by the finalized Newton model builder. up_axis: Up axis for the Newton model builder. - simplify_meshes: Whether to run convex-hull mesh approximation. + simplify_meshes: Whether to run convex-hull mesh approximation. If + ``None``, read from the active :class:`NewtonCfg`. commit_to_manager: Whether :meth:`replicate` should publish the builder to :class:`NewtonManager`. """ self.stage = stage self.device = device self.up_axis = up_axis + if simplify_meshes is None: + from isaaclab_newton.physics import NewtonCfg + + cfg = PhysicsManager._cfg + simplify_meshes = cfg.simplify_meshes if isinstance(cfg, NewtonCfg) else True self.simplify_meshes = simplify_meshes self.commit_to_manager = commit_to_manager self._queue: list[ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi index 176973cd5400..46cc26e9318f 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/physics/__init__.pyi @@ -9,6 +9,8 @@ __all__ = [ "HydroelasticSDFCfg", "NewtonKaminoManager", "KaminoSolverCfg", + "NewtonMPMManager", + "MPMSolverCfg", "NewtonMJWarpManager", "MJWarpSolverCfg", "NewtonCfg", @@ -26,6 +28,8 @@ from .kamino_manager import NewtonKaminoManager from .kamino_manager_cfg import KaminoSolverCfg from .mjwarp_manager import NewtonMJWarpManager from .mjwarp_manager_cfg import MJWarpSolverCfg +from .mpm_manager import NewtonMPMManager +from .mpm_manager_cfg import MPMSolverCfg from .newton_collision_cfg import HydroelasticSDFCfg, NewtonCollisionPipelineCfg from .newton_manager import NewtonManager from .newton_manager_cfg import ( diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py new file mode 100644 index 000000000000..43ae838fa223 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implicit MPM Newton manager.""" + +from __future__ import annotations + +import warp as wp +from newton import BodyFlags, Contacts, Control, Model, ModelBuilder, State +from newton.solvers import SolverImplicitMPM +from warp.fem import TemporaryStore + +from .mpm_manager_cfg import MPMSolverCfg +from .newton_manager import NewtonManager + + +def _make_solver_config(solver_cfg: MPMSolverCfg) -> SolverImplicitMPM.Config: + """Build Newton's implicit MPM solver config from Isaac Lab's cfg.""" + return SolverImplicitMPM.Config( + max_iterations=solver_cfg.max_iterations, + tolerance=solver_cfg.tolerance, + solver=solver_cfg.solver, + warmstart_mode=solver_cfg.warmstart_mode, + collider_velocity_mode=solver_cfg.collider_velocity_mode, + voxel_size=solver_cfg.voxel_size, + grid_type=solver_cfg.grid_type, + grid_padding=solver_cfg.grid_padding, + max_active_cell_count=solver_cfg.max_active_cell_count, + transfer_scheme=solver_cfg.transfer_scheme, + integration_scheme=solver_cfg.integration_scheme, + critical_fraction=solver_cfg.critical_fraction, + air_drag=solver_cfg.air_drag, + collider_normal_from_sdf_gradient=solver_cfg.collider_normal_from_sdf_gradient, + collider_basis=solver_cfg.collider_basis, + strain_basis=solver_cfg.strain_basis, + velocity_basis=solver_cfg.velocity_basis, + ) + + +class NewtonMPMManager(NewtonManager): + """:class:`NewtonManager` specialization for Newton's implicit MPM solver. + + MPM advances particle materials in-place and treats rigid geometry as + colliders, so it does not consume Newton's rigid-body collision pipeline + and steps with a single :class:`State`. + """ + + _project_outside_colliders: bool = False + """Whether :meth:`_step_solver` projects particles out of colliders each substep. + + Set from :attr:`MPMSolverCfg.project_outside_colliders` in + :meth:`_build_solver` and read in :meth:`_step_solver`. + """ + + @classmethod + def _register_builder_attributes(cls, builder: ModelBuilder) -> None: + """Register the particle custom attributes required by :class:`SolverImplicitMPM`. + + Implicit MPM materials are configured per-particle through Newton + custom attributes (``mpm:young_modulus``, ``mpm:viscosity``, ...). + These must be present on the builder *before* particles are added so + that ``add_particles(custom_attributes=...)`` succeeds and so that + ``builder.finalize()`` allocates the matching model arrays. + + Idempotent: ``has_custom_attribute`` guards against re-registration + when the hook is invoked multiple times (e.g. once via + :meth:`create_builder` and again via :meth:`start_simulation`). + """ + if not builder.has_custom_attribute("mpm:young_modulus"): + SolverImplicitMPM.register_custom_attributes(builder) + + @classmethod + def _prepare_builder_for_finalize(cls, builder: ModelBuilder) -> None: + """Normalize kinematic rigid bodies before MPM solver construction. + + Newton's implicit MPM solver treats positive-mass body colliders as + finite-mass colliders. Isaac Lab kinematic assets can import with a + computed mass, so clear mass and inertia for kinematic bodies to match + Newton's direct-builder MPM examples. + """ + kinematic_flag = int(BodyFlags.KINEMATIC) + for body_id, flags in enumerate(builder.body_flags): + if int(flags) & kinematic_flag: + builder.body_mass[body_id] = 0.0 + builder.body_inv_mass[body_id] = 0.0 + builder.body_inertia[body_id] = wp.mat33() + builder.body_inv_inertia[body_id] = wp.mat33() + + @classmethod + def _build_solver(cls, model: Model, solver_cfg: MPMSolverCfg) -> None: + """Construct :class:`SolverImplicitMPM` and populate the base-class slots. + + MPM steps in-place on a single :class:`State` and runs collision + handling internally, so it neither double-buffers state nor drives + Newton's :class:`CollisionPipeline`. + + Args: + model: Finalized Newton model the solver should run on. + solver_cfg: Implicit MPM solver configuration. + """ + NewtonManager._solver = SolverImplicitMPM( + model, + _make_solver_config(solver_cfg), + temporary_store=TemporaryStore(), + ) + NewtonManager._use_single_state = True + NewtonManager._needs_collision_pipeline = False + NewtonManager._needs_fk_before_step = True + cls._project_outside_colliders = solver_cfg.project_outside_colliders + + @classmethod + def _supports_cuda_graph_capture(cls) -> bool: + """Return ``True`` only for fixed-grid MPM. + + Sparse and dense grids reallocate as particles move, which is not + capturable in a CUDA graph; the fixed grid keeps a static topology. + """ + return cls._solver.grid_type == "fixed" + + @classmethod + def _step_solver( + cls, state_0: State, state_1: State, control: Control, contacts: Contacts | None, substep_dt: float + ) -> None: + """Run one implicit MPM substep, optionally projecting particles out of colliders. + + The implicit solve already resolves colliders at the grid level. When + :attr:`MPMSolverCfg.project_outside_colliders` is set, the manager also + runs ``project_outside`` after the step (as in Newton's MPM examples) to + hard-project particles out of collider interiors. The flag is evaluated + when the step is first run, so the chosen branch is baked into any + captured CUDA graph. + """ + cls._solver.step(state_0, state_1, control, contacts, substep_dt) + if cls._project_outside_colliders: + cls._solver.project_outside(state_1, state_1, substep_dt) + + @classmethod + def _solver_specific_clear(cls) -> None: + """Reset MPM-specific class state on teardown. + + :meth:`_build_solver` sets :attr:`_project_outside_colliders` from the + active config. Resetting it here keeps a teardown-only :meth:`clear` + (without a follow-up rebuild) from leaving a stale value on the class, + mirroring how :meth:`NewtonManager.clear` resets the base-class flags. + """ + cls._project_outside_colliders = False diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py new file mode 100644 index 000000000000..9b8c49489b49 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py @@ -0,0 +1,110 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for Newton's implicit MPM solver.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +from isaaclab.utils.configclass import configclass + +from .newton_manager_cfg import NewtonSolverCfg + +if TYPE_CHECKING: + from isaaclab_newton.physics import NewtonManager + + +@configclass +class MPMSolverCfg(NewtonSolverCfg): + """Configuration for Newton's implicit Material Point Method (MPM) solver. + + The implicit MPM solver advances particle materials and treats rigid geometry + as colliders. It is not a rigid-body or articulation dynamics solver. + """ + + class_type: type[NewtonManager] | str = "{DIR}.mpm_manager:NewtonMPMManager" + """Manager class for the implicit MPM solver.""" + + solver_type: str = "implicit_mpm" + """Solver type. Can be "implicit_mpm".""" + + # numerics + max_iterations: int = 250 + """Maximum number of iterations for the rheology solver.""" + + tolerance: float = 1.0e-4 + """Tolerance for the rheology solver.""" + + solver: str | tuple[str, ...] = "auto" + """Rheology solver, or an ordered warm-start sequence of solvers. + + ``"auto"`` lets Newton pick the solver from the velocity basis (``"gs"`` for + ``Q1``, ``"gs-batched"`` for ``B2``/``B3``). Other accepted values include + ``"gauss-seidel"``, ``"jacobi"``, ``"cg"``, ``"cr"``, and ``"gmres"``; pass a + tuple such as ``("cr", "gs")`` to warm-start solvers left-to-right. + """ + + warmstart_mode: Literal["none", "auto", "particles", "grid", "smoothed"] = "auto" + """Warm-start mode for the rheology solver.""" + + collider_velocity_mode: Literal["forward", "backward", "instantaneous", "finite_difference"] = "forward" + """Collider velocity computation mode.""" + + # grid + voxel_size: float = 0.1 + """Size of the MPM grid voxels [m].""" + + grid_type: Literal["sparse", "dense", "fixed"] = "sparse" + """Type of grid to use.""" + + grid_padding: int = 0 + """Number of empty cells to add around particles when allocating the grid.""" + + max_active_cell_count: int = -1 + """Maximum active cell count for dense-grid active subsets. ``-1`` means unlimited.""" + + transfer_scheme: Literal["apic", "pic"] = "apic" + """Particle-grid transfer scheme.""" + + integration_scheme: Literal["pic", "gimp"] = "pic" + """Integration scheme controlling shape-function support.""" + + # material / background + critical_fraction: float = 0.0 + """Dimensionless fraction under which the yield surface collapses.""" + + air_drag: float = 1.0 + """Numerical drag for background air.""" + + # experimental + collider_normal_from_sdf_gradient: bool = False + """Whether collider normals are computed from SDF gradients rather than closest points.""" + + collider_basis: str = "S2" + """Collider basis function, such as ``"S2"`` or ``"Q1"``.""" + + strain_basis: str = "P0" + """Strain basis function, such as ``"P0"``, ``"P1d"``, ``"Q1"``, or ``"Q1d"``.""" + + velocity_basis: str = "Q1" + """Velocity basis function, such as ``"Q1"``, ``"B2"``, or ``"B3"``.""" + + # collision handling (applied by the Isaac Lab manager, not the Newton solver config) + project_outside_colliders: bool = False + """Whether to hard-project particles out of collider interiors after each substep. + + When ``True``, :class:`~isaaclab_newton.physics.NewtonMPMManager` calls + :meth:`SolverImplicitMPM.project_outside` immediately after every solver + substep: it applies a Coulomb response and pushes particles that drifted into + a collider back onto its surface. The implicit solve already resolves + colliders at the grid level; this is the particle-level correction that stops + material from slowly settling inside colliders, mirroring Newton's MPM + examples. Leave it ``False`` for collider-free scenes to skip a per-substep + projection pass over every particle. + + This is a manager-level stepping option and is intentionally **not** part of + ``SolverImplicitMPM.Config``. + """ diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index bc1e04c57cf2..7c51d83c076d 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -11,7 +11,8 @@ import ctypes import logging from abc import abstractmethod -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING import warp as wp @@ -44,6 +45,8 @@ from .newton_manager_cfg import NewtonCfg, NewtonShapeCfg if TYPE_CHECKING: + from pxr import Usd + from isaaclab.sim.simulation_context import SimulationContext from isaaclab_newton.actuators import NewtonActuatorAdapter @@ -104,6 +107,17 @@ def _sync_particle_points( fabric_points[i][j] = wp.transform_point(inv_world_matrix, particle_q[offset + j]) +@dataclass +class _ParticleVisualPrim: + """A ``UsdGeom.Points`` prim mirroring a slice of Newton's particle state.""" + + points_attr: Usd.Attribute + offset: int + count: int + sync_frequency: int + frames_since_sync: int + + @wp.kernel(enable_backward=False) def _or_reset_masks_from_mask( env_mask: wp.array(dtype=wp.bool), @@ -177,7 +191,8 @@ class NewtonManager(PhysicsManager): state, contacts/collision pipeline, sensors, replication, and CUDA-graph orchestration. Concrete subclasses (one per solver) implement :meth:`_build_solver` and - may extend :meth:`_initialize_contacts`, :meth:`_step_solver`, + may extend :meth:`_initialize_contacts`, :meth:`_prepare_builder_for_finalize`, + :meth:`_step_solver`, :meth:`_supports_cuda_graph_capture`, :meth:`_solver_specific_clear`, and :meth:`_log_solver_debug`. Subclasses are selected via :attr:`NewtonSolverCfg.class_type`, which @@ -219,6 +234,7 @@ class NewtonManager(PhysicsManager): # Collision and contacts _contacts: Contacts | None = None _needs_collision_pipeline: bool = False + _needs_fk_before_step: bool = False _collision_pipeline = None _collision_cfg: NewtonCollisionPipelineCfg | None = None _newton_contact_sensors: dict = {} # Maps sensor_key to NewtonContactSensor @@ -251,6 +267,7 @@ class NewtonManager(PhysicsManager): _particles_dirty: bool = False _newton_particle_offset_attr = "newton:particleOffset" _newton_particle_count_attr = "newton:particleCount" + _particle_visual_prims: dict[str, _ParticleVisualPrim] = {} # cubric GPU transform hierarchy (replaces CPU update_world_xforms) _cubric = None @@ -271,6 +288,7 @@ class NewtonManager(PhysicsManager): # Views list for assets to register their views _views: list = [] + _mpm_object_registry: list = [] # CL: Cloning / Replication logic # TODO: These attributes support cloning-specific logic and should be moved into a cloner class @@ -445,48 +463,80 @@ def sync_transforms_to_usd(cls) -> None: @classmethod def sync_particles_to_usd(cls) -> None: - """Write Newton particle_q to Fabric mesh point arrays for Kit viewport rendering. + """Write Newton particle positions to USD/Fabric for Kit viewport rendering. - For each deformable body whose mesh prim carries a ``newton:particleOffset`` - attribute, this function copies the corresponding slice of ``state_0.particle_q`` - into the Fabric ``points`` array so the Kit viewport reflects the current - deformation. + Two prim families are synced from ``state_0.particle_q``: - No-op when there is no ``_usdrt_stage``, no simulation state, or no - deformable bodies registered. + * Fabric mesh prims tagged with ``newton:particleOffset`` / + ``newton:particleCount`` (deformable visual meshes) receive + local-frame points on the GPU via :meth:`_sync_fabric_mesh_particles`. + * ``UsdGeom.Points`` prims registered through + :meth:`register_particle_visual_prim` (MPM particle clouds) receive + world-frame points via :meth:`_sync_particle_points_prims`. + + No-op when there is no particle state or nothing changed since the + last sync. """ - if cls._usdrt_stage is None or cls._state_0 is None or cls._state_0.particle_q is None: - return - if not cls._particles_dirty: + if not cls._particles_dirty or cls._state_0 is None or cls._state_0.particle_q is None: return - pq = cls._state_0.particle_q try: - import usdrt + cls._sync_fabric_mesh_particles() + NewtonManager._particles_dirty = cls._sync_particle_points_prims() + except Exception: + logger.exception("[NewtonManager] sync_particles_to_usd FAILED") - selection = cls._usdrt_stage.SelectPrims( - require_attrs=[ - (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.ReadWrite), - (usdrt.Sdf.ValueTypeNames.UInt, cls._newton_particle_offset_attr, usdrt.Usd.Access.Read), - (usdrt.Sdf.ValueTypeNames.UInt, cls._newton_particle_count_attr, usdrt.Usd.Access.Read), - (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), - ], - device=str(PhysicsManager._device), - ) - if selection.GetCount() == 0: - return - fabric_points = wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f) - fabric_offsets = wp.fabricarray(data=selection, attrib=cls._newton_particle_offset_attr) - fabric_counts = wp.fabricarray(data=selection, attrib=cls._newton_particle_count_attr) - fabric_world_matrices = wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix") - wp.launch( - _sync_particle_points, - dim=selection.GetCount(), - inputs=[fabric_points, fabric_world_matrices, fabric_offsets, fabric_counts, pq], - device=PhysicsManager._device, - ) - NewtonManager._particles_dirty = False - except Exception as exc: - logger.debug("[sync_particles_to_usd] %s", exc) + @classmethod + def _sync_fabric_mesh_particles(cls) -> None: + """Write ``state_0.particle_q`` into Fabric mesh point arrays as local-frame points.""" + if cls._usdrt_stage is None: + return + import usdrt # noqa: PLC0415 + + selection = cls._usdrt_stage.SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.ReadWrite), + (usdrt.Sdf.ValueTypeNames.UInt, cls._newton_particle_offset_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.UInt, cls._newton_particle_count_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ], + device=str(PhysicsManager._device), + ) + if selection.GetCount() == 0: + return + wp.launch( + _sync_particle_points, + dim=selection.GetCount(), + inputs=[ + wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f), + wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix"), + wp.fabricarray(data=selection, attrib=cls._newton_particle_offset_attr), + wp.fabricarray(data=selection, attrib=cls._newton_particle_count_attr), + cls._state_0.particle_q, + ], + device=PhysicsManager._device, + ) + + @classmethod + def _sync_particle_points_prims(cls) -> bool: + """Write registered ``UsdGeom.Points`` prims; return ``True`` while throttled prims remain.""" + if not cls._particle_visual_prims: + return False + + due = [] + for record in cls._particle_visual_prims.values(): + record.frames_since_sync += 1 + if record.frames_since_sync >= record.sync_frequency: + record.frames_since_sync = 0 + due.append(record) + if due: + from pxr import Sdf, Vt # noqa: PLC0415 + + particle_q = cls._state_0.particle_q.numpy() + with Sdf.ChangeBlock(): + for record in due: + points = particle_q[record.offset : record.offset + record.count] + record.points_attr.Set(Vt.Vec3fArray.FromNumpy(points)) + return len(due) < len(cls._particle_visual_prims) @classmethod def _mark_transforms_dirty(cls) -> None: @@ -516,6 +566,29 @@ def _mark_state_dirty(cls) -> None: cls._mark_transforms_dirty() cls._mark_particles_dirty() + @classmethod + def register_particle_visual_prim( + cls, prim_path: str, particle_offset: int, particle_count: int, sync_frequency: int = 1 + ) -> None: + """Register a ``UsdGeom.Points`` prim whose points mirror a slice of Newton's particle state. + + Args: + prim_path: Stage path of an existing ``UsdGeom.Points`` prim. + particle_offset: First index of the prim's slice in ``state.particle_q``. + particle_count: Number of particles in the slice. + sync_frequency: Sync the prim every N dirty render frames. + """ + from pxr import UsdGeom # noqa: PLC0415 + + prim = get_current_stage().GetPrimAtPath(prim_path) + NewtonManager._particle_visual_prims[prim_path] = _ParticleVisualPrim( + points_attr=UsdGeom.Points(prim).GetPointsAttr(), + offset=int(particle_offset), + count=int(particle_count), + sync_frequency=int(sync_frequency), + frames_since_sync=int(sync_frequency), + ) + @classmethod def step(cls) -> None: """Step the physics simulation. @@ -564,11 +637,13 @@ def step(cls) -> None: else: logger.warning("Newton deferred CUDA graph capture failed; using eager execution") - # Ensure body_q is up-to-date before collision detection. - # After env resets, joint_q is written but body_q (used by - # broadphase/narrowphase) is stale until FK runs. + # Ensure body_q is up-to-date before solvers read rigid transforms. + # After env resets or kinematic root writes, joint_q is written but + # body_q is stale until FK runs. Collision-based solvers need this for + # broadphase/narrowphase; collider-based solvers such as MPM need it + # for their internal collider queries. # Only runs FK for dirtied articulations via the accumulated mask. - if cls._needs_collision_pipeline: + if cls._needs_collision_pipeline or cls._needs_fk_before_step: eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, cls._fk_reset_mask) # Zero both masks after consumption @@ -602,6 +677,8 @@ def step(cls) -> None: if cls._usdrt_stage is not None: cls._mark_state_dirty() + elif cls._particle_visual_prims: + cls._mark_particles_dirty() # Launch solver-specific debug logging after stepping. cls._log_solver_debug() @@ -663,6 +740,7 @@ def clear(cls): NewtonManager._control = None NewtonManager._contacts = None NewtonManager._needs_collision_pipeline = False + NewtonManager._needs_fk_before_step = False NewtonManager._collision_pipeline = None NewtonManager._collision_cfg = None NewtonManager._newton_contact_sensors = {} @@ -686,6 +764,8 @@ def clear(cls): NewtonManager._usdrt_stage = None NewtonManager._transforms_dirty = False NewtonManager._particles_dirty = False + NewtonManager._particle_visual_prims = {} + NewtonManager._mpm_object_registry = [] NewtonManager._up_axis = "Z" NewtonManager._scene_data = None NewtonManager._scene_data_mapping = None @@ -723,6 +803,7 @@ def create_builder(cls, up_axis: str | None = None, **kwargs) -> ModelBuilder: New builder with up-axis and per-shape defaults (gap, margin) applied. """ builder = ModelBuilder(up_axis=up_axis or cls._up_axis, **kwargs) + cls._register_builder_attributes(builder) # Resolve which NewtonShapeCfg to apply: user override if active config # is NewtonCfg, else the wrapper's own defaults so callers from non-Newton # contexts (tests, early construction) still get the rough-terrain margin. @@ -731,6 +812,29 @@ def create_builder(cls, up_axis: str | None = None, **kwargs) -> ModelBuilder: checked_apply(shape_cfg, builder.default_shape_cfg) return builder + @classmethod + def _register_builder_attributes(cls, builder: ModelBuilder) -> None: + """Subclass hook to register solver-specific custom attributes on *builder*. + + Override in solver subclasses (e.g. :class:`NewtonMPMManager`) that need + Newton-side particle, shape, or body custom attributes registered before + the builder is finalized. The default implementation is a no-op so + solvers without custom attributes do not need to override it. + + Implementations should be **idempotent** — the same builder may be + passed multiple times across :meth:`create_builder`, + :meth:`instantiate_builder_from_stage`, and :meth:`start_simulation`. + """ + + @classmethod + def _prepare_builder_for_finalize(cls, builder: ModelBuilder) -> None: + """Subclass hook to normalize *builder* before model finalization. + + Override in solver subclasses that need to adapt imported or replicated + builder data before :meth:`ModelBuilder.finalize` allocates model arrays. + The default implementation is a no-op. + """ + @classmethod def cl_register_site(cls, body_pattern: str | None, xform: wp.transform, *, per_world: bool = False) -> str: """Register a site request for injection into prototypes before replication. @@ -932,6 +1036,8 @@ def invalidate_fk( index. Shape ``(world_count, count_per_world)``. Obtained from ``ArticulationView.articulation_ids``. """ + cls._mark_transforms_dirty() + if cls._world_reset_mask is None or cls._fk_reset_mask is None: return @@ -969,6 +1075,7 @@ def start_simulation(cls) -> None: # Create builder from USD stage if not provided if cls._builder is None: cls.instantiate_builder_from_stage() + cls._register_builder_attributes(cls._builder) logger.info("Dispatching MODEL_INIT callbacks") cls.dispatch_event(PhysicsEvent.MODEL_INIT) @@ -984,6 +1091,7 @@ def start_simulation(cls) -> None: if cls._pending_extended_state_attributes: cls._builder.request_state_attributes(*cls._pending_extended_state_attributes) NewtonManager._pending_extended_state_attributes = set() + cls._prepare_builder_for_finalize(cls._builder) with Timer(name="newton_finalize_builder", msg="Finalize builder took:"): NewtonManager._model = cls._builder.finalize(device=device) cls._model.set_gravity(cls._gravity_vector) @@ -1031,9 +1139,16 @@ def start_simulation(cls) -> None: ) NewtonManager._initialize_fabric_body_prims(cls._usdrt_stage, fabric_hierarchy, usdrt, body_bindings) + NewtonManager._initialize_fabric_particle_prims( + cls._usdrt_stage, + fabric_hierarchy, + usdrt, + NewtonManager._particle_visual_prims, + ) - cls._mark_transforms_dirty() + cls._mark_state_dirty() cls.sync_transforms_to_usd() + cls.sync_particles_to_usd() @staticmethod def _initialize_fabric_body_prims(stage, fabric_hierarchy, usdrt, body_bindings: Sequence[tuple[str, int]]) -> None: @@ -1057,6 +1172,18 @@ def _initialize_fabric_body_prims(stage, fabric_hierarchy, usdrt, body_bindings: fabric_hierarchy.update_world_xforms() + @staticmethod + def _initialize_fabric_particle_prims(stage, fabric_hierarchy, usdrt, prim_paths: Iterable[str]) -> None: + """Initialize Fabric world matrices for point prims used by particle sync.""" + prim_paths = tuple(prim_paths) + for prim_path in prim_paths: + prim = stage.GetPrimAtPath(prim_path) + if prim.IsValid(): + usdrt.Rt.Xformable(prim).SetWorldXformFromUsd() + + if prim_paths: + fabric_hierarchy.update_world_xforms() + @classmethod def instantiate_builder_from_stage(cls): """Create builder from USD stage. @@ -1084,7 +1211,7 @@ def instantiate_builder_from_stage(cls): env_paths.append((int(m.group(1)), child.GetPath().pathString)) env_paths.sort(key=lambda x: x[0]) - builder = ModelBuilder(up_axis=up_axis) + builder = cls.create_builder(up_axis=up_axis) schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] @@ -1092,6 +1219,10 @@ def instantiate_builder_from_stage(cls): # No env Xforms — flat loading builder.add_usd(stage, schema_resolvers=schema_resolvers) NewtonManager._world_xforms = [wp.transform()] + if cls._mpm_object_registry: + from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder + + add_registered_mpm_objects_to_builder(builder, 0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) else: # Load everything except the env subtrees (ground plane, lights, etc.) ignore_paths = [path for _, path in env_paths] @@ -1099,7 +1230,7 @@ def instantiate_builder_from_stage(cls): # Build a prototype from the first env (all envs assumed identical) _, proto_path = env_paths[0] - proto = ModelBuilder(up_axis=up_axis) + proto = cls.create_builder(up_axis=up_axis) proto.add_usd( stage, root_path=proto_path, @@ -1142,6 +1273,10 @@ def instantiate_builder_from_stage(cls): local_site_map[label] = [[] for _ in range(num_worlds)] for proto_shape_idx in proto_shape_indices: local_site_map[label][col].append(offset + proto_shape_idx) + if cls._mpm_object_registry: + from isaaclab_newton.assets.mpm_object.mpm_object import add_registered_mpm_objects_to_builder + + add_registered_mpm_objects_to_builder(builder, col, list(pos), list(quat)) builder.end_world() NewtonManager._cl_site_index_map = { @@ -1327,6 +1462,15 @@ def _capture_or_defer_graph(cls) -> None: return use_cuda_graph = cfg.use_cuda_graph and "cuda" in device + if use_cuda_graph and not cls._supports_cuda_graph_capture(): + NewtonManager._graph = None + NewtonManager._graph_capture_pending = False + logger.warning( + "%s does not support CUDA graph capture for the current solver configuration; using eager execution.", + cls.__name__, + ) + return + if use_cuda_graph: with Timer(name="newton_cuda_graph", msg="CUDA graph took:"): if cls._usdrt_stage is None: @@ -1353,6 +1497,11 @@ def _capture_or_defer_graph(cls) -> None: else: NewtonManager._graph = None + @classmethod + def _supports_cuda_graph_capture(cls) -> bool: + """Return whether the active solver configuration supports CUDA graph capture.""" + return True + @classmethod def _capture_relaxed_graph(cls, device: str): """Capture Newton physics (only) as a CUDA graph, RTX-compatible. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py index feb9b28db386..fcf956cce4cd 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager_cfg.py @@ -126,6 +126,9 @@ class NewtonCfg(PhysicsCfg): - :class:`XPBDSolverCfg` (always), - :class:`FeatherstoneSolverCfg` (always). + :class:`~isaaclab_newton.physics.MPMSolverCfg` does not use this pipeline; + implicit MPM treats rigid geometry as colliders internally. + If ``None`` (default), a pipeline with ``broad_phase="explicit"`` is created automatically. Set this to a :class:`NewtonCollisionPipelineCfg` to customize parameters such as broad phase algorithm, contact limits, or hydroelastic mode. @@ -144,6 +147,13 @@ class NewtonCfg(PhysicsCfg): :class:`NewtonShapeCfg` for the declared fields. """ + simplify_meshes: bool = True + """Whether Newton replication simplifies mesh colliders to convex hulls. + + Keep this enabled for most rigid-body scenes. Disable it when exact triangle + meshes are intentional, for example thin or hollow MPM colliders. + """ + def __post_init__(self): # NewtonCfg.class_type is auto-derived from solver_cfg.class_type. # Refuse a user-set value: setting both is ambiguous and was diff --git a/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi index 699fc15ad4df..27677fbd6dbf 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/sim/__init__.pyi @@ -8,6 +8,10 @@ __all__ = [ "NewtonDeformableBodyMaterialCfg", "NewtonDeformableMaterialCfg", "NewtonSurfaceDeformableBodyMaterialCfg", + "MPMGridCfg", + "MPMParticleMaterialCfg", + "MPMParticleSpawnerCfg", + "MPMPointsCfg", "schemas", "spawners", "views", @@ -20,3 +24,4 @@ from .spawners.materials import ( NewtonDeformableMaterialCfg, NewtonSurfaceDeformableBodyMaterialCfg, ) +from .spawners.mpm import MPMGridCfg, MPMParticleMaterialCfg, MPMParticleSpawnerCfg, MPMPointsCfg diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sim/spawners/__init__.pyi index ad99e610b437..460b0a0c28cc 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/spawners/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/__init__.pyi @@ -4,7 +4,8 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "mpm", "materials", ] -from . import materials +from . import materials, mpm diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.py new file mode 100644 index 000000000000..b228122d86b0 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Newton MPM particle spawner utilities.""" + +from isaaclab.utils.module import lazy_export + +lazy_export() diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.pyi new file mode 100644 index 000000000000..a85d97f37875 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/__init__.pyi @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "MPMGridCfg", + "MPMParticleMaterialCfg", + "MPMParticleSpawnerCfg", + "MPMPointsCfg", + "create_mpm_particle_visualization", + "emit_mpm_particles", + "spawn_mpm_particles", +] + +from .mpm import emit_mpm_particles, spawn_mpm_particles +from .mpm_cfg import MPMGridCfg, MPMParticleMaterialCfg, MPMParticleSpawnerCfg, MPMPointsCfg +from .visualization import create_mpm_particle_visualization diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm.py new file mode 100644 index 000000000000..571f8dafa1c0 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import warp as wp + +from isaaclab.sim.utils import clone, create_prim + +if TYPE_CHECKING: + from pxr import Usd + +from .mpm_cfg import MPMGridCfg, MPMParticleMaterialCfg, MPMParticleSpawnerCfg, MPMPointsCfg + + +@clone +def spawn_mpm_particles( + prim_path: str, + cfg: MPMParticleSpawnerCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Create a lightweight placeholder prim for a Newton MPM particle object. + + MPM particles are inserted directly into the Newton model builder during + Newton replication. The USD prim exists so Isaac Lab's normal asset spawning + and clone-planning machinery can reason about the scene entity. + """ + return create_prim(prim_path, prim_type="Xform", translation=translation, orientation=orientation) + + +def _material_custom_attributes(material: MPMParticleMaterialCfg) -> dict[str, float]: + """Map material cfg values to Newton ``mpm:*`` custom attributes for ``add_particles``. + + ``density`` is intentionally absent: it is consumed by the particle + generators to derive per-particle mass, not forwarded to the solver. + """ + return { + "mpm:young_modulus": float(material.young_modulus), + "mpm:poisson_ratio": float(material.poisson_ratio), + "mpm:viscosity": float(material.viscosity), + "mpm:friction": float(material.friction), + "mpm:damping": float(material.damping), + "mpm:yield_pressure": float(material.yield_pressure), + "mpm:tensile_yield_ratio": float(material.tensile_yield_ratio), + "mpm:yield_stress": float(material.yield_stress), + "mpm:hardening": float(material.hardening), + "mpm:dilatancy": float(material.dilatancy), + } + + +def emit_mpm_particles( + builder, + cfg: MPMParticleSpawnerCfg, + *, + position: tuple[float, float, float], + orientation: tuple[float, float, float, float], +) -> None: + """Emit particles described by ``cfg`` into a Newton ``ModelBuilder``.""" + if isinstance(cfg, MPMGridCfg): + _emit_grid(builder, cfg, position=position, orientation=orientation) + elif isinstance(cfg, MPMPointsCfg): + _emit_points(builder, cfg, position=position, orientation=orientation) + else: + raise TypeError(f"Unsupported MPM particle spawner config type: {type(cfg).__name__}") + + +def _emit_grid( + builder, + cfg: MPMGridCfg, + *, + position: tuple[float, float, float], + orientation: tuple[float, float, float, float], +) -> None: + lower = np.asarray(cfg.lower, dtype=np.float32) + upper = np.asarray(cfg.upper, dtype=np.float32) + extent = upper - lower + if np.any(extent <= 0.0): + raise ValueError(f"MPMGridCfg upper corner must be greater than lower corner. Got {cfg.lower=} {cfg.upper=}.") + if cfg.voxel_size <= 0.0: + raise ValueError(f"MPMGridCfg voxel_size must be positive. Got {cfg.voxel_size}.") + if cfg.particles_per_cell <= 0.0: + raise ValueError(f"MPMGridCfg particles_per_cell must be positive. Got {cfg.particles_per_cell}.") + + resolution = np.maximum(np.ceil(cfg.particles_per_cell * extent / cfg.voxel_size), 1).astype(np.int32) + cell_size = extent / resolution + cell_volume = float(np.prod(cell_size)) + mass = float(cfg.mass) if cfg.mass is not None else cell_volume * float(cfg.material.density) + radius = float(cfg.radius) if cfg.radius is not None else 0.5 * float(np.max(cell_size)) + + world_pos = _transform_point(lower, position, orientation) + builder.add_particle_grid( + pos=wp.vec3(*world_pos.tolist()), + rot=wp.quat(*orientation), + vel=wp.vec3(0.0, 0.0, 0.0), + dim_x=int(resolution[0]) + 1, + dim_y=int(resolution[1]) + 1, + dim_z=int(resolution[2]) + 1, + cell_x=float(cell_size[0]), + cell_y=float(cell_size[1]), + cell_z=float(cell_size[2]), + mass=mass, + jitter=float(cfg.jitter), + radius_mean=radius, + custom_attributes=_material_custom_attributes(cfg.material), + ) + + +def _emit_points( + builder, + cfg: MPMPointsCfg, + *, + position: tuple[float, float, float], + orientation: tuple[float, float, float, float], +) -> None: + points = np.asarray(cfg.positions, dtype=np.float32) + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError(f"MPMPointsCfg positions must have shape (N, 3). Got {points.shape}.") + if points.shape[0] == 0: + raise ValueError("MPMPointsCfg positions must contain at least one particle.") + + velocities = np.zeros_like(points) if cfg.velocities is None else np.asarray(cfg.velocities, dtype=np.float32) + if velocities.shape != points.shape: + raise ValueError(f"MPMPointsCfg velocities must match positions shape {points.shape}. Got {velocities.shape}.") + + world_points = _transform_points(points, position, orientation) + world_velocities = _rotate_vectors(velocities, orientation) + + builder.add_particles( + pos=world_points.tolist(), + vel=world_velocities.tolist(), + mass=_expand_scalar_or_sequence(cfg.mass, points.shape[0], "mass"), + radius=_expand_scalar_or_sequence(cfg.radius, points.shape[0], "radius"), + custom_attributes=_material_custom_attributes(cfg.material), + ) + + +def _expand_scalar_or_sequence(value: float | Sequence[float], count: int, name: str) -> list[float]: + if isinstance(value, (int, float)): + return [float(value)] * count + if len(value) != count: + raise ValueError(f"MPMPointsCfg {name} must be scalar or have one value per particle. Got {len(value)} values.") + return [float(v) for v in value] + + +def _transform_point( + point: np.ndarray, + position: tuple[float, float, float], + orientation: tuple[float, float, float, float], +) -> np.ndarray: + return _transform_points(point.reshape(1, 3), position, orientation)[0] + + +def _transform_points( + points: np.ndarray, + position: tuple[float, float, float], + orientation: tuple[float, float, float, float], +) -> np.ndarray: + return _rotate_vectors(points, orientation) + np.asarray(position, dtype=np.float32) + + +def _rotate_vectors(vectors: np.ndarray, orientation: tuple[float, float, float, float]) -> np.ndarray: + q = np.asarray(orientation, dtype=np.float32) + q_vec = q[:3] + q_w = q[3] + uv = np.cross(q_vec, vectors) + uuv = np.cross(q_vec, uv) + return vectors + 2.0 * (q_w * uv + uuv) diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py new file mode 100644 index 000000000000..4bbeb8fd8c39 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py @@ -0,0 +1,116 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import MISSING + +from isaaclab.sim.spawners.spawner_cfg import SpawnerCfg +from isaaclab.utils.configclass import configclass + + +@configclass +class MPMParticleMaterialCfg: + """Per-particle material values consumed by Newton's implicit MPM solver. + + This is a lightweight value config. It does not create or bind a USD material + prim; values are forwarded to Newton as ``mpm:*`` custom attributes when + particles are added to the model builder. + + The defaults model a dry sand-like granular material. + """ + + density: float = 1000.0 + """Particle material density [kg/m^3], used to derive particle mass when a generator does not override it.""" + + young_modulus: float = 1.0e15 + """Young's modulus [Pa].""" + + poisson_ratio: float = 0.3 + """Poisson ratio.""" + + viscosity: float = 0.0 + """Viscosity coefficient.""" + + friction: float = 0.68 + """Particle friction coefficient.""" + + damping: float = 0.0 + """Material damping.""" + + yield_pressure: float = 1.0e12 + """Pressure at which the material yields.""" + + tensile_yield_ratio: float = 0.0 + """Tensile/compressive yield ratio.""" + + yield_stress: float = 0.0 + """Von-Mises yield stress.""" + + hardening: float = 0.0 + """Plastic hardening coefficient.""" + + dilatancy: float = 0.0 + """Granular dilatancy coefficient.""" + + +@configclass +class MPMParticleSpawnerCfg(SpawnerCfg): + """Base configuration for declarative Newton MPM particle generation.""" + + func: Callable | str = "{DIR}.mpm:spawn_mpm_particles" + + material: MPMParticleMaterialCfg = MPMParticleMaterialCfg() + """Material values applied to generated particles.""" + + visual_color: Sequence[float] = (0.7, 0.6, 0.4) + """Display color for Kit particle visualization.""" + + visual_update_frequency: int = 1 + """Kit particle visualization update frequency in render frames.""" + + +@configclass +class MPMGridCfg(MPMParticleSpawnerCfg): + """Generate a regular MPM particle lattice inside an axis-aligned local box.""" + + lower: Sequence[float] = MISSING + """Lower local-space corner of the particle box [m].""" + + upper: Sequence[float] = MISSING + """Upper local-space corner of the particle box [m].""" + + voxel_size: float = MISSING + """Target MPM voxel size [m], used with :attr:`particles_per_cell` to choose lattice resolution.""" + + particles_per_cell: float = 1.0 + """Particle lattice density relative to the MPM grid resolution.""" + + jitter: float = 0.0 + """Newton particle-grid jitter value [m].""" + + mass: float | None = None + """Per-particle mass [kg]. If ``None``, mass is derived from cell volume and material density.""" + + radius: float | None = None + """Particle radius [m]. If ``None``, radius is half the largest generated cell size.""" + + +@configclass +class MPMPointsCfg(MPMParticleSpawnerCfg): + """Generate MPM particles from explicit local-space point positions.""" + + positions: Sequence[Sequence[float]] = MISSING + """Local-space particle positions [m].""" + + velocities: Sequence[Sequence[float]] | None = None + """Optional local-space particle velocities [m/s]. If ``None``, velocities are zero.""" + + mass: float | Sequence[float] = 1.0 + """Particle mass values [kg], either scalar or one value per particle.""" + + radius: float | Sequence[float] = 0.01 + """Particle radius values [m], either scalar or one value per particle.""" diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py new file mode 100644 index 000000000000..b63bda8e31d2 --- /dev/null +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py @@ -0,0 +1,53 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + + +def create_mpm_particle_visualization( + prim_path: str, + positions: np.ndarray, + widths: np.ndarray, + color: Sequence[float], +) -> list[str]: + """Create one ``UsdGeom.Points`` prim per environment for Kit MPM particle rendering. + + The created prims are static USD containers: per-frame position updates are + handled by :meth:`isaaclab_newton.physics.NewtonManager.sync_particles_to_usd` + for prims registered via + :meth:`isaaclab_newton.physics.NewtonManager.register_particle_visual_prim`. + + Args: + prim_path: Base prim path; one ``Points`` prim is created per environment + at ``{prim_path}/env_{idx}``. + positions: Initial world-frame particle positions [m], shape + ``(num_envs, particles_per_env, 3)``. + widths: Particle display widths (diameters) [m], one per particle. + color: RGB display color of the particles. + + Returns: + The created ``Points`` prim paths, one per environment. + """ + from pxr import Gf, Sdf, UsdGeom, Vt # noqa: PLC0415 + + import isaaclab.sim as sim_utils + + stage = sim_utils.get_current_stage() + prim_paths = [f"{prim_path}/env_{env_idx}" for env_idx in range(positions.shape[0])] + points_prims = [UsdGeom.Points.Define(stage, path) for path in prim_paths] + + widths_vt = Vt.FloatArray.FromNumpy(np.ascontiguousarray(widths, dtype=np.float32)) + color_vt = Vt.Vec3fArray([Gf.Vec3f(*(float(value) for value in color))]) + with Sdf.ChangeBlock(): + for env_idx, points in enumerate(points_prims): + points.GetPointsAttr().Set(Vt.Vec3fArray.FromNumpy(positions[env_idx])) + points.CreateWidthsAttr(widths_vt) + points.CreateDisplayColorAttr(color_vt) + + return prim_paths diff --git a/source/isaaclab_newton/test/assets/test_mpm_object.py b/source/isaaclab_newton/test/assets/test_mpm_object.py new file mode 100644 index 000000000000..646cc6e8dfc5 --- /dev/null +++ b/source/isaaclab_newton/test/assets/test_mpm_object.py @@ -0,0 +1,269 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import math + +import numpy as np +import pytest +import torch + +newton = pytest.importorskip("newton") + +from isaaclab_newton.assets.mpm_object import MPMObject, MPMObjectCfg +from isaaclab_newton.assets.mpm_object.mpm_object import MPMObjectRegistryEntry, add_mpm_entry_to_builder +from isaaclab_newton.physics import MPMSolverCfg, NewtonCfg, NewtonMPMManager +from isaaclab_newton.sim.spawners.mpm import MPMGridCfg, MPMParticleMaterialCfg, MPMPointsCfg + +from isaaclab.assets import RigidObjectCfg +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.sim import SimulationCfg, build_simulation_context +from isaaclab.utils.configclass import configclass + + +def test_mpm_particle_material_emits_custom_attributes(): + """MPM materials are value cfgs forwarded as Newton custom attributes, not USD material spawners.""" + from isaaclab_newton.sim.spawners.mpm.mpm import _material_custom_attributes + + attrs = _material_custom_attributes(MPMParticleMaterialCfg(viscosity=0.1)) + + assert attrs["mpm:friction"] == pytest.approx(0.68) + assert attrs["mpm:viscosity"] == pytest.approx(0.1) + assert "density" not in attrs + + +def test_mpm_object_cfg_resolves_asset_class(): + cfg = MPMObjectCfg( + prim_path="/World/envs/env_.*/Sand", + spawn=MPMGridCfg(lower=(0.0, 0.0, 0.0), upper=(0.1, 0.1, 0.1), voxel_size=0.1), + ) + + assert cfg.class_type.__name__ == MPMObject.__name__ + + +def test_mpm_grid_emission_records_constant_offsets_per_env(): + builder = newton.ModelBuilder() + NewtonMPMManager._register_builder_attributes(builder) + + cfg = MPMObjectCfg( + prim_path="/World/envs/env_.*/Sand", + spawn=MPMGridCfg( + lower=(0.0, 0.0, 0.0), + upper=(0.1, 0.1, 0.1), + voxel_size=0.1, + particles_per_cell=1.0, + jitter=0.0, + ), + ) + entry = MPMObjectRegistryEntry(cfg) + + add_mpm_entry_to_builder(builder, entry, 0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) + add_mpm_entry_to_builder(builder, entry, 1, [1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) + + assert entry.particles_per_object == 8 + assert entry.particle_offsets == [0, 8] + assert builder.particle_count == 16 + + +def test_mpm_points_emission_records_constant_offsets_per_env(): + builder = newton.ModelBuilder() + NewtonMPMManager._register_builder_attributes(builder) + + cfg = MPMObjectCfg( + prim_path="/World/envs/env_.*/Fluid", + spawn=MPMPointsCfg( + positions=((0.0, 0.0, 0.0), (0.0, 0.0, 0.1), (0.0, 0.1, 0.0)), + velocities=((0.0, 0.0, 0.0), (0.0, 0.0, 0.1), (0.0, 0.1, 0.0)), + mass=0.01, + radius=0.02, + material=MPMParticleMaterialCfg(viscosity=0.1, friction=0.0), + ), + ) + entry = MPMObjectRegistryEntry(cfg) + + add_mpm_entry_to_builder(builder, entry, 0, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]) + add_mpm_entry_to_builder(builder, entry, 1, [0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]) + + assert entry.particles_per_object == 3 + assert entry.particle_offsets == [0, 3] + assert builder.particle_count == 6 + + +def test_mpm_object_initializes_from_interactive_scene(): + @configclass + class MPMSceneCfg(InteractiveSceneCfg): + media = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Sand", + spawn=MPMGridCfg(lower=(0.0, 0.0, 0.0), upper=(0.1, 0.1, 0.1), voxel_size=0.1), + ) + + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg(solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), use_cuda_graph=False), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + scene = InteractiveScene(MPMSceneCfg(num_envs=2, env_spacing=1.0)) + sim.reset() + + media = scene["media"] + assert media.num_instances == 2 + assert media.particles_per_object == 8 + assert media.data.particle_pos_w.torch.shape == (2, 8, 3) + + default_state = media.data.default_particle_state_w.torch.clone() + shifted_state = default_state[0:1].clone() + shifted_state[..., 2] += 0.05 + + media.write_particle_state_to_sim_index( + shifted_state, + env_ids=torch.tensor([0], device=sim.device, dtype=torch.int32), + ) + torch.testing.assert_close(media.data.particle_state_w.torch[0:1], shifted_state) + + media.reset(env_ids=[0]) + torch.testing.assert_close(media.data.particle_state_w.torch[0], default_state[0]) + + +def test_mpm_solver_refreshes_kinematic_rigid_body_transforms(): + import isaaclab.sim as sim_utils # noqa: PLC0415 + + @configclass + class MPMSceneCfg(InteractiveSceneCfg): + collider = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/KinematicBox", + spawn=sim_utils.CuboidCfg( + size=(0.1, 0.1, 0.1), + rigid_props=sim_utils.NewtonRigidBodyPropertiesCfg( + rigid_body_enabled=True, + kinematic_enabled=True, + disable_gravity=True, + ), + collision_props=sim_utils.NewtonCollisionPropertiesCfg(collision_enabled=True), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)), + ) + media = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Sand", + spawn=MPMGridCfg(lower=(-0.05, -0.05, 0.3), upper=(0.05, 0.05, 0.4), voxel_size=0.05), + ) + + sim_cfg = SimulationCfg( + dt=1.0 / 60.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg(solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), use_cuda_graph=False), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + scene = InteractiveScene(MPMSceneCfg(num_envs=1, env_spacing=0.0)) + sim.reset() + + collider = scene["collider"] + angle = 0.5 + root_pose = torch.tensor( + [[0.1, 0.0, 0.25, 0.0, math.sin(0.5 * angle), 0.0, math.cos(0.5 * angle)]], + dtype=torch.float32, + device=collider.device, + ) + collider.write_root_link_pose_to_sim_index(root_pose=root_pose) + sim.step(render=False) + + body_labels = list(NewtonMPMManager.get_model().body_label) + body_idx = body_labels.index("/World/envs/env_0/KinematicBox") + body_q = NewtonMPMManager.get_state_0().body_q.numpy()[body_idx] + + np.testing.assert_allclose(body_q, root_pose.detach().cpu().numpy()[0], rtol=1.0e-5, atol=1.0e-6) + + +def test_mpm_object_creates_kit_points_when_kit_visualizer_requested(monkeypatch): + @configclass + class MPMSceneCfg(InteractiveSceneCfg): + media = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Sand", + spawn=MPMGridCfg( + lower=(0.0, 0.0, 0.0), + upper=(0.1, 0.1, 0.1), + voxel_size=0.1, + visual_color=(0.1, 0.2, 0.3), + ), + ) + + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg(solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), use_cuda_graph=False), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + monkeypatch.setattr(sim, "resolve_visualizer_types", lambda: ["kit"]) + scene = InteractiveScene(MPMSceneCfg(num_envs=2, env_spacing=1.0)) + sim.reset() + + from pxr import UsdGeom # noqa: PLC0415 + + media = scene["media"] + records = NewtonMPMManager._particle_visual_prims + assert len(records) == media.num_instances + + for env_idx, (prim_path, record) in enumerate(sorted(records.items())): + assert record.offset == media._recorded_particle_offsets[env_idx] + assert record.count == media.particles_per_object + assert record.sync_frequency == 1 + + points_prim = media.stage.GetPrimAtPath(prim_path) + assert points_prim.IsValid() + points = UsdGeom.Points(points_prim) + assert len(points.GetPointsAttr().Get()) == media.particles_per_object + assert len(points.GetWidthsAttr().Get()) == media.particles_per_object + assert tuple(points.GetDisplayColorAttr().Get()[0]) == pytest.approx((0.1, 0.2, 0.3)) + + +def test_mpm_kit_points_follow_particle_state(monkeypatch): + @configclass + class MPMSceneCfg(InteractiveSceneCfg): + media = MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Sand", + spawn=MPMGridCfg( + lower=(0.0, 0.0, 0.1), + upper=(0.1, 0.1, 0.2), + voxel_size=0.05, + visual_color=(0.1, 0.2, 0.3), + ), + ) + + sim_cfg = SimulationCfg( + dt=1.0 / 60.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg(solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), use_cuda_graph=False), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + monkeypatch.setattr(sim, "resolve_visualizer_types", lambda: ["kit"]) + scene = InteractiveScene(MPMSceneCfg(num_envs=1, env_spacing=0.0)) + sim.reset() + + from pxr import UsdGeom # noqa: PLC0415 + + media = scene["media"] + prim_path = next(iter(NewtonMPMManager._particle_visual_prims)) + points = UsdGeom.Points(media.stage.GetPrimAtPath(prim_path)) + points_before = np.asarray(points.GetPointsAttr().Get(), dtype=np.float32) + + for _ in range(3): + sim.step(render=False) + scene.update(sim.get_physics_dt()) + sim.render() + + points_after = np.asarray(points.GetPointsAttr().Get(), dtype=np.float32) + particle_pos = media.data.particle_pos_w.torch.detach().cpu().numpy()[0] + + assert np.max(np.abs(points_after - points_before)) > 0.0 + np.testing.assert_allclose(points_after, particle_pos, rtol=1.0e-5, atol=1.0e-6) diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index 8458f466632c..991a5f63da89 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -25,22 +25,29 @@ from __future__ import annotations +from types import SimpleNamespace + +import numpy as np import pytest +import warp as wp from isaaclab_newton.physics import ( FeatherstoneSolverCfg, KaminoSolverCfg, MJWarpSolverCfg, + MPMSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg, NewtonFeatherstoneManager, NewtonKaminoManager, NewtonManager, NewtonMJWarpManager, + NewtonMPMManager, NewtonSolverCfg, NewtonXPBDManager, XPBDSolverCfg, ) -from newton.solvers import SolverFeatherstone, SolverKamino, SolverMuJoCo, SolverXPBD +from isaaclab_newton.physics.mpm_manager import _make_solver_config +from newton.solvers import SolverFeatherstone, SolverImplicitMPM, SolverKamino, SolverMuJoCo, SolverXPBD from isaaclab.sim import SimulationCfg, build_simulation_context @@ -99,6 +106,14 @@ True, id="kamino_newton_pipeline", ), + pytest.param( + lambda: MPMSolverCfg(max_iterations=2, voxel_size=0.05), + NewtonMPMManager, + SolverImplicitMPM, + True, + False, + id="implicit_mpm", + ), ] @@ -157,13 +172,268 @@ def test_newton_cfg_collision_decimation_warning(num_substeps, collision_decimat assert cfg.collision_decimation == collision_decimation +def test_mpm_solver_cfg_maps_only_newton_solver_fields(): + """MPM config forwarding ignores Isaac Lab metadata fields explicitly.""" + + solver_cfg = MPMSolverCfg( + max_iterations=7, + voxel_size=0.04, + solver_type="isaaclab_metadata_should_not_forward", + ) + + newton_cfg = _make_solver_config(solver_cfg) + + assert newton_cfg.max_iterations == 7 + assert newton_cfg.voxel_size == 0.04 + assert not hasattr(newton_cfg, "class_type") + assert not hasattr(newton_cfg, "solver_type") + # Manager-level stepping option must not leak into the Newton solver config. + assert not hasattr(newton_cfg, "project_outside_colliders") + + +# Tuples of ``(field_name, non_default_value)`` covering every solver-tunable +# field on :class:`MPMSolverCfg`. Each entry exercises the implementation-side +# SolverImplicitMPM.Config construction so a Newton field rename or accidental +# drop is caught here instead of silently producing wrong-physics runs. +_MPM_FIELD_VALUES = [ + ("max_iterations", 13), + ("tolerance", 5.0e-5), + ("solver", "gauss-seidel"), + ("warmstart_mode", "particles"), + ("collider_velocity_mode", "backward"), + ("voxel_size", 0.0375), + ("grid_type", "dense"), + ("grid_padding", 4), + ("max_active_cell_count", 1024), + ("transfer_scheme", "pic"), + ("integration_scheme", "gimp"), + ("critical_fraction", 0.25), + ("air_drag", 0.5), + ("collider_normal_from_sdf_gradient", True), + ("collider_basis", "Q1"), + ("strain_basis", "P1d"), + ("velocity_basis", "B2"), +] + + +@pytest.mark.parametrize("field_name, value", _MPM_FIELD_VALUES) +def test_mpm_solver_cfg_forwards_every_solver_field(field_name, value): + """Every tunable MPM cfg field round-trips into ``SolverImplicitMPM.Config``. + + Guards against MPM manager construction dropping or mis-naming a field if + Newton's config surface changes. + """ + solver_cfg = MPMSolverCfg(**{field_name: value}) + newton_cfg = _make_solver_config(solver_cfg) + assert hasattr(newton_cfg, field_name), ( + f"{field_name!r} disappeared from SolverImplicitMPM.Config — MPMSolverCfg needs to drop or rename it." + ) + assert getattr(newton_cfg, field_name) == value + + +def test_mpm_register_builder_attributes_is_idempotent(): + """The MPM custom-attribute hook is a no-op when attributes are already registered.""" + import newton + + builder = newton.ModelBuilder() + assert not builder.has_custom_attribute("mpm:young_modulus") + + NewtonMPMManager._register_builder_attributes(builder) + assert builder.has_custom_attribute("mpm:young_modulus") + + # Second call must be a no-op (no exceptions, attribute still present). + NewtonMPMManager._register_builder_attributes(builder) + assert builder.has_custom_attribute("mpm:young_modulus") + + +def test_mpm_prepare_builder_makes_kinematic_bodies_massless(): + """Kinematic bodies must be massless so MPM treats them as kinematic colliders.""" + import newton + + builder = newton.ModelBuilder() + kinematic_body = builder.add_body( + mass=0.35, + inertia=wp.mat33(1.0), + is_kinematic=True, + label="kinematic_collider", + ) + dynamic_body = builder.add_body( + mass=1.2, + inertia=wp.mat33(2.0), + is_kinematic=False, + label="dynamic_body", + ) + + NewtonMPMManager._prepare_builder_for_finalize(builder) + + assert builder.body_flags[kinematic_body] & int(newton.BodyFlags.KINEMATIC) + assert builder.body_mass[kinematic_body] == 0.0 + assert builder.body_inv_mass[kinematic_body] == 0.0 + assert np.allclose(np.array(builder.body_inertia[kinematic_body]), 0.0) + assert np.allclose(np.array(builder.body_inv_inertia[kinematic_body]), 0.0) + + assert builder.body_mass[dynamic_body] == pytest.approx(1.2) + assert builder.body_inv_mass[dynamic_body] == pytest.approx(1.0 / 1.2) + assert np.allclose(np.array(builder.body_inertia[dynamic_body]), 2.0) + + +def test_active_manager_create_builder_registers_mpm_attributes(): + """The active MPM manager registers solver-specific builder attributes.""" + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg(solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), use_cuda_graph=False), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + builder = sim.physics_manager.create_builder() + + assert builder.has_custom_attribute("mpm:young_modulus") + + +def test_mpm_end_to_end_with_particle_custom_attributes(): + """End-to-end MPM step using ``add_particles(custom_attributes=...)`` — the production path.""" + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg( + solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05), + use_cuda_graph=False, + ), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + builder = sim.physics_manager.create_builder() + # MPM custom attrs must exist on the builder before particles use them. + assert builder.has_custom_attribute("mpm:young_modulus") + + positions = [(0.0, 0.0, 0.10), (0.05, 0.0, 0.10), (0.0, 0.05, 0.10)] + builder.add_particles( + pos=positions, + vel=[(0.0, 0.0, 0.0)] * len(positions), + mass=[0.01] * len(positions), + radius=[0.02] * len(positions), + custom_attributes={ + "mpm:viscosity": 50.0, + "mpm:friction": 0.0, + "mpm:tensile_yield_ratio": 1.0, + "mpm:yield_pressure": 1.0e15, + "mpm:yield_stress": 0.0, + "mpm:young_modulus": 1.0e15, + "mpm:damping": 0.0, + }, + ) + NewtonManager.set_builder(builder) + + sim.reset() + assert isinstance(NewtonManager._solver, SolverImplicitMPM) + sim.step(render=False) + + +@pytest.mark.parametrize("project_outside", [True, False]) +def test_mpm_project_outside_colliders_gates_projection(project_outside): + """``project_outside_colliders`` controls whether ``project_outside`` runs per substep. + + Wraps the solver's ``project_outside`` with a counter after ``sim.reset()`` + (``use_cuda_graph=False`` keeps the Python callable on the step path) and + runs one tick. The call count is positive only when the flag is set. + """ + sim_cfg = SimulationCfg( + dt=1.0 / 120.0, + device="cuda:0", + gravity=(0.0, 0.0, -9.81), + physics=NewtonCfg( + solver_cfg=MPMSolverCfg(max_iterations=2, voxel_size=0.05, project_outside_colliders=project_outside), + use_cuda_graph=False, + ), + ) + + with build_simulation_context(sim_cfg=sim_cfg) as sim: + builder = sim.physics_manager.create_builder() + builder.add_particles( + pos=[(0.0, 0.0, 0.10), (0.05, 0.0, 0.10), (0.0, 0.05, 0.10)], + vel=[(0.0, 0.0, 0.0)] * 3, + mass=[0.01] * 3, + radius=[0.02] * 3, + custom_attributes={ + "mpm:viscosity": 50.0, + "mpm:friction": 0.0, + "mpm:tensile_yield_ratio": 1.0, + "mpm:yield_pressure": 1.0e15, + "mpm:yield_stress": 0.0, + "mpm:young_modulus": 1.0e15, + "mpm:damping": 0.0, + }, + ) + NewtonManager.set_builder(builder) + sim.reset() + + calls = {"n": 0} + original_project = NewtonManager._solver.project_outside + + def counting_project(*args, **kwargs): + calls["n"] += 1 + return original_project(*args, **kwargs) + + NewtonManager._solver.project_outside = counting_project + try: + sim.step(render=False) + finally: + NewtonManager._solver.project_outside = original_project + + if project_outside: + assert calls["n"] >= 1 + else: + assert calls["n"] == 0 + + +@pytest.mark.parametrize( + "grid_type, expected", + [ + ("fixed", True), + ("sparse", False), + ("dense", False), + ], +) +def test_mpm_cuda_graph_capture_supports_only_fixed_grid(monkeypatch, grid_type, expected): + """Newton implicit MPM is CUDA-graph capturable only with a fixed grid.""" + + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type=grid_type), raising=False) + + assert NewtonMPMManager._supports_cuda_graph_capture() is expected + + +def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): + """Sparse/dense MPM should not enter a CUDA graph capture window.""" + from isaaclab.physics import PhysicsManager + + monkeypatch.setattr( + PhysicsManager, + "_cfg", + NewtonCfg(solver_cfg=MPMSolverCfg(grid_type="sparse"), use_cuda_graph=True), + raising=False, + ) + monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type="sparse"), raising=False) + monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", True, raising=False) + + NewtonMPMManager._capture_or_defer_graph() + + assert NewtonManager._graph is None + assert NewtonManager._graph_capture_pending is False + + # --------------------------------------------------------------------------- # Manager class hierarchy and factory contracts # --------------------------------------------------------------------------- @pytest.mark.parametrize( - "manager", [NewtonMJWarpManager, NewtonXPBDManager, NewtonFeatherstoneManager, NewtonKaminoManager] + "manager", + [NewtonMJWarpManager, NewtonXPBDManager, NewtonFeatherstoneManager, NewtonKaminoManager, NewtonMPMManager], ) def test_subclass_of_newton_manager(manager): """All concrete managers inherit from :class:`NewtonManager`.""" @@ -179,7 +449,8 @@ def test_abstract_build_solver_raises(): @pytest.mark.parametrize( - "manager", [NewtonMJWarpManager, NewtonXPBDManager, NewtonFeatherstoneManager, NewtonKaminoManager] + "manager", + [NewtonMJWarpManager, NewtonXPBDManager, NewtonFeatherstoneManager, NewtonKaminoManager, NewtonMPMManager], ) def test_manager_name_starts_with_newton(manager): """The ``"newton"`` prefix is required by :class:`InteractiveScene` and the @@ -215,12 +486,15 @@ def test_initialize_solver_populates_canonical_state( working regardless of which leaf is active. This test is the regression guard for that contract. - The builder is pre-populated with a minimal one-body / one-joint scene - (instead of relying on a USD stage) for two reasons: + The builder is pre-populated directly (instead of relying on a USD stage) + with either a minimal particle grid for MPM or a one-body / one-joint scene + for rigid/articulation solvers: - 1. :class:`SolverMuJoCo` requires at least one joint to convert the model + 1. :class:`SolverImplicitMPM` requires particles and MPM custom attributes + registered on the builder before particle creation. + 2. :class:`SolverMuJoCo` requires at least one joint to convert the model to MJCF; a ground-plane-only scene fails MJCF conversion. - 2. Pre-populating ``NewtonManager._builder`` causes + 3. Pre-populating ``NewtonManager._builder`` causes :meth:`NewtonManager.start_simulation` to skip :meth:`instantiate_builder_from_stage`, so the test does not depend on USD asset packages. @@ -240,11 +514,28 @@ def test_initialize_solver_populates_canonical_state( assert resolved_manager.__name__ == expected_manager.__name__ assert resolved_manager.__name__.lower().startswith("newton") - # Pre-populate the builder with a minimal scene so MJCF conversion has - # something to work with. - builder = NewtonManager.create_builder() - body = builder.add_body(mass=1.0) - builder.add_joint_revolute(parent=-1, child=body, axis=(0, 0, 1)) + builder = resolved_manager.create_builder() + if expected_solver_cls is SolverImplicitMPM: + assert builder.has_custom_attribute("mpm:young_modulus") + builder.add_particle_grid( + pos=wp.vec3(-0.05, -0.05, 0.10), + rot=wp.quat_identity(), + vel=wp.vec3(0.0), + dim_x=2, + dim_y=2, + dim_z=2, + cell_x=0.05, + cell_y=0.05, + cell_z=0.05, + mass=0.01, + jitter=0.0, + radius_mean=0.02, + ) + else: + # Pre-populate the builder with a minimal scene so MJCF conversion has + # something to work with. + body = builder.add_body(mass=1.0) + builder.add_joint_revolute(parent=-1, child=body, axis=(0, 0, 1)) NewtonManager.set_builder(builder) # Force resolution and bring up the solver. @@ -258,8 +549,8 @@ def test_initialize_solver_populates_canonical_state( # ``_contacts`` is allocated whichever way contacts are handled # (MuJoCo internal buffer or Newton pipeline output). - # Kamino with internal contacts does not currently set NewtonManager._contacts. - if expected_solver_cls is not SolverKamino: + # Kamino with internal contacts and MPM do not currently set NewtonManager._contacts. + if expected_solver_cls not in (SolverKamino, SolverImplicitMPM): assert NewtonManager._contacts is not None # One step should not raise — proves the dispatch wiring lines up @@ -288,7 +579,7 @@ def test_mjwarp_internal_contacts_with_collision_cfg_raises(): ) with build_simulation_context(sim_cfg=sim_cfg) as sim: - builder = NewtonManager.create_builder() + builder = sim.physics_manager.create_builder() body = builder.add_body(mass=1.0) builder.add_joint_revolute(parent=-1, child=body, axis=(0, 0, 1)) NewtonManager.set_builder(builder) @@ -331,7 +622,7 @@ def test_collision_decimation_invokes_mid_loop_collide(num_substeps, collision_d ) with build_simulation_context(sim_cfg=sim_cfg) as sim: - builder = NewtonManager.create_builder() + builder = sim.physics_manager.create_builder() body = builder.add_body(mass=1.0) builder.add_joint_free(child=body) builder.add_shape_sphere(body=body, radius=0.05) diff --git a/source/isaaclab_newton/test/sim/test_mpm_spawners.py b/source/isaaclab_newton/test/sim/test_mpm_spawners.py new file mode 100644 index 000000000000..d87995aedf21 --- /dev/null +++ b/source/isaaclab_newton/test/sim/test_mpm_spawners.py @@ -0,0 +1,101 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import subprocess +import sys +import textwrap + + +def test_mpm_config_imports_do_not_load_pxr(): + code = textwrap.dedent( + """ + import sys + + from isaaclab_newton.assets.mpm_object import MPMObjectCfg + from isaaclab_newton.sim.spawners.mpm import MPMGridCfg, MPMParticleMaterialCfg, MPMPointsCfg + + MPMObjectCfg( + prim_path="/World/envs/env_.*/Sand", + spawn=MPMGridCfg( + lower=(0.0, 0.0, 0.0), + upper=(0.1, 0.1, 0.1), + voxel_size=0.1, + material=MPMParticleMaterialCfg(), + ), + ) + + loaded_pxr_modules = [module for module in sys.modules if module == "pxr" or module.startswith("pxr.")] + if loaded_pxr_modules: + raise SystemExit("pxr loaded before SimulationApp: " + ", ".join(loaded_pxr_modules[:20])) + """ + ) + + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stdout + result.stderr + + +def test_granular_demo_does_not_load_pxr_before_app_launcher(): + code = textwrap.dedent( + """ + import sys + + sys.argv = ["newton_mpm_granular.py", "--max-steps", "1", "--visualizer", "kit"] + + import scripts.demos.mpm.newton_mpm_granular as demo + from isaaclab.app.app_launcher import AppLauncher + + def stop_before_simulation_app(self): + loaded_pxr_modules = [module for module in sys.modules if module == "pxr" or module.startswith("pxr.")] + if loaded_pxr_modules: + raise SystemExit("pxr loaded before SimulationApp: " + ", ".join(loaded_pxr_modules[:20])) + raise RuntimeError("stop before SimulationApp") + + AppLauncher._create_app = stop_before_simulation_app + + try: + demo.main() + except RuntimeError as exc: + if str(exc) != "stop before SimulationApp": + raise + """ + ) + + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stdout + result.stderr + + +def test_particle_pour_demo_does_not_load_pxr_before_app_launcher(): + code = textwrap.dedent( + """ + import sys + + sys.argv = ["particle_pour.py", "--max-steps", "1", "--visualizer", "kit"] + + import scripts.demos.mpm.particle_pour as demo + from isaaclab.app.app_launcher import AppLauncher + + def stop_before_simulation_app(self): + loaded_pxr_modules = [module for module in sys.modules if module == "pxr" or module.startswith("pxr.")] + if loaded_pxr_modules: + raise SystemExit("pxr loaded before SimulationApp: " + ", ".join(loaded_pxr_modules[:20])) + raise RuntimeError("stop before SimulationApp") + + AppLauncher._create_app = stop_before_simulation_app + + try: + demo.main() + except RuntimeError as exc: + if str(exc) != "stop before SimulationApp": + raise + """ + ) + + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stdout + result.stderr diff --git a/source/isaaclab_visualizers/changelog.d/max-newton-mpm-manager.rst b/source/isaaclab_visualizers/changelog.d/max-newton-mpm-manager.rst new file mode 100644 index 000000000000..46e13bfcc061 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/max-newton-mpm-manager.rst @@ -0,0 +1,8 @@ +Changed +^^^^^^^ + +* :class:`~isaaclab_visualizers.newton.NewtonVisualizer` now skips Newton's + per-frame active-particle compaction (two device-to-host reads per render) + when an MPM model's static particle flags are all active, and re-uploads the + particle color buffer only when the point count grows or the configured color + changes. diff --git a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py index 9a41ad20fd26..a10e5ca81935 100644 --- a/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py +++ b/source/isaaclab_visualizers/isaaclab_visualizers/newton/newton_visualizer.py @@ -80,6 +80,8 @@ def __init__( self._particle_color_buffer: wp.array | None = None self._particle_color_buffer_count = 0 self._particle_color_buffer_value: tuple[float, float, float] | None = None + self._mpm_particle_flags_cache_key: tuple[int, int, int] | None = None + self._mpm_particles_all_active = False try: self.register_ui_callback(self._render_training_controls, position="side") @@ -170,6 +172,18 @@ def _particle_color_array(self, count: int) -> wp.array: self._particle_color_buffer_value = color return self._particle_color_buffer + def _particle_color_update_array(self, name: str, count: int) -> wp.array | None: + """Return particle colors only when Newton needs the GL color buffer refreshed.""" + obj = self.objects.get(name) + capacity = obj.num_instances if obj is not None else 0 + if ( + obj is None + or count > capacity + or self._particle_color_buffer_value != self._coerce_color3(self.particle_color) + ): + return self._particle_color_array(max(count, capacity)) + return None + def log_points(self, name, points, radii=None, colors=None, hidden=False): """Apply configured model-particle appearance while preserving Newton's point logging. @@ -180,9 +194,50 @@ def log_points(self, name, points, radii=None, colors=None, hidden=False): if name != "/model/particles" or points is None or self.particle_color is None: return super().log_points(name, points, radii, colors, hidden) - colors = self._particle_color_array(len(points)) + colors = self._particle_color_update_array(name, len(points)) return super().log_points(name, points, radii, colors, hidden) + def _all_mpm_particles_active(self) -> bool: + """Return whether an MPM model's static particle flags are all active.""" + model = self.model + if model is None or getattr(model, "mpm", None) is None or not model.particle_count: + return False + if model.particle_flags is None: + return False + + cache_key = (id(model), id(model.particle_flags), int(model.particle_count)) + if self._mpm_particle_flags_cache_key != cache_key: + import newton as nt + + flags = model.particle_flags.numpy()[: model.particle_count] + self._mpm_particles_all_active = bool(((flags & int(nt.ParticleFlags.ACTIVE)) != 0).all()) + self._mpm_particle_flags_cache_key = cache_key + return self._mpm_particles_all_active + + def _log_particles(self, state): + """Log MPM particles without per-frame active-flag compaction when all particles are active. + + Newton's base implementation stream-compacts active particles every + frame, which costs two device-to-host reads per render. MPM particle + flags are static, so when they are all active the compaction is skipped + and ``state.particle_q`` is logged directly. + """ + if not self._all_mpm_particles_active(): + super()._log_particles(state) + return + + colors = None + if self.model_changed and self.particle_color is None: + colors = wp.full(shape=len(state.particle_q), value=wp.vec3(0.7, 0.6, 0.4), device=self.device) + + self.log_points( + name="/model/particles", + points=state.particle_q, + radii=self.model.particle_radius, + colors=colors, + hidden=not self.show_particles, + ) + def _color_edit3_compat(self, imgui, label: str, color): """ # Handle imgui.color_edit3 API differences between bindings. @@ -523,8 +578,6 @@ def step(self, dt: float) -> None: self._state = NewtonManager.get_state(self._scene_data_provider) return - self._state = NewtonManager.get_state(self._scene_data_provider) - update_frequency = self._viewer._update_frequency if self._viewer else self._update_frequency if self._step_counter % update_frequency != 0: return @@ -533,6 +586,7 @@ def step(self, dt: float) -> None: try: if not self._viewer.is_paused(): + self._state = NewtonManager.get_state(self._scene_data_provider) self._viewer.begin_frame(self._sim_time) try: if self._state is not None: diff --git a/source/isaaclab_visualizers/test/test_newton_adapter.py b/source/isaaclab_visualizers/test/test_newton_adapter.py index 760f7ec0f670..fd93a75ff940 100644 --- a/source/isaaclab_visualizers/test/test_newton_adapter.py +++ b/source/isaaclab_visualizers/test/test_newton_adapter.py @@ -153,6 +153,8 @@ def test_newton_viewer_particle_color_override(monkeypatch): viewer = NewtonViewerGL.__new__(NewtonViewerGL) viewer.device = "cpu" + viewer.objects = {} + viewer.model_changed = False viewer.particle_color = (0.1, 0.2, 0.3) viewer._particle_color_buffer = None viewer._particle_color_buffer_count = 0 @@ -175,6 +177,30 @@ def _log_points(self, name, points, radii=None, colors=None, hidden=False): np.testing.assert_allclose(colors.numpy()[0], np.array([0.1, 0.2, 0.3], dtype=np.float32), rtol=1.0e-6) +def test_newton_viewer_particle_color_override_reuses_existing_color_buffer(monkeypatch): + from newton.viewer import ViewerGL + + viewer = NewtonViewerGL.__new__(NewtonViewerGL) + viewer.device = "cpu" + viewer.model_changed = False + viewer.particle_color = (0.1, 0.2, 0.3) + viewer._particle_color_buffer = wp.zeros(4, dtype=wp.vec3, device="cpu") + viewer._particle_color_buffer_count = 4 + viewer._particle_color_buffer_value = (0.1, 0.2, 0.3) + viewer.objects = {"/model/particles": SimpleNamespace(num_instances=4)} + points = wp.zeros(4, dtype=wp.vec3, device="cpu") + calls = [] + + def _log_points(self, name, points, radii=None, colors=None, hidden=False): + calls.append((name, colors)) + + monkeypatch.setattr(ViewerGL, "log_points", _log_points) + + viewer.log_points("/model/particles", points, colors=object()) + + assert calls[-1] == ("/model/particles", None) + + def test_newton_viewer_particle_color_override_leaves_other_points_unchanged(monkeypatch): from newton.viewer import ViewerGL @@ -194,6 +220,60 @@ def _log_points(self, name, points, radii=None, colors=None, hidden=False): assert calls[-1] == ("/debug/points", original_colors) +def test_newton_viewer_fast_paths_all_active_mpm_particles(monkeypatch): + import newton as nt + + viewer = NewtonViewerGL.__new__(NewtonViewerGL) + viewer.device = "cpu" + viewer.model_changed = False + viewer.particle_color = None + viewer.show_particles = True + viewer._mpm_particle_flags_cache_key = None + viewer._mpm_particles_all_active = False + viewer.model = SimpleNamespace( + mpm=object(), + particle_count=3, + particle_flags=wp.array([int(nt.ParticleFlags.ACTIVE)] * 3, dtype=wp.int32, device="cpu"), + particle_radius=wp.ones(3, dtype=wp.float32, device="cpu"), + ) + state = SimpleNamespace(particle_q=wp.zeros(3, dtype=wp.vec3, device="cpu")) + calls = [] + + def _log_points(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(NewtonViewerGL, "log_points", _log_points) + + viewer._log_particles(state) + + assert calls[-1]["name"] == "/model/particles" + assert calls[-1]["points"] is state.particle_q + assert calls[-1]["radii"] is viewer.model.particle_radius + assert calls[-1]["hidden"] is False + + +def test_newton_viewer_inactive_mpm_particles_use_newton_filter(monkeypatch): + import newton as nt + from newton.viewer import ViewerGL + + viewer = NewtonViewerGL.__new__(NewtonViewerGL) + viewer._mpm_particle_flags_cache_key = None + viewer._mpm_particles_all_active = False + viewer.model = SimpleNamespace( + mpm=object(), + particle_count=2, + particle_flags=wp.array([int(nt.ParticleFlags.ACTIVE), 0], dtype=wp.int32, device="cpu"), + ) + state = object() + fallback_calls = [] + + monkeypatch.setattr(ViewerGL, "_log_particles", lambda self, state: fallback_calls.append(state)) + + viewer._log_particles(state) + + assert fallback_calls == [state] + + class _BodyQ: shape = (1,)