Fix MjSpec zip includes, deep-copy attach assets, and WASM lifetimes

python/mujoco/__init__.py

@@ -18,12 +18,18 @@
 import ctypes.util
 import os
 import platform
+import re
 import subprocess
 from typing import Any, IO, Union, Sequence
 from typing_extensions import TypeAlias
 import warnings
 import zipfile
 
+# Matches <include file="..."/> so from_zip can keep included XMLs.
+_INCLUDE_FILE_RE = re.compile(
+    br"""<\s*include\b[^>]*\bfile\s*=\s*["']([^"']+)["']""",
+    re.IGNORECASE,
+)
 # Extend the path to enable multiple directories to contribute to the same
 # package. Without this line, the `mujoco-mjx` package would not be able to
 # be discovered by import. For more information, see: https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
@@ -133,32 +139,65 @@ def from_zip(file: Union[str, IO[bytes]]) -> _specs.MjSpec:
     An MjSpec object.
   """
   assets = {}
-  xml_string = None
+  xml_files = {}
   if isinstance(file, str):
     file = open(file, 'rb')
   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:
+      if zip_info.filename.endswith(os.path.sep):
+        continue
+      with zip_file.open(zip_info.filename) as f:
+        contents = f.read()
+      if zip_info.filename.endswith('.xml'):
+        xml_files[zip_info.filename] = contents
+      else:
+        assets[zip_info.filename] = contents
+
+  if not xml_files:
     raise ValueError('No XML file found in zip file.')
 
+  # Prefer an XML that is not itself referenced by another archive member via
+  # <include file="..."/>. Remaining XML members are passed as includes.
+  referenced = set()
+  for path, contents in xml_files.items():
+    base_dir = os.path.dirname(path)
+    for match in _INCLUDE_FILE_RE.finditer(contents):
+      ref = match.group(1).decode('utf-8')
+      joined = (
+          os.path.normpath(os.path.join(base_dir, ref))
+          if base_dir
+          else os.path.normpath(ref)
+      )
+      referenced.add(joined)
+      referenced.add(ref)
+      referenced.add(os.path.basename(ref))
+
+  candidates = [
+      path
+      for path in xml_files
+      if path not in referenced and os.path.basename(path) not in referenced
+  ]
+  if not candidates:
+    candidates = list(xml_files.keys())
+  root_path = sorted(candidates, key=lambda path: (path.count('/'), path))[0]
+  xml_string = xml_files.pop(root_path)
+  xml_dir = os.path.dirname(root_path)
+
+  include = {}
+  for path, contents in xml_files.items():
+    include[os.path.relpath(path, xml_dir) if xml_dir else path] = contents
+
   relative_assets = {}
   for key, value in assets.items():
-    new_key = os.path.relpath(key, xml_dir)
-    relative_assets[new_key] = value
-  assets = relative_assets
+    relative_assets[os.path.relpath(key, xml_dir) if xml_dir else key] = value
 
-  return _specs.MjSpec.from_string(xml_string, assets=assets)
+  return _specs.MjSpec.from_string(
+      xml_string,
+      include=include or None,
+      assets=relative_assets or None,
+  )
 
 
 class _MjBindModel:

python/mujoco/specs.cc

@@ -22,6 +22,7 @@
 #include <string>
 #include <string_view>  // IWYU pragma: keep
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>       // IWYU pragma: keep
 
@@ -667,15 +668,17 @@ PYBIND11_MODULE(_specs, m, pybind11::mod_gil_not_used()) {
             throw pybind11::value_error(mjs_getError(self.ptr));
           }
         }
-        // add prefix and suffix to the assets keys
+        // Prefix asset dict keys on the parent.
         std::string pre(p);
         std::string suf(s);
+        std::unordered_set<std::string> child_asset_files;
         for (const auto& asset : child.assets) {
+          std::string original_name = asset.first.cast<std::string>();
+          child_asset_files.insert(original_name);
           std::string asset_name =
-              addPrefixAndSuffix(asset.first.cast<std::string>(), pre, suf);
+              addPrefixAndSuffix(original_name, pre, suf);
           if (self.assets.contains(asset_name) && !self.override_assets) {
-            throw pybind11::value_error("Asset " +
-                                        asset.first.cast<std::string>() +
+            throw pybind11::value_error("Asset " + original_name +
                                         " already exists in parent spec.");
           }
           self.assets[py::str(asset_name)] = asset.second;
@@ -704,22 +707,29 @@ PYBIND11_MODULE(_specs, m, pybind11::mod_gil_not_used()) {
               "asset dict might result in missing assets when attaching again.",
               1);
         }
+        // Rename mesh/texture file fields on the PARENT after attach.
+        // With copy_during_attach, mjs_attach copies child elements first; the
+        // parent copies still reference the original filenames, while asset
+        // keys are prefixed. Mutating the child would leave the parent broken
+        // and would also poison a deep-copied child for later attaches.
         if (child_use_asset_dict) {
-          while (mesh) {
-            std::string file = mjs_getString(mjs_asMesh(mesh)->file);
-            if (!file.empty()) {
+          parent_mesh = mjs_firstElement(self.ptr, mjOBJ_MESH);
+          while (parent_mesh) {
+            std::string file = mjs_getString(mjs_asMesh(parent_mesh)->file);
+            if (!file.empty() && child_asset_files.count(file)) {
               std::string mesh_file = addPrefixAndSuffix(file, pre, suf);
-              mjs_setString(mjs_asMesh(mesh)->file, mesh_file.c_str());
+              mjs_setString(mjs_asMesh(parent_mesh)->file, mesh_file.c_str());
             }
-            mesh = mjs_nextElement(child.ptr, mesh);
+            parent_mesh = mjs_nextElement(self.ptr, parent_mesh);
           }
-          while (tex) {
-            std::string file = mjs_getString(mjs_asTexture(tex)->file);
-            if (!file.empty()) {
+          parent_tex = mjs_firstElement(self.ptr, mjOBJ_TEXTURE);
+          while (parent_tex) {
+            std::string file = mjs_getString(mjs_asTexture(parent_tex)->file);
+            if (!file.empty() && child_asset_files.count(file)) {
               std::string tex_file = addPrefixAndSuffix(file, pre, suf);
-              mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str());
+              mjs_setString(mjs_asTexture(parent_tex)->file, tex_file.c_str());
             }
-            tex = mjs_nextElement(child.ptr, tex);
+            parent_tex = mjs_nextElement(self.ptr, parent_tex);
           }
         }
         child.parent = &self;

python/mujoco/specs_test.py

@@ -16,6 +16,7 @@
 
 import gc
 import inspect
+import io
 import math
 import os
 import textwrap
@@ -1922,6 +1923,86 @@ def test_from_zip(self):
         string_spec.compile()
         self.assertEqual(spec.to_xml(), string_spec.to_xml())
 
+  def test_from_zip_with_includes(self):
+    """XML includes inside a zip must round-trip through from_zip."""
+    root_xml = textwrap.dedent("""
+      <mujoco model="root">
+        <include file="included.xml"/>
+      </mujoco>
+    """).encode("utf-8")
+    included_xml = textwrap.dedent("""
+      <mujoco>
+        <worldbody>
+          <body name="included_body">
+            <geom type="box" size="1 1 1"/>
+          </body>
+        </worldbody>
+      </mujoco>
+    """).encode("utf-8")
+
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as zf:
+      zf.writestr("included.xml", included_xml)
+      zf.writestr("root.xml", root_xml)
+    buf.seek(0)
+
+    spec = mujoco.MjSpec.from_zip(buf)
+    model = spec.compile()
+    self.assertEqual(model.nbody, 2)
+    self.assertIsNotNone(spec.body("included_body"))
+
+  def test_deepcopy_attach_renames_parent_asset_files(self):
+    """Deep-copy attach must prefix mesh files on the parent, not the child."""
+    mesh_obj = b"v 0 0 0\nv 1 0 0\nv 0 1 0\nf 1 2 3\n"
+    child = mujoco.MjSpec.from_string(
+        textwrap.dedent("""
+          <mujoco>
+            <asset>
+              <mesh name="tri" file="tri.obj"/>
+            </asset>
+            <worldbody>
+              <body name="c">
+                <geom type="mesh" mesh="tri"/>
+              </body>
+            </worldbody>
+          </mujoco>
+        """),
+        assets={"tri.obj": mesh_obj},
+    )
+    parent = mujoco.MjSpec.from_string(
+        textwrap.dedent("""
+          <mujoco>
+            <worldbody>
+              <frame name="mount"/>
+            </worldbody>
+          </mujoco>
+        """)
+    )
+    parent.copy_during_attach = True
+    parent.attach(child, frame="mount", prefix="child_")
+
+    self.assertIn("child_tri.obj", parent.assets)
+    self.assertEqual([mesh.file for mesh in parent.meshes], ["child_tri.obj"])
+    # Child remains reusable for another deep-copy attach.
+    self.assertEqual([mesh.file for mesh in child.meshes], ["tri.obj"])
+
+    model = parent.compile()
+    self.assertGreaterEqual(model.nmesh, 1)
+
+    parent2 = mujoco.MjSpec.from_string(
+        textwrap.dedent("""
+          <mujoco>
+            <worldbody>
+              <frame name="mount2"/>
+            </worldbody>
+          </mujoco>
+        """)
+    )
+    parent2.copy_during_attach = True
+    parent2.attach(child, frame="mount2", prefix="again_")
+    self.assertEqual([mesh.file for mesh in parent2.meshes], ["again_tri.obj"])
+    parent2.compile()
+
   def test_rangefinder_sensor(self):
     """Test rangefinder sensor with mjSpec, iterative model building."""
     # Raydata field enum values for dataspec bitfield

wasm/codegen/generated/bindings.cc

@@ -1149,6 +1149,7 @@ std::vector<MjContact> MjData::contact() const {
 
 MjvScene::MjvScene() {
   owned_ = true;
+  model = nullptr;
   ptr_ = new mjvScene;
   mjv_defaultScene(ptr_);
   mjv_makeScene(nullptr, ptr_, 0);
@@ -1178,6 +1179,9 @@ void MjvScene::set(mjvScene* ptr) { ptr_ = ptr; }
 
 // Taken from the python mujoco bindings code for MjvScene Wrapper
 int MjvScene::GetSumFlexFaces() const {
+  if (!model) {
+    return 0;
+  }
   int nflexface = 0;
   int flexfacenum = 0;
   for (int f = 0; f < model->nflex; f++) {

wasm/codegen/generated/bindings.h

@@ -7456,10 +7456,18 @@ struct MjvScene {
         emscripten::typed_memory_view(ptr_->nflex, ptr_->flexfaceused));
   }
   emscripten::val flexedge() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<int*>(nullptr)));
+    }
     return emscripten::val(
         emscripten::typed_memory_view(2 * model->nflexedge, ptr_->flexedge));
   }
   emscripten::val flexvert() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(
         emscripten::typed_memory_view(3 * model->nflexvert, ptr_->flexvert));
   }
@@ -7476,22 +7484,42 @@ struct MjvScene {
         emscripten::typed_memory_view(ptr_->nskin, ptr_->skinvertnum));
   }
   emscripten::val skinvert() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(
         emscripten::typed_memory_view(3 * model->nskinvert, ptr_->skinvert));
   }
   emscripten::val skinnormal() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(
         emscripten::typed_memory_view(3 * model->nskinvert, ptr_->skinnormal));
   }
   emscripten::val flexface() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(emscripten::typed_memory_view(
         9 * MjvScene::GetSumFlexFaces(), ptr_->flexface));
   }
   emscripten::val flexnormal() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(emscripten::typed_memory_view(
         9 * MjvScene::GetSumFlexFaces(), ptr_->flexnormal));
   }
   emscripten::val flextexcoord() const {
+    if (!model) {
+      return emscripten::val(
+          emscripten::typed_memory_view(0, static_cast<float*>(nullptr)));
+    }
     return emscripten::val(emscripten::typed_memory_view(
         6 * MjvScene::GetSumFlexFaces(), ptr_->flextexcoord));
   }
@@ -7618,14 +7646,22 @@ struct MjvScene {
   bool owned_ = false;
 
  public:
-  mjModel* model;
+  mjModel* model = nullptr;
   std::vector<MjvLight> lights;
   std::vector<MjvGLCamera> camera;
 };
 
 struct MjVFS {
   MjVFS() : ptr_(new mjVFS) { mj_defaultVFS(ptr_); }
-  ~MjVFS() { mj_deleteVFS(ptr_); }
+  ~MjVFS() {
+    if (ptr_) {
+      mj_deleteVFS(ptr_);
+      delete ptr_;
+      ptr_ = nullptr;
+    }
+  }
+  MjVFS(const MjVFS&) = delete;
+  MjVFS& operator=(const MjVFS&) = delete;
   void AddBuffer(const std::string& name, const emscripten::val& buffer) {
     std::vector<uint8_t> vec = emscripten::vecFromJSArray<uint8_t>(buffer);
     int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size());

wasm/codegen/templates/bindings.cc

@@ -205,6 +205,7 @@ std::vector<MjContact> MjData::contact() const {
 
 MjvScene::MjvScene() {
   owned_ = true;
+  model = nullptr;
   ptr_ = new mjvScene;
   mjv_defaultScene(ptr_);
   mjv_makeScene(nullptr, ptr_, 0);
@@ -234,6 +235,9 @@ void MjvScene::set(mjvScene* ptr) { ptr_ = ptr; }
 
 // Taken from the python mujoco bindings code for MjvScene Wrapper
 int MjvScene::GetSumFlexFaces() const {
+  if (!model) {
+    return 0;
+  }
   int nflexface = 0;
   int flexfacenum = 0;
   for (int f = 0; f < model->nflex; f++) {