Fix Implicit MPM reset masks

Adapt coupled and standalone Implicit MPM reset masks to the solver's
world-mask contract, and skip unsupported selective shared-world resets.
Add regression tests and changelog fragments.

Author: huidongc

source/isaaclab_contrib/changelog.d/mpm-coupled-reset-mask.rst

@@ -0,0 +1,11 @@
+Fixed
+^^^^^
+
+* Fixed coupled Implicit MPM environment resets raising
+  ``ValueError: world_mask has shape ...`` or
+  ``RuntimeError: Masked reset cannot selectively clear grid-backed warm
+  starts`` when :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state`
+  forwarded Isaac Lab's ``(world_count,)`` mask. MPM entry ``reset`` now receives
+  the ``(world_count + 1,)`` mask required by
+  :meth:`newton.solvers.SolverImplicitMPM.reset` (or skips selective shared-world
+  resets), while MJWarp entries keep the original parent mask.

source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py

@@ -18,9 +18,14 @@
     NewtonCollisionPipelineCfg,
     NewtonSolverCfg,
 )
-from isaaclab_newton.physics.mpm_manager import NewtonMPMManager
+from isaaclab_newton.physics.mpm_manager import (
+    NewtonMPMManager,
+    adapt_world_mask_for_implicit_mpm,
+    should_skip_implicit_mpm_masked_reset,
+)
 from isaaclab_newton.physics.newton_manager import NewtonManager
 from newton import CollisionPipeline, Model, ModelBuilder, ShapeFlags
+from newton.solvers import SolverImplicitMPM
 from newton.solvers.experimental.coupled import SolverCoupled, SolverCoupledADMM, SolverCoupledProxy
 
 from isaaclab.physics import PhysicsManager
@@ -114,6 +119,54 @@ def _build_solver(cls, model: Model, solver_cfg: CouplerCfg) -> None:
         NewtonManager._supports_contact_sensors = False
         NewtonManager._needs_collision_pipeline = needs_collision_pipeline
         NewtonManager._supports_rigid_body_force_input = True
+        cls._install_implicit_mpm_reset_mask_adapter(
+            NewtonManager._solver,
+            [entry.config.name for entry in resolved_entries],
+        )
+
+    @classmethod
+    def _install_implicit_mpm_reset_mask_adapter(cls, solver: object, entry_names: list[str]) -> None:
+        """Adapt coupled-reset masks for Implicit MPM's world-mask contract.
+
+        :class:`~newton.solvers.experimental.coupled.SolverCoupled` forwards the
+        parent ``(world_count,)`` mask to every entry. MJWarp expects that shape,
+        while :class:`~newton.solvers.SolverImplicitMPM` expects one extra bit for
+        global world ``-1``. Shared multi-world Implicit MPM
+        (``separate_worlds=False``) also rejects selective masks when clearing
+        grid-backed warm starts; those resets are skipped. Wrapping only MPM
+        entry ``reset`` methods keeps the coupled parent API unchanged.
+        """
+        # SolverCoupled exposes sub-solvers via ``solver(name)`` (not ``entry_solver``).
+        get_entry_solver = getattr(solver, "solver", None)
+        if not callable(get_entry_solver):
+            return
+
+        for entry_name in entry_names:
+            try:
+                entry_solver = get_entry_solver(entry_name)
+            except (KeyError, TypeError, ValueError):
+                continue
+            if not isinstance(entry_solver, SolverImplicitMPM):
+                continue
+            original_reset = entry_solver.reset
+
+            def _reset(
+                state,
+                world_mask=None,
+                flags=None,
+                *,
+                _original_reset=original_reset,
+                _entry_solver=entry_solver,
+            ):
+                if should_skip_implicit_mpm_masked_reset(_entry_solver, world_mask):
+                    return None
+                return _original_reset(
+                    state,
+                    world_mask=adapt_world_mask_for_implicit_mpm(_entry_solver, world_mask),
+                    flags=flags,
+                )
+
+            entry_solver.reset = _reset  # type: ignore[method-assign]
 
     @classmethod
     def _validate_config(cls, solver_cfg: CouplerCfg) -> None:

source/isaaclab_contrib/test/coupling/test_coupler.py

@@ -864,3 +864,41 @@ def test_admm_build_auto_detects_symmetric_contact_pairs_by_default(monkeypatch)
     cfg.contact_pairs = []
     solver = NewtonCouplerManager._build_admm_coupled_solver(model, entries, cfg)
     assert list(solver.coupling.contact_pairs) == []
+
+
+def test_install_implicit_mpm_reset_mask_adapter_pads_entry_masks(monkeypatch):
+    """Coupled MPM entry resets pad full masks and skip selective shared-world ones."""
+    import warp as wp
+
+    recorded: list[list[bool] | None] = []
+
+    class _FakeImplicitMPM:
+        def __init__(self):
+            self.model = SimpleNamespace(world_count=4)
+            self._separate_worlds = False
+
+        def reset(self, state, world_mask=None, flags=None):
+            del state, flags
+            recorded.append(None if world_mask is None else world_mask.numpy().tolist())
+
+    monkeypatch.setattr(coupler, "SolverImplicitMPM", _FakeImplicitMPM)
+    mpm_solver = _FakeImplicitMPM()
+
+    class _FakeCoupled:
+        def solver(self, name: str):
+            assert name == "mpm"
+            return mpm_solver
+
+    NewtonCouplerManager._install_implicit_mpm_reset_mask_adapter(_FakeCoupled(), ["mpm"])
+
+    # Shared multi-world selective masks are skipped (no call into original reset).
+    mpm_solver.reset(object(), world_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu"))
+    assert recorded == []
+
+    mpm_solver.reset(object(), world_mask=wp.array([True, True, True, True], dtype=wp.bool, device="cpu"))
+    assert recorded == [None]
+
+    # With separate worlds, selective masks are forwarded as (N+1,).
+    mpm_solver._separate_worlds = True
+    mpm_solver.reset(object(), world_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu"))
+    assert recorded == [None, [True, False, True, False, False]]

source/isaaclab_newton/changelog.d/mpm-world-mask-shape.rst

@@ -0,0 +1,8 @@
+Fixed
+^^^^^
+
+* Fixed :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state` under
+  Implicit MPM rejecting Isaac Lab ``(world_count,)`` masks. The manager now
+  expands them to the ``(world_count + 1,)`` shape required by
+  :meth:`newton.solvers.SolverImplicitMPM.reset`, and skips selective masks when
+  shared multi-world Implicit MPM cannot clear grid-backed warm starts per world.

