From 70ac76bb4bb8f9c89d134c9049c6eb9dadbbe8d5 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 19 Jul 2024 04:33:35 -0700 Subject: [PATCH] Add assets to MjSpec.compile() PiperOrigin-RevId: 653965318 Change-Id: Ib9ec1c70a7241f7ee547b799a08cc58e7d4bace7 --- doc/changelog.rst | 5 ++ python/mujoco/specs.cc | 111 ++++++++++++++++++++++++++++++------ python/mujoco/specs_test.py | 20 +++++++ python/mujoco/structs.cc | 14 ++++- python/mujoco/structs.h | 2 +- 5 files changed, 130 insertions(+), 22 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 6173007f..88918ff0 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,11 @@ MJX 1. Added more fields to ``mjx.Model`` and ``mjx.Data`` for further compatibility with the corresponding MuJoCo structs. +Python bindings +^^^^^^^^^^^^^^^ +2. Added support for asset dictionary argument in ``mujoco.spec.from_file``, ``mujoco.spec.from_string`` and + ``mujoco.spec.compile``. + Version 3.2.0 (Jul 15, 2024) ---------------------------- diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 872761c7..70bbc18f 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -16,6 +16,7 @@ #include // IWYU pragma: keep #include #include +#include #include #include // IWYU pragma: keep #include // IWYU pragma: keep @@ -159,6 +160,30 @@ PYBIND11_MODULE(_specs, m) { mjSpec.def("compile", [mjmodel_from_spec_ptr](MjSpec& self) { return mjmodel_from_spec_ptr(reinterpret_cast(self.ptr)); }); + mjSpec.def( + "compile", + [mjmodel_from_spec_ptr](MjSpec& self, py::dict& assets) -> py::object { + mjVFS vfs; + mj_defaultVFS(&vfs); + for (auto item : assets) { + std::string buffer = py::cast(item.second); + mj_addBufferVFS(&vfs, py::cast(item.first).c_str(), + buffer.c_str(), buffer.size()); + }; + auto model = + mjmodel_from_spec_ptr(reinterpret_cast(self.ptr), + reinterpret_cast(&vfs)); + mj_deleteVFS(&vfs); + return model; + }, R"mydelimiter( + Compiles the spec and returns the compiled model. + + Parameters + ---------- + assets : dict, optional + A dictionary of assets to be used by the spec. The keys are asset names + and the values are asset contents. + )mydelimiter"); mjSpec.def( "copy_back", [](MjSpec& self, raw::MjModel& model) { @@ -177,24 +202,74 @@ PYBIND11_MODULE(_specs, m) { } return std::string(buf.get()); }); - mjSpec.def("from_file", [](MjSpec& self, std::string& filename) -> void { - std::array err; - err[0] = '\0'; - mj_deleteSpec(self.ptr); - self.ptr = mj_parseXML(filename.c_str(), 0, err.data(), err.size()); - if (!self.ptr) { - throw FatalError(std::string(err.data())); - } - }); - mjSpec.def("from_string", [](MjSpec& self, std::string& xml) -> void { - std::array err; - err[0] = '\0'; - mj_deleteSpec(self.ptr); - self.ptr = mj_parseXMLString(xml.c_str(), 0, err.data(), err.size()); - if (!self.ptr) { - throw FatalError(std::string(err.data())); - } - }); + 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())); + } + }, + py::arg("filename"), + py::arg("assets") = py::none(), R"mydelimiter( + Creates a spec from an XML file. + + Parameters + ---------- + filename : str + Path to the XML file. + assets : dict, optional + A dictionary of assets to be used by the spec. The keys are asset names + and the values are asset contents. + )mydelimiter"); + 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())); + } + }, + py::arg("xml"), + py::arg("assets") = py::none(), R"mydelimiter( + Creates a spec from an XML string. + + Parameters + ---------- + xml : str + XML string. + assets : dict, optional + A dictionary of assets to be used by the spec. The keys are asset names + and the values are asset contents. + )mydelimiter"); mjSpec.def( "add_default", [](MjSpec* spec, std::string& classname, diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index 0f36653c..0459b013 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -249,6 +249,26 @@ class SpecsTest(absltest.TestCase): geom = spec.worldbody.next_geom(geom) i += 1 + def test_assets(self): + cube = """ + 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""" + spec = mujoco.MjSpec() + mesh = spec.add_mesh() + mesh.name = 'cube' + mesh.file = 'cube.obj' + geom = spec.worldbody.add_geom() + geom.type = mujoco.mjtGeom.mjGEOM_MESH + geom.meshname = 'cube' + model = spec.compile({'cube.obj': cube}) + self.assertEqual(model.nmeshvert, 8) + if __name__ == '__main__': absltest.main() diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index c70e28cc..fd8443d9 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -423,8 +423,9 @@ MjModelWrapper MjModelWrapper::LoadXML( return MjModelWrapper(model); } -MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec) { - auto m = mj_compile(spec, nullptr); +MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec, + const mjVFS* vfs) { + auto m = mj_compile(spec, vfs); if (!m || mjs_isWarning(spec)) { throw py::value_error(mjs_getError(spec)); } @@ -1572,7 +1573,14 @@ R"(Loads an MjModel from an XML string and an optional assets dictionary.)")); mjModel.def_static( "_from_spec_ptr", [](uintptr_t addr) { return MjModelWrapper::CompileSpec( - reinterpret_cast(addr)); + reinterpret_cast(addr), + nullptr); + }); + mjModel.def_static( + "_from_spec_ptr", [](uintptr_t addr, uintptr_t vfs) { + return MjModelWrapper::CompileSpec( + reinterpret_cast(addr), + reinterpret_cast(vfs)); }); mjModel.def_static( "from_xml_path", &MjModelWrapper::LoadXMLFile, diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h index 70e399ce..dcbfed9d 100644 --- a/python/mujoco/structs.h +++ b/python/mujoco/structs.h @@ -489,7 +489,7 @@ class MjWrapper : public WrapperBase { const std::optional< std::unordered_map>& assets); - static MjWrapper CompileSpec(raw::MjSpec* spec); + static MjWrapper CompileSpec(raw::MjSpec* spec, const mjVFS* vfs); static constexpr char kFromRawPointer[] = "__MUJOCO_STRUCTS_MJMODELWRAPPER_LOOKUP";