From 22a38d37a364570139ae101c312f06363cbae815 Mon Sep 17 00:00:00 2001 From: camevor Date: Fri, 4 Sep 2026 15:06:08 +0200 Subject: [PATCH 1/3] Newton-side cloth replication --- .../changelog.d/newton-cloth-replication.rst | 5 + .../deformable/deformable_object.py | 102 ++++++++++--- .../test_deformable_builder_hooks.py | 81 +++++++++++ .../changelog.d/newton-cloth-replication.rst | 5 + .../cloner/newton_clone_utils.py | 136 +++++++++++++----- 5 files changed, 274 insertions(+), 55 deletions(-) create mode 100644 source/isaaclab_contrib/changelog.d/newton-cloth-replication.rst create mode 100644 source/isaaclab_newton/changelog.d/newton-cloth-replication.rst diff --git a/source/isaaclab_contrib/changelog.d/newton-cloth-replication.rst b/source/isaaclab_contrib/changelog.d/newton-cloth-replication.rst new file mode 100644 index 000000000000..9c6555b8b71e --- /dev/null +++ b/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. diff --git a/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py b/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py index 05f6725e344b..b970267d0d1e 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py +++ b/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,65 @@ 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 whether Newton replication can apply every relative particle transform. + + ``ModelBuilder.replicate`` offsets ``particle_q`` by each transform's translation and drops + its rotation. + """ + if not SimulationManager._deformable_registry: + return True + rotations = xforms[:, 3:] + return bool( + np.allclose(rotations[:, :3], 0.0, atol=1.0e-6, rtol=0.0) + and np.allclose(np.abs(rotations[:, 3]), 1.0, 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)." + ) + for entry, source_offset, particle_count in replicated_entries: + entry.particles_per_body = particle_count + entry.particle_offsets = [ + destination_particle_base + world * source_particle_stride + source_offset + for world in range(num_worlds) + ] + + 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: diff --git a/source/isaaclab_contrib/test/deformable/test_deformable_builder_hooks.py b/source/isaaclab_contrib/test/deformable/test_deformable_builder_hooks.py index a0d4800c95af..539803199355 100644 --- a/source/isaaclab_contrib/test/deformable/test_deformable_builder_hooks.py +++ b/source/isaaclab_contrib/test/deformable/test_deformable_builder_hooks.py @@ -6,9 +6,13 @@ 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 @@ -16,6 +20,7 @@ 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() diff --git a/source/isaaclab_newton/changelog.d/newton-cloth-replication.rst b/source/isaaclab_newton/changelog.d/newton-cloth-replication.rst new file mode 100644 index 000000000000..7e0b3b9afd7b --- /dev/null +++ b/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. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py index cd0a605c2129..038c5dc7b04c 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py @@ -5,6 +5,7 @@ from __future__ import annotations +import logging from collections.abc import Callable, Sequence from typing import Any @@ -19,6 +20,8 @@ from isaaclab_newton.renderers.visual_material import import_builder_visual_material_paths +logger = logging.getLogger(__name__) + def _has_visible_non_collision_geometry(stage: Usd.Stage, prim_path: str) -> bool: """Return whether a prim hierarchy contains visible geometry without collision.""" @@ -229,6 +232,51 @@ def _rebase_labels(builder: ModelBuilder, source: str, destination: str) -> str: return prefix +def _collect_replication_preparers( + hooks: Sequence[Callable[[ModelBuilder, int, np.ndarray, np.ndarray], None]], + xforms: np.ndarray, +) -> list[Callable[..., Callable[[], None]]] | None: + """Replication preparers for every hook, or ``None`` if any hook needs the per-world loop. + + Args: + hooks: Per-world builder hooks registered for this replication. + xforms: Per-world transforms relative to world 0, shape ``[num_worlds, 7]``, position [m] + followed by a quaternion in xyzw order. + + Returns: + One preparer per hook, in registration order, or ``None`` if any hook does not support + replication or declines these transforms. + + Raises: + TypeError: If a hook carries one opt-in attribute as a callable but not the other. + """ + opted_in = [] + for hook in hooks: + can_replicate = getattr(hook, "_can_replicate_builder", None) + prepare = getattr(hook, "_prepare_builder_replication", None) + # Checked for every hook before any is consulted, so a half-implemented hook is rejected + # regardless of registration order. + if callable(can_replicate) != callable(prepare): + missing = "_prepare_builder_replication" if callable(can_replicate) else "_can_replicate_builder" + raise TypeError( + f"Newton world-builder hook '{getattr(hook, '__qualname__', repr(hook))}' opts into replication" + f" but has no callable '{missing}'." + ) + opted_in.append((hook, can_replicate, prepare)) + + preparers = [] + for hook, can_replicate, prepare in opted_in: + name = getattr(hook, "__qualname__", repr(hook)) + if not callable(can_replicate): + logger.debug("Hook '%s' does not support replication; building each Newton world separately.", name) + return None + if not can_replicate(xforms): + logger.debug("Hook '%s' declined replication; building each Newton world separately.", name) + return None + preparers.append(prepare) + return preparers + + def replicate_builder_mapping( builder: ModelBuilder, sources: Sequence[str], @@ -243,7 +291,13 @@ def replicate_builder_mapping( env_root_sites: dict[str, wp.transform] | None = None, per_world_builder_hooks: Sequence[Callable[[ModelBuilder, int, np.ndarray, np.ndarray], None]] = (), ) -> tuple[dict[str, list[list[int]]], list[wp.transform], list[tuple[str, int]]]: - """Replicate source builders, naming homogeneous copies at their destinations.""" + """Replicate source builders, naming homogeneous copies at their destinations. + + A hook joins the batched path by carrying ``_can_replicate_builder(xforms)`` and + ``_prepare_builder_replication(builder, source_builder, num_worlds, source_position, + source_rotation)``, the latter returning a callback invoked after replication. One hook + that declines sends them all back to the per-world loop. + """ source_site_indices = source_site_indices or {} env_root_sites = env_root_sites or {} num_worlds = mapping.shape[1] @@ -253,51 +307,57 @@ def replicate_builder_mapping( xforms_np = np.concatenate((positions, quaternions), axis=1) world_xforms = [wp.transform(*row) for row in xforms_np] - can_batch = ( + if ( len(sources) == 1 and mapping.shape[0] == 1 and num_worlds > 0 and bool(mapping.all()) - and not per_world_builder_hooks and bool(destinations) and env_ids is not None - ) - if can_batch: + ): source_builder = source_builders[sources[0]] - - # Inject env-root sites into the source so replicate() copies them. Prefixed - # by world_xforms[0] so R_w = world_xform_w * inv(world_xform_0) lands each - # copy at world_xform_w * xform. - site_local_indices: dict[str, list[int]] = {} - for label, xform in env_root_sites.items(): - idx = source_builder.add_site(body=-1, xform=wp.transform_multiply(world_xforms[0], xform), label=label) - site_local_indices.setdefault(label, []).append(idx) - for label, indices in source_site_indices.get(id(source_builder), {}).items(): - site_local_indices.setdefault(label, []).extend(indices) - - # Site index after replicate: base_shape + world * stride + source_local_index. - base_shape = builder.shape_count - stride = source_builder.shape_count - source_xform_inv = _invert_xform(xforms_np[0]) - xforms = _compose_world_xforms(positions, quaternions, source_xform_inv) - - label_groups = _label_groups(source_builder) - original_labels = {name: list(labels) for name, labels in label_groups.items()} - try: - prefix = _rebase_labels(source_builder, sources[0], destinations[0]) - prefixes = [prefix.format(int(env_id)) for env_id in env_ids] - builder.replicate(source_builder, num_worlds, xforms=xforms, label_prefixes=prefixes) - finally: - for name, labels in original_labels.items(): - label_groups[name][:] = labels - - for label, local_indices in site_local_indices.items(): - local_site_map[label] = [ - [base_shape + world * stride + local for local in local_indices] for world in range(num_worlds) + xforms = _compose_world_xforms(positions, quaternions, _invert_xform(xforms_np[0])) + hook_preparers = _collect_replication_preparers(per_world_builder_hooks, xforms) + + if hook_preparers is not None: + # Inject env-root sites into the source so replicate() copies them. Prefixed + # by world_xforms[0] so R_w = world_xform_w * inv(world_xform_0) lands each + # copy at world_xform_w * xform. + site_local_indices: dict[str, list[int]] = {} + for label, xform in env_root_sites.items(): + idx = source_builder.add_site(body=-1, xform=wp.transform_multiply(world_xforms[0], xform), label=label) + site_local_indices.setdefault(label, []).append(idx) + for label, indices in source_site_indices.get(id(source_builder), {}).items(): + site_local_indices.setdefault(label, []).extend(indices) + + finish_callbacks = [ + prepare(builder, source_builder, num_worlds, xforms_np[0, :3].copy(), xforms_np[0, 3:].copy()) + for prepare in hook_preparers ] - - bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping, skip_entity_labels=True) - return local_site_map, world_xforms, bindings + # Site index after replicate: base_shape + world * stride + source_local_index. Read + # after the site injection and the preparers, both of which can add shapes. + base_shape = builder.shape_count + stride = source_builder.shape_count + + label_groups = _label_groups(source_builder) + original_labels = {name: list(labels) for name, labels in label_groups.items()} + try: + prefix = _rebase_labels(source_builder, sources[0], destinations[0]) + prefixes = [prefix.format(int(env_id)) for env_id in env_ids] + builder.replicate(source_builder, num_worlds, xforms=xforms, label_prefixes=prefixes) + finally: + for name, labels in original_labels.items(): + label_groups[name][:] = labels + for finish in finish_callbacks: + finish() + + for label, local_indices in site_local_indices.items(): + local_site_map[label] = [ + [base_shape + world * stride + local for local in local_indices] for world in range(num_worlds) + ] + + bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping, skip_entity_labels=True) + return local_site_map, world_xforms, bindings source_world_indices = mapping.argmax(axis=1) From c5bb1601e825654307b1d587911be1020ef29210 Mon Sep 17 00:00:00 2001 From: camevor Date: Fri, 4 Sep 2026 17:49:23 +0200 Subject: [PATCH 2/3] Harden batched Newton replication --- .../cloner/newton_clone_utils.py | 38 ++++++---- .../test/cloner/test_rename_builder_labels.py | 74 +++++++++++++++++++ 2 files changed, 96 insertions(+), 16 deletions(-) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py index 038c5dc7b04c..f994fc48adc2 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py @@ -246,26 +246,11 @@ def _collect_replication_preparers( Returns: One preparer per hook, in registration order, or ``None`` if any hook does not support replication or declines these transforms. - - Raises: - TypeError: If a hook carries one opt-in attribute as a callable but not the other. """ - opted_in = [] + preparers = [] for hook in hooks: can_replicate = getattr(hook, "_can_replicate_builder", None) prepare = getattr(hook, "_prepare_builder_replication", None) - # Checked for every hook before any is consulted, so a half-implemented hook is rejected - # regardless of registration order. - if callable(can_replicate) != callable(prepare): - missing = "_prepare_builder_replication" if callable(can_replicate) else "_can_replicate_builder" - raise TypeError( - f"Newton world-builder hook '{getattr(hook, '__qualname__', repr(hook))}' opts into replication" - f" but has no callable '{missing}'." - ) - opted_in.append((hook, can_replicate, prepare)) - - preparers = [] - for hook, can_replicate, prepare in opted_in: name = getattr(hook, "__qualname__", repr(hook)) if not callable(can_replicate): logger.debug("Hook '%s' does not support replication; building each Newton world separately.", name) @@ -277,6 +262,21 @@ def _collect_replication_preparers( return preparers +def _validate_replication_hook_attributes( + hooks: Sequence[Callable[[ModelBuilder, int, np.ndarray, np.ndarray], None]], +) -> None: + """Reject hooks that define only half of the replication opt-in contract.""" + for hook in hooks: + can_replicate = getattr(hook, "_can_replicate_builder", None) + prepare = getattr(hook, "_prepare_builder_replication", None) + if callable(can_replicate) != callable(prepare): + missing = "_prepare_builder_replication" if callable(can_replicate) else "_can_replicate_builder" + raise TypeError( + f"Newton world-builder hook '{getattr(hook, '__qualname__', repr(hook))}' opts into replication" + f" but has no callable '{missing}'." + ) + + def replicate_builder_mapping( builder: ModelBuilder, sources: Sequence[str], @@ -306,6 +306,7 @@ def replicate_builder_mapping( quaternions = quaternions.astype(np.float32, copy=False) xforms_np = np.concatenate((positions, quaternions), axis=1) world_xforms = [wp.transform(*row) for row in xforms_np] + _validate_replication_hook_attributes(per_world_builder_hooks) if ( len(sources) == 1 @@ -330,6 +331,11 @@ def replicate_builder_mapping( for label, indices in source_site_indices.get(id(source_builder), {}).items(): site_local_indices.setdefault(label, []).extend(indices) + # Hook-added particles inherit the aggregate builder's model-wide velocity limit. + # A source that already owns particles keeps Newton's source-wins behavior. + if source_builder.particle_count == 0: + source_builder.particle_max_velocity = builder.particle_max_velocity + finish_callbacks = [ prepare(builder, source_builder, num_worlds, xforms_np[0, :3].copy(), xforms_np[0, 3:].copy()) for prepare in hook_preparers diff --git a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py index 0d9c84703bbc..e2aebc1fe908 100644 --- a/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py +++ b/source/isaaclab_newton/test/cloner/test_rename_builder_labels.py @@ -556,6 +556,80 @@ def hook(builder, *_): [f"/World/envs/env_{env_id}/Robot/{label}" for env_id in env_ids for label in ("base", "hook")], ) + def test_half_replication_hook_is_rejected_on_non_homogeneous_plan(self): + def hook(*_): + pass + + hook._can_replicate_builder = lambda _: True + sources = (f"{self._SRC}/a", f"{self._SRC}/b") + with self.assertRaisesRegex(TypeError, "has no callable '_prepare_builder_replication'"): + replicate_builder_mapping( + newton.ModelBuilder(), + sources, + np.eye(2, dtype=np.bool_), + np.zeros((2, 3), dtype=np.float32), + np.array([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=np.float32), + {source: newton.ModelBuilder() for source in sources}, + destinations=(f"{self._ENV}/a", f"{self._ENV}/b"), + env_ids=np.arange(2, dtype=np.int64), + per_world_builder_hooks=(hook,), + ) + + def test_hook_added_particles_inherit_destination_max_velocity(self): + source = newton.ModelBuilder() + builder = newton.ModelBuilder() + builder.particle_max_velocity = 17.0 + + def hook(*_): + pass + + def prepare(_builder, source_builder, *_): + self.assertEqual(source_builder.particle_max_velocity, 17.0) + source_builder.add_particle(pos=(0.0, 0.0, 0.0), vel=(0.0, 0.0, 0.0), mass=1.0, radius=0.1) + return lambda: None + + hook._can_replicate_builder = lambda _: True + hook._prepare_builder_replication = prepare + replicate_builder_mapping( + builder, + (self._SRC,), + np.ones((1, 2), dtype=np.bool_), + np.zeros((2, 3), dtype=np.float32), + np.array([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=np.float32), + {self._SRC: source}, + destinations=(f"{self._ENV}/Robot",), + env_ids=np.arange(2, dtype=np.int64), + per_world_builder_hooks=(hook,), + ) + self.assertEqual(source.particle_max_velocity, 17.0) + self.assertEqual(builder.particle_max_velocity, 17.0) + + def test_particle_source_keeps_own_max_velocity_with_replication_hook(self): + source = newton.ModelBuilder() + source.add_particle(pos=(0.0, 0.0, 0.0), vel=(0.0, 0.0, 0.0), mass=1.0, radius=0.1) + source.particle_max_velocity = 23.0 + builder = newton.ModelBuilder() + builder.particle_max_velocity = 17.0 + + def hook(*_): + pass + + hook._can_replicate_builder = lambda _: True + hook._prepare_builder_replication = lambda *_: lambda: None + replicate_builder_mapping( + builder, + (self._SRC,), + np.ones((1, 2), dtype=np.bool_), + np.zeros((2, 3), dtype=np.float32), + np.array([[0.0, 0.0, 0.0, 1.0]] * 2, dtype=np.float32), + {self._SRC: source}, + destinations=(f"{self._ENV}/Robot",), + env_ids=np.arange(2, dtype=np.int64), + per_world_builder_hooks=(hook,), + ) + self.assertEqual(source.particle_max_velocity, 23.0) + self.assertEqual(builder.particle_max_velocity, 23.0) + class TestRootJointNaming(unittest.TestCase): """The importer leaves a floating base's root joint unnamed; every other entity is named.""" From 49533ea662ae8f335cad71f25713d1f36556d6f9 Mon Sep 17 00:00:00 2001 From: camevor Date: Tue, 8 Sep 2026 10:50:47 +0200 Subject: [PATCH 3/3] Address PR comments --- .../deformable/deformable_object.py | 21 ++++++++----------- .../cloner/newton_clone_utils.py | 7 ++++--- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py b/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py index b970267d0d1e..d5182aa02fa0 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py +++ b/source/isaaclab_contrib/isaaclab_contrib/deformable/deformable_object.py @@ -191,18 +191,15 @@ def add_registered_deformables_to_builder( def _can_replicate_registered_deformables(xforms: np.ndarray) -> bool: - """Return whether Newton replication can apply every relative particle transform. + """Return False if any registered deformables prevent replication. - ``ModelBuilder.replicate`` offsets ``particle_q`` by each transform's translation and drops - its rotation. + 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 - rotations = xforms[:, 3:] - return bool( - np.allclose(rotations[:, :3], 0.0, atol=1.0e-6, rtol=0.0) - and np.allclose(np.abs(rotations[:, 3]), 1.0, atol=1.0e-6, rtol=0.0) - ) + 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( @@ -235,12 +232,12 @@ def finish_replication() -> None: 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 = [ - destination_particle_base + world * source_particle_stride + source_offset - for world in range(num_worlds) - ] + entry.particle_offsets = (world_particle_offsets + source_offset).tolist() return finish_replication diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py index f994fc48adc2..24f17712f70a 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/newton_clone_utils.py @@ -357,10 +357,11 @@ def replicate_builder_mapping( for finish in finish_callbacks: finish() + world_shape_offsets = base_shape + np.arange(num_worlds, dtype=np.int64) * stride for label, local_indices in site_local_indices.items(): - local_site_map[label] = [ - [base_shape + world * stride + local for local in local_indices] for world in range(num_worlds) - ] + local_site_map[label] = ( + world_shape_offsets[:, None] + np.asarray(local_indices, dtype=np.int64) + ).tolist() bindings = rename_builder_labels(builder, sources, destinations, env_ids, mapping, skip_entity_labels=True) return local_site_map, world_xforms, bindings