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

Fix bug in frame saving. Fixes #1802

Browse Source 
PiperOrigin-RevId: 651696068
Change-Id: Ie1e2e1de4e83c42f2b5ebcc3afd0540b9afe7bbd
This commit is contained in:
Alessio Quaglino
2024-07-12 01:55:35 -07:00
committed by Copybara-Service
parent 3d1d1d0718
commit 2f46dec669
5 changed files with 175 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/xml/xml_native_writer.cc
+80 -31
View File
@@ -867,7 +867,7 @@ string mjXWriter::Write(char *error, size_t error_sz) {
Extension(root);
Custom(root);
Asset(root);
Body(InsertEnd(root, "worldbody"), model->GetWorld());
Body(InsertEnd(root, "worldbody"), model->GetWorld(), nullptr);
Contact(root);
Deformable(root);
Equality(root);
@@ -1526,14 +1526,17 @@ XMLElement* mjXWriter::OneFrame(XMLElement* elem, mjCFrame* frame) {
XMLElement* frame_elem = InsertEnd(elem, "frame");
WriteAttrTxt(frame_elem, "name", frame->name);
WriteAttrTxt(frame_elem, "childclass", frame->classname);
if (frame->classname != "main") {
WriteAttrTxt(frame_elem, "childclass", frame->classname);
}
return frame_elem;
}
// recursive body and frame writer
void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclass) {
void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame,
std::string_view childclass) {
double unitq[4] = {1, 0, 0, 0};
if (!body) {
@@ -1541,7 +1544,7 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
}
// write body attributes and inertial
else if (body!=model->GetWorld()) {
else if (!frame && body!=model->GetWorld()) {
WriteAttrTxt(elem, "name", body->name);
if (childclass != body->classname && body->classname != "main") {
WriteAttrTxt(elem, "childclass", body->classname);
@@ -1574,57 +1577,67 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
}
}
// write joints
for (int i=0; i<body->joints.size(); i++) {
XMLElement* celem = OneFrame(elem, body->joints[i]->frame);
// joints in this frame
for (int i = 0; i < body->joints.size(); i++) {
if (body->joints[i]->frame != frame) {
continue;
}
std::string classname = body->joints[i]->frame && !body->joints[i]->frame->classname.empty()
? body->joints[i]->frame->classname
: body->classname;
OneJoint(InsertEnd(celem, "joint"), body->joints[i],
OneJoint(InsertEnd(elem, "joint"), body->joints[i],
model->def_map[body->joints[i]->classname],
classname.empty() ? childclass : classname);
}
// write geoms
for (int i=0; i<body->geoms.size(); i++) {
XMLElement* celem = OneFrame(elem, body->geoms[i]->frame);
// geoms in this frame
for (int i = 0; i < body->geoms.size(); i++) {
if (body->geoms[i]->frame != frame) {
continue;
}
std::string classname = body->geoms[i]->frame && !body->geoms[i]->frame->classname.empty()
? body->geoms[i]->frame->classname
: body->classname;
OneGeom(InsertEnd(celem, "geom"), body->geoms[i],
OneGeom(InsertEnd(elem, "geom"), body->geoms[i],
model->def_map[body->geoms[i]->classname],
classname.empty() ? childclass : classname);
}
// write sites
for (int i=0; i<body->sites.size(); i++) {
XMLElement* celem = OneFrame(elem, body->sites[i]->frame);
// sites in this frame
for (int i = 0; i < body->sites.size(); i++) {
if (body->sites[i]->frame != frame) {
continue;
}
std::string classname = body->sites[i]->frame && !body->sites[i]->frame->classname.empty()
? body->sites[i]->frame->classname
: body->classname;
OneSite(InsertEnd(celem, "site"), body->sites[i],
OneSite(InsertEnd(elem, "site"), body->sites[i],
model->def_map[body->sites[i]->classname],
classname.empty() ? childclass : classname);
}
// write cameras
for (int i=0; i<body->cameras.size(); i++) {
XMLElement* celem = OneFrame(elem, body->cameras[i]->frame);
// cameras in this frame
for (int i = 0; i < body->cameras.size(); i++) {
if (body->cameras[i]->frame != frame) {
continue;
}
std::string classname = body->cameras[i]->frame && !body->cameras[i]->frame->classname.empty()
? body->cameras[i]->frame->classname
: body->classname;
OneCamera(InsertEnd(celem, "camera"), body->cameras[i],
OneCamera(InsertEnd(elem, "camera"), body->cameras[i],
model->def_map[body->cameras[i]->classname],
classname.empty() ? childclass : classname);
}
// write lights
for (int i=0; i<body->lights.size(); i++) {
XMLElement* celem = OneFrame(elem, body->lights[i]->frame);
// lights in this frame
for (int i = 0; i < body->lights.size(); i++) {
if (body->lights[i]->frame != frame) {
continue;
}
std::string classname = body->lights[i]->frame && !body->lights[i]->frame->classname.empty()
? body->lights[i]->frame->classname
: body->classname;
OneLight(InsertEnd(celem, "light"), body->lights[i],
OneLight(InsertEnd(elem, "light"), body->lights[i],
model->def_map[body->lights[i]->classname],
classname.empty() ? childclass : classname);
}
@@ -1634,13 +1647,49 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
OnePlugin(InsertEnd(elem, "plugin"), &body->plugin);
}
// write child bodies recursively
for (int i=0; i<body->bodies.size(); i++) {
XMLElement* celem = OneFrame(elem, body->bodies[i]->frame);
std::string classname = body->bodies[i]->frame && !body->bodies[i]->frame->classname.empty()
? body->bodies[i]->frame->classname
: body->classname;
Body(InsertEnd(celem, "body"), body->bodies[i], classname.empty() ? childclass : classname);
// write children recursively
int i = 0, j = 0;
while (i < body->bodies.size() || body->bodies.empty()) {
mjCFrame* bframe = body->bodies.empty() ? nullptr : body->bodies[i]->frame;
// write body if its frame matches the current frame, avoid access if there are no bodies
if (bframe == frame && !body->bodies.empty()) {
std::string classname = bframe && !bframe->classname.empty()
? bframe->classname
: body->classname;
Body(InsertEnd(elem, "body"), body->bodies[i], nullptr,
classname.empty() ? childclass : classname);
}
i++;
// do not go to frames until we reach a body with a frame or we are done with bodies
if (!bframe && i < body->bodies.size()) {
continue;
}
// loop over the remaining frames in the current body
while (j < body->frames.size()) {
mjCFrame* fframe = body->frames[j++];
// write frame if its frame matches the current frame
if (fframe->frame == frame) {
std::string classname = fframe && !fframe->classname.empty()
? fframe->classname
: body->classname;
Body(OneFrame(elem, fframe), body, fframe, childclass);
}
// stop if we reached the frame of the current child body, ignore if there are no bodies
if (bframe && bframe == fframe) {
break;
}
}
// if there are no bodies, we only want to run the loop once
if (body->bodies.empty()) {
break;
}
}
}
src/xml/xml_native_writer.h
+1 -1
View File
@@ -59,7 +59,7 @@ class mjXWriter : public mjXBase {
void Keyframe(tinyxml2::XMLElement* root); // keyframe section
// body/world section
void Body(tinyxml2::XMLElement* elem, mjCBody* body, std::string_view childclass = "");
void Body(tinyxml2::XMLElement* elem, mjCBody* body, mjCFrame* frame, std::string_view childclass = "");
// single element writers, used in defaults and main body
void OneFlex(tinyxml2::XMLElement* elem, const mjCFlex* pflex);
test/user/testdata/frame1.xml
Vendored
+24
View File
@@ -0,0 +1,24 @@
<mujoco model="load save load test model with 2 frames">
<worldbody>
<frame name="frame1" pos=".1 .2 .3">
<geom name="geom1" type="sphere" size=".1"/>
<body name="root1">
<joint name="joint1"/>
<geom type="sphere" size=".1"/>
</body>
<body name="root2">
<frame name="frame2" pos=".01 .02 .03">
<geom type="sphere" size=".1"/>
<body name="root22">
<geom type="sphere" size=".1"/>
</body>
</frame><!-- end frame2 -->
</body>
</frame> <!-- end frame1 -->
</worldbody>
</mujoco>
test/user/testdata/frame2.xml
Vendored
+31
View File
@@ -0,0 +1,31 @@
<mujoco model="load save load test model with 2 frames">
<worldbody>
<frame name="frame1" pos=".1 .2 .3">
<geom name="geom1" type="sphere" size=".1"/>
<body name="root1">
<geom type="sphere" size=".1"/>
</body>
<frame name="frame2" pos=".01 .02 .03">
<geom name="geom2" type="sphere" size=".1"/>
<body name="root2">
<geom type="sphere" size=".1"/>
<body name="root22">
<geom type="sphere" size=".1"/>
</body>
</body>
<geom name="geom3" type="sphere" size=".1"/>
<body name="root3">
<geom type="sphere" size=".1"/>
<body name="root5">
<geom type="sphere" size=".1"/>
</body>
</body>
</frame> <!-- end frame1 -->
</frame><!-- end frame2 -->
</worldbody>
</mujoco>
test/xml/xml_native_writer_test.cc
+39 -10
View File
@@ -726,7 +726,7 @@ TEST_F(XMLWriterTest, WritesActuatorDefaults) {
TEST_F(XMLWriterTest, WritesFrameDefaults) {
static constexpr char xml[] = R"(
<mujoco>
<mujoco model="test">
<default>
<default class="dframe">
<geom size=".1"/>
@@ -741,24 +741,53 @@ TEST_F(XMLWriterTest, WritesFrameDefaults) {
<body name="body">
<frame pos="0 1 0" name="f2" childclass="dframe">
<geom pos="0 1 0"/>
<body pos="1 0 0">
<geom pos="0 0 1"/>
</body>
<frame pos="0 1 0" name="f3">
<frame pos="0 1 0" name="f4">
<body pos="1 0 0">
<geom pos="0 0 1"/>
</body>
</frame>
</frame>
</frame>
</body>
</worldbody>
</mujoco>
)";
static constexpr char xml_expected[] = R"(<mujoco model="test">
<compiler angle="radian"/>
<default>
<default class="dframe">
<geom size="0.1 0 0"/>
</default>
</default>
<worldbody>
<body name="body">
<frame name="f2" childclass="dframe">
<geom pos="0 2 0"/>
<frame name="f3" childclass="dframe">
<frame name="f4" childclass="dframe">
<body pos="1 3 0">
<geom pos="0 0 1"/>
</body>
</frame>
</frame>
</frame>
</body>
<frame name="f1">
<geom size="0.5" quat="0.906308 0 0 0.422618"/>
</frame>
</worldbody>
</mujoco>
)";
std::array<char, 1024> error;
mjModel* model = LoadModelFromString(xml, error.data(), error.size());
EXPECT_THAT(model, NotNull()) << error.data();
std::string saved_xml = SaveAndReadXml(model);
EXPECT_THAT(saved_xml, HasSubstr("frame name=\"f1\""));
EXPECT_THAT(saved_xml, HasSubstr("frame name=\"f2\" childclass=\"dframe\">"));
EXPECT_THAT(saved_xml, HasSubstr("geom pos=\"0 2 0\""));
EXPECT_THAT(saved_xml, HasSubstr("body pos=\"1 1 0\""));
EXPECT_THAT(saved_xml, HasSubstr("geom pos=\"0 0 1\""));
EXPECT_THAT(saved_xml, Not(HasSubstr("<frame>")));
EXPECT_STREQ(saved_xml.c_str(), xml_expected);
mj_deleteModel(model);
}