Skip to content
Skild AISkild AISteel

Shapes onto sand

Coupling

A dozen rigid shapes dropped onto MPM material with two-way coupling.

0.00 s
The steel_implicit_mpm preset (the default), recordedon an NVIDIA RTX PRO 6000 Blackwell.

Run it

Terminal window
cd examples
# the default preset (steel_implicit_mpm), then the others
uv run shelf play shapes_onto_sand
uv run shelf play shapes_onto_sand --solver mujoco_implicit_mpm
# change solver settings live
uv run shelf tune shapes_onto_sand

Source: examples/shelf/09_coupling/shapes_onto_sand/ · ported from newton/examples/mpm/example_mpm_twoway_coupling.py @ skild-ai/newton 7565a26b1

Presets

PresetSolverStatusChecked
steel_implicit_mpmdefaultcoupledworks2026-11-04, newton 7565a26b1, NVIDIA GeForce RTX 3090
mujoco_implicit_mpmcoupledworks2026-10-28, newton 7565a26b1, NVIDIA RTX A6000

Solvers with no row are untested on this example. The full settings are inpresets.toml.

Code

examples/shelf/09_coupling/shapes_onto_sand/example.py
# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers
# SPDX-License-Identifier: Apache-2.0
# Modified for Steel Shelf: preset-driven execution and shared runner integration.


###########################################################################
# Example MPM 2-Way Coupling
#
# A simple scene spawning a dozen rigid shapes above a plane. The shapes
# fall and collide using the MuJoCo solver. Demonstrates basic builder APIs
# and the standard example structure.
#
# Command: python -m shelf run shapes_onto_sand
#
###########################################################################


from __future__ import annotations


import numpy as np
import warp as wp


import newton
import newton.examples
from newton.solvers import SolverImplicitMPM
from common import presets




@wp.kernel
def compute_body_forces(
    dt: float,
    collider_ids: wp.array[int],
    collider_impulses: wp.array[wp.vec3],
    collider_impulse_pos: wp.array[wp.vec3],
    body_ids: wp.array[int],
    body_q: wp.array[wp.transform],
    body_com: wp.array[wp.vec3],
    body_f: wp.array[wp.spatial_vector],
):
    """Compute forces applied by sand to rigid bodies.


    Sum the impulses applied on each mpm grid node and convert to
    forces and torques at the body's center of mass.
    """


    i = wp.tid()


    cid = collider_ids[i]
    if cid >= 0 and cid < body_ids.shape[0]:
        body_index = body_ids[cid]
        if body_index == -1:
            return


        f_world = collider_impulses[i] / dt


        X_wb = body_q[body_index]
        X_com = body_com[body_index]
        r = collider_impulse_pos[i] - wp.transform_point(X_wb, X_com)
        wp.atomic_add(body_f, body_index, wp.spatial_vector(f_world, wp.cross(r, f_world)))




@wp.kernel
def subtract_body_force(
    dt: float,
    body_q: wp.array[wp.transform],
    body_qd: wp.array[wp.spatial_vector],
    body_f: wp.array[wp.spatial_vector],
    body_inv_inertia: wp.array[wp.mat33],
    body_inv_mass: wp.array[float],
    body_q_res: wp.array[wp.transform],
    body_qd_res: wp.array[wp.spatial_vector],
):
    """Update the rigid bodies velocity to remove the forces applied by sand at the last step.


    This is necessary to compute the total impulses that are required to enforce the complementarity-based
    frictional contact boundary conditions.
    """


    body_id = wp.tid()


    # Remove previously applied force
    f = body_f[body_id]
    delta_v = dt * body_inv_mass[body_id] * wp.spatial_top(f)
    r = wp.transform_get_rotation(body_q[body_id])


    delta_w = dt * wp.quat_rotate(r, body_inv_inertia[body_id] * wp.quat_rotate_inv(r, wp.spatial_bottom(f)))


    body_q_res[body_id] = body_q[body_id]
    body_qd_res[body_id] = body_qd[body_id] - wp.spatial_vector(delta_v, delta_w)




