diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 37fc9e02fcd..80555641c91 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -151,6 +151,7 @@ Guidelines for modifications: * Muhong Guo * Narendra Dahile * Neel Anand Jawale +* NeoZng * Nicola Loi * Nicholas Blauch * Nicolas Moenne-Loccoz diff --git a/source/isaaclab_newton/changelog.d/neozng-mjwarp-usd-joint-properties.rst b/source/isaaclab_newton/changelog.d/neozng-mjwarp-usd-joint-properties.rst new file mode 100644 index 00000000000..f144c98aab3 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/neozng-mjwarp-usd-joint-properties.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed MJWarp USD import dropping ``mjc:frictionloss`` by enabling the MuJoCo + schema resolver for MJWarp stage imports. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index 8d0f201ffe3..e76f20aef20 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -14,7 +14,6 @@ import torch import warp as wp from newton import ModelBuilder -from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx from pxr import Usd @@ -119,8 +118,8 @@ def _build_newton_builder_from_mapping( quaternions = torch.zeros((mapping.size(1), 4), device=mapping.device, dtype=torch.float32) quaternions[:, 3] = 1.0 - schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] manager_cls = PhysicsManager._sim.physics_manager + schema_resolvers = manager_cls._get_usd_import_schema_resolvers() builder = manager_cls.create_builder(up_axis=up_axis) import_paths = (PhysicsManager._sim.cfg.physics_prim_path, *global_paths) diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py index ea1ab276970..572271a7895 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mjwarp_manager.py @@ -11,8 +11,9 @@ import numpy as np import warp as wp -from newton import Contacts, Model +from newton import Contacts, Model, ModelBuilder from newton.solvers import SolverMuJoCo +from newton.usd import SchemaResolver, SchemaResolverMjc from isaaclab.physics import PhysicsManager @@ -33,6 +34,16 @@ class NewtonMJWarpManager(NewtonManager): _builder_attribute_solvers = (SolverMuJoCo,) + @classmethod + def _register_builder_attributes(cls, builder: ModelBuilder) -> None: + """Register the MuJoCo custom attributes consumed during USD import.""" + SolverMuJoCo.register_custom_attributes(builder) + + @classmethod + def _get_usd_import_schema_resolvers(cls) -> list[SchemaResolver]: + """Add MuJoCo schema support after the generic Newton and PhysX schemas.""" + return [*super()._get_usd_import_schema_resolvers(), SchemaResolverMjc()] + @classmethod def _create_solver(cls, model: Model, solver_cfg: MJWarpSolverCfg) -> SolverMuJoCo: """Construct the configured MuJoCo Warp solver.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index 0765df0e370..4a227483c9e 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -76,7 +76,7 @@ def _paused_gc(): from newton.sensors import SensorFrameTransform from newton.sensors import SensorIMU as NewtonSensorIMU from newton.solvers import SolverBase, SolverKamino -from newton.usd import SchemaResolverNewton, SchemaResolverPhysx +from newton.usd import SchemaResolver, SchemaResolverNewton, SchemaResolverPhysx from pxr import Usd, UsdGeom @@ -1846,6 +1846,16 @@ def _get_usd_import_ignore_paths(cls) -> list[str]: """Return solver-specific prim paths excluded from USD import.""" return [] + @classmethod + def _get_usd_import_schema_resolvers(cls) -> list[SchemaResolver]: + """Return the schema resolvers used to import the stage into Newton. + + Solver managers may extend this list when they consume solver-specific + USD attributes. Resolver order defines precedence when multiple schemas + author the same Newton model property. + """ + return [SchemaResolverNewton(), SchemaResolverPhysx()] + @classmethod def instantiate_builder_from_stage(cls): """Create builder from USD stage. @@ -1875,7 +1885,7 @@ def instantiate_builder_from_stage(cls): builder = cls.create_builder(up_axis=up_axis) - schema_resolvers = [SchemaResolverNewton(), SchemaResolverPhysx()] + schema_resolvers = cls._get_usd_import_schema_resolvers() # NOTE: None of the add_usd calls below pass joint_ordering or # bodies_follow_joint_ordering, so the live articulation's native diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index cc52bd3492e..684c6aa0236 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -31,6 +31,7 @@ import isaaclab_newton.physics.newton_manager as newton_manager_module import numpy as np import pytest +import torch import warp as wp from isaaclab_newton.assets.articulation import articulation as articulation_module from isaaclab_newton.physics import ( @@ -699,6 +700,67 @@ def test_active_manager_create_builder_registers_mpm_attributes(): assert builder.has_custom_attribute("mpm:young_modulus") +@pytest.mark.parametrize("import_path", ["clone", "standalone"]) +def test_mjwarp_production_imports_mujoco_joint_properties(monkeypatch, import_path): + """MJWarp production imports preserve MuJoCo friction loss and passive damping.""" + from isaaclab_newton.cloner.replicate import _build_newton_builder_from_mapping + + from pxr import Sdf, Usd, UsdGeom, UsdPhysics + + stage = Usd.Stage.CreateInMemory() + UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z) + UsdGeom.SetStageMetersPerUnit(stage, 1.0) + + root_path = "/Sources/robot" if import_path == "clone" else "/World/robot" + root = UsdGeom.Cube.Define(stage, root_path).GetPrim() + UsdPhysics.RigidBodyAPI.Apply(root) + UsdPhysics.ArticulationRootAPI.Apply(root) + + child_path = f"{root_path}/child" + child = UsdGeom.Cube.Define(stage, child_path).GetPrim() + UsdPhysics.RigidBodyAPI.Apply(child) + + joint = UsdPhysics.RevoluteJoint.Define(stage, f"{child_path}/joint") + joint.CreateAxisAttr().Set("Z") + joint.CreateBody0Rel().SetTargets([root_path]) + joint.CreateBody1Rel().SetTargets([child_path]) + joint.GetPrim().CreateAttribute("mjc:frictionloss", Sdf.ValueTypeNames.Double, True).Set(0.11) + joint.GetPrim().CreateAttribute("mjc:damping", Sdf.ValueTypeNames.Double, True).Set(0.23) + + monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(physics_manager=NewtonMJWarpManager)) + monkeypatch.setattr(PhysicsManager, "_cfg", NewtonCfg(solver_cfg=MJWarpSolverCfg())) + monkeypatch.setattr(PhysicsManager, "_device", "cpu") + monkeypatch.setattr(NewtonManager, "_builder", None) + monkeypatch.setattr(NewtonManager, "_deformable_registry", []) + monkeypatch.setattr(NewtonManager, "_cl_pending_sites", {}) + monkeypatch.setattr(NewtonManager, "_per_world_builder_hooks", []) + monkeypatch.setattr(NewtonManager, "_world_xforms", None) + + if import_path == "clone": + builder, _, _, _, _ = _build_newton_builder_from_mapping( + stage=stage, + sources=(root_path,), + destinations=("/World/envs/env_{}",), + env_ids=torch.tensor([0], dtype=torch.int32), + mapping=torch.ones((1, 1), dtype=torch.bool), + load_visual_shapes=False, + ) + else: + monkeypatch.setattr(newton_manager_module, "get_current_stage", lambda: stage) + monkeypatch.setattr( + newton_manager_module, "_restore_visible_colliders_without_visual_shapes", lambda *args: None + ) + monkeypatch.setattr(newton_manager_module, "replace_newton_builder_shape_colors", lambda *args: None) + monkeypatch.setattr(newton_manager_module, "import_builder_visual_material_paths", lambda *args: None) + NewtonMJWarpManager.instantiate_builder_from_stage() + builder = NewtonManager._builder + + model = builder.finalize(device="cpu") + + assert model.joint_friction.numpy()[-1] == pytest.approx(0.11) + assert model.joint_damping.numpy()[-1] == pytest.approx(0.23) + + def test_mpm_end_to_end_with_particle_custom_attributes(): """End-to-end MPM step using ``add_particles(custom_attributes=...)`` — the production path.""" sim_cfg = SimulationCfg(