Add assets to MjSpec.compile()

PiperOrigin-RevId: 653965318
Change-Id: Ib9ec1c70a7241f7ee547b799a08cc58e7d4bace7

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)
 ----------------------------
 

python/mujoco/specs.cc

@@ -16,6 +16,7 @@
 #include <cstddef>  // IWYU pragma: keep
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <string_view>  // IWYU pragma: keep
 #include <vector>       // 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<uintptr_t>(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<std::string>(item.second);
+          mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
+                          buffer.c_str(), buffer.size());
+        };
+        auto model =
+            mjmodel_from_spec_ptr(reinterpret_cast<uintptr_t>(self.ptr),
+                                  reinterpret_cast<uintptr_t>(&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 {
+  mjSpec.def(
-    std::array<char, 1024> err;
+      "from_file",
-    err[0] = '\0';
+      [](MjSpec& self, std::string& filename,
-    mj_deleteSpec(self.ptr);
+         std::optional<py::dict>& assets) -> void {
-    self.ptr = mj_parseXML(filename.c_str(), 0, err.data(), err.size());
+        mjVFS vfs;
-    if (!self.ptr) {
+        mj_defaultVFS(&vfs);
-      throw FatalError(std::string(err.data()));
+        if (assets.has_value()) {
-    }
+          for (auto item : assets.value()) {
-  });
+            std::string buffer = py::cast<std::string>(item.second);
-  mjSpec.def("from_string", [](MjSpec& self, std::string& xml) -> void {
+            mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
-    std::array<char, 1024> err;
+                            buffer.c_str(), buffer.size());
-    err[0] = '\0';
+          };
-    mj_deleteSpec(self.ptr);
+        }
-    self.ptr = mj_parseXMLString(xml.c_str(), 0, err.data(), err.size());
+        std::array<char, 1024> err;
-    if (!self.ptr) {
+        err[0] = '\0';
-      throw FatalError(std::string(err.data()));
+        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<py::dict>& assets) -> void {
+        mjVFS vfs;
+        mj_defaultVFS(&vfs);
+        if (assets.has_value()) {
+          for (auto item : assets.value()) {
+            std::string buffer = py::cast<std::string>(item.second);
+            mj_addBufferVFS(&vfs, py::cast<std::string>(item.first).c_str(),
+                            buffer.c_str(), buffer.size());
+          };
+        }
+        std::array<char, 1024> 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,

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()

python/mujoco/structs.cc

@@ -423,8 +423,9 @@ MjModelWrapper MjModelWrapper::LoadXML(
   return MjModelWrapper(model);
 }
 
-MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec) {
+MjModelWrapper MjModelWrapper::CompileSpec(raw::MjSpec* spec,
-  auto m = mj_compile(spec, nullptr);
+                                           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<raw::MjSpec*>(addr));
+            reinterpret_cast<raw::MjSpec*>(addr),
+            nullptr);
+      });
+  mjModel.def_static(
+      "_from_spec_ptr", [](uintptr_t addr, uintptr_t vfs) {
+        return MjModelWrapper::CompileSpec(
+            reinterpret_cast<raw::MjSpec*>(addr),
+            reinterpret_cast<mjVFS*>(vfs));
       });
   mjModel.def_static(
       "from_xml_path", &MjModelWrapper::LoadXMLFile,

python/mujoco/structs.h

@@ -489,7 +489,7 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
       const std::optional<
           std::unordered_map<std::string, pybind11::bytes>>& assets);
 
-  static MjWrapper CompileSpec(raw::MjSpec* spec);
+  static MjWrapper CompileSpec(raw::MjSpec* spec, const mjVFS* vfs);
 
   static constexpr char kFromRawPointer[] =
       "__MUJOCO_STRUCTS_MJMODELWRAPPER_LOOKUP";