SolverMuJoCo: custom attributes for MJCF/MJC schema options

newton/_src/sim/builder.py

@@ -1806,6 +1806,10 @@ def transform_mul(a, b):
             # Determine the offset based on frequency
             if attr.frequency == ModelAttributeFrequency.ONCE:
                 offset = 0
+            elif attr.frequency == ModelAttributeFrequency.WORLD:
+                # For WORLD frequency, use current_world as the offset
+                # This remaps local index 0 to the actual world index
+                offset = self.current_world if self.current_world >= 0 else 0
             elif attr.frequency == ModelAttributeFrequency.BODY:
                 offset = start_body_idx
             elif attr.frequency == ModelAttributeFrequency.SHAPE:
@@ -5809,6 +5813,8 @@ def finalize(self, device: Devicelike | None = None, requires_grad: bool = False
                 # determine count by frequency
                 if frequency == ModelAttributeFrequency.ONCE:
                     count = 1
+                elif frequency == ModelAttributeFrequency.WORLD:
+                    count = m.num_worlds
                 elif frequency == ModelAttributeFrequency.BODY:
                     count = m.body_count
                 elif frequency == ModelAttributeFrequency.SHAPE:

newton/_src/sim/model.py

@@ -68,6 +68,8 @@ class ModelAttributeFrequency(IntEnum):
     """Attribute frequency follows the number of shapes (see :attr:`~newton.Model.shape_count`)."""
     ARTICULATION = 6
     """Attribute frequency follows the number of articulations (see :attr:`~newton.Model.articulation_count`)."""
+    WORLD = 7
+    """Attribute frequency follows the number of worlds (see :attr:`~newton.Model.num_worlds`)."""
 
 
 class AttributeNamespace:

newton/_src/solvers/mujoco/kernels.py

@@ -897,6 +897,25 @@ def repeat_array_kernel(
     dst[tid] = src[src_idx]
 
 
+@wp.kernel
+def update_solver_options_kernel(
+    newton_impratio: wp.array(dtype=float),
+    # outputs
+    opt_impratio: wp.array(dtype=float),
+):
+    """Update per-world solver options from Newton model.
+
+    Args:
+        newton_impratio: Per-world impratio values from Newton model (None if overridden)
+        opt_impratio: MuJoCo Warp opt.impratio array to update (shape: nworld)
+    """
+    worldid = wp.tid()
+
+    # Only update if Newton array exists (None means overridden or not available)
+    if newton_impratio:
+        opt_impratio[worldid] = newton_impratio[worldid]
+
+
 @wp.kernel
 def update_axis_properties_kernel(
     mjc_actuator_to_newton_axis: wp.array2d(dtype=wp.int32),

newton/_src/solvers/mujoco/solver_mujoco.py

@@ -64,6 +64,7 @@
     update_mocap_transforms_kernel,
     update_model_properties_kernel,
     update_shape_mappings_kernel,
+    update_solver_options_kernel,
 )
 
 if TYPE_CHECKING:
@@ -285,6 +286,18 @@ def register_custom_attributes(cls, builder: ModelBuilder) -> None:
                 mjcf_attribute_name="actuatorgravcomp",
             )
         )
