Skip to content

Commit 3ad2892

Browse files
committed
[Newton] Avoid repeated model startup work (#7295)
## Summary This is the model/articulation part of the scoped Newton startup work. It keeps startup work with the component that owns it: - each Newton manager declares only the custom builder schema used by its active solver; - articulation target modes are resolved for the prototype and copied to its replicas; - the base physics manager owns the articulation-view registry; - articulations create and register their view, while joint-wrench sensors reuse it or create it when used independently; - root expressions remain regular expressions instead of taking a lossy regex-to-glob round trip. The active-solver and reusable-view findings originated in Chris's #7269. This PR isolates those model/articulation changes so #7269 can remain the contact/raycast change under Chris's PR. This replaces the cloner's unconditional MuJoCo + Kamino registration and the duplicate joint-wrench view. There is no view scan, manager-specific cache API, compatibility fallback, or duplicate registry. The production diff against the current base is `+57/-60` (net `-3`). The explicit-global-ownership part is #7292. The contact/raycast part remains in #7269. ## Startup benchmark RTX 5090, CUDA device 1, 4096 environments, three fresh processes per revision/task. Values are median end-to-end startup wall time; raw runs are included below. Base: `86cf66651bd`. PR: `947869af173`. | Task | Base | PR | Change | |---|---:|---:|---:| | `Isaac-Cartpole` | 8.316 s | 7.655 s | -8.0% | | `Isaac-Velocity-Rough-UnitreeGo2` | 17.525 s | 14.844 s | -15.3% | | `Isaac-Lift-KukaAllegro-Camera` | 41.691 s | 36.832 s | -11.7% | Raw totals: - Cartpole base: 10.595, 8.190, 8.316 s; PR: 7.655, 7.674, 7.469 s. - Go2 rough base: 17.568, 17.525, 16.426 s; PR: 14.852, 14.844, 14.724 s. - Kuka camera base: 47.590, 41.691, 38.724 s; PR: 36.832, 36.666, 37.032 s. The measured `env_creation` medians improve from 6.473 to 5.801 s for Cartpole, 15.558 to 12.863 s for Go2 rough, and 37.481 to 32.578 s for Kuka camera. ## Test plan - `247 passed` across the physics-manager lifecycle, Newton cloner, manager abstraction, coupled-manager, and joint-wrench reuse tests. - `test_rename_builder_labels.py`: `17 passed` after removing obsolete solver-registration mocks. - Ruff check and format pass on all changed Python files. - The three 4096-environment benchmark tasks provide end-to-end Newton MJWarp startup coverage. (cherry picked from commit c4a2759)
1 parent c780f65 commit 3ad2892

24 files changed

Lines changed: 168 additions & 67 deletions

File tree

source/isaaclab/changelog.d/ooctipus-newton-active-solver-attributes.skip

Whitespace-only changes.

source/isaaclab/isaaclab/physics/physics_manager.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ class PhysicsManager(ABC):
8686
_sim_time: ClassVar[float] = 0.0
8787
_callbacks: ClassVar[dict[int, tuple[Any, Callable, int, str | None, Any]]] = {}
8888
_callback_id: ClassVar[int] = 0
89+
views: ClassVar[dict[tuple[type, str], Any]] = {}
8990

9091
@classmethod
9192
def _prepare_stage_creation(cls) -> None:
@@ -457,6 +458,7 @@ def close(cls) -> None:
457458
cls.clear_callbacks()
458459
finally:
459460
if is_active_manager:
461+
PhysicsManager.views.clear()
460462
PhysicsManager._sim = None
461463
PhysicsManager._cfg = None
462464
PhysicsManager._sim_time = 0.0

source/isaaclab/test/sim/test_physics_manager_lifecycle.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ class TestManager(PhysicsManager):
2626
monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(physics_manager=TestManager))
2727
monkeypatch.setattr(PhysicsManager, "_cfg", object())
2828
monkeypatch.setattr(PhysicsManager, "_sim_time", 1.0)
29+
monkeypatch.setattr(PhysicsManager, "views", {(TestManager, "/World/Robot"): object()})
2930