class Example:
    def __init__(self, viewer, args):
        # setup simulation parameters first
        self.preset = presets.load(__file__)
        self.fps = self.preset.fps
        self.frame_dt = 1.0 / self.fps
        self.sim_time = 0.0
        self.sim_substeps = self.preset.substeps
        self.sim_dt = self.frame_dt / self.sim_substeps


        self.viewer = viewer


        # setup rigid-body model builder
        builder = newton.ModelBuilder()
        builder.default_shape_cfg.mu = 0.5
        self._emit_rigid_bodies(builder)


        # add ground plane
        builder.add_ground_plane()


        # setup sand model builder
        sand_builder = newton.ModelBuilder()
        # Optional cap on sand advection speed [m/s], set per preset (see presets.toml).
        max_velocity = self.preset.model.get("sand_particle_max_velocity")
        if max_velocity is not None:
            sand_builder.particle_max_velocity = float(max_velocity)


        # Register MPM custom attributes before adding particles
        SolverImplicitMPM.register_custom_attributes(sand_builder)


        voxel_size = 0.05  # 5 cm
        self._emit_particles(sand_builder, voxel_size)


        # finalize models
        self.model = builder.finalize()
        self.sand_model = sand_builder.finalize()
        self.voxel_size = voxel_size


        # simulation state
        self.state_0 = self.model.state()
        self.state_1 = self.model.state()


        self.sand_state_0 = self.sand_model.state()
        self.sand_state_0.body_q = wp.empty_like(self.state_0.body_q)
        self.sand_state_0.body_qd = wp.empty_like(self.state_0.body_qd)
        self.sand_state_0.body_f = wp.empty_like(self.state_0.body_f)


        self.control = self.model.control()


        self.collision_pipeline = newton.CollisionPipeline(self.model)
        self.contacts = self.collision_pipeline.contacts()


        # viewer
        self.viewer.set_model(self.model)
        if hasattr(self.viewer, "register_ui_callback"):
            self.viewer.register_ui_callback(self.render_ui, position="side")
        self.viewer.show_particles = True
        self.show_impulses = False


        # not required for MuJoCo, but required for other solvers
        newton.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, self.state_0)


        # Additional buffers for tracking two-way coupling forces
        max_nodes = 1 << 20
        self.collider_impulses = wp.zeros(max_nodes, dtype=wp.vec3, device=self.model.device)
        self.collider_impulse_pos = wp.zeros(max_nodes, dtype=wp.vec3, device=self.model.device)
        self.collider_impulse_ids = wp.full(max_nodes, value=-1, dtype=int, device=self.model.device)


        self.build_solver()


        # per-body forces and torques applied by sand to rigid bodies
        self.body_sand_forces = wp.zeros_like(self.state_0.body_f)


        self.particle_render_colors = wp.full(
            self.sand_model.particle_count, value=wp.vec3(0.7, 0.6, 0.4), dtype=wp.vec3, device=self.sand_model.device
        )


        self.capture()


    def build_solver(self):
        """Build the solver and everything bound to its instance; the tuner calls this again to swap it."""
        # setup mpm solver
        mpm_options = self.preset.entry("implicit_mpm").make_config(self.sand_model)
        mpm_options.voxel_size = self.voxel_size


        self.mpm_solver = SolverImplicitMPM(self.sand_model, config=mpm_options)
        # read colliders from the RB model rather than the sand model
        self.mpm_solver.setup_collider(model=self.model)


        # setup rigid-body solver (MuJoCo or Steel, both on Newton's contacts)
        rigid = self.preset.name.split("_")[0]
        if rigid == "steel":
            # Steel sizes its contact scratch from the model before it sees a contact buffer
            self.model.rigid_contact_max = self.contacts.rigid_contact_max
        self.solver = self.preset.entry(rigid).make_solver(self.model)


        # Read the new solver's pre-step impulses: a sand state that was already stepped still
        # carries the previous solver instance's grid fields, which must not be read here.
        self.collect_collider_impulses(None)


        # map from collider index to body index
        self.collider_body_id = self.mpm_solver.collider_body_index


    def capture(self):
        with wp.ScopedCapture() as capture:
            self.simulate()
        self.graph = capture.graph


    def simulate(self):
        for substep in range(self.sim_substeps):
            self.state_0.clear_forces()


            wp.launch(
                compute_body_forces,
                dim=self.collider_impulse_ids.shape[0],
                inputs=[
                    self.frame_dt,
                    self.collider_impulse_ids,
                    self.collider_impulses,
                    self.collider_impulse_pos,
                    self.collider_body_id,
                    self.state_0.body_q,
                    self.model.body_com,
                    self.state_0.body_f,
                ],
            )
            # saved applied force to subtract later on
            self.body_sand_forces.assign(self.state_0.body_f)


            # apply forces to the model
            self.viewer.apply_forces(self.state_0)


            self.collision_pipeline.collide(self.state_0, self.contacts)
            self.solver.step(self.state_0, self.state_1, self.control, self.contacts, self.sim_dt)


            # swap states
            if self.sim_substeps % 2 == 1 and substep == self.sim_substeps - 1:
                # an odd count under graph capture: copy the last result into state_0 instead of swapping,
                # so every replay of the captured graph starts from the newest state
                self.state_0.assign(self.state_1)
            else:
                self.state_0, self.state_1 = self.state_1, self.state_0


        self.simulate_sand()


    def collect_collider_impulses(self, state):
        collider_impulses, collider_impulse_pos, collider_impulse_ids = self.mpm_solver.collect_collider_impulses(
            state
        )
        self.collider_impulse_ids.fill_(-1)
        n_colliders = min(collider_impulses.shape[0], self.collider_impulses.shape[0])
        self.collider_impulses[:n_colliders].assign(collider_impulses[:n_colliders])
        self.collider_impulse_pos[:n_colliders].assign(collider_impulse_pos[:n_colliders])
        self.collider_impulse_ids[:n_colliders].assign(collider_impulse_ids[:n_colliders])


    def simulate_sand(self):
        # Subtract previously applied impulses from body velocities


        if self.sand_state_0.body_q is not None:
            wp.launch(
                subtract_body_force,
                dim=self.sand_state_0.body_q.shape,
                inputs=[
                    self.frame_dt,
                    self.state_0.body_q,
                    self.state_0.body_qd,
                    self.body_sand_forces,
                    self.model.body_inv_inertia,
                    self.model.body_inv_mass,
                    self.sand_state_0.body_q,
                    self.sand_state_0.body_qd,
                ],
            )


        self.mpm_solver.step(self.sand_state_0, self.sand_state_0, contacts=None, control=None, dt=self.frame_dt)


        # Save impulses to apply back to rigid bodies
        self.collect_collider_impulses(self.sand_state_0)


    def step(self):
        if self.graph:
            wp.capture_launch(self.graph)
        else:
            self.simulate()


        self.sim_time += self.frame_dt


    def test_final(self):
        newton.examples.test_body_state(
            self.model,
            self.state_0,
            "all bodies are above the sand",
            lambda q, qd: q[2] > 0.45,
        )
        voxel_size = self.mpm_solver.voxel_size
        newton.examples.test_particle_state(
            self.sand_state_0,
            "all particles are above the ground",
            lambda q, qd: q[2] > -voxel_size,
        )


    def render(self):
        self.viewer.begin_frame(self.sim_time)
        self.viewer.log_state(self.state_0)
        self.viewer.log_contacts(self.contacts, self.state_0)


        self.viewer.log_points(
            "/sand",
            points=self.sand_state_0.particle_q,
            radii=self.sand_model.particle_radius,
            colors=self.particle_render_colors,
            hidden=not self.viewer.show_particles,
        )


        if self.show_impulses:
            impulses, pos, _cid = self.mpm_solver.collect_collider_impulses(self.sand_state_0)
            self.viewer.log_lines(
                "/impulses",
                starts=pos,
                ends=pos + impulses,
                colors=wp.full(pos.shape[0], value=wp.vec3(1.0, 0.0, 0.0), dtype=wp.vec3),
            )
        else:
            self.viewer.log_lines("/impulses", None, None, None)


        self.viewer.end_frame()


    def render_ui(self, imgui):
        _changed, self.show_impulses = imgui.checkbox("Show Impulses", self.show_impulses)


    def _emit_rigid_bodies(self, builder: newton.ModelBuilder):
        # z height to drop shapes from
        drop_z = 2.0


        # layout: spawn shapes near the same XY so they collide/stack
        offsets_xy = [
            (0.00, 0.00),
            (0.10, 0.00),
            (-0.10, 0.00),
            (0.00, 0.10),
            (0.00, -0.10),
            (0.10, 0.10),
            (-0.10, 0.10),
            (0.10, -0.10),
            (-0.10, -0.10),
            (0.15, 0.00),
            (-0.15, 0.00),
            (0.00, 0.15),
        ]
        offset_index = 0
        z_index = 0
        z_separation = 0.6  # vertical spacing to avoid initial overlap


        # generate a few boxes with varying sizes
        # boxes = [(0.45, 0.35, 0.25)]  # (hx, hy, hz)
        boxes = [
            (0.25, 0.35, 0.25),
            (0.25, 0.25, 0.25),
            (0.3, 0.2, 0.2),
            (0.25, 0.35, 0.25),
            (0.25, 0.25, 0.25),
            (0.3, 0.2, 0.2),
        ]  # (hx, hy, hz)
        for box in boxes:
            (hx, hy, hz) = box


            ox, oy = offsets_xy[offset_index % len(offsets_xy)]
            offset_index += 1
            pz = drop_z + float(z_index) * z_separation
            z_index += 1
            body = builder.add_body(
                xform=wp.transform(p=wp.vec3(float(ox), float(oy), pz), q=wp.normalize(wp.quatf(0.0, 0.0, 0.0, 1.0))),
                mass=75.0,
            )
            builder.add_shape_box(body, hx=float(hx), hy=float(hy), hz=float(hz))


    def _emit_particles(self, sand_builder: newton.ModelBuilder, voxel_size: float):
        # ------------------------------------------
        # Add sand bed (2m x 2m x 0.5m) above ground
        # ------------------------------------------


        particles_per_cell = 3.0
        density = 2500.0


        bed_lo = np.array([-1.0, -1.0, 0.0])
        bed_hi = np.array([1.0, 1.0, 0.5])
        bed_res = np.array(np.ceil(particles_per_cell * (bed_hi - bed_lo) / voxel_size), dtype=int)


        cell_size = (bed_hi - bed_lo) / bed_res
        cell_volume = np.prod(cell_size)
        radius = float(np.max(cell_size) * 0.5)
        mass = float(np.prod(cell_volume) * density)


        sand_builder.add_particle_grid(
            pos=wp.vec3(bed_lo),
            rot=wp.quat_identity(),
            vel=wp.vec3(0.0),
            dim_x=bed_res[0] + 1,
            dim_y=bed_res[1] + 1,
            dim_z=bed_res[2] + 1,
            cell_x=cell_size[0],
            cell_y=cell_size[1],
            cell_z=cell_size[2],
            mass=mass,
            jitter=2.0 * radius,
            radius_mean=radius,
            custom_attributes={"mpm:friction": 0.75},
        )




if __name__ == "__main__":
    # Parse arguments and initialize viewer
    viewer, args = newton.examples.init()


    # Create example and run
    newton.examples.run(Example(viewer, args), args)

example.py on GitHub

Steel docs and tutorials are licensed under Apache-2.0. Third-party notices. "Skild", "Skild AI" and the Skild logo are trademarks of Skild AI and are not licensed under Apache-2.0.