chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow to attach an entire mjSpec in mjs_attach.

Browse Source 
PiperOrigin-RevId: 742737713
Change-Id: Ie172568b8c35232e1f13d75aebee20346e76e95b
This commit is contained in:
Alessio Quaglino
2025-04-01 09:50:42 -07:00
committed by Copybara-Service
parent e5912c3ce4
commit b839fe79e1
4 changed files with 63 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/XMLreference.rst
+3 -2
View File
@@ -3716,8 +3716,9 @@ all attachments will appear in the saved XML file.
.. _body-attach-body:
:at:`body`: :at-val:`string, required`
Name of the body in the sub-model to attach here. The body and its subtree will be attached.
:at:`body`: :at-val:`string, optional`
Name of the body in the sub-model to attach here. The body and its subtree will be attached. If this attribute is not
specified, the contents of the world body will be attached in a new :ref:`frame<body-frame>`.
.. _body-attach-prefix:
python/mujoco/specs.cc
+2 -31
View File
@@ -269,16 +269,6 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
return list; // list of pointers, so they can be copied
}
void SetFrame(raw::MjsBody* body, mjtObj objtype, raw::MjsFrame* frame) {
mjsElement* el = mjs_firstChild(body, objtype, 0);
while (el) {
if (frame->element != el && mjs_getFrame(el) == nullptr) {
mjs_setFrame(el, frame);
}
el = mjs_nextChild(body, el, 0);
}
}
PYBIND11_MODULE(_specs, m) {
auto structs_m = py::module::import("mujoco._structs");
py::function mjmodel_from_raw_ptr =
@@ -526,18 +516,6 @@ PYBIND11_MODULE(_specs, m) {
throw pybind11::value_error(
"Only one of frame or site can be specified.");
}
auto worldbody = mjs_findBody(child.ptr, "world");
if (!worldbody) {
throw pybind11::value_error("Child does not have a world body.");
}
auto worldframe = mjs_addFrame(worldbody, nullptr);
SetFrame(worldbody, mjOBJ_BODY, worldframe);
SetFrame(worldbody, mjOBJ_SITE, worldframe);
SetFrame(worldbody, mjOBJ_FRAME, worldframe);
SetFrame(worldbody, mjOBJ_JOINT, worldframe);
SetFrame(worldbody, mjOBJ_GEOM, worldframe);
SetFrame(worldbody, mjOBJ_LIGHT, worldframe);
SetFrame(worldbody, mjOBJ_CAMERA, worldframe);
const char* p = prefix.has_value() ? prefix.value().c_str() : "";
const char* s = suffix.has_value() ? suffix.value().c_str() : "";
raw::MjsElement* attached_frame = nullptr;
@@ -556,18 +534,11 @@ PYBIND11_MODULE(_specs, m) {
throw pybind11::value_error(
"Frame spec does not match parent spec.");
}
raw::MjsBody* parent_body = mjs_getParent(frame_ptr->element);
if (!parent_body) {
throw pybind11::value_error("Frame does not have a parent body.");
}
attached_frame =
mjs_attach(parent_body->element, worldframe->element, p, s);
mjs_attach(frame_ptr->element, child.ptr->element, p, s);
if (!attached_frame) {
throw pybind11::value_error(mjs_getError(self.ptr));
}
if (mjs_setFrame(attached_frame, frame_ptr) != 0) {
throw pybind11::value_error(mjs_getError(self.ptr));
}
}
if (site.has_value()) {
raw::MjsSite* site_ptr = nullptr;
@@ -585,7 +556,7 @@ PYBIND11_MODULE(_specs, m) {
"Site spec does not match parent spec.");
}
attached_frame =
mjs_attach(site_ptr->element, worldframe->element, p, s);
mjs_attach(site_ptr->element, child.ptr->element, p, s);
if (!attached_frame) {
throw pybind11::value_error(mjs_getError(self.ptr));
}
src/user/user_api.cc
+44 -1
View File
@@ -119,6 +119,18 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) {
}
// set frame for all elements of a body
static void SetFrame(mjsBody* body, mjtObj objtype, mjsFrame* frame) {
mjsElement* el = mjs_firstChild(body, objtype, 0);
while (el) {
if (frame->element != el && mjs_getFrame(el) == nullptr) {
mjs_setFrame(el, frame);
}
el = mjs_nextChild(body, el, 0);
}
}
// attach body to a frame of the parent
static mjsElement* attachBody(mjCFrame* parent, const mjCBody* child,
@@ -207,13 +219,44 @@ mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
return nullptr;
}
mjCModel* model = static_cast<mjCModel*>(mjs_getSpec(parent)->element);
if (child->elemtype == mjOBJ_MODEL) {
mjCModel* child_model = static_cast<mjCModel*>((mjsElement*)child);
mjsBody* worldbody = mjs_findBody(&child_model->spec, "world");
if (!worldbody) {
model->SetError(mjCError(0, "Child does not have a world body."));
return nullptr;
}
mjsFrame* worldframe = mjs_addFrame(worldbody, nullptr);
SetFrame(worldbody, mjOBJ_BODY, worldframe);
SetFrame(worldbody, mjOBJ_SITE, worldframe);
SetFrame(worldbody, mjOBJ_FRAME, worldframe);
SetFrame(worldbody, mjOBJ_JOINT, worldframe);
SetFrame(worldbody, mjOBJ_GEOM, worldframe);
SetFrame(worldbody, mjOBJ_LIGHT, worldframe);
SetFrame(worldbody, mjOBJ_CAMERA, worldframe);
child = worldframe->element;
}
switch (parent->elemtype) {
case mjOBJ_FRAME:
if (child->elemtype == mjOBJ_BODY) {
return attachBody(static_cast<mjCFrame*>(parent),
static_cast<const mjCBody*>(child), prefix, suffix);
} else if (child->elemtype == mjOBJ_FRAME) {
mjsBody* parent_body = mjs_getParent(parent);
if (!parent_body) {
model->SetError(mjCError(0, "Frame does not have a parent body."));
return nullptr;
}
mjCFrame* frame = static_cast<mjCFrame*>(parent);
mjsElement* attached_frame =
attachFrame(static_cast<mjCBody*>(parent_body->element),
static_cast<const mjCFrame*>(child), prefix, suffix);
if (mjs_setFrame(attached_frame, &frame->spec)) {
return nullptr;
}
return attached_frame;
} else {
model->SetError(mjCError(0, "child element is not a body"));
model->SetError(mjCError(0, "child element is not a body or frame"));
return nullptr;
}
case mjOBJ_BODY:
src/xml/xml_native_reader.cc
+14 -8
View File
@@ -3627,27 +3627,33 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
else if (name == "attach") {
string model_name, body_name, prefix;
ReadAttrTxt(elem, "model", model_name, /*required=*/true);
ReadAttrTxt(elem, "body", body_name, /*required=*/true);
ReadAttrTxt(elem, "body", body_name, /*required=*/false);
ReadAttrTxt(elem, "prefix", prefix, /*required=*/true);
mjsBody* child = mjs_findBody(spec, (prefix+body_name).c_str());
mjsBody* child_body = mjs_findBody(spec, (prefix+body_name).c_str());
mjsFrame* pframe = frame ? frame : mjs_addFrame(body, nullptr);
if (!child) {
if (!child_body) {
mjSpec* asset = mjs_findSpec(spec, model_name.c_str());
if (!asset) {
throw mjXError(elem, "could not find model '%s'", model_name.c_str());
}
child = mjs_findBody(asset, body_name.c_str());
if (!child) {
throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str());
mjsElement* child;
if (body_name.empty()) {
child = asset->element;
} else {
child_body = mjs_findBody(asset, body_name.c_str());
if (!child_body) {
throw mjXError(elem, "could not find body '%s''%s'", body_name.c_str());
}
child = child_body->element;
}
if (!mjs_attach(pframe->element, child->element, prefix.c_str(), "")) {
if (!mjs_attach(pframe->element, child, prefix.c_str(), "")) {
throw mjXError(elem, mjs_getError(spec));
}
} else {
// only set frame to existing body
if (mjs_setFrame(child->element, pframe)) {
if (mjs_setFrame(child_body->element, pframe)) {
throw mjXError(elem, mjs_getError(spec));
}
}