[Newton] Speed up deformable cloning using Newton replication

source/isaaclab_contrib/changelog.d/newton-cloth-replication.rst

@@ -0,0 +1,5 @@
+Changed
+^^^^^^^
+
+* Changed Newton cloth and soft-body replication to add each registered deformable mesh once when environments
+  differ only by translation, reducing startup time for large scenes. No migration is required.

source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -39,10 +39,9 @@
 class DeformableRegistryEntry:
     """Entry in the deformable body registry.
 
-    Registered by :class:`DeformableObject` during ``__init__``, consumed by
-    the Newton clone context inside the per-world ``begin_world``/``end_world`` loop.
-    After replication, ``particle_offsets`` and ``particles_per_body`` are filled in
-    so the asset can bind to the correct particle ranges.
+    Registered by :class:`DeformableObject` during ``__init__`` and consumed when the Newton
+    clone context builds the worlds. After builder construction, ``particle_offsets`` and
+    ``particles_per_body`` identify the particle ranges bound to the asset.
     """
 
     prim_path: str
@@ -65,7 +64,6 @@ class DeformableRegistryEntry:
     k_mu: float = 1e5
     k_lambda: float = 1e5
     k_damp: float = 0.0
-    # Filled by the Newton clone context:
     particle_offsets: list[int] = field(default_factory=list)
     particles_per_body: int = 0
 
@@ -99,12 +97,32 @@ def add_deformable_entry_to_builder(
         env_position: World position [x, y, z] [m] for this environment.
         env_rotation: World orientation as quaternion ``(x, y, z, w)`` for this environment.
     """
+    # Cleared before the mesh is added, so a build that fails partway cannot leave offsets from
+    # an earlier build behind: they would have the expected length and pass every later check.
     if env_idx == 0:
         entry.particle_offsets.clear()
         entry.particles_per_body = 0
 
-    before_count = getattr(builder, "particle_count", 0)
+    before_count, particle_count = _add_deformable_entry_geometry(builder, entry, env_position, env_rotation)
 
+    if env_idx == 0:
+        entry.particles_per_body = particle_count
+    elif entry.particles_per_body != particle_count:
+        raise RuntimeError(
+            f"Deformable body '{entry.prim_path}' produced {particle_count} particles in env {env_idx}, "
+            f"but env 0 produced {entry.particles_per_body}."
+        )
+    entry.particle_offsets.append(before_count)
+
+
+def _add_deformable_entry_geometry(
+    builder,
+    entry: DeformableRegistryEntry,
+    env_position: Sequence[float],
+    env_rotation: Sequence[float],
+) -> tuple[int, int]:
+    """Add one deformable mesh and return its particle offset and count."""
+    before_count = getattr(builder, "particle_count", 0)
     env_pos = wp.vec3(float(env_position[0]), float(env_position[1]), float(env_position[2]))
     env_rot = wp.quat(
         float(env_rotation[0]),
@@ -158,16 +176,7 @@ def add_deformable_entry_to_builder(
         )
 
     after_count = getattr(builder, "particle_count", 0)
-    delta = after_count - before_count
-
-    entry.particle_offsets.append(before_count)
-    if env_idx == 0:
-        entry.particles_per_body = delta
-    elif entry.particles_per_body != delta:
-        raise RuntimeError(
-            f"Deformable body '{entry.prim_path}' produced {delta} particles in env {env_idx}, "
-            f"but env 0 produced {entry.particles_per_body}."
-        )
+    return before_count, after_count - before_count
 
 
 def add_registered_deformables_to_builder(
@@ -181,6 +190,62 @@ def add_registered_deformables_to_builder(
         add_deformable_entry_to_builder(builder, entry, world_idx, env_position, env_rotation)
 
 
+def _can_replicate_registered_deformables(xforms: np.ndarray) -> bool:
+    """Return False if any registered deformables prevent replication.
+
+    Every rotation must be identity because ``ModelBuilder.replicate`` only translates particles.
+    """
+    # An empty registry is compatible because the hook adds no particles
+    if not SimulationManager._deformable_registry:
+        return True
+    identity_rotation = np.array((0.0, 0.0, 0.0, 1.0), dtype=xforms.dtype)
+    return bool(np.allclose(np.abs(xforms[:, 3:]), identity_rotation, atol=1.0e-6, rtol=0.0))
+
+
+def _prepare_registered_deformable_replication(
+    builder,
+    source_builder,
+    num_worlds: int,
+    source_position: Sequence[float],
+    source_rotation: Sequence[float],
+) -> Callable[[], None]:
+    """Add registered deformables once to the source builder and return the offset binding."""
+    destination_particle_base = getattr(builder, "particle_count", 0)
+    # Cleared first, for the same reason as in ``add_deformable_entry_to_builder``.
+    for entry in SimulationManager._deformable_registry:
+        entry.particle_offsets.clear()
+        entry.particles_per_body = 0
+    replicated_entries = [
+        (entry, *_add_deformable_entry_geometry(source_builder, entry, source_position, source_rotation))
+        for entry in SimulationManager._deformable_registry
+    ]
+
+    def finish_replication() -> None:
+        source_particle_stride = getattr(source_builder, "particle_count", 0)
+        expected_particle_count = destination_particle_base + num_worlds * source_particle_stride
+        replicated_particle_count = getattr(builder, "particle_count", 0)
+        # The offsets below assume replication appended ``num_worlds`` contiguous copies of the
+        # source, so that layout is verified rather than trusted.
+        if replicated_particle_count != expected_particle_count:
+            raise RuntimeError(
+                f"Builder holds {replicated_particle_count} particles after replication, but deformable"
+                f" particle offsets assume {expected_particle_count} ({num_worlds} worlds of"
+                f" {source_particle_stride} particles after {destination_particle_base} existing particles)."
+            )
+        world_particle_offsets = (
+            destination_particle_base + np.arange(num_worlds, dtype=np.int64) * source_particle_stride
+        )
+        for entry, source_offset, particle_count in replicated_entries:
+            entry.particles_per_body = particle_count
+            entry.particle_offsets = (world_particle_offsets + source_offset).tolist()
+
+    return finish_replication
+
+
+add_registered_deformables_to_builder._can_replicate_builder = _can_replicate_registered_deformables
+add_registered_deformables_to_builder._prepare_builder_replication = _prepare_registered_deformable_replication
+
+
 def setup_registered_deformable_fabric_sync(manager_cls: type[SimulationManager]) -> None:
     """Bind registered deformable visual meshes to their Newton particle slices in Fabric."""
     if manager_cls._clone_physics_only or not manager_cls._deformable_registry:

source/isaaclab_contrib/test/deformable/test_deformable_builder_hooks.py

@@ -6,16 +6,21 @@
 import math
 import sys
 from types import SimpleNamespace
+from unittest import mock
 
+import newton
+import numpy as np
 import pytest
 import warp as wp
+from isaaclab_newton.cloner.newton_clone_utils import replicate_builder_mapping
 from isaaclab_newton.physics import NewtonManager
 from isaaclab_newton.sim.spawners.materials import NewtonDeformableMaterialCfg
 
 from isaaclab_contrib.deformable import DeformableObject
 from isaaclab_contrib.deformable.deformable_object import (
     DeformableRegistryEntry,
     add_deformable_entry_to_builder,
+    add_registered_deformables_to_builder,
     setup_registered_deformable_fabric_sync,
 )
 
@@ -131,6 +136,82 @@ def test_builder_hook_resets_entry_offsets_on_first_environment():
     assert entry.particles_per_body == 3
 
 
+def _replicate_deformable_worlds(builder, source, positions, quaternions):
+    """Replicate ``source`` into one world per row of ``positions``, running the deformable hook."""
+    positions = np.asarray(positions, dtype=np.float32)
+    return replicate_builder_mapping(
+        builder,
+        ("/World/envs/env_0",),
+        np.ones((1, len(positions)), dtype=bool),
+        positions,
+        np.asarray(quaternions, dtype=np.float32),
+        {"/World/envs/env_0": source},
+        ("/World/envs/env_{}",),
+        np.arange(len(positions)),
+        per_world_builder_hooks=(add_registered_deformables_to_builder,),
+    )
+
+
+def test_homogeneous_deformables_are_added_once_and_replicated(monkeypatch):
+    """Test that translation-only worlds replicate each source deformable mesh once."""
+    entries = [_make_surface_entry(), _make_surface_entry()]
+    entries[1].init_pos = (0.0, 2.0, 0.0)
+    monkeypatch.setattr(NewtonManager, "_deformable_registry", entries)
+    source = newton.ModelBuilder()
+    builder = newton.ModelBuilder()
+    builder.add_particle(pos=(-1.0, -2.0, -3.0), vel=(0.0, 0.0, 0.0), mass=1.0, radius=0.1)
+    positions = np.asarray([[2.0, 3.0, 0.0], [5.0, 3.0, 0.0], [2.0, 7.0, 0.0]], dtype=np.float32)
+
+    with (
+        mock.patch.object(source, "add_cloth_mesh", wraps=source.add_cloth_mesh) as add_cloth_mesh,
+        mock.patch.object(builder, "replicate", wraps=builder.replicate) as replicate,
+    ):
+        _replicate_deformable_worlds(builder, source, positions, [[0.0, 0.0, 0.0, 1.0]] * 3)
+
+    assert add_cloth_mesh.call_count == len(entries)
+    replicate.assert_called_once()
+    assert builder.world_count == 3
+    assert [entry.particle_offsets for entry in entries] == [[1, 7, 13], [4, 10, 16]]
+    assert [entry.particles_per_body for entry in entries] == [3, 3]
+    np.testing.assert_array_equal(np.asarray(builder.tri_indices), np.arange(1, 19).reshape(-1, 3))
+    # Both entries carry the same mesh rotated 90 deg about z, the second offset 2 m along y.
+    local_particles = np.asarray(
+        [[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 3.0, 0.0], [-1.0, 2.0, 0.0]]
+    )
+    expected_particles = np.concatenate(
+        ([[-1.0, -2.0, -3.0]], (positions[:, None, :] + local_particles).reshape(-1, 3))
+    )
+    np.testing.assert_allclose(np.asarray(builder.particle_q), expected_particles, atol=1.0e-6)
+
+
+def test_rotated_deformable_worlds_keep_per_world_fallback(monkeypatch):
+    """Test that worlds with differing rotations fall back to per-world construction."""
+    entry = _make_surface_entry()
+    monkeypatch.setattr(NewtonManager, "_deformable_registry", [entry])
+    source = newton.ModelBuilder()
+    builder = newton.ModelBuilder()
+    half_sqrt = math.sqrt(0.5)
+
+    with mock.patch.object(builder, "replicate", wraps=builder.replicate) as replicate:
+        _replicate_deformable_worlds(
+            builder,
+            source,
+            [[0.0, 0.0, 0.0], [4.0, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, half_sqrt, half_sqrt]],
+        )
+
+    replicate.assert_not_called()
+    assert source.particle_count == 0
+    assert builder.world_count == 2
+    assert entry.particle_offsets == [0, 3]
+    assert entry.particles_per_body == 3
+    np.testing.assert_allclose(
+        np.asarray(builder.particle_q)[3:],
+        [[4.0, 1.0, 0.0], [3.0, 1.0, 0.0], [4.0, 0.0, 0.0]],
+        atol=1.0e-6,
+    )
+
+
 def test_fabric_particle_sync_skips_missing_fabric_prim(monkeypatch):
     """Test that missing Fabric prims are skipped before attributes are authored."""
     entry = _make_surface_entry()

source/isaaclab_newton/changelog.d/newton-cloth-replication.rst

@@ -0,0 +1,5 @@
+Changed
+^^^^^^^
+
+* Allowed compatible per-world Newton builder hooks to use homogeneous source-builder replication instead of the
+  per-world construction loop. No migration is required.