Add mjs_attachToSite for attaching a body to a site.

PiperOrigin-RevId: 684432183
Change-Id: I7100bebfb6521f9dbfbb07f31490b2b9298bf18c
This commit is contained in:
Alessio Quaglino
2024-10-10 07:11:16 -07:00
committed by Copybara-Service
parent 694bdebc19
commit cdb167f275
10 changed files with 180 additions and 52 deletions
+9 -9
View File
@@ -3813,6 +3813,15 @@ mjs_attachFrame
Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
.. _mjs_attachToSite:
mjs_attachToSite
~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_attachToSite
Attach child body to a parent site, return the attached body if success or NULL otherwise.
.. _mjs_detachBody:
mjs_detachBody
@@ -4126,15 +4135,6 @@ mjs_getSpec
Get spec from body.
.. _mjs_getSpecFromFrame:
mjs_getSpecFromFrame
~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getSpecFromFrame
Get spec from frame.
.. _mjs_findBody:
mjs_findBody
+3 -2
View File
@@ -3564,6 +3564,8 @@ mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
const char* prefix, const char* suffix);
mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix);
int mjs_detachBody(mjSpec* s, mjsBody* b);
mjsBody* mjs_addBody(mjsBody* body, mjsDefault* def);
mjsSite* mjs_addSite(mjsBody* body, mjsDefault* def);
@@ -3596,8 +3598,7 @@ mjsHField* mjs_addHField(mjSpec* s);
mjsSkin* mjs_addSkin(mjSpec* s);
mjsTexture* mjs_addTexture(mjSpec* s);
mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
mjSpec* mjs_getSpec(mjsBody* body);
mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
mjSpec* mjs_getSpec(mjsElement* element);
mjsBody* mjs_findBody(mjSpec* s, const char* name);
mjsElement* mjs_findElement(mjSpec* s, mjtObj type, const char* name);
mjsBody* mjs_findChild(mjsBody* body, const char* name);
+5 -4
View File
@@ -1415,6 +1415,10 @@ MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
// Attach child body to a parent site, return the attached body if success or NULL otherwise.
MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix);
// Detach body from mjSpec, remove all references and delete the body, return 0 on success.
MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b);
@@ -1524,10 +1528,7 @@ MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
//---------------------------------- Find and get utilities ----------------------------------------
// Get spec from body.
MJAPI mjSpec* mjs_getSpec(mjsBody* body);
// Get spec from frame.
MJAPI mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
MJAPI mjSpec* mjs_getSpec(mjsElement* element);
// Find body in spec by name.
MJAPI mjsBody* mjs_findBody(mjSpec* s, const char* name);
+36 -18
View File
@@ -9030,6 +9030,40 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Attach child frame to a parent body, return the attached frame if success or NULL otherwise.', # pylint: disable=line-too-long
)),
('mjs_attachToSite',
FunctionDecl(
name='mjs_attachToSite',
return_type=PointerType(
inner_type=ValueType(name='mjsBody'),
),
parameters=(
FunctionParameterDecl(
name='parent',
type=PointerType(
inner_type=ValueType(name='mjsSite'),
),
),
FunctionParameterDecl(
name='child',
type=PointerType(
inner_type=ValueType(name='mjsBody', is_const=True),
),
),
FunctionParameterDecl(
name='prefix',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='suffix',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
),
doc='Attach child body to a parent site, return the attached body if success or NULL otherwise.', # pylint: disable=line-too-long
)),
('mjs_detachBody',
FunctionDecl(
name='mjs_detachBody',
@@ -9674,30 +9708,14 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
parameters=(
FunctionParameterDecl(
name='body',
name='element',
type=PointerType(
inner_type=ValueType(name='mjsBody'),
inner_type=ValueType(name='mjsElement'),
),
),
),
doc='Get spec from body.',
)),
('mjs_getSpecFromFrame',
FunctionDecl(
name='mjs_getSpecFromFrame',
return_type=PointerType(
inner_type=ValueType(name='mjSpec'),
),
parameters=(
FunctionParameterDecl(
name='frame',
type=PointerType(
inner_type=ValueType(name='mjsFrame'),
),
),
),
doc='Get spec from frame.',
)),
('mjs_findBody',
FunctionDecl(
name='mjs_findBody',
+19 -4
View File
@@ -426,7 +426,7 @@ PYBIND11_MODULE(_specs, m) {
[](raw::MjsBody& self, mjtObj objtype) -> py::list {
py::list list;
raw::MjsElement* el = mjs_firstChild(&self, objtype, true);
std::string error = mjs_getError(mjs_getSpec(&self));
std::string error = mjs_getError(mjs_getSpec(self.element));
if (!el && !error.empty()) {
throw pybind11::value_error(error);
}
@@ -557,7 +557,9 @@ PYBIND11_MODULE(_specs, m) {
py::return_value_policy::reference_internal);
mjsBody.def(
"spec",
[](raw::MjsBody& self) -> raw::MjSpec* { return mjs_getSpec(&self); },
[](raw::MjsBody& self) -> raw::MjSpec* {
return mjs_getSpec(self.element);
},
py::return_value_policy::reference_internal);
mjsBody.def(
"attach_frame",
@@ -566,7 +568,7 @@ PYBIND11_MODULE(_specs, m) {
auto new_frame =
mjs_attachFrame(&self, &frame, prefix.c_str(), suffix.c_str());
if (!new_frame) {
throw pybind11::value_error(mjs_getError(mjs_getSpec(&self)));
throw pybind11::value_error(mjs_getError(mjs_getSpec(self.element)));
}
return new_frame;
},
@@ -587,7 +589,7 @@ PYBIND11_MODULE(_specs, m) {
mjs_attachBody(&self, &body, prefix.c_str(), suffix.c_str());
if (!new_body) {
throw pybind11::value_error(
mjs_getError(mjs_getSpecFromFrame(&self)));
mjs_getError(mjs_getSpec(self.element)));
}
return new_body;
},
@@ -643,6 +645,19 @@ PYBIND11_MODULE(_specs, m) {
return mjs_getDefault(self.element);
},
py::return_value_policy::reference_internal);
mjsSite.def(
"attach",
[](raw::MjsSite& self, raw::MjsBody& body, std::string& prefix,
std::string& suffix) -> raw::MjsBody* {
auto new_body =
mjs_attachToSite(&self, &body, prefix.c_str(), suffix.c_str());
if (!new_body) {
throw pybind11::value_error(
mjs_getError(mjs_getSpec(self.element)));
}
return new_body;
},
py::return_value_policy::reference_internal);
// ============================= MJSCAMERA ===================================
mjsCamera.def_property_readonly("id", [](raw::MjsCamera& self) -> int {
+9
View File
@@ -857,6 +857,15 @@ class SpecsTest(absltest.TestCase):
' attribute',
)
def test_attach_body_to_site(self):
child = mujoco.MjSpec()
parent = mujoco.MjSpec()
site = parent.worldbody.add_site(pos=[1, 2, 3])
body = child.worldbody.add_body()
self.assertIsNotNone(site.attach(body, '', ''))
model = parent.compile()
np.testing.assert_array_equal(model.body_pos[1], [1, 2, 3])
if __name__ == '__main__':
absltest.main()
+26 -9
View File
@@ -29,6 +29,7 @@
#include "user/user_model.h"
#include "user/user_objects.h"
#include "user/user_cache.h"
#include "user/user_util.h"
namespace {
@@ -160,6 +161,29 @@ mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
// attach child body to a parent site
mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix) {
if (!parent) {
mju_error("parent site is null");
return nullptr;
}
mjCSite* site = static_cast<mjCSite*>(parent->element);
mjCBody* body = site->Body();
mjCFrame* frame = body->AddFrame(site->frame);
frame->spec.pos[0] = site->spec.pos[0];
frame->spec.pos[1] = site->spec.pos[1];
frame->spec.pos[2] = site->spec.pos[2];
frame->spec.quat[0] = site->spec.quat[0];
frame->spec.quat[1] = site->spec.quat[1];
frame->spec.quat[2] = site->spec.quat[2];
frame->spec.quat[3] = site->spec.quat[3];
frame->SetParent(body);
return mjs_attachBody(&frame->spec, child, prefix, suffix);
}
// get error message from model
const char* mjs_getError(mjSpec* s) {
mjCModel* modelC = static_cast<mjCModel*>(s->element);
@@ -521,15 +545,8 @@ mjsDefault* mjs_addDefault(mjSpec* s, const char* classname, const mjsDefault* p
// get spec from body
mjSpec* mjs_getSpec(mjsBody* body) {
return &(static_cast<mjCBody*>(body->element)->model->spec);
}
// get spec from frame
mjSpec* mjs_getSpecFromFrame(mjsFrame* frame) {
return &(static_cast<mjCFrame*>(frame->element)->model->spec);
mjSpec* mjs_getSpec(mjsElement* element) {
return &(static_cast<mjCBase*>(element)->model->spec);
}
+7 -6
View File
@@ -71,11 +71,15 @@ MJAPI int mjs_activatePlugin(mjSpec* s, const char* name);
// Attach child body to a parent frame, return the attached body if success or NULL otherwise.
MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
const char* prefix, const char* suffix);
const char* prefix, const char* suffix);
// Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
const char* prefix, const char* suffix);
// Attach child body to a parent site, return the attached body if success or NULL otherwise.
MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix);
// Detach body from mjSpec, remove all references and delete the body, return 0 on success.
MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b);
@@ -186,10 +190,7 @@ MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def);
//---------------------------------- Find/get utilities --------------------------------------------
// Get spec from body.
MJAPI mjSpec* mjs_getSpec(mjsBody* body);
// Get spec from frame.
MJAPI mjSpec* mjs_getSpecFromFrame(mjsFrame* frame);
MJAPI mjSpec* mjs_getSpec(mjsElement* element);
// Find spec (model asset) by name.
MJAPI mjSpec* mjs_findSpec(mjSpec* spec, const char* name);
+3
View File
@@ -587,6 +587,9 @@ class mjCSite : public mjCSite_, private mjsSite {
mjsSite spec; // variables set by user
// site's body
mjCBody* Body() const { return body; }
// use strings from mjCBase rather than mjStrings from mjsSite
using mjCBase::name;
using mjCBase::info;
+63
View File
@@ -1079,6 +1079,69 @@ TEST_F(MujocoTest, DetachBody) {
TestDetachBody(/*compile=*/true);
}
TEST_F(MujocoTest, AttachToSite) {
std::array<char, 1000> er;
mjtNum tol = 0;
std::string field = "";
static constexpr char xml_parent[] = R"(
<mujoco>
<worldbody>
<site name="site" pos="1 0 0" quat="0 1 0 0"/>
</worldbody>
</mujoco>)";
static constexpr char xml_child[] = R"(
<mujoco>
<worldbody>
<body name="sphere">
<joint type="slide"/>
<geom size=".1"/>
</body>
</worldbody>
</mujoco>)";
static constexpr char xml_result[] = R"(
<mujoco>
<worldbody>
<site name="site" pos="1 0 0" quat="0 1 0 0"/>
<frame pos="1 0 0" quat="0 1 0 0">
<body name="attached-sphere-1">
<joint type="slide"/>
<geom size=".1"/>
</body>
</frame>
</worldbody>
</mujoco>)";
mjSpec* parent = mj_parseXMLString(xml_parent, 0, er.data(), er.size());
EXPECT_THAT(parent, NotNull()) << er.data();
mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size());
EXPECT_THAT(child, NotNull()) << er.data();
mjsBody* world = mjs_findBody(parent, "world");
EXPECT_THAT(world, NotNull());
mjsSite* site = mjs_asSite(mjs_firstChild(world, mjOBJ_SITE, 0));
EXPECT_THAT(site, NotNull());
mjsBody* body = mjs_findBody(child, "sphere");
EXPECT_THAT(body, NotNull());
mjsBody* attached = mjs_attachToSite(site, body, "attached-", "-1");
EXPECT_THAT(attached, NotNull());
mjModel* model = mj_compile(parent, 0);
EXPECT_THAT(model, NotNull());
mjModel* expected = LoadModelFromString(xml_result, er.data(), er.size());
EXPECT_THAT(expected, NotNull()) << er.data();
EXPECT_LE(CompareModel(model, expected, field), tol)
<< "Expected and attached models are different!\n"
<< "Different field: " << field << '\n';
mj_deleteSpec(parent);
mj_deleteSpec(child);
mj_deleteModel(model);
mj_deleteModel(expected);
}
TEST_F(MujocoTest, PreserveState) {
std::array<char, 1000> er;
std::string field = "";