From 07e7417d26d2192c6f3aea405509f2f40d90c0e6 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Thu, 14 Aug 2025 03:48:13 -0700 Subject: [PATCH] Add prefix to asset paths when appending specs. PiperOrigin-RevId: 794966246 Change-Id: I4825a11cc3234e6f403340dd997742cea0e2183b --- python/mujoco/specs.cc | 82 +++++++++++++++++++++++++++++-------- python/mujoco/specs_test.py | 40 +++++++++--------- 2 files changed, 85 insertions(+), 37 deletions(-) diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 8be8a34c..324a6207 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -192,6 +192,29 @@ static std::string addSuffixBeforeExtension(const std::string& original_path, } } +static std::string addPrefixToFileName(const std::string& original_path, + const std::string& prefix_to_add) { + // Find the position of the last slash + size_t slash_pos = original_path.find_last_of("/\\"); + + // Check if a slash was found + if (slash_pos != std::string::npos) { + std::string new_path = original_path; + new_path.insert(slash_pos + 1, prefix_to_add); + return new_path; + } else { + // No slash found, just prepend + return prefix_to_add + original_path; + } +} + +static std::string addPrefixAndSuffix(const std::string& original_path, + const std::string& prefix_to_add, + const std::string& suffix_to_add) { + std::string prefixed_path = addPrefixToFileName(original_path, prefix_to_add); + return addSuffixBeforeExtension(prefixed_path, suffix_to_add); +} + PYBIND11_MODULE(_specs, m) { auto structs_m = py::module::import("mujoco._structs"); py::function mjmodel_from_raw_ptr = @@ -487,11 +510,12 @@ PYBIND11_MODULE(_specs, m) { throw pybind11::value_error(mjs_getError(self.ptr)); } } - // append the assets adding the suffix to their path - std::string suff(s); + // add prefix and suffix to the assets keys + std::string pre(p); + std::string suf(s); for (const auto& asset : child.assets) { std::string asset_name = - addSuffixBeforeExtension(asset.first.cast(), suff); + addPrefixAndSuffix(asset.first.cast(), pre, suf); if (self.assets.contains(asset_name) && !self.override_assets) { throw pybind11::value_error("Asset " + asset.first.cast() + @@ -500,22 +524,46 @@ PYBIND11_MODULE(_specs, m) { self.assets[py::str(asset_name)] = asset.second; } raw::MjsElement* mesh = mjs_firstElement(child.ptr, mjOBJ_MESH); - while (mesh) { - std::string file = mjs_getString(mjs_asMesh(mesh)->file); - if (!file.empty()) { - std::string mesh_file = addSuffixBeforeExtension(file, suff); - mjs_setString(mjs_asMesh(mesh)->file, mesh_file.c_str()); - } - mesh = mjs_nextElement(child.ptr, mesh); - } raw::MjsElement* tex = mjs_firstElement(child.ptr, mjOBJ_TEXTURE); - while (tex) { - std::string file = mjs_getString(mjs_asTexture(tex)->file); - if (!file.empty()) { - std::string tex_file = addSuffixBeforeExtension(file, suff); - mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str()); + raw::MjsElement* parent_mesh = mjs_firstElement(self.ptr, mjOBJ_MESH); + raw::MjsElement* parent_tex = mjs_firstElement(self.ptr, mjOBJ_TEXTURE); + bool child_has_assets = mesh || tex; + bool parent_has_assets = parent_mesh || parent_tex; + bool child_use_asset_dict = !child.assets.empty(); + bool parent_use_asset_dict = !self.assets.empty(); + if (!child_use_asset_dict && child_has_assets && + parent_use_asset_dict) { + PyErr_WarnEx( + PyExc_Warning, + "Attaching a child without asset dict to a parent with an " + "asset dict might result in missing assets when attaching again.", + 1); + } + if (!parent_use_asset_dict && parent_has_assets && + child_use_asset_dict) { + PyErr_WarnEx( + PyExc_Warning, + "Attaching a child with asset dict to a parent without an " + "asset dict might result in missing assets when attaching again.", + 1); + } + if (child_use_asset_dict) { + while (mesh) { + std::string file = mjs_getString(mjs_asMesh(mesh)->file); + if (!file.empty()) { + std::string mesh_file = addPrefixAndSuffix(file, pre, suf); + mjs_setString(mjs_asMesh(mesh)->file, mesh_file.c_str()); + } + mesh = mjs_nextElement(child.ptr, mesh); + } + while (tex) { + std::string file = mjs_getString(mjs_asTexture(tex)->file); + if (!file.empty()) { + std::string tex_file = addPrefixAndSuffix(file, pre, suf); + mjs_setString(mjs_asTexture(tex)->file, tex_file.c_str()); + } + tex = mjs_nextElement(child.ptr, tex); } - tex = mjs_nextElement(child.ptr, tex); } child.parent = &self; return mjs_asFrame(attached_frame); diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index e1827128..6f08cbfb 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -1047,7 +1047,7 @@ class SpecsTest(absltest.TestCase): def test_attach_to_site(self): parent = mujoco.MjSpec() - parent.assets = {'cube.obj': 'cube_content'} + parent.assets = {'path/cube.obj': 'cube_content'} site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1]) site.name = 'site' @@ -1063,11 +1063,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') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') # Attach entire spec to site and compile again. child2 = mujoco.MjSpec() - child2.assets = {'cube2.obj': 'cube2_content'} + child2.assets = {'path/cube2.obj': 'cube2_content'} body2 = child2.worldbody.add_body(name='body') self.assertIsNotNone(parent.attach(child2, site=site, suffix='-child2')) self.assertIsNotNone(child2.worldbody) @@ -1080,14 +1080,14 @@ 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-child2.obj'], 'cube2_content') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') + self.assertEqual(parent.assets['path/cube2-child2.obj'], 'cube2_content') # Attach another spec to site (referenced by name) and compile again. child3 = mujoco.MjSpec() - child3.assets = {'cube3.obj': 'cube3_content'} + child3.assets = {'path/cube3.obj': 'cube3_content'} body3 = child3.worldbody.add_body(name='body') - self.assertIsNotNone(parent.attach(child3, site='site', suffix='-child3')) + self.assertIsNotNone(parent.attach(child3, site='site', prefix='child3-')) self.assertIsNotNone(child3.worldbody) self.assertEqual(child3.parent, parent) body3.pos = [-2, -2, -2] @@ -1100,9 +1100,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-child2.obj'], 'cube2_content') - self.assertEqual(parent.assets['cube3-child3.obj'], 'cube3_content') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') + self.assertEqual(parent.assets['path/cube2-child2.obj'], 'cube2_content') + self.assertEqual(parent.assets['path/child3-cube3.obj'], 'cube3_content') # Fail to attach to a site that does not exist. child4 = mujoco.MjSpec() @@ -1127,13 +1127,13 @@ class SpecsTest(absltest.TestCase): def test_attach_to_frame(self): parent = mujoco.MjSpec() - parent.assets = {'cube.obj': 'cube_content'} + parent.assets = {'path/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'} + child1.assets = {'path/cube1.obj': 'cube1_content'} body1 = child1.worldbody.add_body() self.assertIs(body1, frame.attach_body(body1, prefix='_')) self.assertIsNotNone(child1.worldbody) @@ -1143,11 +1143,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') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') # Attach entire spec to frame and compile again. child2 = mujoco.MjSpec() - child2.assets = {'cube2.obj': 'cube2_content'} + child2.assets = {'path/cube2.obj': 'cube2_content'} body2 = child2.worldbody.add_body(name='body') body2.set_frame(child2.worldbody.add_frame(pos=[-1, -1, 1])) self.assertIsNotNone(parent.attach(child2, frame=frame, suffix='-child')) @@ -1161,15 +1161,15 @@ class SpecsTest(absltest.TestCase): np.testing.assert_array_equal(model2.body_pos[2], [3, 4, 3]) 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-child.obj'], 'cube2_content') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') + self.assertEqual(parent.assets['path/cube2-child.obj'], 'cube2_content') # Attach another spec to frame (referenced by name) and compile again. child3 = mujoco.MjSpec() - child3.assets = {'cube2.obj': 'new_cube2_content'} + child3.assets = {'path/cube2.obj': 'new_content'} body3 = child3.worldbody.add_body(name='body') body3.set_frame(child3.worldbody.add_frame(pos=[-1, -1, 1])) - self.assertIsNotNone(parent.attach(child3, frame='frame', suffix='-child3')) + self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-')) self.assertIsNotNone(child3.worldbody) self.assertEqual(child3.parent, parent) body3.pos = [-2, -2, -2] @@ -1182,8 +1182,8 @@ 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-child3.obj'], 'new_cube2_content') + self.assertEqual(parent.assets['path/cube.obj'], 'cube_content') + self.assertEqual(parent.assets['path/child3-cube2.obj'], 'new_content') # Fail to attach to a frame that does not exist. child4 = mujoco.MjSpec()