+        # Solver options (frequency WORLD for per-world values)
+        builder.add_custom_attribute(
+            ModelBuilder.CustomAttribute(
+                name="impratio",
+                frequency=ModelAttributeFrequency.WORLD,
+                assignment=ModelAttributeAssignment.MODEL,
+                dtype=wp.float32,
+                default=1.0,
+                namespace="mujoco",
+                usd_attribute_name="mjc:option:impratio",
+            )
+        )
 
     def __init__(
         self,
@@ -300,7 +313,7 @@ def __init__(
         solver: int | str = "cg",
         integrator: int | str = "implicitfast",
         cone: int | str = "pyramidal",
-        impratio: float = 1.0,
+        impratio: float | None = None,
         use_mujoco_cpu: bool = False,
         disable_contacts: bool = False,
         default_actuator_gear: float | None = None,
@@ -326,7 +339,7 @@ def __init__(
             solver (int | str): Solver type. Can be "cg" or "newton", or their corresponding MuJoCo integer constants.
             integrator (int | str): Integrator type. Can be "euler", "rk4", or "implicitfast", or their corresponding MuJoCo integer constants.
             cone (int | str): The type of contact friction cone. Can be "pyramidal", "elliptic", or their corresponding MuJoCo integer constants.
-            impratio (float): Frictional-to-normal constraint impedance ratio.
+            impratio (float | None): Frictional-to-normal constraint impedance ratio. If None, uses the value from model custom attribute ``mujoco:impratio`` or MuJoCo's default (1.0).
             use_mujoco_cpu (bool): If True, use the MuJoCo-C CPU backend instead of `mujoco_warp`.
             disable_contacts (bool): If True, disable contact computation in MuJoCo.
             register_collision_groups (bool): If True, register collision groups from the Newton model in MuJoCo.
@@ -894,7 +907,7 @@ def _convert_to_mjc(
         integrator: int | str = "implicitfast",
         disableflags: int = 0,
         disable_contacts: bool = False,
-        impratio: float = 1.0,
+        impratio: float | None = None,
         tolerance: float = 1e-6,
         ls_tolerance: float = 0.01,
         cone: int | str = "pyramidal",
@@ -996,6 +1009,14 @@ def fill_arr_from_dict(arr: nparray, d: dict[int, Any]):
             else:
                 arr[tuple(keys.T)] = vals
 
+        # Resolve solver options: constructor arg > custom attribute > MuJoCo default
+        # Track if overridden (tile() handles expansion) vs needs per-world update from Newton model
+        impratio_overridden = impratio is not None
+        if impratio is None:
+            mujoco_attrs = getattr(model, "mujoco", None)
+            if mujoco_attrs and hasattr(mujoco_attrs, "impratio"):
+                impratio = float(mujoco_attrs.impratio.numpy()[0])
+
         spec = mujoco.MjSpec()
         spec.option.disableflags = disableflags
         spec.option.gravity = np.array([*model.gravity.numpy()[0]])
@@ -1004,7 +1025,8 @@ def fill_arr_from_dict(arr: nparray, d: dict[int, Any]):
         spec.option.iterations = iterations
         spec.option.ls_iterations = ls_iterations
         spec.option.cone = cone
-        spec.option.impratio = impratio
+        if impratio is not None:
+            spec.option.impratio = impratio
         spec.option.tolerance = tolerance
         spec.option.ls_tolerance = ls_tolerance
         spec.option.jacobian = mujoco.mjtJacobian.mjJAC_AUTO
@@ -1905,6 +1927,9 @@ def add_geoms(newton_body_id: int):
             # expand model fields that can be expanded:
             self.expand_model_fields(self.mjw_model, nworld)
 
+            # update solver options from Newton model (only if not overridden by constructor)
+            self.update_solver_options(impratio_overridden=impratio_overridden)
+
             # so far we have only defined the first world,
             # now complete the data from the Newton model
             self.notify_model_changed(SolverNotifyFlags.ALL)
@@ -1987,6 +2012,20 @@ def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
             # "mat_rgba",
         }
 
+        # Solver option fields to expand (nested in mj_model.opt)
+        opt_fields_to_expand = {
+            # "timestep",
+            "impratio",
+            # "tolerance",
+            # "ls_tolerance",
+            # "ccd_tolerance",
+            # "density",
+            # "viscosity",
+            # "gravity",
+            # "wind",
+            # "magnetic",
+        }
+
         def tile(x: wp.array):
             # Create new array with same shape but first dim multiplied by nworld
             new_shape = list(x.shape)
@@ -2014,6 +2053,41 @@ def tile(x: wp.array):
                 array = getattr(mj_model, field)
                 setattr(mj_model, field, tile(array))
 
+        for field in opt_fields_to_expand:
+            if hasattr(mj_model.opt, field):
+                array = getattr(mj_model.opt, field)
+                setattr(mj_model.opt, field, tile(array))
+
+    def update_solver_options(self, impratio_overridden: bool = False):
+        """Update solver options from Newton model to MuJoCo Warp.
+
+        Copies per-world values from Newton custom attributes to the MuJoCo Warp model.
+        If a value was overridden by constructor, tile() already handled expansion so we skip it.
+
+        Args:
+            impratio_overridden: If True, impratio was set by constructor and tile() handled it.
+        """
+        mujoco_attrs = getattr(self.model, "mujoco", None)
+        nworld = self.model.num_worlds
+
+        # Get Newton arrays - pass None if overridden or not available (kernel checks for None)
+        if not impratio_overridden and mujoco_attrs and hasattr(mujoco_attrs, "impratio"):
+            newton_impratio = mujoco_attrs.impratio
+        else:
+            newton_impratio = None
+
+        # Skip kernel if all options are None (add more checks here as options are added)
+        all_none = newton_impratio is None  # and other_option is None and ...
+        if all_none:
+            return
+
+        wp.launch(
+            update_solver_options_kernel,
+            dim=nworld,
+            inputs=[newton_impratio],
+            outputs=[self.mjw_model.opt.impratio],
+        )
+
     def update_model_inertial_properties(self):
         if self.model.body_count == 0:
             return

newton/_src/utils/import_mjcf.py

@@ -143,6 +143,18 @@ def parse_mjcf(
     else:
         mesh_dir = "."
 
+    # Parse MJCF option tag for ONCE and WORLD frequency custom attributes (solver options)
+    # WORLD frequency attributes use index 0 here; they get remapped during add_world()
+    builder_custom_attr_option: list[ModelBuilder.CustomAttribute] = builder.get_custom_attributes_by_frequency(
+        [ModelAttributeFrequency.ONCE, ModelAttributeFrequency.WORLD]
+    )
+    option_elem = root.find("option")
+    if option_elem is not None and builder_custom_attr_option:
+        option_attrs = parse_custom_attributes(option_elem.attrib, builder_custom_attr_option, "mjcf")
+        for key, value in option_attrs.items():
+            if key in builder.custom_attributes:
+                builder.custom_attributes[key].values[0] = value
+
     mesh_assets = {}
     for asset in root.findall("asset"):
         for mesh in asset.findall("mesh"):

newton/_src/utils/import_usd.py

@@ -956,13 +956,21 @@ def define_joint_targets(dof, joint_desc):
     )
 
     if physics_scene_prim is not None:
-        # Extract custom attributes for model (ONCE) frequency from the PhysicsScene prim
+        # Extract custom attributes for model (ONCE and WORLD frequency) from the PhysicsScene prim
+        # WORLD frequency attributes use index 0 here; they get remapped during add_world()
         builder_custom_attr_model: list[ModelBuilder.CustomAttribute] = [
-            attr for attr in builder.custom_attributes.values() if attr.frequency == ModelAttributeFrequency.ONCE
+            attr
+            for attr in builder.custom_attributes.values()
+            if attr.frequency in (ModelAttributeFrequency.ONCE, ModelAttributeFrequency.WORLD)
         ]
         scene_custom_attrs = usd.get_custom_attribute_values(physics_scene_prim, builder_custom_attr_model)
         scene_attributes.update(scene_custom_attrs)
 
+        # Set values on builder's custom attributes (similar to MJCF parsing)
+        for key, value in scene_custom_attrs.items():
+            if key in builder.custom_attributes:
+                builder.custom_attributes[key].values[0] = value
+
     joint_descriptions = {}
     # maps from joint prim path to joint index in builder
     path_joint_map: dict[str, int] = {}

newton/tests/test_import_mjcf.py

@@ -1385,6 +1385,82 @@ def test_geom_solimp_parsing(self):
             for i, (a, e) in enumerate(zip(actual, expected, strict=False)):
                 self.assertAlmostEqual(a, e, places=4, msg=f"geom_solimp[{shape_idx}][{i}] should be {e}, got {a}")
 
+    def test_option_impratio_parsing(self):
+        """Test parsing of impratio from MJCF option tag."""
+        mjcf = """<?xml version="1.0" ?>
+<mujoco>
+    <option impratio="1.5"/>
+    <worldbody>
+        <body name="body1" pos="0 0 1">
+            <joint type="hinge" axis="0 0 1"/>
+            <geom type="sphere" size="0.1"/>
+        </body>
+    </worldbody>
+</mujoco>
+"""
+        builder = newton.ModelBuilder()
+        SolverMuJoCo.register_custom_attributes(builder)
+        builder.add_mjcf(mjcf)
+        model = builder.finalize()
+
+        self.assertTrue(hasattr(model, "mujoco"))
+        self.assertTrue(hasattr(model.mujoco, "impratio"))
+
+        impratio = model.mujoco.impratio.numpy()
+
+        # Single world should have single value
+        self.assertEqual(len(impratio), 1)
+        self.assertAlmostEqual(impratio[0], 1.5, places=4)
+
+    def test_option_impratio_per_world(self):
+        """Test that impratio is correctly remapped per world when merging builders."""
+        # Robot A with impratio=1.5
+        robot_a = newton.ModelBuilder()
+        SolverMuJoCo.register_custom_attributes(robot_a)
+        robot_a.add_mjcf("""
+<mujoco>
+    <option impratio="1.5"/>
+    <worldbody>
+        <body name="a" pos="0 0 1">
+            <joint type="hinge" axis="0 0 1"/>
+            <geom type="sphere" size="0.1"/>
+        </body>
+    </worldbody>
+</mujoco>
+""")
+
+        # Robot B with impratio=2.0
+        robot_b = newton.ModelBuilder()
+        SolverMuJoCo.register_custom_attributes(robot_b)
+        robot_b.add_mjcf("""
+<mujoco>
+    <option impratio="2.0"/>
+    <worldbody>
+        <body name="b" pos="0 0 1">
+            <joint type="hinge" axis="0 0 1"/>
+            <geom type="sphere" size="0.1"/>
+        </body>
+    </worldbody>
+</mujoco>
+""")
+
+        # Merge into main builder
+        main = newton.ModelBuilder()
+        SolverMuJoCo.register_custom_attributes(main)
+        main.add_world(robot_a)
+        main.add_world(robot_b)
+        model = main.finalize()
+
+        self.assertTrue(hasattr(model, "mujoco"))
+        self.assertTrue(hasattr(model.mujoco, "impratio"))
+
+        impratio = model.mujoco.impratio.numpy()
+
+        # Should have 2 worlds with different impratio values
+        self.assertEqual(len(impratio), 2)
+        self.assertAlmostEqual(impratio[0], 1.5, places=4, msg="World 0 should have impratio=1.5")
+        self.assertAlmostEqual(impratio[1], 2.0, places=4, msg="World 1 should have impratio=2.0")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)