Update pymjcf schema to support SDF geoms in worldbody.

PiperOrigin-RevId: 597803513 Change-Id: I4feb79575f82e0a98fb92c28331cb7f749cf9078
google-deepmind · Jan 12, 2024 · 241851f · 241851f
1 parent 60c32e1
commit 241851f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/dm_control/mjcf/physics_test.py b/dm_control/mjcf/physics_test.py
@@ -608,6 +608,9 @@ def test_plugins_elasticity(self):
 
   def test_plugins_sdf(self):
     root = mjcf.RootElement()
+    root.option.sdf_iterations = 10
+    root.option.sdf_initpoints = 40
+
     extension = root.extension.add('plugin', plugin='mujoco.sdf.torus')
     instance = extension.add('instance', name='torus')
     instance.add('config', key='radius1', value='0.35')
@@ -617,10 +620,16 @@ def test_plugins_sdf(self):
     mesh = root.asset.add('mesh', name='torus')
     mesh.add('plugin', instance='torus')
 
+    # Test we can add SDF geom to the worldbody.
+    worldbody_geom = root.worldbody.add(
+        'geom', type='sdf', mesh='torus', rgba=[.2, .2, .8, 1])
+    worldbody_geom.add('plugin', instance='torus')
+
+    # Test we can add SDF geom to a body.
     body = root.worldbody.add('body', pos=[-1, 0, 3.8])
     body.add('freejoint')
-    geom = body.add('geom', type='sdf', mesh='torus', rgba=[.2, .2, .8, 1])
-    geom.add('plugin', instance='torus')
+    body_geom = body.add('geom', type='sdf', mesh='torus', rgba=[.2, .2, .8, 1])
+    body_geom.add('plugin', instance='torus')
 
     physics = mjcf.Physics.from_mjcf_model(root)
     physics.step()

diff --git a/dm_control/mjcf/schema.xml b/dm_control/mjcf/schema.xml
@@ -68,6 +68,8 @@
         <attribute name="noslip_tolerance" type="float" conflict_allowed="true" conflict_behavior="min"/>
         <attribute name="mpr_iterations" type="int" conflict_allowed="true" conflict_behavior="max"/>
         <attribute name="mpr_tolerance" type="float" conflict_allowed="true" conflict_behavior="min"/>
+        <attribute name="sdf_iterations" type="int" conflict_allowed="true" conflict_behavior="max"/>
+        <attribute name="sdf_initpoints" type="int" conflict_allowed="true" conflict_behavior="max"/>
       </attributes>
       <children>
         <element name="flag">
@@ -1014,6 +1016,23 @@
             <attribute name="fluidcoef" type="array" array_type="float" array_size="5"/>
             <attribute name="user" type="array" array_type="float"/>
           </attributes>
+          <children>
+            <element name="plugin" repeated="true" namespace="geom">
+              <attributes>
+                <attribute name="name" type="string"/>
+                <attribute name="plugin" type="string"/>
+                <attribute name="instance" type="reference" reference_namespace="plugin"/>
+              </attributes>
+              <children>
+                <element name="config" repeated="true">
+                  <attributes>
+                    <attribute name="key" type="string"/>
+                    <attribute name="value" type="string"/>
+                  </attributes>
+                </element>
+              </children>
+            </element>
+          </children>
         </element>
         <element name="site" repeated="true">
           <attributes>