Add asset argument to MjSpec constructors.

PiperOrigin-RevId: 720628689
Change-Id: I189aa34973e58b9fb87951e34cb923d94e081a55

python/mujoco/specs.cc

@@ -71,17 +71,25 @@ using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
 
 struct MjSpec {
   MjSpec() : ptr(mj_makeSpec()) {}
-  MjSpec(raw::MjSpec* ptr) : ptr(ptr) {}
+  MjSpec(raw::MjSpec* ptr, const py::dict& assets_ = {}) : ptr(ptr) {
+    for (const auto [key, value] : assets_) {
+      assets[key] = value;
+    }
+  }
 
   // copy constructor and assignment
   MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {
     override_assets = other.override_assets;
-    assets = other.assets;
+    for (const auto [key, value] : other.assets) {
+      assets[key] = value;
+    }
   }
   MjSpec& operator=(const MjSpec& other) {
     override_assets = other.override_assets;
     ptr = mj_copySpec(other.ptr);
-    assets = other.assets;
+    for (const auto [key, value] : other.assets) {
+      assets[key] = value;
+    }
     return *this;
   }
 
@@ -89,14 +97,18 @@ struct MjSpec {
   MjSpec(MjSpec&& other) : ptr(other.ptr) {
     override_assets = other.override_assets;
     other.ptr = nullptr;
-    assets = other.assets;
+    for (const auto [key, value] : other.assets) {
+      assets[key] = value;
+    }
     other.assets.clear();
   }
   MjSpec& operator=(MjSpec&& other) {
     override_assets = other.override_assets;
     ptr = other.ptr;
     other.ptr = nullptr;
-    assets = other.assets;
+    for (const auto [key, value] : other.assets) {
+      assets[key] = value;
+    }
     other.assets.clear();
     return *this;
   }
@@ -263,8 +275,8 @@ PYBIND11_MODULE(_specs, m) {
   mjSpec.def_static(
       "from_file",
       [](std::string& filename,
-         std::optional<std::unordered_map<std::string, py::bytes>>& include)
+         std::optional<std::unordered_map<std::string, py::bytes>>& include,
-          -> MjSpec {
+         std::optional<py::dict>& assets) -> MjSpec {
         const auto files = _impl::ConvertAssetsDict(include);
         raw::MjSpec* spec;
         {
@@ -280,9 +292,13 @@ PYBIND11_MODULE(_specs, m) {
             throw py::value_error(error);
           }
         }
+        if (assets.has_value()) {
+          return MjSpec(spec, assets.value());
+        }
         return MjSpec(spec);
       },
-      py::arg("filename"), py::arg("include") = py::none(), R"mydelimiter(
+      py::arg("filename"), py::arg("include") = py::none(),
+      py::arg("assets") = py::none(), R"mydelimiter(
     Creates a spec from an XML file.
 
     Parameters
@@ -292,13 +308,16 @@ PYBIND11_MODULE(_specs, m) {
     include : dict, optional
         A dictionary of xml files included by the model. The keys are file names
         and the values are file contents.
+    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",
       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<std::unordered_map<std::string, py::bytes>>& include,
-          -> MjSpec {
+         std::optional<py::dict>& assets) -> MjSpec {
         auto files = _impl::ConvertAssetsDict(include);
         raw::MjSpec* spec;
         {
@@ -324,9 +343,13 @@ PYBIND11_MODULE(_specs, m) {
             throw py::value_error(error);
           }
         }
+        if (assets.has_value()) {
+          return MjSpec(spec, assets.value());
+        }
         return MjSpec(spec);
       },
-      py::arg("xml"), py::arg("include") = py::none(), R"mydelimiter(
+      py::arg("xml"), py::arg("include") = py::none(),
+      py::arg("assets") = py::none(), R"mydelimiter(
     Creates a spec from an XML string.
 
     Parameters
@@ -336,6 +359,9 @@ PYBIND11_MODULE(_specs, m) {
     include : dict, optional
         A dictionary of xml files included by the model. The keys are file names
         and the values are file contents.
+    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",
       py::return_value_policy::move);
   mjSpec.def("recompile", [mjmodel_mjdata_from_spec_ptr](

python/mujoco/specs_test.py

@@ -749,6 +749,22 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model.nmeshvert, 8)
     self.assertEqual(spec.assets['cube.obj'], cube)
 
+    xml = """
+    <mujoco model="test">
+      <asset>
+        <mesh name="cube" file="cube.obj"/>
+      </asset>
+      <worldbody>
+        <geom mesh="cube"/>
+      </worldbody>
+    </mujoco>
+    """
+    assets = {'cube.obj': cube}
+    spec = mujoco.MjSpec.from_string(xml, assets=assets)
+    model = spec.compile()
+    self.assertEqual(model.nmeshvert, 8)
+    self.assertEqual(spec.assets['cube.obj'], cube)
+
   def test_include(self):
     included_xml = """
       <mujoco>