Add support for geom contact margin

newton/_src/solvers/mujoco/kernels.py

@@ -1214,6 +1214,7 @@ def update_geom_properties_kernel(
     shape_kd: wp.array(dtype=float),
     shape_size: wp.array(dtype=wp.vec3f),
     shape_transform: wp.array(dtype=wp.transform),
+    shape_contact_margin: wp.array(dtype=float),
     mjc_geom_to_newton_shape: wp.array2d(dtype=wp.int32),
     geom_type: wp.array(dtype=int),
     GEOM_TYPE_MESH: int,
@@ -1235,6 +1236,7 @@ def update_geom_properties_kernel(
     geom_solimp: wp.array2d(dtype=vec5),
     geom_solmix: wp.array2d(dtype=float),
     geom_gap: wp.array2d(dtype=float),
+    geom_margin: wp.array2d(dtype=float),
 ):
     """Update MuJoCo geom properties from Newton shape properties.
 
@@ -1281,6 +1283,9 @@ def update_geom_properties_kernel(
     if shape_geom_gap:
         geom_gap[world, geom_idx] = shape_geom_gap[shape_idx]
 
+    # update geom margin
+    geom_margin[world, geom_idx] = shape_contact_margin[shape_idx]
+
     # update size
     geom_size[world, geom_idx] = shape_size[shape_idx]
 

newton/_src/solvers/mujoco/solver_mujoco.py

@@ -1120,6 +1120,7 @@ def fill_arr_from_dict(arr: nparray, d: dict[int, Any]):
         shape_mu = model.shape_material_mu.numpy()
         shape_torsional_friction = model.shape_material_torsional_friction.numpy()
         shape_rolling_friction = model.shape_material_rolling_friction.numpy()
+        shape_contact_margin = model.shape_contact_margin.numpy()
 
         # retrieve MuJoCo-specific attributes
         mujoco_attrs = getattr(model, "mujoco", None)
@@ -1377,6 +1378,8 @@ def add_geoms(newton_body_id: int):
                     geom_params["solmix"] = shape_geom_solmix[shape]
                 if shape_geom_gap is not None:
                     geom_params["gap"] = shape_geom_gap[shape]
+                if shape_contact_margin is not None:
+                    geom_params["margin"] = shape_contact_margin[shape]
 
                 body.add_geom(**geom_params)
                 # store the geom name instead of assuming index
@@ -2019,8 +2022,8 @@ def expand_model_fields(self, mj_model: MjWarpModel, nworld: int):
             "geom_pos",
             "geom_quat",
             "geom_friction",
-            # "geom_margin",
             "geom_gap",
+            "geom_margin",
             # "geom_rgba",
             # "site_pos",
             # "site_quat",
@@ -2288,6 +2291,7 @@ def update_geom_properties(self):
                 self.model.shape_material_kd,
                 self.model.shape_scale,
                 self.model.shape_transform,
+                self.model.shape_contact_margin,
                 self.mjc_geom_to_newton_shape,
                 self.mjw_model.geom_type,
                 self._mujoco.mjtGeom.mjGEOM_MESH,
@@ -2310,6 +2314,7 @@ def update_geom_properties(self):
                 self.mjw_model.geom_solimp,
                 self.mjw_model.geom_solmix,
                 self.mjw_model.geom_gap,
+                self.mjw_model.geom_margin,
             ],
             device=self.model.device,
         )

newton/_src/utils/import_mjcf.py

@@ -333,6 +333,9 @@ def parse_shapes(defaults, body_name, link, geoms, density, visible=True, just_v
                 if len(friction_values) >= 3:
                     shape_cfg.rolling_friction = float(friction_values[2])
 
+            # Parse MJCF margin (contact margin for collision detection)
+            shape_cfg.contact_margin = parse_float(geom_attrib, "margin", shape_cfg.contact_margin)
+
             custom_attributes = parse_custom_attributes(geom_attrib, builder_custom_attr_shape, parsing_mode="mjcf")
             shape_kwargs = {
                 "key": geom_name,

newton/_src/utils/import_usd.py

@@ -1374,6 +1374,12 @@ def warn_invalid_desc(path, descriptor) -> bool:
                         thickness=usd.get_float_with_fallback(
                             prim_and_scene, "newton:contact_thickness", builder.default_shape_cfg.thickness
                         ),
+                        contact_margin=R.get_value(
+                            prim,
+                            prim_type=PrimType.SHAPE,
+                            key="rigid_contact_margin",
+                            default=builder.rigid_contact_margin,
+                        ),
                         mu=material.dynamicFriction,
                         restitution=material.restitution,
                         density=body_density.get(body_path, default_shape_density),

newton/examples/robot/example_robot_allegro_hand.py

@@ -93,6 +93,7 @@ def __init__(self, viewer, num_worlds=4):
         newton.solvers.SolverMuJoCo.register_custom_attributes(allegro_hand)
         allegro_hand.default_shape_cfg.ke = 1.0e3
         allegro_hand.default_shape_cfg.kd = 1.0e2
+        allegro_hand.rigid_contact_margin = 0.0
 
         asset_path = newton.utils.download_asset("wonik_allegro")
         asset_file = str(asset_path / "usd" / "allegro_left_hand_with_cube.usda")

newton/tests/test_anymal_reset.py

@@ -36,6 +36,7 @@ def setUp(self):
 
     def _setup_simulation(self, cone_type):
         builder = newton.ModelBuilder(up_axis=newton.Axis.Z)
+        builder.rigid_contact_margin = 0.0
         builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig(
             armature=0.06,
             limit_ke=1.0e2,

newton/tests/test_import_mjcf.py

@@ -1497,6 +1497,49 @@ def test_geom_gap_parsing(self):
             actual = geom_gap[shape_idx].tolist()
             self.assertAlmostEqual(actual, expected, places=4)
 
+    def test_geom_margin_parsing(self):
+        """Test that geom_margin attribute is parsed correctly from MJCF."""
+        mjcf = """<?xml version="1.0" ?>
+<mujoco>
+    <worldbody>
+        <body name="body1">
+            <freejoint/>
+            <geom type="box" size="0.1 0.1 0.1" margin="0.5"/>
+        </body>
+        <body name="body2">
+            <freejoint/>
+            <geom type="sphere" size="0.05"/>
+        </body>
+        <body name="body3">
+            <freejoint/>
+            <geom type="capsule" size="0.05 0.1" margin="0.8"/>
+        </body>
+    </worldbody>
+</mujoco>
+"""
+
+        builder = newton.ModelBuilder()
+        builder.rigid_contact_margin = 0.17  # Note: we need to set this before calling add_mjcf()
+        builder.add_mjcf(mjcf)
+        model = builder.finalize()
+
+        # shape_contact_margin is a built-in Newton attribute, not a custom mujoco attribute
+        self.assertTrue(hasattr(model, "shape_contact_margin"), "Model should have shape_contact_margin attribute")
+
+        geom_margin = model.shape_contact_margin.numpy()
+        self.assertEqual(model.shape_count, 3, "Should have 3 shapes")
+
+        # Expected values: shape 0 has margin=0.5, shape 1 has default margin=0.17 (builder.rigid_contact_margin), shape 2 has margin=0.8
+        expected_values = {
+            0: 0.5,
+            1: 0.17,  # default taken from builder.rigid_contact_margin. Note that Mujoco defaults to 0 but we don't adopt that convention.
+            2: 0.8,
+        }
+
+        for shape_idx, expected in expected_values.items():
+            actual = geom_margin[shape_idx]
+            self.assertAlmostEqual(actual, expected, places=4)
+
     def test_default_inheritance(self):
         """Test nested default class inheritanc."""
         mjcf_content = """<?xml version="1.0" ?>

newton/tests/test_import_usd.py

@@ -1425,6 +1425,126 @@ def floats_match(arr, expected, tol=1e-4):
         self.assertTrue(found_default, f"Expected default gap {expected_default} not found in model")
         self.assertTrue(found_explicit_2, f"Expected gap {expected_explicit_2} not found in model")
 
+    @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+    def test_geom_margin_parsing(self):
+        """Test that geom_margin attribute is parsed correctly from USD."""
+        from pxr import Usd  # noqa: PLC0415
+
+        from newton._src.usd.schemas import SchemaResolverMjc  # noqa: PLC0415
+
+        usd_content = """#usda 1.0
+(
+    upAxis = "Z"
+)
+
+def PhysicsScene "physicsScene"
+{
+}
+
+def Xform "Body1" (
+    prepend apiSchemas = ["PhysicsRigidBodyAPI", "PhysicsArticulationRootAPI"]
+)
+{
+    double3 xformOp:translate = (0, 0, 1)
+    uniform token[] xformOpOrder = ["xformOp:translate"]
+
+    def Cube "Collision1" (
+        prepend apiSchemas = ["PhysicsCollisionAPI"]
+    )
+    {
+        double size = 1.0
+        # MuJoCo margin attribute (1 float)
+        double mjc:margin = 0.82
+    }
+}
+
+def Xform "Body2" (
+    prepend apiSchemas = ["PhysicsRigidBodyAPI"]
+)
+{
+    double3 xformOp:translate = (1, 0, 1)
+    uniform token[] xformOpOrder = ["xformOp:translate"]
+
+    def Sphere "Collision2" (
+        prepend apiSchemas = ["PhysicsCollisionAPI"]
+    )
+    {
+        double radius = 1.0
+        # No geom_margin - should use defaults
+    }
+}
+
+def PhysicsRevoluteJoint "Joint1"
+{
+    rel physics:body0 = </Body1>
+    rel physics:body1 = </Body2>
+    point3f physics:localPos0 = (0, 0, 0)
+    point3f physics:localPos1 = (0, 0, 0)
+    quatf physics:localRot0 = (1, 0, 0, 0)
+    quatf physics:localRot1 = (1, 0, 0, 0)
+    token physics:axis = "Z"
+}
+
+def Xform "Body3" (
+    prepend apiSchemas = ["PhysicsRigidBodyAPI"]
+)
+{
+    double3 xformOp:translate = (2, 0, 1)
+    uniform token[] xformOpOrder = ["xformOp:translate"]
+
+    def Capsule "Collision3" (
+        prepend apiSchemas = ["PhysicsCollisionAPI"]
+    )
+    {
+        double radius = 0.05
+        double height = 0.2
+        # Different margin values
+        double mjc:margin = 0.71
+    }
+}
+
+def PhysicsRevoluteJoint "Joint2"
+{
+    rel physics:body0 = </Body2>
+    rel physics:body1 = </Body3>
+    point3f physics:localPos0 = (0, 0, 0)
+    point3f physics:localPos1 = (0, 0, 0)
+    quatf physics:localRot0 = (1, 0, 0, 0)
+    quatf physics:localRot1 = (1, 0, 0, 0)
+    token physics:axis = "Y"
+}
+"""
+        stage = Usd.Stage.CreateInMemory()
+        stage.GetRootLayer().ImportFromString(usd_content)
+
+        builder = newton.ModelBuilder()
+        builder.rigid_contact_margin = 0.2
+        builder.add_usd(stage, schema_resolvers=[SchemaResolverMjc()])
+        model = builder.finalize()
+
+        # shape_contact_margin is a built-in Newton attribute, not a custom mujoco attribute
+        self.assertTrue(hasattr(model, "shape_contact_margin"), "Model should have shape_contact_margin attribute")
+
+        geom_margin = model.shape_contact_margin.numpy()
+        self.assertEqual(model.shape_count, 3, "Should have 3 shapes")
+
+        def floats_match(arr, expected, tol=1e-4):
+            return abs(arr - expected) < tol
+
+        # Check that we have shapes with expected values
+        expected_explicit_1 = 0.82
+        expected_default = 0.0  # default Note that we default to 0.0 when parsing USD and that builder.rigid_contact_margin plays no role.
+        expected_explicit_2 = 0.71
+
+        # Find shapes matching each expected value
+        found_explicit_1 = any(floats_match(geom_margin[i], expected_explicit_1) for i in range(model.shape_count))
+        found_default = any(floats_match(geom_margin[i], expected_default) for i in range(model.shape_count))
+        found_explicit_2 = any(floats_match(geom_margin[i], expected_explicit_2) for i in range(model.shape_count))
+
+        self.assertTrue(found_explicit_1, f"Expected geom margin {expected_explicit_1} not found in model")
+        self.assertTrue(found_default, f"Expected default geom margin {expected_default} not found in model")
+        self.assertTrue(found_explicit_2, f"Expected geom margin {expected_explicit_2} not found in model")
+
 
 class TestImportSampleAssets(unittest.TestCase):
     def verify_usdphysics_parser(self, file, model, compare_min_max_coords, floating):

newton/tests/test_mujoco_solver.py

@@ -2594,6 +2594,100 @@ def test_geom_solmix_conversion_and_update(self):
                     msg=f"Updated geom_solmix mismatch for shape {shape_idx} in world {world_idx}",
                 )
 