3031
TestManager.register_callback(
3132
lambda _payload: events.append("first"),
@@ -71,6 +72,7 @@ def failing_listener(_payload):
7172
assert PhysicsManager._sim is None
7273
assert PhysicsManager._cfg is None
7374
assert PhysicsManager._sim_time == 0.0
75+
assert PhysicsManager.views == {}
7476

7577

7678
def test_close_surfaces_stop_errors_stored_by_safe_callback_invoke(monkeypatch):
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Fixed
2+
^^^^^
3+
4+
* Registered every configured child solver's builder attributes for coupled Newton models.

source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -188,17 +188,15 @@ def _validate_resolved_entries(
188188
def _register_builder_attributes(cls, builder: ModelBuilder) -> None:
189189
"""Register custom attributes required by nested coupled entries."""
190190
super()._register_builder_attributes(builder)
191-
solver_cfg = getattr(PhysicsManager._cfg, "solver_cfg", None)
192-
if any(isinstance(entry.solver_cfg, MPMSolverCfg) for entry in getattr(solver_cfg, "entries", ())):
193-
NewtonMPMManager._register_builder_attributes(builder)
191+
for entry in PhysicsManager._cfg.solver_cfg.entries:
192+
entry.solver_cfg.class_type._register_builder_attributes(builder)
194193

195194
@classmethod
196195
def _prepare_builder_for_finalize(cls, builder: ModelBuilder) -> None:
197196
"""Normalize kinematic colliders when a coupled entry uses implicit MPM."""
198197
super()._prepare_builder_for_finalize(builder)
199-
solver_cfg = getattr(PhysicsManager._cfg, "solver_cfg", None)
200-
if any(isinstance(entry.solver_cfg, MPMSolverCfg) for entry in getattr(solver_cfg, "entries", ())):
201-
NewtonMPMManager._prepare_builder_for_finalize(builder)
198+
for entry in PhysicsManager._cfg.solver_cfg.entries:
199+
entry.solver_cfg.class_type._prepare_builder_for_finalize(builder)
202200

203201
@classmethod
204202
def _initialize_contacts(cls) -> None:

source/isaaclab_contrib/isaaclab_contrib/custom_coupling/coupled_mjwarp_vbd_manager.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
from newton import Contacts, Control, Model, State
1414
from newton.solvers import SolverBase, SolverMuJoCo, SolverVBD
1515

16+
from isaaclab.physics import PhysicsManager
17+
1618
from .kernels import _kernel_body_particle_reaction
1719
from .newton_manager_cfg import CoupledMJWarpVBDSolverCfg
1820

@@ -27,12 +29,11 @@ class NewtonCoupledMJWarpVBDManager(NewtonVBDManager):
2729
_rigid_solver: SolverMuJoCo | None = None
2830
_soft_solver: SolverVBD | None = None
2931
_coupling_mode: str | None = None
32+
_builder_attribute_solvers = (SolverMuJoCo,)
3033

3134
@classmethod
3235
def step(cls) -> None:
3336
"""Step the physics simulation."""
34-
from isaaclab.physics import PhysicsManager
35-
3637
sim = PhysicsManager._sim
3738
if sim is None or not sim.is_playing():
3839
return

source/isaaclab_contrib/test/coupling/test_coupler.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
XPBDSolverCfg,
3131
)
3232
from isaaclab_newton.physics.newton_manager import NewtonManager
33-
from newton import ShapeFlags
33+
from newton import ModelBuilder, ShapeFlags
3434
from newton.solvers.experimental.coupled import SolverCoupledADMM, SolverCoupledProxy
3535

3636
from isaaclab_contrib.coupling import (
@@ -579,6 +579,23 @@ def test_mpm_entry_reuses_builder_lifecycle_hooks(monkeypatch):
579579
assert events == [("register", builder), ("finalize", builder)]
580580

581581

582+
def test_nested_solvers_register_their_builder_attributes(monkeypatch):
583+
"""A coupled model declares the schemas consumed by each configured child solver."""
584+
builder = ModelBuilder()
585+
solver_cfg = CouplerProxyCfg(
586+
entries=[
587+
CouplerEntryCfg(name="rigid", solver_cfg=MJWarpSolverCfg()),
588+
CouplerEntryCfg(name="media", solver_cfg=MPMSolverCfg()),
589+
]
590+
)
591+
monkeypatch.setattr(coupler.PhysicsManager, "_cfg", SimpleNamespace(solver_cfg=solver_cfg))
592+
593+
NewtonCouplerManager._register_builder_attributes(builder)
594+
595+
assert builder.has_custom_attribute("mujoco:condim")
596+
assert builder.has_custom_attribute("mpm:young_modulus")
597+
598+
582599
def test_contact_initialization_prepares_coupled_solver_buffers(monkeypatch):
583600
"""Entry-local contact buffers are allocated before graph capture."""
584601
events: list[tuple[str, object | None]] = []

source/isaaclab_contrib/test/custom_coupling/test_manager.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,29 @@
55

66
"""Unit tests for the custom coupling manager."""
77

8+
from types import SimpleNamespace
89
from unittest.mock import MagicMock
910

1011
import pytest
1112
from isaaclab_newton.physics import MJWarpSolverCfg, VBDSolverCfg
13+
from newton import ModelBuilder
1214

1315
import isaaclab_contrib.custom_coupling.coupled_mjwarp_vbd_manager as manager_module
1416
from isaaclab_contrib.custom_coupling.coupled_mjwarp_vbd_manager import NewtonCoupledMJWarpVBDManager
1517
from isaaclab_contrib.custom_coupling.newton_manager_cfg import CoupledMJWarpVBDSolverCfg
1618

1719

20+
def test_register_builder_attributes_includes_nested_solvers(monkeypatch: pytest.MonkeyPatch) -> None:
21+
"""The custom coupled manager delegates builder setup to both configured children."""
22+
cfg = CoupledMJWarpVBDSolverCfg()
23+
monkeypatch.setattr(manager_module.PhysicsManager, "_cfg", SimpleNamespace(solver_cfg=cfg))
24+
builder = ModelBuilder()
25+
26+
NewtonCoupledMJWarpVBDManager._register_builder_attributes(builder)
27+
28+
assert builder.has_custom_attribute("mujoco:condim")
29+
30+
1831
def test_reset_forwards_to_both_subsolvers(monkeypatch: pytest.MonkeyPatch) -> None:
1932
"""Reset the real sub-solvers instead of the dummy solver slot."""
2033
rigid_solver = MagicMock()
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Changed
2+
^^^^^^^
3+
4+
* Registered builder attributes only for the active Newton solver instead of importing and allocating inactive
5+
solver data.
6+
* Reused target-mode resolution across identical articulation clones and one canonical articulation view between
7+
each articulation and its joint-wrench sensor.

source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from __future__ import annotations
1010

1111
import logging
12+
import re
1213
import warnings
1314
from collections.abc import Sequence
1415
from typing import TYPE_CHECKING
@@ -27,7 +28,7 @@
2728
from isaaclab.assets.articulation import ordering_kernels
2829
from isaaclab.assets.articulation.base_articulation import BaseArticulation
2930
from isaaclab.physics import PhysicsEvent
30-
from isaaclab.sim.utils.queries import path_expr_to_glob, resolve_matching_prims_from_source
31+
from isaaclab.sim.utils.queries import resolve_matching_prims_from_source
3132
from isaaclab.utils.string import resolve_matching_names, resolve_matching_names_values
3233
from isaaclab.utils.version import get_isaac_sim_version, has_kit
3334
from isaaclab.utils.warp import ProxyArray
@@ -84,10 +85,11 @@ def has_articulation_root_api(prim) -> bool:
8485

8586
def _configure_builder_joint_target_modes(builder, cfg: ArticulationCfg) -> None:
8687
"""Resolve configured actuator gains into Newton builder target modes before finalization."""
87-
root_prim_path_regex = path_expr_to_glob(_resolve_articulation_root_prim_path_expr(cfg)).replace("*", ".*")
88+
root_prim_path_regex = _resolve_articulation_root_prim_path_expr(cfg)
8889
articulation_ids, _ = resolve_matching_names(
8990
root_prim_path_regex, builder.articulation_label, raise_when_no_match=False
9091
)
92+
source_dof_ids = None
9193
for articulation_id in articulation_ids:
9294
joint_start = builder.articulation_start[articulation_id]
9395
joint_end = builder.articulation_end[articulation_id]
@@ -107,6 +109,11 @@ def _configure_builder_joint_target_modes(builder, cfg: ArticulationCfg) -> None
107109
dof_ids.append(dof_id)
108110
dof_names.append(joint_name if dof_end - dof_start == 1 else f"{joint_name}:{axis_index}")
109111

112+
if source_dof_ids is not None:
113+
for source_dof_id, dof_id in zip(source_dof_ids, dof_ids, strict=True):
114+
builder.joint_target_mode[dof_id] = builder.joint_target_mode[source_dof_id]
115+
continue
116+
110117
for actuator_cfg in cfg.actuators.values():
111118
matched_indices, matched_names = resolve_matching_names(
112119
actuator_cfg.joint_names_expr, dof_names, raise_when_no_match=False
@@ -130,6 +137,7 @@ def _configure_builder_joint_target_modes(builder, cfg: ArticulationCfg) -> None
130137
if _is_implicit_actuator_cfg(actuator_cfg)
131138
else JointTargetMode.EFFORT
132139
)
140+
source_dof_ids = dof_ids
133141

134142

135143
class Articulation(BaseArticulation):
@@ -3279,19 +3287,14 @@ def write_spatial_tendon_properties_to_sim_mask(
32793287
"""
32803288

32813289
def _initialize_impl(self):
3282-
# obtain global simulation view
3283-
self._physics_sim_view = SimulationManager.get_physics_sim_view()
3284-
32853290
root_prim_path_expr = _resolve_articulation_root_prim_path_expr(self.cfg)
32863291
# -- articulation
3287-
self._root_view = ArticulationView(
3292+
self._root_view = SimulationManager.views[SimulationManager, root_prim_path_expr] = ArticulationView(
32883293
SimulationManager.get_model(),
3289-
path_expr_to_glob(root_prim_path_expr),
3294+
re.compile(root_prim_path_expr),
32903295
verbose=False,
32913296
exclude_joint_types=[JointType.FREE, JointType.FIXED],
32923297
)
3293-
# Register view with Newton manager so sensors (e.g. FrameTransformer) can find it.
3294-
SimulationManager.get_physics_sim_view().append(self._root_view)
32953298

32963299
# container for data access
32973300
self._data = ArticulationData(self.root_view, self.device)

0 commit comments

Comments
 (0)