Merge child asset dictionary into parent spec during attach().

PiperOrigin-RevId: 717939728 Change-Id: I54f8ac87062cf19f55a16af37b777701545ed865
google-deepmind · Jan 21, 2025 · 1412a29 · 1412a29
1 parent 4bb6aeb
commit 1412a29
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
@@ -552,6 +552,14 @@ PYBIND11_MODULE(_specs, m) {
         if (!attached_world) {
           throw pybind11::value_error(mjs_getError(self.ptr));
         }
+        for (const auto& asset : child.assets) {
+          if (self.assets.contains(asset.first)) {
+            throw pybind11::value_error("Asset " +
+                                        asset.first.cast<std::string>() +
+                                        " already exists in parent spec.");
+          }
+          self.assets[asset.first] = asset.second;
+        }
         return mjs_bodyToFrame(&attached_world);
       },
       py::arg("child"), py::arg("prefix") = py::none(),

diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
@@ -938,11 +938,13 @@ def test_attach_units(self):
 
   def test_attach_to_site(self):
     parent = mujoco.MjSpec()
+    parent.assets = {'cube.obj': 'cube_content'}
     site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1])
     site.name = 'site'
 
     # Attach body to site and compile.
     child1 = mujoco.MjSpec()
+    child1.assets = {'cube1.obj': 'cube1_content'}
     body1 = child1.worldbody.add_body()
     self.assertIs(body1, site.attach_body(body1, prefix='_'))
     body1.pos = [1, 1, 1]
@@ -951,9 +953,11 @@ def test_attach_to_site(self):
     self.assertEqual(model1.nbody, 2)
     np.testing.assert_array_equal(model1.body_pos[1], [0, 1, 4])
     np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
+    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
 
     # Attach entire spec to site and compile again.
     child2 = mujoco.MjSpec()
+    child2.assets = {'cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child2, site=site, prefix='child2-'))
     body2.pos = [-1, -1, -1]
@@ -964,9 +968,12 @@ def test_attach_to_site(self):
     np.testing.assert_array_equal(model2.body_pos[2], [2, 3, 2])
     np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1])
+    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
+    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
 
     # Attach another spec to site (referenced by name) and compile again.
     child3 = mujoco.MjSpec()
+    child3.assets = {'cube3.obj': 'cube3_content'}
     body3 = child3.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-'))
     body3.pos = [-2, -2, -2]
@@ -979,6 +986,9 @@ def test_attach_to_site(self):
     np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1])
     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')
 
     # Fail to attach to a site that does not exist.
     child4 = mujoco.MjSpec()
@@ -994,11 +1004,13 @@ def test_body_to_frame(self):
 
   def test_attach_to_frame(self):
     parent = mujoco.MjSpec()
+    parent.assets = {'cube.obj': 'cube_content'}
     frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1])
     frame.name = 'frame'
 
     # Attach body to frame and compile.
     child1 = mujoco.MjSpec()
+    child1.assets = {'cube1.obj': 'cube1_content'}
     body1 = child1.worldbody.add_body()
     self.assertIs(body1, frame.attach_body(body1, prefix='_'))
     body1.pos = [1, 1, 1]
@@ -1007,9 +1019,11 @@ def test_attach_to_frame(self):
     self.assertEqual(model1.nbody, 2)
     np.testing.assert_array_equal(model1.body_pos[1], [0, 1, 4])
     np.testing.assert_array_equal(model1.body_quat[1], [0, 0, 0, 1])
+    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
 
     # Attach entire spec to frame and compile again.
     child2 = mujoco.MjSpec()
+    child2.assets = {'cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child2, frame=frame, prefix='child-'))
     body2.pos = [-1, -1, -1]
@@ -1020,9 +1034,12 @@ def test_attach_to_frame(self):
     np.testing.assert_array_equal(model2.body_pos[2], [2, 3, 2])
     np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1])
+    self.assertEqual(parent.assets['cube.obj'], 'cube_content')
+    self.assertEqual(parent.assets['cube2.obj'], 'cube2_content')
 
     # Attach another spec to frame (referenced by name) and compile again.
     child3 = mujoco.MjSpec()
+    child3.assets = {'cube3.obj': 'cube3_content'}
     body3 = child3.worldbody.add_body(name='body')
     self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
     body3.pos = [-2, -2, -2]
@@ -1035,6 +1052,9 @@ def test_attach_to_frame(self):
     np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1])
     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')
 
     # Fail to attach to a frame that does not exist.
     child4 = mujoco.MjSpec()