Use newton-usd-schemas

Author: andrewkaufman

mujoco_usd_converter/_impl/cli.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import mujoco
+import newton_usd_schemas
 import usdex.core
 from pxr import Tf, Usd
 
@@ -47,6 +48,7 @@ def run() -> int:
     Tf.Status(f"USD Version: {Usd.GetVersion()}")
     Tf.Status(f"USDEX Version: {usdex.core.version()}")
     Tf.Status(f"MuJoCo Version: {mujoco.__version__}")
+    Tf.Status(f"Newton USD Schemas Version: {newton_usd_schemas.__version__}")
 
     try:
         converter = Converter(layer_structure=not args.no_layer_structure, scene=not args.no_physics_scene, comment=args.comment)

mujoco_usd_converter/_impl/convert.py

@@ -5,6 +5,9 @@
 from dataclasses import dataclass
 
 import mujoco
+
+# importing registers the Newton schema plugin with USD
+import newton_usd_schemas  # noqa: F401
 import usdex.core
 from pxr import Sdf, Tf, Usd, UsdGeom, UsdPhysics
 

mujoco_usd_converter/_impl/joint.py

@@ -67,18 +67,21 @@ def convert_joints(parent: Usd.Prim, body: mujoco.MjsBody, data: ConversionData)
 
         data.references[Tokens.Physics][joint.name] = joint_prim.GetPrim()
 
-        apply_mjc_joint_api(joint_prim.GetPrim(), joint)
+        convert_joint_attributes(joint_prim.GetPrim(), joint)
 
 
-def apply_mjc_joint_api(prim: Usd.Prim, joint: mujoco.MjsJoint):
+def convert_joint_attributes(prim: Usd.Prim, joint: mujoco.MjsJoint):
+    prim.ApplyAPI(Usd.SchemaRegistry.GetSchemaTypeName("NewtonJointAPI"))
     prim.ApplyAPI(Usd.SchemaRegistry.GetSchemaTypeName("MjcPhysicsJointAPI"))
 
+    set_schema_attribute(prim, "newton:armature", joint.armature)
+    set_schema_attribute(prim, "mjc:armature", joint.armature)  # FUTURE: Remove once MuJoCo supports Newton schemas
+
     limited_token = mj_limited_to_token(joint.actfrclimited)
     set_schema_attribute(prim, "mjc:actuatorfrclimited", limited_token)
     set_schema_attribute(prim, "mjc:actuatorfrcrange:min", joint.actfrcrange[0])
     set_schema_attribute(prim, "mjc:actuatorfrcrange:max", joint.actfrcrange[1])
     set_schema_attribute(prim, "mjc:actuatorgravcomp", bool(joint.actgravcomp))
-    set_schema_attribute(prim, "mjc:armature", joint.armature)
     set_schema_attribute(prim, "mjc:damping", joint.damping)
     set_schema_attribute(prim, "mjc:frictionloss", joint.frictionloss)
     set_schema_attribute(prim, "mjc:group", joint.group)

pyproject.toml

@@ -9,6 +9,7 @@ requires-python = ">=3.10,<3.13"
 dependencies = [
     "mujoco==3.3.5",
     "usd-exchange>=2.1.0a3", # locked to USD 25.05
+    "newton-usd-schemas>=0.1.0",
     "numpy-stl>=3.2",
     "tinyobjloader>=2.0.0rc13",
 ]
@@ -39,6 +40,9 @@ path = "mujoco_usd_converter/_version.py"
 package = true
 python-preference = "only-managed"
 
+[tool.uv.sources]
+newton-usd-schemas = { path = "../newton-usd-schemas", editable = true }
+
 # Dev Tools
 
 [tool.poe.tasks.test]

tests/testJoints.py

@@ -342,6 +342,33 @@ def test_no_autolimits(self):
         self.assertFalse(unlimited_joint.GetLowerLimitAttr().HasAuthoredValue())
         self.assertFalse(unlimited_joint.GetUpperLimitAttr().HasAuthoredValue())
 
+    def test_newton_schema(self):
+        model = pathlib.Path("./tests/data/joint_attributes.xml")
+        asset: Sdf.AssetPath = mujoco_usd_converter.Converter().convert(model, self.tmpDir())
+        stage: Usd.Stage = Usd.Stage.Open(asset.path)
+        self.assertIsValidUsd(stage)
+
+        custom_joint: Usd.Prim = stage.GetPrimAtPath("/joint_attributes/Geometry/body1/custom_joint")
+        self.assertTrue(custom_joint.IsValid())
+        self.assertTrue(custom_joint.IsA(UsdPhysics.RevoluteJoint))
+        self.assertTrue(custom_joint.HasAPI(Usd.SchemaRegistry.GetAPISchemaTypeName("NewtonJointAPI")))
+
+        # Check that all Newton properties are authored
+        for property in custom_joint.GetPropertiesInNamespace("newton"):
+            self.assertTrue(property.HasAuthoredValue(), f"Property {property.GetName()} is not authored")
+
+        # Check that all attributes are authored correctly
+        self.assertTrue(custom_joint.GetAttribute("newton:armature").HasAuthoredValue())
+        self.assertAlmostEqual(custom_joint.GetAttribute("newton:armature").Get(), 0.1)
+
+        # A joint with explicitly authored default values does not need to author any values in USD
+        default_joint: Usd.Prim = stage.GetPrimAtPath("/joint_attributes/Geometry/body2/default_joint")
+        self.assertTrue(default_joint.IsValid())
+        self.assertTrue(default_joint.IsA(UsdPhysics.PrismaticJoint))
+        self.assertTrue(default_joint.HasAPI(Usd.SchemaRegistry.GetAPISchemaTypeName("NewtonJointAPI")))
+        self.assertFalse(default_joint.GetAttribute("newton:armature").HasAuthoredValue())
+        self.assertAlmostEqual(default_joint.GetAttribute("newton:armature").Get(), 0.0)
+
     def test_mjc_schema(self):
         # Test that all joint attributes are authored correctly
         model = pathlib.Path("./tests/data/joint_attributes.xml")