From 723b8b1ea60237c6ce2ef4a47f19050288bdcd7d Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Fri, 1 May 2026 09:33:43 -0700 Subject: [PATCH] Introduce explicit VFS in Python MjVfs This change introduces MjVFS, an explicit object to mirror the C mjVFS. This is meant to replace the MjSpec.assets dict. This latter while convenient guides users towards harmful authoring patterns with respect to data duplication and spec attachment workflows. Making VFS management explicit should encourage better memory usage and allow us to make better compile time optimizations. From this change, `spec.assets` is deprecated. However we will temporarily support backwards compatibility due to the wide spread usage. An error will be thrown if calling code tries to use a spec that uses both the assets dict and the new MjVfs. PiperOrigin-RevId: 908772050 Change-Id: I77c6d0369307fc300c954768fed17401261e18da --- doc/changelog.rst | 13 ++ doc/python.rst | 58 ++++++- python/mujoco/specs.cc | 109 ++++++++++++- python/mujoco/specs_wrapper.cc | 60 +++---- python/mujoco/specs_wrapper.h | 2 +- python/mujoco/structs.cc | 18 +-- python/mujoco/structs.h | 12 +- python/mujoco/structs_wrappers.cc | 108 ++++++++++--- python/mujoco/vfs.h | 46 ++++++ python/mujoco/vfs_test.py | 251 ++++++++++++++++++++++++++++++ 10 files changed, 599 insertions(+), 78 deletions(-) create mode 100644 python/mujoco/vfs.h create mode 100644 python/mujoco/vfs_test.py diff --git a/doc/changelog.rst b/doc/changelog.rst index eb2b05f2..26f1f591 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -55,6 +55,19 @@ Bug fixes than the parent spec. This prevents the origin of the parent spec to affect the resolution of asset paths in the child spec. +Python +^^^^^^ + +- Added ``mujoco.MjVfs`` Python binding to interact with the Virtual File System directly from Python. + See :ref:`Virtual File System ` for usage details. + + .. warning:: + The previous way of passing assets via a dictionary mapping asset names to bytes is **deprecated** and will be + removed in an upcoming release. You cannot specify both the ``assets`` dictionary and the ``vfs`` argument at the same + time. ``MjVfs`` should be used as a drop-in replacement. + + + Version 3.7.0 (April 14, 2026) ------------------------------ diff --git a/doc/python.rst b/doc/python.rst index 95beb3ff..963d366d 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -275,6 +275,8 @@ that create a new :ref:`mjModel` instance: ``mujoco.MjModel.from_xml_string``, ` functions accept the path to either an XML or MJB model file. All three functions optionally accept a Python dictionary which is converted into a MuJoCo :ref:`Virtualfilesystem` for use during model compilation. + + .. _PyFunctions: Functions @@ -514,11 +516,53 @@ The ``MjSpec`` object wraps the :ref:`mjSpec` struct and can be constructed in t Note the ``from_string()`` and ``from_file()`` methods can only be called at construction time. +.. _PyVFS: + Assets ^^^^^^ -All three methods take in an optional argument called ``assets`` which is used to resolve asset references in the XML. -This argument is a dictionary that maps asset name (string) to asset data (bytes), as demonstrated below: +MuJoCo optionally uses a :ref:`Virtual File System ` (VFS) to load assets (like meshes and textures) +from memory. Some :ref:`decoders` may also choose to leverage the VFS as a way to load assets on +demand, such as when addressing files in an archive format. This requires the same VFS to be used when parsing and +compiling a spec (and all attached specs) into a model. + +The Python bindings provide the ``mujoco.MjVfs`` as a wrapper around the :ref:`mjVFS` C struct. + +``MjVfs`` supports the context manager protocol, which ensures that resources are properly freed when leaving the block: + +.. code-block:: python + + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = b"" + spec = mujoco.MjSpec.from_string("model.xml", vfs=vfs) + spec.compile(vfs=vfs) + +You can also create an instance directly and call ``close()`` when done: + +.. code-block:: python + + vfs = mujoco.MjVfs() + vfs["model.xml"] = some_xml_string.encode("utf-8") + spec = mujoco.MjSpec.from_file("model.xml", vfs=vfs) + spec.compile(vfs=vfs) + vfs.close() + +The ``MjVfs`` object supports dictionary-like operations to manage buffers: + +- ``vfs["name"] = data``: Adds a buffer to the VFS. ``data`` must be of type ``bytes``. +- ``del vfs["name"]``: Deletes a file from the VFS. +- ``"name" in vfs``: Checks if a file exists in the VFS. + +The static factory functions ``mujoco.MjModel.from_xml_string``, ``mujoco.MjModel.from_xml_path``, +``mujoco.MjSpec.from_string`` and ``mujoco.MjSpec.from_file`` accept an optional ``vfs`` argument. Additionally, the +``spec.compile()`` function also accepts an optional ``vfs`` argument. + +.. warning:: + The previous way of passing assets via a dictionary mapping asset names to bytes is **deprecated** and will be + removed in the next release. You cannot specify both the ``assets`` dictionary and the ``vfs`` argument at the same + time. ``MjVfs`` should be used as a drop-in replacement. + +For reference, the deprecated ``assets`` dictionary approach looked like this: .. code-block:: python @@ -526,6 +570,12 @@ This argument is a dictionary that maps asset name (string) to asset data (bytes spec = mujoco.MjSpec.from_string(xml_referencing_image_png, assets=assets) model = spec.compile() + # Or + + spec = mujoco.MjSpec.from_string(xml_referencing_image_png) + spec.assets = {'image.png': b'image_data'} + model = spec.compile() + Save to XML ----------- @@ -566,8 +616,8 @@ It is possible to combine multiple specs by using attachments. The following opt the reference to a frame, which is the attached worldbody transformed into a frame. The site must belong to the child spec. Prefix and suffix can also be specified as keyword arguments. - Attach a child spec to a frame in the parent spec: ``parent_spec.attach(child_spec, frame=frame_name_or_obj)``, - returns the reference to a frame, which is the attached worldbody transformed into a frame. The frame must belong to - the child spec. Prefix and suffix can also be specified as keyword arguments. + returns the reference to a frame, which is the attached worldbody transformed into a frame. The frame must + belong to the child spec. Prefix and suffix can also be specified as keyword arguments. The default behavior of attaching is to not copy, so all the child references (except for the worldbody) are still valid in the parent and therefore modifying the child will modify the parent. This is not true for the attach diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 2ba20431..c5a60b52 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -18,6 +18,7 @@ #include #include #include +#include #include #include // IWYU pragma: keep #include @@ -33,6 +34,7 @@ #include "specs_wrapper.h" // IWYU pragma: keep #include "raw.h" #include "structs.h" // IWYU pragma: keep +#include "vfs.h" #include #include #include @@ -264,6 +266,49 @@ PYBIND11_MODULE(_specs, m) { DefineArray(m, "MjFloatVec"); DefineArray(m, "MjIntVec"); + // ============================= MJVFS ===================================== + py::class_(m, "MjVfs") + .def(py::init<>()) + .def("close", &MjVfs::Close) + .def("__enter__", [](MjVfs& self) -> MjVfs& { return self; }) + .def("__exit__", + [](MjVfs& self, py::object, py::object, py::object) { + self.Close(); + }) + .def("__setitem__", + [](MjVfs& self, const std::string& name, py::bytes data) { + if (!self.is_open()) { + throw std::runtime_error("VFS is closed"); + } + std::string_view buffer = data; + const int err = mj_addBufferVFS( + self.get(), name.c_str(), buffer.data(), buffer.size()); + if (err == 2) { + throw py::value_error( + "Repeated file name in VFS: " + name); + } else if (err) { + throw py::value_error( + "Failed to add buffer to VFS: " + name); + } + }) + .def("__delitem__", + [](MjVfs& self, const std::string& name) { + if (!self.is_open()) { + throw std::runtime_error("VFS is closed"); + } + if (mj_deleteFileVFS(self.get(), name.c_str())) { + throw py::key_error(name); + } + }) + .def("__contains__", + [](MjVfs& self, const std::string& name) { + if (!self.is_open()) { + throw std::runtime_error("VFS is closed"); + } + return mj_containsBufferVFS(self.get(), name.c_str()) == 1; + }); + + // ============================= MJSPEC ===================================== mjSpec.def(py::init<>()); mjSpec.def_property_readonly( @@ -272,7 +317,26 @@ PYBIND11_MODULE(_specs, m) { "from_file", [](std::string& filename, std::optional>& include, - std::optional& assets) -> MjSpec { + std::optional& assets, + MjVfs* vfs) -> MjSpec { + if (vfs && (include.has_value() || assets.has_value())) { + throw py::value_error( + "Cannot specify both 'vfs' and 'include'/'assets'."); + } + if (vfs) { + raw::MjSpec* spec; + { + py::gil_scoped_release no_gil; + char error[1024]; + spec = InterceptMjErrors(mj_parse)( + filename.c_str(), nullptr, vfs->get(), + error, sizeof(error)); + if (!spec) { + throw py::value_error(error); + } + } + return MjSpec(spec); + } const auto files = _impl::ConvertAssetsDict(include); raw::MjSpec* spec; { @@ -294,7 +358,8 @@ PYBIND11_MODULE(_specs, m) { return MjSpec(spec); }, py::arg("filename"), py::arg("include") = py::none(), - py::arg("assets") = py::none(), R"mydelimiter( + py::arg("assets") = py::none(), py::arg("vfs") = py::none(), + R"mydelimiter( Creates a spec from an XML file. Parameters @@ -307,13 +372,34 @@ PYBIND11_MODULE(_specs, m) { assets : dict, optional A dictionary of assets to be used by the spec. The keys are asset names and the values are asset contents. + vfs : MjVfs, optional + A VFS to use for resolving includes and assets. Cannot be used with + include or assets. )mydelimiter", py::return_value_policy::move); mjSpec.def_static( "from_string", [](std::string& xml, std::optional>& include, - std::optional& assets) -> MjSpec { + std::optional& assets, + MjVfs* vfs) -> MjSpec { + if (vfs && (include.has_value() || assets.has_value())) { + throw py::value_error( + "Cannot specify both 'vfs' and 'include'/'assets'."); + } + if (vfs) { + raw::MjSpec* spec; + { + py::gil_scoped_release no_gil; + char error[1024]; + spec = InterceptMjErrors(mj_parseXMLString)( + xml.c_str(), vfs->get(), error, sizeof(error)); + if (!spec) { + throw py::value_error(error); + } + } + return MjSpec(spec); + } auto files = _impl::ConvertAssetsDict(include); raw::MjSpec* spec; { @@ -345,7 +431,8 @@ PYBIND11_MODULE(_specs, m) { return MjSpec(spec); }, py::arg("xml"), py::arg("include") = py::none(), - py::arg("assets") = py::none(), R"mydelimiter( + py::arg("assets") = py::none(), py::arg("vfs") = py::none(), + R"mydelimiter( Creates a spec from an XML string. Parameters @@ -358,6 +445,9 @@ PYBIND11_MODULE(_specs, m) { assets : dict, optional A dictionary of assets to be used by the spec. The keys are asset names and the values are asset contents. + vfs : MjVfs, optional + A VFS to use for resolving includes and assets. Cannot be used with + include or assets. )mydelimiter", py::return_value_policy::move); mjSpec.def("recompile", [mjmodel_mjdata_from_spec_ptr]( @@ -391,9 +481,14 @@ PYBIND11_MODULE(_specs, m) { return mjs_findDefault(self.ptr, classname.c_str()); }, py::return_value_policy::reference_internal); - mjSpec.def("compile", [mjmodel_from_raw_ptr](MjSpec& self) -> py::object { - return mjmodel_from_raw_ptr(reinterpret_cast(self.Compile())); - }); + mjSpec.def("compile", + [mjmodel_from_raw_ptr](MjSpec& self, + std::optional vfs) -> py::object { + mjVFS* vfs_ptr = vfs.has_value() ? (*vfs)->get() : nullptr; + return mjmodel_from_raw_ptr( + reinterpret_cast(self.Compile(vfs_ptr))); + }, + py::arg("vfs") = py::none()); mjSpec.def_property( "assets", [](MjSpec& self) -> py::dict { diff --git a/python/mujoco/specs_wrapper.cc b/python/mujoco/specs_wrapper.cc index f1ba1188..31d52e41 100644 --- a/python/mujoco/specs_wrapper.cc +++ b/python/mujoco/specs_wrapper.cc @@ -15,6 +15,7 @@ #include "specs_wrapper.h" #include // IWYU pragma: keep +#include #include #include // IWYU pragma: keep #include // IWYU pragma: keep @@ -87,44 +88,45 @@ MjSpec& MjSpec::operator=(MjSpec&& other) { MjSpec::~MjSpec() { mj_deleteSpec(ptr); } -raw::MjModel* MjSpec::Compile() { - if (assets.empty()) { - raw::MjModel* m; - { - // Release GIL before calling mj_compile which may spawn threads - py::gil_scoped_release no_gil; - m = mj_compile(ptr, 0); - } - if (!m || mjs_isWarning(ptr)) { - throw py::value_error(mjs_getError(ptr)); - } - return m; +raw::MjModel* MjSpec::Compile(mjVFS* vfs) { + if (vfs != nullptr && !assets.empty()) { + throw py::value_error("Cannot specify both 'vfs' and 'assets'."); } - mjVFS vfs; - mj_defaultVFS(&vfs); - for (const auto& asset : assets) { - std::string buffer_name = py::cast(asset.first).c_str(); - std::string buffer = py::cast(asset.second); - const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( - &vfs, buffer_name.c_str(), buffer.c_str(), buffer.size()); - if (vfs_error) { - mj_deleteVFS(&vfs); - if (vfs_error == 2) { - throw py::value_error("Repeated file name in assets dict: " + - buffer_name); - } else { - throw py::value_error("Asset failed to load: " + buffer_name); + + std::optional local_vfs; + if (vfs == nullptr) { + vfs = &local_vfs.emplace(); + mj_defaultVFS(vfs); + + for (const auto& asset : assets) { + std::string buffer_name = py::cast(asset.first).c_str(); + std::string buffer = py::cast(asset.second); + const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( + vfs, buffer_name.c_str(), buffer.c_str(), buffer.size()); + if (vfs_error) { + mj_deleteVFS(vfs); + if (vfs_error == 2) { + throw py::value_error("Repeated file name in assets dict: " + + buffer_name); + } else { + throw py::value_error("Asset failed to load: " + buffer_name); + } } } } raw::MjModel* m; { - // Release GIL before calling mj_compile which may spawn threads py::gil_scoped_release no_gil; - m = mj_compile(ptr, &vfs); + m = mj_compile(ptr, vfs); } - mj_deleteVFS(&vfs); + + if (local_vfs.has_value()) { + // vfs points at local_vfs value. + mj_deleteVFS(vfs); + local_vfs = std::nullopt; + } + if (!m || mjs_isWarning(ptr)) { throw py::value_error(mjs_getError(ptr)); } diff --git a/python/mujoco/specs_wrapper.h b/python/mujoco/specs_wrapper.h index 2ebaac60..66c1adab 100644 --- a/python/mujoco/specs_wrapper.h +++ b/python/mujoco/specs_wrapper.h @@ -35,7 +35,7 @@ struct MjSpec { MjSpec& operator=(MjSpec&& other); ~MjSpec(); - raw::MjModel* Compile(); + raw::MjModel* Compile(mjVFS* vfs = nullptr); raw::MjSpec* ptr; py::dict assets; diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 1c9e04d9..79840742 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -41,6 +41,7 @@ #include #include #include +#include "vfs.h" namespace mujoco::python::_impl { @@ -301,24 +302,21 @@ PYBIND11_MODULE(_structs, m) { // ==================== MJMODEL ============================================== py::class_ mjModel(m, "MjModel"); mjModel.def_static( - "from_xml_string", &MjModelWrapper::LoadXML, py::arg("xml"), - py::arg_v("assets", py::none()), + "from_xml_string", MjModelWrapper::LoadXML, py::arg("xml"), + py::arg_v("assets", py::none()), py::arg("vfs") = py::none(), py::doc( - R"(Loads an MjModel from an XML string and an optional assets dictionary.)")); + R"(Loads an MjModel from an XML string and optional assets dict or VFS.)")); mjModel.def_static("_from_model_ptr", [](uintptr_t addr) { return MjModelWrapper::WrapRawModel(reinterpret_cast(addr)); }); mjModel.def_static( - "from_xml_path", &MjModelWrapper::LoadXMLFile, py::arg("filename"), - py::arg_v("assets", py::none()), + "from_xml_path", MjModelWrapper::LoadXMLFile, py::arg("filename"), + py::arg_v("assets", py::none()), py::arg("vfs") = py::none(), py::doc( - R"(Loads an MjModel from an XML file and an optional assets dictionary. - -The filename for the XML can also refer to a key in the assets dictionary. -This is useful for example when the XML is not available as a file on disk.)")); + R"(Loads an MjModel from an XML file and optional assets dict or VFS.)")); mjModel.def_static( "from_binary_path", &MjModelWrapper::LoadBinaryFile, py::arg("filename"), - py::arg_v("assets", py::none()), + py::arg_v("assets", py::none()), py::arg("vfs") = py::none(), py::doc( R"(Loads an MjModel from an MJB file and an optional assets dictionary. diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 5c1e70cb..03080618 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -44,6 +44,9 @@ namespace py = ::pybind11; namespace mujoco::python { + +class MjVfs; + namespace _impl { struct VfsAsset { @@ -501,17 +504,20 @@ class MjWrapper : public WrapperBase { static MjWrapper LoadXMLFile( const std::string& filename, const std::optional< - std::unordered_map>& assets); + std::unordered_map>& assets, + MjVfs* vfs = nullptr); static MjWrapper LoadBinaryFile( const std::string& filename, const std::optional< - std::unordered_map>& assets); + std::unordered_map>& assets, + MjVfs* vfs = nullptr); static MjWrapper LoadXML( const std::string& xml, const std::optional< - std::unordered_map>& assets); + std::unordered_map>& assets, + MjVfs* vfs = nullptr); static MjWrapper WrapRawModel(raw::MjModel* m); diff --git a/python/mujoco/structs_wrappers.cc b/python/mujoco/structs_wrappers.cc index ebfa907b..0f46d811 100644 --- a/python/mujoco/structs_wrappers.cc +++ b/python/mujoco/structs_wrappers.cc @@ -19,6 +19,7 @@ #include #include // NOLINT(build/c++11) #include +#include #include #include #include @@ -38,6 +39,7 @@ #include "raw.h" #include "serialization.h" #include "structs.h" +#include "vfs.h" #include #include #include @@ -79,6 +81,18 @@ constexpr auto XArrayShapeImpl(const std::string_view dim1_str) { inline std::size_t NConMax(const mjData* d) { return d->narena / sizeof(mjContact); } +template +struct Cleanup final { + Callback clean_func; + Cleanup(Callback callback) : clean_func(std::move(callback)) {} + ~Cleanup() { clean_func(); } +}; + +// `Cleanup c = /* callback */;` +// +// C++17 type deduction API for creating an instance of `Cleanup` +template +Cleanup(Callback callback) -> Cleanup; } // namespace @@ -295,18 +309,29 @@ MjModelWrapper::~MjWrapper() { template static raw::MjModel* LoadModelFileImpl(const std::string& filename, const std::vector& assets, + mjVFS* vfs, LoadFunc&& loadfunc) { - mjVFS vfs; - mjVFS* vfs_ptr = nullptr; + if (!assets.empty() && vfs != nullptr) { + throw py::value_error("Cannot specify both 'assets' and 'vfs'."); + } + + std::optional local_vfs; + Cleanup vfs_cleanup = [&]() { + if (local_vfs.has_value()) { + mj_deleteVFS(vfs); + }; + }; + if (!assets.empty()) { - mj_defaultVFS(&vfs); - vfs_ptr = &vfs; + vfs = &local_vfs.emplace(); + mj_defaultVFS(vfs); + for (const auto& asset : assets) { std::string buffer_name = StripPath(asset.name); const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( - vfs_ptr, buffer_name.c_str(), asset.content, asset.content_size); + vfs, buffer_name.c_str(), asset.content, + asset.content_size); if (vfs_error) { - mj_deleteVFS(vfs_ptr); if (vfs_error == 2) { throw py::value_error("Repeated file name in assets dict: " + buffer_name); @@ -317,28 +342,34 @@ static raw::MjModel* LoadModelFileImpl(const std::string& filename, } } - raw::MjModel* model = loadfunc(filename.c_str(), vfs_ptr); - mj_deleteVFS(vfs_ptr); + raw::MjModel* model = loadfunc(filename.c_str(), vfs); if (model && !model->buffer) { mj_deleteModel(model); model = nullptr; } + return model; } MjModelWrapper MjModelWrapper::LoadXMLFile( const std::string& filename, - const std::optional>& assets) { + const std::optional>& assets, + MjVfs* vfs) { + if (assets.has_value() && vfs != nullptr) { + throw py::value_error("Cannot specify both 'assets' and 'vfs'."); + } + const auto converted_assets = ConvertAssetsDict(assets); raw::MjModel* model; { py::gil_scoped_release no_gil; char error[1024]; - model = LoadModelFileImpl(filename, converted_assets, - [&error](const char* filename, const mjVFS* vfs) { - return InterceptMjErrors(mj_loadXML)( - filename, vfs, error, sizeof(error)); - }); + model = LoadModelFileImpl( + filename, converted_assets, vfs ? vfs->get() : nullptr, + [&error](const char* filename, const mjVFS* vfs) { + return InterceptMjErrors(mj_loadXML)( + filename, vfs, error, sizeof(error)); + }); if (!model) { throw py::value_error(error); } @@ -348,12 +379,18 @@ MjModelWrapper MjModelWrapper::LoadXMLFile( MjModelWrapper MjModelWrapper::LoadBinaryFile( const std::string& filename, - const std::optional>& assets) { + const std::optional>& assets, + MjVfs* vfs) { + if (assets.has_value() && vfs != nullptr) { + throw py::value_error("Cannot specify both 'assets' and 'vfs'."); + } + const auto converted_assets = ConvertAssetsDict(assets); raw::MjModel* model; { py::gil_scoped_release no_gil; model = LoadModelFileImpl(filename, converted_assets, + vfs ? vfs->get() : nullptr, InterceptMjErrors(mj_loadModel)); if (!model) { throw py::value_error("mj_loadModel: failed to load from mjb"); @@ -364,22 +401,44 @@ MjModelWrapper MjModelWrapper::LoadBinaryFile( MjModelWrapper MjModelWrapper::LoadXML( const std::string& xml, - const std::optional>& assets) { + const std::optional>& assets, + MjVfs* vfs) { + if (assets.has_value() && vfs != nullptr) { + throw py::value_error("Cannot specify both 'assets' and 'vfs'."); + } + auto converted_assets = ConvertAssetsDict(assets); raw::MjModel* model; { py::gil_scoped_release no_gil; - std::string model_filename = "model_.xml"; - if (assets.has_value()) { - while (assets->find(model_filename) != assets->end()) { - model_filename = - model_filename.substr(0, model_filename.size() - 4) + "_.xml"; + std::string model_identifier = "model_.xml"; + bool file_added = false; + Cleanup file_cleanup = [&]() { + if (file_added) { + mj_deleteFileVFS(vfs->get(), model_identifier.c_str()); } + }; + if (vfs != nullptr) { + while (mj_containsBufferVFS(vfs->get(), model_identifier.c_str())) { + model_identifier = + model_identifier.substr(0, model_identifier.size() - 4) + "_.xml"; + } + + mj_addBufferVFS(vfs->get(), model_identifier.c_str(), xml.c_str(), + xml.length()); + file_added = true; + } else { + while (assets.has_value() && + assets->find(model_identifier) != assets->end()) { + model_identifier = + model_identifier.substr(0, model_identifier.size() - 4) + "_.xml"; + } + converted_assets.emplace_back(model_identifier.c_str(), xml.c_str(), + xml.length()); } - converted_assets.emplace_back(model_filename.c_str(), xml.c_str(), - xml.length()); char error[1024]; - model = LoadModelFileImpl(model_filename, converted_assets, + model = LoadModelFileImpl(model_identifier, converted_assets, + vfs ? vfs->get() : nullptr, [&error](const char* filename, const mjVFS* vfs) { return InterceptMjErrors(mj_loadXML)( filename, vfs, error, sizeof(error)); @@ -468,6 +527,7 @@ std::unique_ptr MjModelWrapper::Deserialize( raw::MjModel* model = LoadModelFileImpl( "model.mjb", {{"model.mjb", model_bytes.data(), static_cast(model_size)}}, + nullptr, InterceptMjErrors(mj_loadModel)); if (!model) { throw py::value_error("Invalid serialized mjModel."); diff --git a/python/mujoco/vfs.h b/python/mujoco/vfs.h new file mode 100644 index 00000000..f14ec198 --- /dev/null +++ b/python/mujoco/vfs.h @@ -0,0 +1,46 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_PYTHON_VFS_H_ +#define MUJOCO_PYTHON_VFS_H_ + +#include + +#include + +namespace mujoco::python { + +class MjVfs { + public: + MjVfs() : vfs_(new mjVFS) { mj_defaultVFS(vfs_.get()); } + + void Close() { vfs_.reset(); } + + mjVFS* get() const { return vfs_.get(); } + + bool is_open() const { return vfs_ != nullptr; } + + private: + struct VfsDeleter { + void operator()(mjVFS* vfs) const { + mj_deleteVFS(vfs); + delete vfs; + } + }; + std::unique_ptr vfs_; +}; + +} // namespace mujoco::python + +#endif // MUJOCO_PYTHON_VFS_H_ diff --git a/python/mujoco/vfs_test.py b/python/mujoco/vfs_test.py new file mode 100644 index 00000000..dff3cfd1 --- /dev/null +++ b/python/mujoco/vfs_test.py @@ -0,0 +1,251 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import textwrap + +from absl.testing import absltest +import mujoco + + +SIMPLE_XML = b"" + +XML_WITH_MESH = textwrap.dedent("""\ + + + + + + + + + + +""").encode() + +BOX_OBJ = textwrap.dedent("""\ + v -1 -1 -1 + v 1 -1 -1 + v 1 1 -1 + v -1 1 -1 + v -1 -1 1 + v 1 -1 1 + v 1 1 1 + v -1 1 1 + f 1 2 3 4 + f 5 8 7 6 + f 1 5 6 2 + f 2 6 7 3 + f 3 7 8 4 + f 4 8 5 1 +""").encode() + + +class VfsLifecycleTest(absltest.TestCase): + + def test_create_and_close(self): + vfs = mujoco.MjVfs() + vfs.close() + + def test_context_manager(self): + with mujoco.MjVfs() as vfs: + self.assertIsNotNone(vfs) + + def test_double_close_is_safe(self): + vfs = mujoco.MjVfs() + vfs.close() + vfs.close() + + def test_operations_after_close_raise(self): + vfs = mujoco.MjVfs() + vfs.close() + with self.assertRaises(RuntimeError): + vfs["model.xml"] = SIMPLE_XML + + +class VfsBufferTest(absltest.TestCase): + + def test_add_buffer(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = SIMPLE_XML + self.assertIn("model.xml", vfs) + + def test_add_duplicate_raises(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = SIMPLE_XML + with self.assertRaises(ValueError): + vfs["model.xml"] = SIMPLE_XML + + def test_delete_buffer(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = SIMPLE_XML + self.assertIn("model.xml", vfs) + del vfs["model.xml"] + self.assertNotIn("model.xml", vfs) + + def test_delete_missing_raises(self): + with mujoco.MjVfs() as vfs: + with self.assertRaises(KeyError): + del vfs["nonexistent"] + + def test_contains_missing(self): + with mujoco.MjVfs() as vfs: + self.assertNotIn("nonexistent", vfs) + + +class VfsCompileTest(absltest.TestCase): + + def test_compile_simple_model(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = SIMPLE_XML + model = mujoco.MjModel.from_xml_path("model.xml", vfs=vfs) + self.assertEqual(model.nq, 0) + + def test_compile_model_with_mesh(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = XML_WITH_MESH + vfs["box.obj"] = BOX_OBJ + model = mujoco.MjModel.from_xml_path("model.xml", vfs=vfs) + self.assertEqual(model.nmesh, 1) + + def test_spec_from_string_with_vfs(self): + with mujoco.MjVfs() as vfs: + vfs["box.obj"] = BOX_OBJ + spec = mujoco.MjSpec.from_string(XML_WITH_MESH.decode(), vfs=vfs) + model = spec.compile(vfs=vfs) + self.assertEqual(model.nmesh, 1) + + def test_from_xml_string_with_vfs(self): + with mujoco.MjVfs() as vfs: + vfs["box.obj"] = BOX_OBJ + model = mujoco.MjModel.from_xml_string( + XML_WITH_MESH.decode(), vfs=vfs + ) + self.assertEqual(model.nmesh, 1) + + def test_spec_compile_with_vfs(self): + with mujoco.MjVfs() as vfs: + vfs["box.obj"] = BOX_OBJ + spec = mujoco.MjSpec.from_string(XML_WITH_MESH.decode()) + model = spec.compile(vfs=vfs) + self.assertEqual(model.nmesh, 1) + + def test_spec_recompile_with_vfs(self): + with mujoco.MjVfs() as vfs: + vfs["box.obj"] = BOX_OBJ + spec = mujoco.MjSpec.from_string(XML_WITH_MESH.decode()) + model1 = spec.compile(vfs=vfs) + self.assertEqual(model1.ngeom, 1) + + body = spec.worldbody.add_body() + body.add_geom(size=[1, 0, 0]) + model2 = spec.compile(vfs=vfs) + self.assertEqual(model2.ngeom, 2) + + def test_long_lived_vfs_without_context(self): + vfs = mujoco.MjVfs() + vfs["box.obj"] = BOX_OBJ + + spec = mujoco.MjSpec.from_string(XML_WITH_MESH.decode()) + model1 = spec.compile(vfs=vfs) + self.assertEqual(model1.nmesh, 1) + + spec.worldbody.add_body().add_geom(size=[1, 0, 0]) + model2 = spec.compile(vfs=vfs) + self.assertEqual(model2.nmesh, 1) + self.assertEqual(model2.ngeom, 2) + + vfs.close() + + def test_vfs_and_assets_raises(self): + with mujoco.MjVfs() as vfs: + vfs["model.xml"] = SIMPLE_XML + with self.assertRaises(ValueError): + mujoco.MjModel.from_xml_string( + SIMPLE_XML.decode(), assets={"a": b"b"}, vfs=vfs + ) + + def test_backward_compat_assets_dict(self): + model = mujoco.MjModel.from_xml_string( + XML_WITH_MESH.decode(), assets={"box.obj": BOX_OBJ} + ) + self.assertEqual(model.nmesh, 1) + + +class VfsAttachTest(absltest.TestCase): + + def test_attach_shared_vfs(self): + child_xml = textwrap.dedent("""\ + + + + + + + + + + + + """) + parent_xml = textwrap.dedent("""\ + + + + + + + + """) + + with mujoco.MjVfs() as vfs: + vfs["box.obj"] = BOX_OBJ + + parent = mujoco.MjSpec.from_string(parent_xml) + child = mujoco.MjSpec.from_string(child_xml) + parent.attach(child, site="mount") + model = parent.compile(vfs=vfs) + self.assertEqual(model.nmesh, 1) + self.assertEmpty(parent.assets) + + def test_attach_no_asset_dict_needed(self): + child_xml = textwrap.dedent("""\ + + + + + + + + + """) + parent_xml = textwrap.dedent("""\ + + + + + + + + """) + + with mujoco.MjVfs() as vfs: + parent = mujoco.MjSpec.from_string(parent_xml) + child = mujoco.MjSpec.from_string(child_xml) + parent.attach(child, site="mount") + model = parent.compile(vfs=vfs) + self.assertGreater(model.nbody, 1) + + +if __name__ == "__main__": + absltest.main()