From 52a0149cd1676170a1987fa84c3cdfa8e23254fd Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Sun, 13 Oct 2024 10:01:26 -0700 Subject: [PATCH] Resizes keyframes before storing them during attach/detach. Improve keyframe tests. PiperOrigin-RevId: 685450257 Change-Id: I97aa129c4d7a66f613290605b85657e8081700b6 --- src/user/user_api.cc | 7 +- src/user/user_model.cc | 164 ++++++++++++++++++++--------- src/user/user_model.h | 12 ++- src/xml/xml_native_reader.cc | 4 +- test/user/user_api_test.cc | 43 ++++++++ test/xml/xml_native_reader_test.cc | 36 ++++--- 6 files changed, 199 insertions(+), 67 deletions(-) diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 8eac4a2a..151c0c2d 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -198,7 +198,12 @@ const char* mjs_getError(mjSpec* s) { int mjs_detachBody(mjSpec* s, mjsBody* b) { mjCModel* model = static_cast(s->element); mjCBody* body = static_cast(b->element); - *model -= *body; + try { + *model -= *body; + } catch (mjCError& e) { + model->SetError(e); + return -1; + } delete body; return 0; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index d8b4256c..420546a0 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -264,11 +264,11 @@ void mjCModel::ResetTreeLists() { // save associated state addresses in related elements -void mjCModel::SaveDofOffsets() { +void mjCModel::SaveDofOffsets(bool computesize) { int qposadr = 0; int dofadr = 0; int actadr = 0; - int nmocap = 0; + int mocapadr = 0; for (auto joint : joints_) { joint->qposadr_ = qposadr; @@ -289,12 +289,19 @@ void mjCModel::SaveDofOffsets() { for (mjCBody* body : bodies_) { if (body->spec.mocap) { - body->mocapid = nmocap; - nmocap++; + body->mocapid = mocapadr++; } else { body->mocapid = -1; } } + + if (computesize) { + nq = qposadr; + nv = dofadr; + na = actadr; + nu = (int)actuators_.size(); + nmocap = mocapadr; + } } template @@ -1556,6 +1563,7 @@ void mjCModel::SetSizes() { ntuple = (int)tuples_.size(); nkey = (int)keys_.size(); nplugin = (int)plugins_.size(); + nq = nv = nu = na = nmocap = 0; // nq, nv for (int i=0; iqpos0, nq); + mjuu_copyvec(body_pos0.data(), m->body_pos, 3*nbody); + mjuu_copyvec(body_quat0.data(), m->body_quat, 4*nbody); } @@ -3101,7 +3113,7 @@ void mjCModel::RestoreState(const std::string& state_name, const mjtNum* pos0, for (unsigned int i=0; imocap) { + if (!body->spec.mocap) { continue; } if (mpos) { @@ -3143,11 +3155,16 @@ void mjCModel::StoreKeyframes() { resetlists = true; } - // do not change the offset computed during compilation in case the user wants to recompile + // do not change compilation quantities in case the user wants to recompile preserving the state if (!compiled) { - SaveDofOffsets(); + SaveDofOffsets(/*computesize=*/true); + qpos0.resize(nq); + body_pos0.resize(3*bodies_.size()); + body_quat0.resize(4*bodies_.size()); + ComputeReference(qpos0, body_pos0, body_quat0); } + // save keyframe info and resize keyframes for (auto& key : keys_) { mjKeyInfo info; info.name = prefix + key->name + suffix; @@ -3159,6 +3176,7 @@ void mjCModel::StoreKeyframes() { info.mpos = !key->spec_mpos_.empty(); info.mquat = !key->spec_mquat_.empty(); key_pending_.push_back(info); + ResizeKeyframe(key, qpos0.data(), body_pos0.data(), body_quat0.data()); SaveState(info.name, key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data(), key->spec_ctrl_.data(), key->spec_mpos_.data(), key->spec_mquat_.data()); @@ -3167,6 +3185,10 @@ void mjCModel::StoreKeyframes() { if (resetlists) { ResetTreeLists(); } + + if (!compiled) { + nq = nv = na = nu = nmocap = 0; + } } @@ -3697,6 +3719,87 @@ void mjCModel::CompileMeshes(const mjVFS* vfs) { +// compute qpos0 +template +void mjCModel::ComputeReference(std::vector& q0, std::vector& bpos, std::vector& bquat) { + int b = 0; + for (auto body : bodies_) { + mjuu_copyvec(bpos.data()+3*b, body->spec.pos, 3); + mjuu_copyvec(bquat.data()+4*b, body->spec.quat, 4); + for (auto joint : body->joints) { + switch (joint->type) { + case mjJNT_FREE: + mjuu_copyvec(q0.data()+joint->qposadr_, body->spec.pos, 3); + mjuu_copyvec(q0.data()+joint->qposadr_+3, body->spec.quat, 4); + break; + + case mjJNT_BALL: + mjuu_setvec(q0.data()+joint->qposadr_, 1, 0, 0, 0); + break; + + case mjJNT_SLIDE: + case mjJNT_HINGE: + q0[joint->qposadr_] = (T)joint->spec.ref; + break; + + default: + throw mjCError(joint, "unknown joint type"); + } + } + b++; + } +} + + + +// resizes a keyframe, filling in missing values +void mjCModel::ResizeKeyframe(mjCKey* key, const mjtNum* qpos0_, + const mjtNum* bpos, const mjtNum* bquat) { + if (!key->spec_qpos_.empty()) { + int nq0 = key->spec_qpos_.size(); + key->spec_qpos_.resize(nq); + for (int i=nq0; ispec_qpos_[i] = (double)qpos0_[i]; + } + } + if (!key->spec_qvel_.empty()) { + key->spec_qvel_.resize(nv); + } + if (!key->spec_act_.empty()) { + key->spec_act_.resize(na); + } + if (!key->spec_ctrl_.empty()) { + key->spec_ctrl_.resize(nu); + } + if (!key->spec_mpos_.empty()) { + int nmocap0 = key->spec_mpos_.size() / 3; + key->spec_mpos_.resize(3*nmocap); + for (unsigned int j = 0; j < bodies_.size(); j++) { + if (bodies_[j]->mocapid < nmocap0) { + continue; + } + int i = bodies_[j]->mocapid; + key->spec_mpos_[3*i+0] = (double)bpos[3*j+0]; + key->spec_mpos_[3*i+1] = (double)bpos[3*j+1]; + key->spec_mpos_[3*i+2] = (double)bpos[3*j+2]; + } + } + if (!key->spec_mquat_.empty()) { + int nmocap0 = key->spec_mquat_.size() / 4; + key->spec_mquat_.resize(4*nmocap); + for (unsigned int j = 0; j < bodies_.size(); j++) { + if (bodies_[j]->mocapid < nmocap0) { + continue; + } + int i = bodies_[j]->mocapid; + key->spec_mquat_[4*i+0] = (double)bquat[4*j+0]; + key->spec_mquat_[4*i+1] = (double)bquat[4*j+1]; + key->spec_mquat_[4*i+2] = (double)bquat[4*j+2]; + key->spec_mquat_[4*i+3] = (double)bquat[4*j+3]; + } + } +} + // convert pending keyframes info to actual keyframes void mjCModel::ResolveKeyframes(const mjModel* m) { if (key_pending_.empty()) { @@ -3707,51 +3810,8 @@ void mjCModel::ResolveKeyframes(const mjModel* m) { SaveDofOffsets(); // resize existing keyframes to the new state, fill in missing default values - for (unsigned int i = 0; i < nkey - key_pending_.size(); i++) { - mjCKey* key = keys_[i]; - if (!key->spec_qpos_.empty()) { - int nq0 = key->spec_qpos_.size(); - key->spec_qpos_.resize(nq); - for (int i=nq0; inq; i++) { - key->spec_qpos_[i] = (double)m->qpos0[i]; - } - } - if (!key->spec_qvel_.empty()) { - key->spec_qvel_.resize(nv); - } - if (!key->spec_act_.empty()) { - key->spec_act_.resize(na); - } - if (!key->spec_ctrl_.empty()) { - key->spec_ctrl_.resize(nu); - } - if (!key->spec_mpos_.empty()) { - int nmocap0 = key->spec_mpos_.size() / 3; - key->spec_mpos_.resize(3*nmocap); - for (unsigned int j = 0; j < bodies_.size(); j++) { - if (bodies_[j]->mocapid < nmocap0) { - continue; - } - int i = bodies_[j]->mocapid; - key->spec_mpos_[3*i+0] = (double)m->body_pos[3*j+0]; - key->spec_mpos_[3*i+1] = (double)m->body_pos[3*j+1]; - key->spec_mpos_[3*i+2] = (double)m->body_pos[3*j+2]; - } - } - if (!key->spec_mquat_.empty()) { - int nmocap0 = key->spec_mquat_.size() / 4; - key->spec_mquat_.resize(4*nmocap); - for (unsigned int j = 0; j < bodies_.size(); j++) { - if (bodies_[j]->mocapid < nmocap0) { - continue; - } - int i = bodies_[j]->mocapid; - key->spec_mquat_[4*i+0] = (double)m->body_quat[4*j+0]; - key->spec_mquat_[4*i+1] = (double)m->body_quat[4*j+1]; - key->spec_mquat_[4*i+2] = (double)m->body_quat[4*j+2]; - key->spec_mquat_[4*i+3] = (double)m->body_quat[4*j+3]; - } - } + for (auto* key : keys_) { + ResizeKeyframe(key, m->qpos0, m->body_pos, m->body_quat); } // create new keyframes, fill in missing default values diff --git a/src/user/user_model.h b/src/user/user_model.h index 78e12a32..d6c6d81a 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -128,6 +128,8 @@ class mjCModel_ : public mjsElement { // save qpos0, to recognize changed key_qpos in write std::vector qpos0; + std::vector body_pos0; + std::vector body_quat0; // variable-size attributes std::string comment_; // comment at top of XML @@ -382,11 +384,19 @@ class mjCModel : public mjCModel_, private mjSpec { void ResetTreeLists(); // save dof offsets in joints and actuators - void SaveDofOffsets(); + void SaveDofOffsets(bool computesize = false); // convert pending keyframes info to actual keyframes void ResolveKeyframes(const mjModel* m); + // resize a keyframe, filling in missing values + void ResizeKeyframe(mjCKey* key, const mjtNum* qpos0_, const mjtNum* bpos, const mjtNum* bquat); + + // compute qpos0 + template + void ComputeReference(std::vector& q0, std::vector& bpos, + std::vector& bquat); + mjListKeyMap ids; // map from object names to ids mjCError errInfo; // last error info std::vector key_pending_; // attached keyframes diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 404420cc..5961592b 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -3617,7 +3617,9 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame, } // delete subtree - mjs_detachBody(spec, subtree); + if (mjs_detachBody(spec, subtree)) { + throw mjXError(elem, mjs_getError(spec)); + } } // body sub-element diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 7c06c6e3..0b71ed0e 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -783,6 +783,10 @@ TEST_F(MujocoTest, AttachDifferent) { + + + + )"; static constexpr char xml_result[] = R"( @@ -837,6 +841,7 @@ TEST_F(MujocoTest, AttachDifferent) { + @@ -905,6 +910,10 @@ TEST_F(MujocoTest, AttachFrame) { + + + + )"; static constexpr char xml_result[] = R"( @@ -959,6 +968,7 @@ TEST_F(MujocoTest, AttachFrame) { + @@ -1356,6 +1366,39 @@ TEST_F(MujocoTest, AttachMocap) { mj_deleteModel(m_expected); } +TEST_F(MujocoTest, ReplicateKeyframe) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, testing::NotNull()) << error.data(); + EXPECT_THAT(m->ngeom, 1); + EXPECT_THAT(m->nbody, 2); + + // check that the keyframe is resized + EXPECT_THAT(m->nkey, 1); + EXPECT_THAT(m->nq, 1); + EXPECT_THAT(m->key_qpos[0], 0); + EXPECT_STREQ(mj_id2name(m, mjOBJ_KEY, 0), "keyframe"); + + mj_deleteModel(m); +} + TEST_F(MujocoTest, AttachUnnamedAssets) { static constexpr char cube[] = R"( v -1 -1 1 diff --git a/test/xml/xml_native_reader_test.cc b/test/xml/xml_native_reader_test.cc index df24d742..505659f9 100644 --- a/test/xml/xml_native_reader_test.cc +++ b/test/xml/xml_native_reader_test.cc @@ -1234,21 +1234,28 @@ TEST_F(XMLReaderTest, ParseReplicate) { + + + + + + + - - - - + + + + )"; @@ -1287,14 +1294,19 @@ TEST_F(XMLReaderTest, ParseReplicate) { } // check that the final pose is correct - int n = 104; - EXPECT_NEAR(m->geom_pos[3*n+0], 0, 1e-8); - EXPECT_NEAR(m->geom_pos[3*n+1], 1, 1e-8); - EXPECT_EQ(m->geom_pos[3*n+2], 0); - EXPECT_NEAR(m->geom_quat[4*n+0], 0, 1e-8); - EXPECT_EQ(m->geom_quat[4*n+1], 0); - EXPECT_EQ(m->geom_quat[4*n+2], 0); - EXPECT_EQ(m->geom_quat[4*n+3], 1); + int n = m->nbody-1; + EXPECT_THAT(m->nbody, 102); + EXPECT_NEAR(m->body_pos[3*n+0], 0, 1e-8); + EXPECT_NEAR(m->body_pos[3*n+1], 1, 1e-8); + EXPECT_EQ(m->body_pos[3*n+2], 0); + EXPECT_NEAR(m->body_quat[4*n+0], 0, 1e-8); + EXPECT_EQ(m->body_quat[4*n+1], 0); + EXPECT_EQ(m->body_quat[4*n+2], 0); + EXPECT_EQ(m->body_quat[4*n+3], 1); + + // check that the pending keyframes are lost while detaching + EXPECT_THAT(m->nkey, 0); + EXPECT_THAT(m->nq, 101); mj_deleteModel(m); }