source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py

@@ -7,15 +7,94 @@
 
 from __future__ import annotations
 
+import numpy as np
 import warp as wp
-from newton import BodyFlags, Contacts, Control, GeoType, Model, ModelBuilder, State
+from newton import BodyFlags, Contacts, Control, GeoType, Model, ModelBuilder, State, StateFlags
 from newton.solvers import SolverImplicitMPM
 from warp.fem import TemporaryStore
 
 from .mpm_manager_cfg import MPMSolverCfg
 from .newton_manager import NewtonManager
 
 
+def adapt_world_mask_for_implicit_mpm(
+    solver: SolverImplicitMPM,
+    world_mask: wp.array | None,
+) -> wp.array | None:
+    """Adapt an Isaac Lab world-reset mask to Implicit MPM's mask contract.
+
+    Isaac Lab and most Newton solvers use a per-world mask of shape
+    ``(world_count,)``. :meth:`SolverImplicitMPM.reset` instead expects
+    ``(world_count + 1,)``, where the trailing entry selects global objects
+    whose world index is ``-1``.
+
+    Args:
+        solver: Implicit MPM solver whose model defines ``world_count``.
+        world_mask: Optional Isaac Lab / Newton mask of shape ``(world_count,)``
+            or an already-adapted Implicit MPM mask of shape
+            ``(world_count + 1,)``.
+
+    Returns:
+        ``None`` when no mask is provided or every world is selected (a full
+        reset, including global world ``-1``). Otherwise a boolean Warp array
+        of shape ``(world_count + 1,)`` with the global bit left ``False``.
+
+    Raises:
+        ValueError: If ``world_mask`` has neither ``(world_count,)`` nor
+            ``(world_count + 1,)`` shape.
+    """
+    if world_mask is None:
+        return None
+
+    local_selected = _local_world_selection(solver, world_mask)
+    if bool(np.all(local_selected)):
+        # Full world selection is equivalent to an unmasked reset and also
+        # covers global (world index -1) particle/collider history.
+        return None
+
+    world_count = int(solver.model.world_count)
+    if tuple(world_mask.shape) == (world_count + 1,):
+        return world_mask
+
+    padded = np.zeros(world_count + 1, dtype=bool)
+    padded[:-1] = local_selected
+    return wp.array(padded, dtype=wp.bool, device=world_mask.device)
+
+
+def should_skip_implicit_mpm_masked_reset(
+    solver: SolverImplicitMPM,
+    world_mask: wp.array | None,
+) -> bool:
+    """Whether a masked Implicit MPM reset must be skipped.
+
+    Shared multi-world Implicit MPM (``separate_worlds=False``) rejects selective
+    masks when clearing grid-backed warm starts. Matching
+    :meth:`NewtonMPMManager._reset_solver_internals`, Isaac Lab skips those
+    resets instead of raising. Full-world selection still proceeds as an
+    unmasked reset.
+    """
+    if world_mask is None:
+        return False
+    if bool(getattr(solver, "_separate_worlds", False)) or int(solver.model.world_count) <= 1:
+        return False
+    return not bool(np.all(_local_world_selection(solver, world_mask)))
+
+
+def _local_world_selection(solver: SolverImplicitMPM, world_mask: wp.array) -> np.ndarray:
+    """Return the per-world selection bits from an Isaac Lab or Implicit MPM mask."""
+    world_count = int(solver.model.world_count)
+    shape = tuple(world_mask.shape)
+    selected = world_mask.numpy()
+    if shape == (world_count + 1,):
+        return selected[:-1]
+    if shape == (world_count,):
+        return selected
+    raise ValueError(
+        f"world_mask has shape {shape}, expected ({world_count},) or ({world_count + 1},) "
+        "for SolverImplicitMPM.reset."
+    )
+
+
 def _make_solver_config(solver_cfg: MPMSolverCfg) -> SolverImplicitMPM.Config:
     """Build Newton's implicit MPM solver config from Isaac Lab's cfg."""
     return SolverImplicitMPM.Config(
@@ -161,6 +240,34 @@ def _reset_solver_internals(cls, world_mask: wp.array | None) -> None:
             world_mask: Per-world reset mask, ignored.
         """
 
+    @classmethod
+    def reset_solver_state(
+        cls,
+        state: State | None = None,
+        world_mask: wp.array(dtype=wp.bool) | None = None,
+        flags: StateFlags | int | None = None,
+    ) -> None:
+        """Reset Implicit MPM history after simulation state is rewritten.
+
+        Expands Isaac Lab's ``(world_count,)`` mask to the
+        ``(world_count + 1,)`` shape required by :meth:`SolverImplicitMPM.reset`
+        before delegating to :meth:`NewtonManager.reset_solver_state`. Shared
+        multi-world selective masks are skipped; see
+        :func:`should_skip_implicit_mpm_masked_reset`.
+        """
+        if not isinstance(cls._solver, SolverImplicitMPM):
+            raise RuntimeError(
+                f"{cls.__name__}.reset_solver_state requires an active SolverImplicitMPM; "
+                f"got {type(cls._solver).__name__}."
+            )
+        if should_skip_implicit_mpm_masked_reset(cls._solver, world_mask):
+            return
+        super().reset_solver_state(
+            state=state,
+            world_mask=adapt_world_mask_for_implicit_mpm(cls._solver, world_mask),
+            flags=flags,
+        )
+
     @classmethod
     def _solver_specific_clear(cls) -> None:
         """Reset MPM-specific class state on teardown.

source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py

@@ -632,6 +632,45 @@ def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch):
     assert NewtonManager._graph_capture_pending is False
 
 
