From c191cfaaf965d18a3113a8af858a0d3c6cd87061 Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Wed, 29 Jul 2026 03:06:31 +0530 Subject: [PATCH] Fix MjSpec zip includes, deep-copy attach assets, and WASM lifetimes. from_zip now selects a root XML and passes sibling XML members as includes; deep-copy attach prefixes mesh/texture files on the parent instead of the child; default MjvScene no longer null-dereferences model accessors; and MjVFS frees its wrapper allocation after mj_deleteVFS. --- python/mujoco/__init__.py | 69 +++++++++++++++++++------ python/mujoco/specs.cc | 38 ++++++++------ python/mujoco/specs_test.py | 81 ++++++++++++++++++++++++++++++ wasm/codegen/generated/bindings.cc | 4 ++ wasm/codegen/generated/bindings.h | 40 ++++++++++++++- wasm/codegen/templates/bindings.cc | 4 ++ wasm/codegen/templates/bindings.h | 40 ++++++++++++++- 7 files changed, 243 insertions(+), 33 deletions(-) diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py index 1147be58ed9..3d2ef6e11c3 100644 --- a/python/mujoco/__init__.py +++ b/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 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 + # . 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: diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 94372f5e876..9346984c0a4 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -22,6 +22,7 @@ #include #include // IWYU pragma: keep #include +#include #include #include // 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 child_asset_files; for (const auto& asset : child.assets) { + std::string original_name = asset.first.cast(); + child_asset_files.insert(original_name); std::string asset_name = - addPrefixAndSuffix(asset.first.cast(), pre, suf); + addPrefixAndSuffix(original_name, pre, suf); if (self.assets.contains(asset_name) && !self.override_assets) { - throw pybind11::value_error("Asset " + - asset.first.cast() + + 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; diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 0b6fa50cbab..5b774ee3baf 100644 --- a/python/mujoco/specs_test.py +++ b/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(""" + + + + """).encode("utf-8") + included_xml = textwrap.dedent(""" + + + + + + + + """).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(""" + + + + + + + + + + + """), + assets={"tri.obj": mesh_obj}, + ) + parent = mujoco.MjSpec.from_string( + textwrap.dedent(""" + + + + + + """) + ) + 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(""" + + + + + + """) + ) + 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 diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 861be173130..7198c20ed41 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -1149,6 +1149,7 @@ std::vector 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++) { diff --git a/wasm/codegen/generated/bindings.h b/wasm/codegen/generated/bindings.h index 48697ecc235..0a7f33047e4 100644 --- a/wasm/codegen/generated/bindings.h +++ b/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(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(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(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(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(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(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(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 lights; std::vector 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 vec = emscripten::vecFromJSArray(buffer); int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size()); diff --git a/wasm/codegen/templates/bindings.cc b/wasm/codegen/templates/bindings.cc index 356f7749fdd..0a8e40a0c6e 100644 --- a/wasm/codegen/templates/bindings.cc +++ b/wasm/codegen/templates/bindings.cc @@ -205,6 +205,7 @@ std::vector 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++) { diff --git a/wasm/codegen/templates/bindings.h b/wasm/codegen/templates/bindings.h index 92e3c771458..f44cb96c797 100644 --- a/wasm/codegen/templates/bindings.h +++ b/wasm/codegen/templates/bindings.h @@ -378,10 +378,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(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(nullptr))); + } return emscripten::val( emscripten::typed_memory_view(3 * model->nflexvert, ptr_->flexvert)); } @@ -398,22 +406,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(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(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(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(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(nullptr))); + } return emscripten::val(emscripten::typed_memory_view( 6 * MjvScene::GetSumFlexFaces(), ptr_->flextexcoord)); } @@ -424,14 +452,22 @@ struct MjvScene { bool owned_ = false; public: - mjModel* model; + mjModel* model = nullptr; std::vector lights; std::vector 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 vec = emscripten::vecFromJSArray(buffer); int result = mj_addBufferVFS(ptr_, name.c_str(), vec.data(), vec.size());