Add the override_assets attribute to MjSpec.

When specified, child asset with repeated names will override parent's, rather than throw an error. Default behavior is to override.

PiperOrigin-RevId: 717979621
Change-Id: I18022dc266619b4ba1a8636d1c60f85b0eb5f862

python/mujoco/specs.cc

@@ -81,9 +81,11 @@ struct MjSpec {
 
   // copy constructor and assignment
   MjSpec(const MjSpec& other) : ptr(mj_copySpec(other.ptr)) {
+    override_assets = other.override_assets;
     assets = other.assets;
   }
   MjSpec& operator=(const MjSpec& other) {
+    override_assets = other.override_assets;
     ptr = mj_copySpec(other.ptr);
     assets = other.assets;
     return *this;
@@ -91,11 +93,13 @@ struct MjSpec {
 
   // move constructor and move assignment
   MjSpec(MjSpec&& other) : ptr(other.ptr) {
+    override_assets = other.override_assets;
     other.ptr = nullptr;
     assets = other.assets;
     other.assets.clear();
   }
   MjSpec& operator=(MjSpec&& other) {
+    override_assets = other.override_assets;
     ptr = other.ptr;
     other.ptr = nullptr;
     assets = other.assets;
@@ -108,6 +112,7 @@ struct MjSpec {
   }
   raw::MjSpec* ptr;
   py::dict assets;
+  bool override_assets = true;
 };
 
 template <typename LoadFunc>
@@ -467,6 +472,14 @@ PYBIND11_MODULE(_specs, m) {
           self.assets[item.first] = item.second;
         };
   }, py::return_value_policy::reference_internal);
+  mjSpec.def_property(
+      "override_assets",
+      [](MjSpec& self) -> bool {
+        return self.override_assets;
+      },
+      [](MjSpec& self, bool override_assets) {
+        self.override_assets = override_assets;
+      });
   mjSpec.def("to_xml", [](MjSpec& self) -> std::string {
     int size = mj_saveXMLString(self.ptr, nullptr, 0, nullptr, 0);
     std::unique_ptr<char[]> buf(new char[size + 1]);
@@ -553,7 +566,7 @@ PYBIND11_MODULE(_specs, m) {
           throw pybind11::value_error(mjs_getError(self.ptr));
         }
         for (const auto& asset : child.assets) {
-          if (self.assets.contains(asset.first)) {
+          if (self.assets.contains(asset.first) && !self.override_assets) {
             throw pybind11::value_error("Asset " +
                                         asset.first.cast<std::string>() +
                                         " already exists in parent spec.");

python/mujoco/specs_test.py

@@ -1039,7 +1039,7 @@ class SpecsTest(absltest.TestCase):
 
     # Attach another spec to frame (referenced by name) and compile again.
     child3 = mujoco.MjSpec()
-    child3.assets = {'cube3.obj': 'cube3_content'}
+    child3.assets = {'cube2.obj': 'new_cube2_content'}
     body3 = child3.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
     body3.pos = [-2, -2, -2]
@@ -1053,8 +1053,7 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1])
     np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
-    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
-    self.assertEqual(parent.assets['cube3.obj'], 'cube3_content')
+    self.assertEqual(parent.assets['cube2.obj'], 'new_cube2_content')
 
     # Fail to attach to a frame that does not exist.
     child4 = mujoco.MjSpec()