+def test_adapt_world_mask_for_implicit_mpm_pads_selective_masks():
+    """Isaac Lab (world_count,) masks expand to Implicit MPM's (world_count + 1,) contract."""
+    from isaaclab_newton.physics.mpm_manager import (
+        adapt_world_mask_for_implicit_mpm,
+        should_skip_implicit_mpm_masked_reset,
+    )
+
+    solver = SimpleNamespace(model=SimpleNamespace(world_count=4), _separate_worlds=False)
+    world_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu")
+
+    assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is True
+    adapted = adapt_world_mask_for_implicit_mpm(solver, world_mask)
+
+    assert adapted is not None
+    assert adapted.numpy().tolist() == [True, False, True, False, False]
+
+    solver._separate_worlds = True
+    assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is False
+
+
+def test_adapt_world_mask_for_implicit_mpm_full_selection_becomes_unmasked():
+    """Selecting every world is equivalent to an unmasked Implicit MPM reset."""
+    from isaaclab_newton.physics.mpm_manager import (
+        adapt_world_mask_for_implicit_mpm,
+        should_skip_implicit_mpm_masked_reset,
+    )
+
+    solver = SimpleNamespace(model=SimpleNamespace(world_count=3), _separate_worlds=False)
+    world_mask = wp.array([True, True, True], dtype=wp.bool, device="cpu")
+
+    assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is False
+    assert adapt_world_mask_for_implicit_mpm(solver, world_mask) is None
+    assert adapt_world_mask_for_implicit_mpm(solver, None) is None
+    assert should_skip_implicit_mpm_masked_reset(solver, None) is False
+
+    empty = wp.array([False, False, False], dtype=wp.bool, device="cpu")
+    assert should_skip_implicit_mpm_masked_reset(solver, empty) is True
+
+
 def test_cuda_graph_capture_uses_simulation_device(monkeypatch):
     """CUDA graph capture should use the simulation device instead of Warp's default device."""
     from isaaclab.physics import PhysicsManager