Update to new newton modelbuilder and experimental coupled solvers

source/isaaclab_newton/isaaclab_newton/physics/coupled_manager.py

@@ -14,11 +14,11 @@
 
 import numpy as np
 from newton import CollisionPipeline, Model
-from newton.solvers.coupled_experimental import (
+from newton.solvers.experimental.coupled import (
     CouplingInterface,
-    SolverAdmmCoupled,
     SolverCoupled,
-    SolverProxyCoupled,
+    SolverCoupledAdmm,
+    SolverCoupledProxy,
 )
 
 from isaaclab.managers import SceneEntityCfg
@@ -44,7 +44,7 @@
     from isaaclab.scene import InteractiveSceneCfg
 
 
-class SolverOneWayCoupled(SolverProxyCoupled):
+class SolverOneWayCoupled(SolverCoupledProxy):
     """Proxy-coupled solver that only transfers source motion to destination proxies.
 
     This intentionally reuses Newton's proxy collision/contact preparation path
@@ -137,10 +137,10 @@ def _build_solver(cls, model: Model, solver_cfg: CoupledSolverCfg) -> None:
         if solver_cfg.coupling_type == "base":
             NewtonManager._solver = SolverCoupled(model=model, entries=entries)
         elif solver_cfg.coupling_type == "proxy":
-            NewtonManager._solver = SolverProxyCoupled(
+            NewtonManager._solver = SolverCoupledProxy(
                 model=model,
                 entries=entries,
-                coupling=SolverProxyCoupled.Config(
+                coupling=SolverCoupledProxy.Config(
                     proxies=[cls._build_proxy(proxy_cfg) for proxy_cfg in solver_cfg.proxy_coupling.proxies],
                     iterations=solver_cfg.proxy_coupling.iterations,
                 ),
@@ -149,13 +149,13 @@ def _build_solver(cls, model: Model, solver_cfg: CoupledSolverCfg) -> None:
             NewtonManager._solver = SolverOneWayCoupled(
                 model=model,
                 entries=entries,
-                coupling=SolverProxyCoupled.Config(
+                coupling=SolverCoupledProxy.Config(
                     proxies=[cls._build_proxy(proxy_cfg) for proxy_cfg in solver_cfg.one_way_coupling.proxies],
                     iterations=1,
                 ),
             )
         elif solver_cfg.coupling_type == "admm":
-            NewtonManager._solver = SolverAdmmCoupled(
+            NewtonManager._solver = SolverCoupledAdmm(
                 model=model,
                 entries=entries,
                 coupling=cls._build_admm(solver_cfg.admm_coupling, entries),
@@ -594,14 +594,14 @@ def _entry_uses_local_collision_pipeline(entry_cfg: CoupledSolverEntryCfg) -> bo
         )
 
     @classmethod
-    def _build_proxy(cls, proxy_cfg: CoupledProxyCfg) -> SolverProxyCoupled.Proxy:
+    def _build_proxy(cls, proxy_cfg: CoupledProxyCfg) -> SolverCoupledProxy.Proxy:
         """Build a Newton proxy mapping from an Isaac Lab proxy cfg."""
         if not proxy_cfg.source or not proxy_cfg.destination:
             raise ValueError("CoupledProxyCfg source and destination must be non-empty.")
         if not proxy_cfg.bodies and not proxy_cfg.particles:
             raise ValueError("CoupledProxyCfg must map at least one body or particle.")
 
-        return SolverProxyCoupled.Proxy(
+        return SolverCoupledProxy.Proxy(
             source=proxy_cfg.source,
             destination=proxy_cfg.destination,
             bodies=list(proxy_cfg.bodies),
@@ -628,21 +628,21 @@ def _build_proxy_mode(mode: str | int) -> str:
     @classmethod
     def _build_admm(
         cls, admm_cfg: AdmmCouplingCfg, entries: list[SolverCoupled.Entry] | None = None
-    ) -> SolverAdmmCoupled.Config:
+    ) -> SolverCoupledAdmm.Config:
         """Build a Newton ADMM coupling config from an Isaac Lab cfg."""
         contact_pairs = [cls._build_admm_contact_pair(pair_cfg) for pair_cfg in admm_cfg.contact_pairs]
         if admm_cfg.auto_contact_pairs:
             if entries is None:
                 raise ValueError("AdmmCouplingCfg.auto_contact_pairs requires coupled solver entries.")
             contact_pairs.extend(
-                SolverAdmmCoupled.auto_detect_contact_pairs(
+                SolverCoupledAdmm.auto_detect_contact_pairs(
                     entries,
                     contact_distance=admm_cfg.auto_contact_distance,
                     detection_margin=admm_cfg.auto_detection_margin,
                 )
             )
 
-        return SolverAdmmCoupled.Config(
+        return SolverCoupledAdmm.Config(
             iterations=admm_cfg.iterations,
             rho=admm_cfg.rho,
             gamma=admm_cfg.gamma,
@@ -655,9 +655,9 @@ def _build_admm(
         )
 
     @staticmethod
-    def _build_admm_contact_pair(pair_cfg: AdmmContactPairCfg) -> SolverAdmmCoupled.ContactPair:
+    def _build_admm_contact_pair(pair_cfg: AdmmContactPairCfg) -> SolverCoupledAdmm.ContactPair:
         """Build a Newton ADMM contact-pair config from an Isaac Lab cfg."""
-        return SolverAdmmCoupled.ContactPair(
+        return SolverCoupledAdmm.ContactPair(
             source=pair_cfg.source,
             destination=pair_cfg.destination,
             contact_distance=pair_cfg.contact_distance,

source/isaaclab_newton/isaaclab_newton/physics/coupled_manager_cfg.py

@@ -166,7 +166,7 @@ class CoupledProxyCfg:
     """Scale factor for proxy mass/inertia in the destination view."""
 
     mode: Literal["lagged", "staggered"] | int = "lagged"
-    """Proxy transfer mode passed to Newton's ``SolverProxyCoupled``."""
+    """Proxy transfer mode passed to Newton's ``SolverCoupledProxy``."""
 
     collision_pipeline_factory: Callable | None = None
     """Optional factory for a proxy-local collision pipeline.
@@ -184,7 +184,7 @@ class ProxyCouplingCfg:
     """Lagged-impulse proxy coupling configuration."""
 
     proxies: list[CoupledProxyCfg] = field(default_factory=list)
-    """Proxy mappings used by ``SolverProxyCoupled``."""
+    """Proxy mappings used by ``SolverCoupledProxy``."""
 
     iterations: int = 1
     """Number of proxy relaxation passes per coupled step."""

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -828,9 +828,9 @@ def _register_solver_custom_attributes(cls, builder: ModelBuilder) -> None:
 
             SolverImplicitMPM.register_custom_attributes(builder)
         if getattr(solver_cfg, "coupling_type", None) == "admm":
-            from newton.solvers.coupled_experimental import SolverAdmmCoupled
+            from newton.solvers.experimental.coupled import SolverCoupledAdmm
 
-            SolverAdmmCoupled.register_custom_attributes(builder)
+            SolverCoupledAdmm.register_custom_attributes(builder)
 
     @classmethod
     def _set_fk_articulation_filter(cls, mask: np.ndarray | list[bool] | None) -> None:

source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py

@@ -57,11 +57,11 @@
 )
 from newton.solvers import SolverFeatherstone, SolverImplicitMPM, SolverKamino, SolverMuJoCo, SolverXPBD
 try:
-    from newton.solvers.coupled_experimental import SolverAdmmCoupled, SolverCoupled, SolverProxyCoupled
+    from newton.solvers.experimental.coupled import SolverCoupled, SolverCoupledAdmm, SolverCoupledProxy
 except ImportError:
-    SolverAdmmCoupled = None
+    SolverCoupledAdmm = None
     SolverCoupled = None
-    SolverProxyCoupled = None
+    SolverCoupledProxy = None
 
 from isaaclab.sim import SimulationCfg, build_simulation_context
 
@@ -174,7 +174,7 @@
             use_collision_pipeline=True,
         ),
         NewtonCoupledManager,
-        SolverProxyCoupled,
+        SolverCoupledProxy,
         False,
         True,
         id="proxy_coupled_xpbd_body_particle",
@@ -232,7 +232,7 @@
             use_collision_pipeline=False,
         ),
         NewtonCoupledManager,
-        SolverAdmmCoupled,
+        SolverCoupledAdmm,
         False,
         False,
         id="admm_coupled_xpbd_body_particle",
@@ -573,8 +573,8 @@ def test_initialize_solver_populates_canonical_state(
                 radius_mean=0.02,
             )
         elif (
-            SolverProxyCoupled is not None
-            and issubclass(expected_solver_cls, SolverProxyCoupled)
+            SolverCoupledProxy is not None
+            and issubclass(expected_solver_cls, SolverCoupledProxy)
         ):
             body = builder.add_body(mass=1.0)
             builder.add_shape_box(body, hx=0.05, hy=0.05, hz=0.05)
@@ -594,7 +594,7 @@ def test_initialize_solver_populates_canonical_state(
                 mass=0.1,
                 radius=0.02,
             )
-        elif SolverAdmmCoupled is not None and expected_solver_cls is SolverAdmmCoupled:
+        elif SolverCoupledAdmm is not None and expected_solver_cls is SolverCoupledAdmm:
             assert builder.has_custom_attribute("coupling:body_particle_attachment_body")
             body = builder.add_body(mass=1.0)
             particle = builder.add_particle(
@@ -603,7 +603,7 @@ def test_initialize_solver_populates_canonical_state(
                 mass=0.1,
                 radius=0.02,
             )
-            SolverAdmmCoupled.add_body_particle_attachment(builder, body, particle, stiffness=10.0)
+            SolverCoupledAdmm.add_body_particle_attachment(builder, body, particle, stiffness=10.0)
         else:
             # Pre-populate the builder with a minimal scene so MJCF conversion has
             # something to work with.
@@ -620,11 +620,11 @@ def test_initialize_solver_populates_canonical_state(
         if SolverCoupled is not None and expected_solver_cls is SolverCoupled:
             assert NewtonCoupledManager.get_entry_solver("rigid") is not None
             assert NewtonCoupledManager.get_entry_solver("particle") is not None
-        if SolverProxyCoupled is not None and issubclass(expected_solver_cls, SolverProxyCoupled):
+        if SolverCoupledProxy is not None and issubclass(expected_solver_cls, SolverCoupledProxy):
             rigid_solver = NewtonCoupledManager.get_entry_solver("rigid")
             assert rigid_solver is not None
             assert NewtonCoupledManager.get_entry_solver("particle") is not None
-        if SolverAdmmCoupled is not None and expected_solver_cls is SolverAdmmCoupled:
+        if SolverCoupledAdmm is not None and expected_solver_cls is SolverCoupledAdmm:
             assert NewtonCoupledManager.get_entry_solver("rigid") is not None
             assert NewtonCoupledManager.get_entry_solver("particle") is not None
         assert NewtonManager._use_single_state is expected_use_single_state

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/waterhose_robot_demo/coupled_builder.py

@@ -593,15 +593,12 @@ def _copy_body_shapes(
         shape_type = int(src_builder.shape_type[shape_id])
         scale = src_builder.shape_scale[shape_id]
         xform = src_builder.shape_transform[shape_id]
-        pos = wp.transform_get_translation(xform)
-        rot = wp.transform_get_rotation(xform)
 
         if shape_type == int(GeoType.SPHERE):
             copied_id = dst_builder.add_shape_sphere(
                 body=dst_body_id,
                 radius=float(scale[0]),
-                pos=pos,
-                rot=rot,
+                xform=xform,
                 cfg=cfg,
             )
         elif shape_type == int(GeoType.BOX):
@@ -610,17 +607,15 @@ def _copy_body_shapes(
                 hx=float(scale[0]),
                 hy=float(scale[1]),
                 hz=float(scale[2]),
-                pos=pos,
-                rot=rot,
+                xform=xform,
                 cfg=cfg,
             )
         elif shape_type == int(GeoType.CAPSULE):
             copied_id = dst_builder.add_shape_capsule(
                 body=dst_body_id,
                 radius=float(scale[0]),
                 half_height=float(scale[1]),
-                pos=pos,
-                rot=rot,
+                xform=xform,
                 cfg=cfg,
             )
         elif shape_type == int(GeoType.MESH):

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/waterhose_robot_demo/simulation.py

@@ -1232,13 +1232,13 @@ def _create_proxy_bodies(self, vbd_builder):
                 shape_xform_data = shape_transform_np[shape_idx]
                 pos = wp.vec3(shape_xform_data[0], shape_xform_data[1], shape_xform_data[2])
                 rot = wp.quat(shape_xform_data[3], shape_xform_data[4], shape_xform_data[5], shape_xform_data[6])
+                shape_xform = wp.transform(p=pos, q=rot)
 
                 if shape_type == GeoType.SPHERE:
                     sid = vbd_builder.add_shape_sphere(
                         body=proxy_body_id,
                         radius=float(shape_scale[0]),
-                        pos=pos,
-                        rot=rot,
+                        xform=shape_xform,
                         cfg=proxy_shape_cfg,
                     )
                     shape_ids.append(int(sid))
@@ -1248,8 +1248,7 @@ def _create_proxy_bodies(self, vbd_builder):
                         hx=float(shape_scale[0]),
                         hy=float(shape_scale[1]),
                         hz=float(shape_scale[2]),
-                        pos=pos,
-                        rot=rot,
+                        xform=shape_xform,
                         cfg=proxy_shape_cfg,
                     )
                     shape_ids.append(int(sid))
@@ -1258,14 +1257,12 @@ def _create_proxy_bodies(self, vbd_builder):
                         body=proxy_body_id,
                         radius=float(shape_scale[0]),
                         half_height=float(shape_scale[1]),
-                        pos=pos,
-                        rot=rot,
+                        xform=shape_xform,
                         cfg=proxy_shape_cfg,
                     )
                     shape_ids.append(int(sid))
                 elif shape_type == GeoType.MESH:
                     shape_source = self.mujoco_model.shape_source[shape_idx]
-                    shape_xform = wp.transform(p=pos, q=rot)
                     sid = vbd_builder.add_shape_mesh(
                         body=proxy_body_id,
                         mesh=shape_source,