Add assetdir functionality to pymjcf.

PiperOrigin-RevId: 618313406
Change-Id: I24d8e277ad15c284ddd2c6a2e8c7ab53b2392068

Author: DeepMind

dm_control/mjcf/attribute.py

@@ -511,24 +511,41 @@ def _get_asset_from_path(self, path):
     _, basename = os.path.split(path)
     filename, extension = os.path.splitext(basename)
 
-    # Look in the dict of pre-loaded assets before checking the filesystem.
-    try:
+    assetdir = None
+    if self._parent.namescope.has_identifier(
+        constants.BASEPATH, constants.ASSETDIR_NAMESPACE
+    ):
+      assetdir = self._parent.namescope.get(
+          constants.BASEPATH, constants.ASSETDIR_NAMESPACE
+      )
+
+    if path in self._parent.namescope.assets:
+      # Look in the dict of pre-loaded assets before checking the filesystem.
       contents = self._parent.namescope.assets[path]
-    except KeyError:
+    else:
       # Construct the full path to the asset file, prefixed by the path to the
       # model directory, and by `meshdir` or `texturedir` if appropriate.
       path_parts = []
       if self._parent.namescope.model_dir:
         path_parts.append(self._parent.namescope.model_dir)
-      try:
-        base_path = self._parent.namescope.get(constants.BASEPATH,
-                                               self._path_namespace)
+
+      if self._parent.namescope.has_identifier(
+          constants.BASEPATH, self._path_namespace
+      ):
+        base_path = self._parent.namescope.get(
+            constants.BASEPATH, self._path_namespace
+        )
         path_parts.append(base_path)
-      except KeyError:
-        pass
+      elif (
+          self._path_namespace
+          in (constants.TEXTUREDIR_NAMESPACE, constants.MESHDIR_NAMESPACE)
+          and assetdir is not None
+      ):
+        path_parts.append(assetdir)
       path_parts.append(path)
       full_path = os.path.join(*path_parts)  # pylint: disable=no-value-for-parameter
       contents = resources.GetResource(full_path)
+
     if self._parent.tag == constants.SKIN:
       return SkinAsset(contents=contents, parent=self._parent,
                        extension=extension, prefix=filename)

dm_control/mjcf/constants.py

@@ -39,6 +39,11 @@
 TENDON = 'tendon'
 WORLDBODY = 'worldbody'
 
+# Path namespaces.
+MESHDIR_NAMESPACE = 'mesh'
+TEXTUREDIR_NAMESPACE = 'texture'
+ASSETDIR_NAMESPACE = 'asset'
+
 MJDATA_TRIGGERS_DIRTY = [
     'qpos', 'qvel', 'act', 'ctrl', 'qfrc_applied', 'xfrc_applied']
 MJMODEL_DOESNT_TRIGGER_DIRTY = [

dm_control/mjcf/export_with_assets_test.py

@@ -27,6 +27,7 @@
 
 _ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'test_assets')
 _TEST_MODEL_WITH_ASSETS = os.path.join(_ASSETS_DIR, 'model_with_assets.xml')
+_TEST_MODEL_WITH_ASSETDIR = os.path.join(_ASSETS_DIR, 'model_with_assetdir.xml')
 _TEST_MODEL_WITHOUT_ASSETS = os.path.join(_ASSETS_DIR, 'lego_brick.xml')
 
 
@@ -42,6 +43,7 @@ class ExportWithAssetsTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
       ('with_assets', _TEST_MODEL_WITH_ASSETS, 'mujoco_with_assets.xml'),
+      ('with_assetdir', _TEST_MODEL_WITH_ASSETDIR, 'mujoco_with_assetdir.xml'),
       ('without_assets', _TEST_MODEL_WITHOUT_ASSETS, 'mujoco.xml'),)
   def test_export_model(self, xml_path, out_xml_name):
     """Save processed MJCF model."""

dm_control/mjcf/test_assets/model_with_assetdir.xml

@@ -0,0 +1,19 @@
+<mujoco model="Textured cube and mesh">
+  <compiler assetdir="meshes" texturedir="textures"/>
+  <asset>
+    <mesh name="cube" file="cube.stl"/>
+    <mesh name="another_cube" file="more_meshes/cube.stl"/>
+    <mesh name="unused_asset_should_not_cause_problems" file="cube.stl"/>
+    <mesh name="cube_msh" file="cube.msh"/>
+    <texture name="texture" file="deepmind.png"/>
+    <material name="mat_texture" texture="texture"/>
+    <hfield name="hill" file="../textures/deepmind.png" size="0.5 0.5 0.5 0.1"/>
+  </asset>
+  <worldbody>
+    <light diffuse=".5 .5 .5" pos="0 0 3" dir="0 0 -1"/>
+    <geom type="mesh" mesh="cube" material="mat_texture"/>
+    <geom type="mesh" mesh="another_cube" material="mat_texture" pos="2.5 0. 0."/>
+    <geom type="mesh" mesh="cube_msh" material="mat_texture" pos="4. 0. 0."/>
+    <geom type="hfield" hfield="hill" pos="1.2 0. 0" rgba="0. 0.9 0. 1" size="40 40 0.1"/>
+  </worldbody>
+</mujoco>