chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Introduce explicit VFS in Python MjVfs

Browse Source 
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
This commit is contained in:
Sam Haves
2026-05-01 09:33:43 -07:00
committed by Copybara-Service
parent 910b3336ed
commit 723b8b1ea6
10 changed files with 599 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/specs.cc
+102 -7
View File
@@ -18,6 +18,7 @@
#include <cstring>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view> // IWYU pragma: keep
#include <unordered_map>
@@ -33,6 +34,7 @@
#include "specs_wrapper.h" // IWYU pragma: keep
#include "raw.h"
#include "structs.h" // IWYU pragma: keep
#include "vfs.h"
#include <pybind11/cast.h>
#include <pybind11/eigen.h>
#include <pybind11/eigen/matrix.h>
@@ -264,6 +266,49 @@ PYBIND11_MODULE(_specs, m) {
DefineArray<float>(m, "MjFloatVec");
DefineArray<int>(m, "MjIntVec");
// ============================= MJVFS =====================================
py::class_<MjVfs>(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<std::unordered_map<std::string, py::bytes>>& include,
std::optional<py::dict>& assets) -> MjSpec {
std::optional<py::dict>& 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<std::unordered_map<std::string, py::bytes>>& include,
std::optional<py::dict>& assets) -> MjSpec {
std::optional<py::dict>& 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<uintptr_t>(self.Compile()));
});
mjSpec.def("compile",
[mjmodel_from_raw_ptr](MjSpec& self,
std::optional<MjVfs*> vfs) -> py::object {
mjVFS* vfs_ptr = vfs.has_value() ? (*vfs)->get() : nullptr;
return mjmodel_from_raw_ptr(
reinterpret_cast<uintptr_t>(self.Compile(vfs_ptr)));
},
py::arg("vfs") = py::none());
mjSpec.def_property(
"assets",
[](MjSpec& self) -> py::dict {
python/mujoco/specs_wrapper.cc
+31 -29
View File
@@ -15,6 +15,7 @@
#include "specs_wrapper.h"
#include <cstddef> // IWYU pragma: keep
#include <optional>
#include <string>
#include <string_view> // IWYU pragma: keep
#include <vector> // 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<std::string>(asset.first).c_str();
std::string buffer = py::cast<std::string>(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<mjVFS> local_vfs;
if (vfs == nullptr) {
vfs = &local_vfs.emplace();
mj_defaultVFS(vfs);
for (const auto& asset : assets) {
std::string buffer_name = py::cast<std::string>(asset.first).c_str();
std::string buffer = py::cast<std::string>(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));
}
python/mujoco/specs_wrapper.h
+1 -1
View File
@@ -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;
python/mujoco/structs.cc
+8 -10
View File
@@ -41,6 +41,7 @@
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
#include "vfs.h"
namespace mujoco::python::_impl {
@@ -301,24 +302,21 @@ PYBIND11_MODULE(_structs, m) {
// ==================== MJMODEL ==============================================
py::class_<MjModelWrapper> 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<raw::MjModel*>(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.
python/mujoco/structs.h
+9 -3
View File
@@ -44,6 +44,9 @@
namespace py = ::pybind11;
namespace mujoco::python {
class MjVfs;
namespace _impl {
struct VfsAsset {
@@ -501,17 +504,20 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
static MjWrapper LoadXMLFile(
const std::string& filename,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
std::unordered_map<std::string, pybind11::bytes>>& assets,
MjVfs* vfs = nullptr);
static MjWrapper LoadBinaryFile(
const std::string& filename,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
std::unordered_map<std::string, pybind11::bytes>>& assets,
MjVfs* vfs = nullptr);
static MjWrapper LoadXML(
const std::string& xml,
const std::optional<
std::unordered_map<std::string, pybind11::bytes>>& assets);
std::unordered_map<std::string, pybind11::bytes>>& assets,
MjVfs* vfs = nullptr);
static MjWrapper WrapRawModel(raw::MjModel* m);
python/mujoco/structs_wrappers.cc
+84 -24
View File
@@ -19,6 +19,7 @@
#include <cstring>
#include <chrono> // NOLINT(build/c++11)
#include <exception>
#include <functional>
#include <ios>
#include <iostream>
#include <memory>
@@ -38,6 +39,7 @@
#include "raw.h"
#include "serialization.h"
#include "structs.h"
#include "vfs.h"
#include <pybind11/cast.h>
#include <pybind11/detail/common.h>
#include <pybind11/numpy.h>
@@ -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 <typename Callback = void()>
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 <typename Callback>
Cleanup(Callback callback) -> Cleanup<Callback>;
} // namespace
@@ -295,18 +309,29 @@ MjModelWrapper::~MjWrapper() {
template <typename LoadFunc>
static raw::MjModel* LoadModelFileImpl(const std::string& filename,
const std::vector<VfsAsset>& 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<mjVFS> 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<std::unordered_map<std::string, py::bytes>>& assets) {
const std::optional<std::unordered_map<std::string, py::bytes>>& 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<std::unordered_map<std::string, py::bytes>>& assets) {
const std::optional<std::unordered_map<std::string, py::bytes>>& 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<std::unordered_map<std::string, py::bytes>>& assets) {
const std::optional<std::unordered_map<std::string, py::bytes>>& 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> MjModelWrapper::Deserialize(
raw::MjModel* model = LoadModelFileImpl(
"model.mjb",
{{"model.mjb", model_bytes.data(), static_cast<std::size_t>(model_size)}},
nullptr,
InterceptMjErrors(mj_loadModel));
if (!model) {
throw py::value_error("Invalid serialized mjModel.");
python/mujoco/vfs.h
+46
View File
@@ -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 <memory>
#include <mujoco/mujoco.h>
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<mjVFS, VfsDeleter> vfs_;
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_VFS_H_
python/mujoco/vfs_test.py
+251
View File
@@ -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"<mujoco/>"
XML_WITH_MESH = textwrap.dedent("""\
<mujoco>
<asset>
<mesh name="box" file="box.obj"/>
</asset>
<worldbody>
<body>
<geom type="mesh" mesh="box"/>
</body>
</worldbody>
</mujoco>
""").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("""\
<mujoco>
<asset>
<mesh name="child_mesh" file="box.obj"/>
</asset>
<worldbody>
<body name="child">
<geom type="mesh" mesh="child_mesh" size="1 1 1"/>
<site name="attachment_site"/>
</body>
</worldbody>
</mujoco>
""")
parent_xml = textwrap.dedent("""\
<mujoco>
<worldbody>
<body name="parent">
<site name="mount"/>
</body>
</worldbody>
</mujoco>
""")
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("""\
<mujoco>
<worldbody>
<body name="child">
<geom size="1"/>
<site name="child_site"/>
</body>
</worldbody>
</mujoco>
""")
parent_xml = textwrap.dedent("""\
<mujoco>
<worldbody>
<body name="parent">
<site name="mount"/>
</body>
</worldbody>
</mujoco>
""")
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()