+    def test_geom_margin_conversion_and_update(self):
+        """Test per-shape geom_margin conversion to MuJoCo and dynamic updates across multiple worlds."""
+
+        # Note: shape_contact_margin is a standard Newton Model attribute, not a MuJoCo custom attribute,
+        # so no register_custom_attributes() call is needed (unlike geom_gap, geom_solmix, etc.)
+        num_worlds = 2
+        template_builder = newton.ModelBuilder()
+        shape_cfg = newton.ModelBuilder.ShapeConfig(density=1000.0)
+
+        # Create bodies with shapes
+        body1 = template_builder.add_link(mass=0.1)
+
+        template_builder.add_shape_box(body=body1, hx=0.1, hy=0.1, hz=0.1, cfg=shape_cfg)
+        joint1 = template_builder.add_joint_free(child=body1)
+
+        body2 = template_builder.add_link(mass=0.1)
+        template_builder.add_shape_sphere(body=body2, radius=0.1, cfg=shape_cfg)
+        joint2 = template_builder.add_joint_revolute(parent=body1, child=body2, axis=(0.0, 0.0, 1.0))
+
+        template_builder.add_articulation([joint1, joint2])
+
+        builder = newton.ModelBuilder()
+        builder.replicate(template_builder, num_worlds)
+        model = builder.finalize()
+
+        # Use per-world iteration to handle potential global shapes correctly
+        shape_world = model.shape_world.numpy()
+        initial_margin = np.zeros(model.shape_count, dtype=np.float32)
+
+        # Set unique margin values per shape and world
+        for world_idx in range(model.num_worlds):
+            world_shape_indices = np.where(shape_world == world_idx)[0]
+            for local_idx, shape_idx in enumerate(world_shape_indices):
+                initial_margin[shape_idx] = 0.4 + local_idx * 0.2 + world_idx * 0.05
+
+        model.shape_contact_margin.assign(wp.array(initial_margin, dtype=wp.float32, device=model.device))
+
+        solver = SolverMuJoCo(model, iterations=1, disable_contacts=True)
+        to_newton_shape_index = solver.mjc_geom_to_newton_shape.numpy()
+        num_geoms = solver.mj_model.ngeom
+
+        # Verify initial conversion
+        geom_margin = solver.mjw_model.geom_margin.numpy()
+        tested_count = 0
+        for world_idx in range(model.num_worlds):
+            for geom_idx in range(num_geoms):
+                shape_idx = to_newton_shape_index[world_idx, geom_idx]
+                if shape_idx < 0:
+                    continue
+
+                tested_count += 1
+                expected_margin = initial_margin[shape_idx]
+                actual_margin = geom_margin[world_idx, geom_idx]
+                self.assertAlmostEqual(
+                    float(actual_margin),
+                    expected_margin,
+                    places=5,
+                    msg=f"Initial geom_margin mismatch for shape {shape_idx} in world {world_idx}, geom {geom_idx}",
+                )
+
+        self.assertGreater(tested_count, 0, "Should have tested at least one shape")
+
+        # Update with different values
+        updated_margin = np.zeros(model.shape_count, dtype=np.float32)
+
+        # Set unique margin values per shape and world
+        for world_idx in range(model.num_worlds):
+            world_shape_indices = np.where(shape_world == world_idx)[0]
+            for local_idx, shape_idx in enumerate(world_shape_indices):
+                updated_margin[shape_idx] = 0.7 + local_idx * 0.03 + world_idx * 0.06
+
+        model.shape_contact_margin.assign(wp.array(updated_margin, dtype=wp.float32, device=model.device))
+
+        solver.notify_model_changed(SolverNotifyFlags.SHAPE_PROPERTIES)
+
+        # Verify updates
+        updated_geom_margin = solver.mjw_model.geom_margin.numpy()
+
+        for world_idx in range(model.num_worlds):
+            for geom_idx in range(num_geoms):
+                shape_idx = to_newton_shape_index[world_idx, geom_idx]
+                if shape_idx < 0:
+                    continue
+
+                expected_margin = updated_margin[shape_idx]
+                actual_margin = updated_geom_margin[world_idx, geom_idx]
+
+                self.assertAlmostEqual(
+                    float(actual_margin),
+                    expected_margin,
+                    places=5,
+                    msg=f"Updated geom_margin mismatch for shape {shape_idx} in world {world_idx}",
+                )
+
 
 class TestMuJoCoSolverEqualityConstraintProperties(TestMuJoCoSolverPropertiesBase):
     def test_eq_solref_conversion_and_update(self):
@@ -3696,6 +3790,7 @@ def test_mocap_body_transform_updates_collision_geoms(self):
         builder = newton.ModelBuilder()
         builder.default_shape_cfg.ke = 1e4
         builder.default_shape_cfg.kd = 1000.0
+        builder.rigid_contact_margin = 0.0
 
         # Create fixed-base (mocap) body at root (at origin)
         # This body will have a FIXED joint to the world, making it a mocap body in MuJoCo