Do not copy the spec during attach.

Use a reference count for managing the memory.

PiperOrigin-RevId: 716169486
Change-Id: Id270c4858c17b9250115e9544d5ea143584e2d5f
This commit is contained in:
Alessio Quaglino
2025-01-16 03:52:31 -08:00
committed by Copybara-Service
parent 6436055c6c
commit c2138c3fb0
19 changed files with 505 additions and 173 deletions
+8
View File
@@ -350,6 +350,14 @@ PYBIND11_MODULE(_specs, m) {
mjSpec.def("copy", [](const MjSpec& self) -> MjSpec {
return MjSpec(self);
});
mjSpec.def_property(
"copy_during_attach",
[](MjSpec& self) {
throw pybind11::value_error("copy_during_attach can only be set.");
},
[](MjSpec& self, bool deepcopy) {
return mjs_setDeepCopy(self.ptr, deepcopy);
});
mjSpec.def_property_readonly(
"worldbody",
[](MjSpec& self) -> raw::MjsBody* {
+24 -16
View File
@@ -918,27 +918,31 @@ class SpecsTest(absltest.TestCase):
model = parent.compile()
np.testing.assert_almost_equal(model.body_quat[1], [1, 0, 0, 0])
def test_attach_body_to_site(self):
child = mujoco.MjSpec()
def test_attach_to_site(self):
parent = mujoco.MjSpec()
site = parent.worldbody.add_site(pos=[1, 2, 3], quat=[0, 0, 0, 1])
body = child.worldbody.add_body()
# Attach body to site and compile.
self.assertIsNotNone(site.attach_body(body, prefix='_'))
child1 = mujoco.MjSpec()
body1 = child1.worldbody.add_body()
self.assertIs(body1, site.attach_body(body1, prefix='_'))
body1.pos = [1, 1, 1]
model1 = parent.compile()
self.assertIsNotNone(model1)
self.assertEqual(model1.nbody, 2)
np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3])
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])
# Attach entire spec to site and compile again.
self.assertIsNotNone(site.attach(child, prefix='child-'))
child2 = mujoco.MjSpec()
body2 = child2.worldbody.add_body(name='body')
self.assertIsNotNone(site.attach(child2, prefix='child-'))
body2.pos = [-1, -1, -1]
model2 = parent.compile()
self.assertIsNotNone(model2)
self.assertEqual(model2.nbody, 3)
np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3])
np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 3])
np.testing.assert_array_equal(model2.body_pos[1], [0, 1, 4])
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])
@@ -949,27 +953,31 @@ class SpecsTest(absltest.TestCase):
frame = body.to_frame()
np.testing.assert_array_equal(frame.pos, [1, 2, 3])
def test_attach_spec_to_frame(self):
child = mujoco.MjSpec()
def test_attach_to_frame(self):
parent = mujoco.MjSpec()
frame = parent.worldbody.add_frame(pos=[1, 2, 3], quat=[0, 0, 0, 1])
body = child.worldbody.add_body()
# Attach body to frame and compile.
self.assertIsNotNone(frame.attach_body(body, prefix='_'))
child1 = mujoco.MjSpec()
body1 = child1.worldbody.add_body()
self.assertIs(body1, frame.attach_body(body1, prefix='_'))
body1.pos = [1, 1, 1]
model1 = parent.compile()
self.assertIsNotNone(model1)
self.assertEqual(model1.nbody, 2)
np.testing.assert_array_equal(model1.body_pos[1], [1, 2, 3])
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])
# Attach entire spec to frame and compile again.
self.assertIsNotNone(frame.attach(child, prefix='child-'))
child2 = mujoco.MjSpec()
body2 = child2.worldbody.add_body(name='body')
self.assertIsNotNone(frame.attach(child2, prefix='child-'))
body2.pos = [-1, -1, -1]
model2 = parent.compile()
self.assertIsNotNone(model2)
self.assertEqual(model2.nbody, 3)
np.testing.assert_array_equal(model2.body_pos[1], [1, 2, 3])
np.testing.assert_array_equal(model2.body_pos[2], [1, 2, 3])
np.testing.assert_array_equal(model2.body_pos[1], [0, 1, 4])
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])