From 5bca7876c444aa3e6d1b5648d2413c8cf0df34bd Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 24 Sep 2024 06:56:27 -0700 Subject: [PATCH] Update spec.from_string() and spec.from_file() to match the respective mjModel functions. PiperOrigin-RevId: 678230403 Change-Id: I7f3d8c5fd44d30b5695eee1f8513cc27f4158d94 --- python/mujoco/specs.cc | 113 ++++++++++++++++++++++++------------ python/mujoco/specs_test.py | 19 ++++++ python/mujoco/structs.cc | 41 ------------- python/mujoco/structs.h | 44 ++++++++++++++ 4 files changed, 138 insertions(+), 79 deletions(-) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 07d52fc3..50aecd44 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -19,6 +19,7 @@ #include #include #include // IWYU pragma: keep +#include #include // IWYU pragma: keep #include @@ -73,6 +74,37 @@ struct MjSpec { raw::MjSpec* ptr; }; +template +static raw::MjSpec* LoadSpecFileImpl( + const std::string& filename, + const std::vector<_impl::VfsAsset>& assets, + LoadFunc&& loadfunc) { + mjVFS vfs; + mjVFS* vfs_ptr = nullptr; + if (!assets.empty()) { + mj_defaultVFS(&vfs); + vfs_ptr = &vfs; + for (const auto& asset : assets) { + std::string buffer_name = _impl::StripPath(asset.name); + const int vfs_error = InterceptMjErrors(mj_addBufferVFS)( + vfs_ptr, 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); + } else { + throw py::value_error("Asset failed to load: " + buffer_name); + } + } + } + } + + raw::MjSpec* spec = loadfunc(filename.c_str(), vfs_ptr); + mj_deleteVFS(vfs_ptr); + return spec; +} + PYBIND11_MODULE(_specs, m) { auto structs_m = py::module::import("mujoco._structs"); py::function mjmodel_from_spec_ptr = @@ -196,27 +228,25 @@ PYBIND11_MODULE(_specs, m) { mjSpec.def( "from_file", [](MjSpec& self, std::string& filename, - std::optional& assets) -> void { - mjVFS vfs; - mj_defaultVFS(&vfs); - if (assets.has_value()) { - for (auto item : assets.value()) { - std::string buffer = py::cast(item.second); - mj_addBufferVFS(&vfs, py::cast(item.first).c_str(), - buffer.c_str(), buffer.size()); - }; - } - std::array err; - err[0] = '\0'; - mj_deleteSpec(self.ptr); - self.ptr = mj_parseXML(filename.c_str(), &vfs, err.data(), err.size()); - mj_deleteVFS(&vfs); - if (!self.ptr) { - throw FatalError(std::string(err.data())); + std::optional>& assets) + -> void { + const auto converted_assets = _impl::ConvertAssetsDict(assets); + { + py::gil_scoped_release no_gil; + char error[1024]; + mj_deleteSpec(self.ptr); + self.ptr = LoadSpecFileImpl( + filename, converted_assets, + [&error](const char* filename, const mjVFS* vfs) { + return InterceptMjErrors(mj_parseXML)( + filename, vfs, error, sizeof(error)); + }); + if (!self.ptr) { + throw py::value_error(error); + } } }, - py::arg("filename"), - py::arg("assets") = py::none(), R"mydelimiter( + py::arg("filename"), py::arg("assets") = py::none(), R"mydelimiter( Creates a spec from an XML file. Parameters @@ -230,27 +260,34 @@ PYBIND11_MODULE(_specs, m) { mjSpec.def( "from_string", [](MjSpec& self, std::string& xml, - std::optional& assets) -> void { - mjVFS vfs; - mj_defaultVFS(&vfs); - if (assets.has_value()) { - for (auto item : assets.value()) { - std::string buffer = py::cast(item.second); - mj_addBufferVFS(&vfs, py::cast(item.first).c_str(), - buffer.c_str(), buffer.size()); - }; - } - std::array err; - err[0] = '\0'; - mj_deleteSpec(self.ptr); - self.ptr = mj_parseXMLString(xml.c_str(), &vfs, err.data(), err.size()); - mj_deleteVFS(&vfs); - if (!self.ptr) { - throw FatalError(std::string(err.data())); + std::optional>& assets) + -> void { + auto converted_assets = _impl::ConvertAssetsDict(assets); + { + 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"; + } + } + converted_assets.emplace_back( + model_filename.c_str(), xml.c_str(), xml.length()); + char error[1024]; + mj_deleteSpec(self.ptr); + self.ptr = LoadSpecFileImpl( + model_filename, converted_assets, + [&error](const char* filename, const mjVFS* vfs) { + return InterceptMjErrors(mj_parseXML)( + filename, vfs, error, sizeof(error)); + }); + if (!self.ptr) { + throw py::value_error(error); + } } }, - py::arg("xml"), - py::arg("assets") = py::none(), R"mydelimiter( + py::arg("xml"), py::arg("assets") = py::none(), R"mydelimiter( Creates a spec from an XML string. Parameters diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 8fed347e..4dcaa5fc 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -611,6 +611,25 @@ class SpecsTest(absltest.TestCase): model = spec.compile({'cube.obj': cube}) self.assertEqual(model.nmeshvert, 8) + def test_include(self): + included_xml = """ + + + + + + + + """ + spec = mujoco.MjSpec() + spec.from_string(textwrap.dedent(""" + + + + """), {'included.xml': included_xml.encode('utf-8')}) + self.assertEqual(spec.worldbody.first_body().first_geom().type, + mujoco.mjtGeom.mjGEOM_BOX) + def test_delete(self): filename = '../../test/testdata/model.xml' diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index fd060fb1..93b9f014 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -87,21 +87,6 @@ inline std::size_t NConMax(const mjData* d) { return d->narena / sizeof(mjContact); } -// strip path prefix from filename and make lowercase -std::string StripPath(const char* name) { - std::string filename(name); - size_t start = filename.find_last_of("/\\"); - - // get name without path - if (start != std::string::npos) { - filename = filename.substr(start + 1, filename.size() - start - 1); - } - - // make lowercase - std::transform(filename.begin(), filename.end(), filename.begin(), - [](unsigned char c) { return std::tolower(c); }); - return filename; -} } // namespace // ==================== MJOPTION =============================================== @@ -318,16 +303,6 @@ MjModelWrapper::~MjWrapper() { } } -namespace { -struct VfsAsset { - VfsAsset(const char* name, const void* content, std::size_t content_size) - : name(name), content(content), content_size(content_size) {} - const char* name; - const void* content; - std::size_t content_size; -}; -} - // Helper function for both LoadXMLFile and LoadBinaryFile. // Creates a temporary MJB from the assets dictionary if one is supplied. template @@ -365,22 +340,6 @@ static raw::MjModel* LoadModelFileImpl( return model; } -// Converts a dict with py::bytes value to a vector of standard C++ types. -// This allows us to release the GIL early. Note that the vector consists only -// of pointers to existing data so no substantial data copies are being made. -static std::vector -ConvertAssetsDict( - const std::optional>& assets) { - std::vector out; - if (assets.has_value()) { - for (const auto& [name, content] : *assets) { - out.emplace_back(name.c_str(), PYBIND11_BYTES_AS_STRING(content.ptr()), - py::len(content)); - } - } - return out; -} - MjModelWrapper MjModelWrapper::LoadXMLFile( const std::string& filename, const std::optional>& assets) { diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index dcbfed9d..da8c283a 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -15,7 +15,9 @@ #ifndef MUJOCO_PYTHON_STRUCTS_H_ #define MUJOCO_PYTHON_STRUCTS_H_ +#include #include +#include #include #include #include @@ -39,9 +41,51 @@ #include #include +namespace py = ::pybind11; + namespace mujoco::python { namespace _impl { +struct VfsAsset { + VfsAsset(const char* name, const void* content, std::size_t content_size) + : name(name), content(content), content_size(content_size) {} + const char* name; + const void* content; + std::size_t content_size; +}; + +// strip path prefix from filename and make lowercase +inline std::string StripPath(const char* name) { + std::string filename(name); + size_t start = filename.find_last_of("/\\"); + + // get name without path + if (start != std::string::npos) { + filename = filename.substr(start + 1, filename.size() - start - 1); + } + + // make lowercase + std::transform(filename.begin(), filename.end(), filename.begin(), + [](unsigned char c) { return std::tolower(c); }); + return filename; +} + + +// Converts a dict with py::bytes value to a vector of standard C++ types. +// This allows us to release the GIL early. Note that the vector consists only +// of pointers to existing data so no substantial data copies are being made. +inline std::vector ConvertAssetsDict( + const std::optional>& assets) { + std::vector out; + if (assets.has_value()) { + for (const auto& [name, content] : *assets) { + out.emplace_back(name.c_str(), PYBIND11_BYTES_AS_STRING(content.ptr()), + py::len(content)); + } + } + return out; +} + template class WrapperBase { public: