Store asset keys relative to the MJCF. This CL fixes cases when loading a spec with "from_zip" that fails on zipped folders that contain a root directory.

PiperOrigin-RevId: 825484553
Change-Id: I79c6f542edbd3ffd167d0b5cecc704c32c8d6824

python/mujoco/__init__.py

@@ -124,15 +124,25 @@ def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec:
   if not zipfile.is_zipfile(file):
     raise ValueError(f'File {file} is not a zip file.')
   with zipfile.ZipFile(file, 'r') as zip_file:
+    xml_dir = None
     for zip_info in zip_file.infolist():
       if not zip_info.filename.endswith(os.path.sep):
         with zip_file.open(zip_info.filename) as f:
           if zip_info.filename.endswith('.xml'):
             xml_string = f.read()
+            xml_dir = os.path.dirname(zip_info.filename)
           else:
             assets[zip_info.filename] = f.read()
+
   if not xml_string:
     raise ValueError('No XML file found in zip file.')
+
+  relative_assets = {}
+  for key, value in assets.items():
+    new_key = os.path.relpath(key, xml_dir)
+    relative_assets[new_key] = value
+  assets = relative_assets
+
   return _specs.MjSpec.from_string(xml_string, assets=assets)
 
 

python/mujoco/specs_test.py

@@ -1615,5 +1615,25 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(wrap_joint2.type, mujoco.mjtWrap.mjWRAP_JOINT)
     self.assertEqual(wrap_joint2.coef, 2.0)
 
+  def test_from_zip(self):
+    """Tests that the assets are correctly parsed from a zip file."""
+    model_path_root = (
+        epath.resource_path("mujoco") / "testdata" / "MJCF_Root.zip"
+    )
+    model_path_no_root = (
+        epath.resource_path("mujoco") / "testdata" / "MJCF_NoRoot.zip"
+    )
+    filenames = [model_path_root.as_posix(), model_path_no_root.as_posix()]
+
+    for filename in filenames:
+      with self.subTest(filename):
+        spec = mujoco.MjSpec.from_zip(filename)
+        spec.compile()
+        assets = spec.assets
+        xml_string = spec.to_xml()
+        string_spec = mujoco.MjSpec.from_string(xml_string, assets=assets)
+        string_spec.compile()
+        self.assertEqual(spec.to_xml(), string_spec.to_xml())
+
 if __name__ == '__main__':
   absltest.main()

python/pyproject.toml

@@ -57,6 +57,7 @@ mujoco = [
     "testdata/*.xml",
     "testdata/*.msh",
     "testdata/*.usda",
+    "testdata/*.zip",
 ]
 
 [project.optional-dependencies]