From 1412a29dace0ff21d2f60379629300ae18dff9df Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Tue, 21 Jan 2025 08:56:25 -0800 Subject: [PATCH] Merge child asset dictionary into parent spec during attach(). PiperOrigin-RevId: 717939728 Change-Id: I54f8ac87062cf19f55a16af37b777701545ed865 --- python/mujoco/specs.cc | 8 ++++++++ python/mujoco/specs_test.py | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 08716d1b..dfad9d93 100644 --- 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() + + " 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 index 6465619b..d4aa345e 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -938,11 +938,13 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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 @@ class SpecsTest(absltest.TestCase): 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()