Harden batched Newton replication

Author: camevor

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

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."""