diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 338f7b7c..645ea9e2 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -3808,7 +3808,6 @@ all attachments will appear in the saved XML file. - An entire model cannot be attached (i.e. including all elements, referenced or not). - All assets from the child model will be copied in, whether they are referenced or not. - Self-attach or circular references are not checked for and will lead to infinite loops. - - :ref:`Keyframes` are not yet supported. When attaching, all keyframes will be deleted. .. _body-attach-model: diff --git a/doc/changelog.rst b/doc/changelog.rst index e48c492b..80294d25 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -7,10 +7,9 @@ Upcoming version (not yet released) General ^^^^^^^ -1. Add :ref:`mjSpec` option for creating a texture from a buffer. +1. Added :ref:`mjSpec` option for creating a texture from a buffer. 2. :ref:`shellinertia ` is now supported by all geom types. -3. Add support for :ref:`attaching` keyframes. Note: this only supports keyframe containing qpos, qvel, - and act. +3. Added support for :ref:`attaching` keyframes. Version 3.2.2 (Aug 8, 2024) --------------------------- diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 7ff36d73..bf3a98c6 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -93,15 +93,17 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) { // recompile spec into existing model and data while preserving the state void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) { mjCModel* modelC = static_cast(s->element); + std::string state_name = "state"; mjtNum time = 0; if (d) { time = d->time; - modelC->SaveState(d->qpos, d->qvel, d->act, d->ctrl); + modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat); } modelC->Compile(vfs, &m); if (d) { modelC->MakeData(m, &d); - modelC->RestoreState(m->qpos0, d->qpos, d->qvel, d->act, d->ctrl); + modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel, + d->act, d->ctrl, d->mocap_pos, d->mocap_quat); d->time = time; } } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 178dbdd7..bfcce1fd 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -109,9 +109,6 @@ mjCModel::mjCModel() { // this class allocated the plugins plugin_owner = true; - - // default state name - state_name_ = "state"; } @@ -216,6 +213,7 @@ void mjCModel::SaveDofOffsets() { int qposadr = 0; int dofadr = 0; int actadr = 0; + int nmocap = 0; for (auto joint : joints_) { joint->qposadr_ = qposadr; @@ -233,6 +231,15 @@ void mjCModel::SaveDofOffsets() { actuator->actadr_ = actuator->actdim_ ? actadr : -1; actadr += actuator->actdim_; } + + for (mjCBody* body : bodies_) { + if (body->spec.mocap) { + body->mocapid = nmocap; + nmocap++; + } else { + body->mocapid = -1; + } + } } @@ -2863,23 +2870,36 @@ void mjCModel::CopyObjects(mjModel* m) { // save the current state template -void mjCModel::SaveState(const T* qpos, const T* qvel, const T* act, const T* ctrl) { +void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act, + const T* ctrl, const T* mpos, const T* mquat) { for (auto joint : joints_) { if (joint->qposadr_ == -1 || joint->dofadr_ == -1) { throw mjCError(NULL, "SaveState: joint %s has no address", joint->name.c_str()); } - if (qpos) mjuu_copyvec(joint->qpos(), qpos + joint->qposadr_, joint->nq()); - if (qvel) mjuu_copyvec(joint->qvel(), qvel + joint->dofadr_, joint->nv()); + if (qpos) mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq()); + if (qvel) mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv()); } for (unsigned int i=0; iactadr_ != -1 && actuator->actdim_ != -1 && act) { - actuator->act().assign(actuator->actdim_, 0); - mjuu_copyvec(actuator->act().data(), act + actuator->actadr_, actuator->actdim_); + actuator->act(state_name).assign(actuator->actdim_, 0); + mjuu_copyvec(actuator->act(state_name).data(), act + actuator->actadr_, actuator->actdim_); } if (ctrl) { - actuator->ctrl() = ctrl[i]; + actuator->ctrl(state_name) = ctrl[i]; + } + } + + for (auto body : bodies_) { + if (!body->spec.mocap) { + continue; + } + if (mpos) { + mjuu_copyvec(body->mpos(state_name), mpos + 3*body->mocapid, 3); + } + if (mquat) { + mjuu_copyvec(body->mquat(state_name), mquat + 4*body->mocapid, 4); } } } @@ -2900,42 +2920,63 @@ void mjCModel::MakeData(const mjModel* m, mjData** dest) { // restore the previous state template -void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act, T* ctrl) { +void mjCModel::RestoreState(const std::string& state_name, const mjtNum* pos0, + const mjtNum* mpos0, const mjtNum* mquat0, T* qpos, + T* qvel, T* act, T* ctrl, T* mpos, T* mquat) { for (auto joint : joints_) { if (qpos) { - if (mjuu_defined(joint->qpos()[0])) { - mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(), joint->nq()); + if (mjuu_defined(joint->qpos(state_name)[0])) { + mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(state_name), joint->nq()); } else { mjuu_copyvec(qpos + joint->qposadr_, pos0 + joint->qposadr_, joint->nq()); } } - if (mjuu_defined(joint->qvel()[0]) && qvel) { - mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(), joint->nv()); + if (mjuu_defined(joint->qvel(state_name)[0]) && qvel) { + mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(state_name), joint->nv()); } } // restore act for (unsigned int i=0; iact().empty() && mjuu_defined(actuator->act()[0]) && act) { - mjuu_copyvec(act + actuator->actadr_, actuator->act().data(), actuator->actdim_); + if (!actuator->act(state_name).empty() && mjuu_defined(actuator->act(state_name)[0]) && act) { + mjuu_copyvec(act + actuator->actadr_, actuator->act(state_name).data(), actuator->actdim_); } if (ctrl) { - ctrl[i] = mjuu_defined(actuator->ctrl()) ? actuator->ctrl() : 0; + ctrl[i] = mjuu_defined(actuator->ctrl(state_name)) ? actuator->ctrl(state_name) : 0; + } + } + + for (unsigned int i=0; imocap) { + continue; + } + if (mpos) { + if (mjuu_defined(body->mpos(state_name)[0])) { + mjuu_copyvec(mpos + 3*body->mocapid, body->mpos(state_name), 3); + } else { + mjuu_copyvec(mpos + 3*body->mocapid, mpos0 + 3*i, 3); + } + } + if (mquat) { + if (mjuu_defined(body->mquat(state_name)[0])) { + mjuu_copyvec(mquat + 4*body->mocapid, body->mquat(state_name), 4); + } else { + mjuu_copyvec(mquat + 4*body->mocapid, mquat0 + 4*i, 4); + } } } } - - // force explicit instantiations -template void mjCModel::SaveState(const mjtNum* qpos, - const mjtNum* qvel, - const mjtNum* act, - const mjtNum* ctrl); +template void mjCModel::SaveState( + const std::string& name, const mjtNum* qpos, const mjtNum* qvel, const mjtNum* act, + const mjtNum* ctrl, const mjtNum* mpos, const mjtNum* mquat); -template void mjCModel::RestoreState(const mjtNum* qpos0, mjtNum* qpos, - mjtNum* qvel, mjtNum* act, mjtNum* ctrl); +template void mjCModel::RestoreState( + const std::string& name, const mjtNum* qpos0, const mjtNum* mpos0, const mjtNum* mquat0, + mjtNum* qpos, mjtNum* qvel, mjtNum* act, mjtNum* ctrl, mjtNum* mpos, mjtNum* mquat); @@ -2957,10 +2998,12 @@ void mjCModel::StoreKeyframes() { info.qvel = !key->spec_qvel_.empty(); info.act = !key->spec_act_.empty(); info.ctrl = !key->spec_ctrl_.empty(); + info.mpos = !key->spec_mpos_.empty(); + info.mquat = !key->spec_mquat_.empty(); key_pending_.push_back(info); - state_name_ = info.name; - SaveState(key->spec_qpos_.data(), key->spec_qvel_.data(), - key->spec_act_.data(), key->spec_ctrl_.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()); } if (resetlists) { @@ -3502,7 +3545,10 @@ void mjCModel::ResolveKeyframes(const mjModel* m) { return; } - // resize non-pending keyframes to the new number of dofs + // store dof offsets in joints and actuators + 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()) { @@ -3521,12 +3567,36 @@ void mjCModel::ResolveKeyframes(const mjModel* m) { 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]; + } + } } - // store dof offsets in joints and actuators - SaveDofOffsets(); - - // copy state stored in joints and actuators to keyframes + // create new keyframes, fill in missing default values for (const auto& info : key_pending_) { mjCKey* key = (mjCKey*)FindObject(mjOBJ_KEY, info.name); key->name = info.name; @@ -3535,9 +3605,12 @@ void mjCModel::ResolveKeyframes(const mjModel* m) { if (info.qvel) key->spec_qvel_.assign(nv, 0); if (info.act) key->spec_act_.assign(na, 0); if (info.ctrl) key->spec_ctrl_.assign(nu, 0); - state_name_ = info.name; - RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(), - key->spec_act_.data(), key->spec_ctrl_.data()); + if (info.mpos) key->spec_mpos_.assign(3*nmocap, 0); + if (info.mquat) key->spec_mquat_.assign(4*nmocap, 0); + RestoreState(info.name, m->qpos0, m->body_pos, m->body_quat, + key->spec_qpos_.data(), key->spec_qvel_.data(), + key->spec_act_.data(), key->spec_ctrl_.data(), + key->spec_mpos_.data(), key->spec_mquat_.data()); } // the attached keyframes have been copied into the model diff --git a/src/user/user_model.h b/src/user/user_model.h index ed7ef8cc..f4ccf839 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -41,6 +41,8 @@ typedef struct mjKeyInfo_ { bool qvel; bool act; bool ctrl; + bool mpos; + bool mquat; } mjKeyInfo; class mjCModel_ : public mjsElement { @@ -284,9 +286,15 @@ class mjCModel : public mjCModel_, private mjSpec { template void DeleteMaterial(std::vector& list, std::string_view name = ""); - // save/restore the current state - template void SaveState(const T* qpos, const T* qvel, const T* act, const T* ctrl); - template void RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act, T* ctrl); + // save the current state + template + void SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act, + const T* ctrl, const T* mpos, const T* mquat); + + // restore the previously saved state + template + void RestoreState(const std::string& state_name, const mjtNum* pos0, const mjtNum* mpos0, + const mjtNum* mquat0, T* qpos, T* qvel, T* act, T* ctrl, T* mpos, T* mquat); // clear existing data void MakeData(const mjModel* m, mjData** dest); @@ -383,7 +391,5 @@ class mjCModel : public mjCModel_, private mjSpec { mjCError errInfo; // last error info bool plugin_owner; // this class allocated the plugins std::vector key_pending_; // attached keyframes - - std::string state_name_; }; #endif // MUJOCO_SRC_USER_USER_MODEL_H_ diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 12f55c1f..3a509b15 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -1411,6 +1411,8 @@ void mjCBody::ForgetKeyframes() const { joint->qpos_.clear(); joint->qvel_.clear(); } + model->FindBody((mjCBody*)this, name)->mpos_.clear(); // this is a hack to avoid const + model->FindBody((mjCBody*)this, name)->mquat_.clear(); // this is a hack to avoid const for (auto body : bodies) { body->ForgetKeyframes(); } @@ -1418,6 +1420,24 @@ void mjCBody::ForgetKeyframes() const { +mjtNum* mjCBody::mpos(const std::string& state_name) { + if (mpos_.find(state_name) == mpos_.end()) { + mpos_[state_name] = {mjNAN, 0, 0}; + } + return mpos_.at(state_name).data(); +} + + + +mjtNum* mjCBody::mquat(const std::string& state_name) { + if (mquat_.find(state_name) == mquat_.end()) { + mquat_[state_name] = {mjNAN, 0, 0, 0}; + } + return mquat_.at(state_name).data(); +} + + + // compiler void mjCBody::Compile(void) { CopyFromSpec(); @@ -1828,20 +1848,20 @@ int mjCJoint::nv(mjtJoint joint_type) { -mjtNum* mjCJoint::qpos() { - if (qpos_.find(model->state_name_) == qpos_.end()) { - qpos_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0, 0}; +mjtNum* mjCJoint::qpos(const std::string& state_name) { + if (qpos_.find(state_name) == qpos_.end()) { + qpos_[state_name] = {mjNAN, 0, 0, 0, 0, 0, 0}; } - return qpos_.at(model->state_name_).data(); + return qpos_.at(state_name).data(); } -mjtNum* mjCJoint::qvel() { - if (qvel_.find(model->state_name_) == qvel_.end()) { - qvel_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0}; +mjtNum* mjCJoint::qvel(const std::string& state_name) { + if (qvel_.find(state_name) == qvel_.end()) { + qvel_[state_name] = {mjNAN, 0, 0, 0, 0, 0}; } - return qvel_.at(model->state_name_).data(); + return qvel_.at(state_name).data(); } @@ -5405,20 +5425,20 @@ bool mjCActuator::is_actlimited() const { return islimited(actlimited, actrange) -std::vector& mjCActuator::act() { - if (act_.find(model->state_name_) == act_.end()) { - act_[model->state_name_] = std::vector(model->nu, mjNAN); +std::vector& mjCActuator::act(const std::string& state_name) { + if (act_.find(state_name) == act_.end()) { + act_[state_name] = std::vector(model->nu, mjNAN); } - return act_.at(model->state_name_); + return act_.at(state_name); } -mjtNum& mjCActuator::ctrl() { - if (ctrl_.find(model->state_name_) == ctrl_.end()) { - ctrl_[model->state_name_] = mjNAN; +mjtNum& mjCActuator::ctrl(const std::string& state_name) { + if (ctrl_.find(state_name) == ctrl_.end()) { + ctrl_[state_name] = mjNAN; } - return ctrl_.at(model->state_name_); + return ctrl_.at(state_name); } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 62401489..c55cdcd0 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -255,6 +255,10 @@ class mjCBody_ : public mjCBase { std::string plugin_instance_name; std::vector userdata_; std::vector spec_userdata_; + + // variables used for temporarily storing the state of the mocap bodies + std::map> mpos_; // saved mocap_pos + std::map> mquat_; // saved mocap_quat }; class mjCBody : public mjCBody_, private mjsBody { @@ -323,6 +327,10 @@ class mjCBody : public mjCBody_, private mjsBody { // reset keyframe references for allowing self-attach void ForgetKeyframes() const; + // get mocap position and quaternion + mjtNum* mpos(const std::string& state_name); + mjtNum* mquat(const std::string& state_name); + private: mjCBody(const mjCBody& other, mjCModel* _model); // copy constructor mjCBody& operator=(const mjCBody& other); // copy assignment @@ -444,8 +452,8 @@ class mjCJoint : public mjCJoint_, private mjsJoint { int nq() const { return nq(spec.type); } int nv() const { return nv(spec.type); } - mjtNum* qpos(); - mjtNum* qvel(); + mjtNum* qpos(const std::string& state_name); + mjtNum* qvel(const std::string& state_name); private: int Compile(void); // compiler; return dofnum @@ -1436,8 +1444,8 @@ class mjCActuator : public mjCActuator_, private mjsActuator { bool is_forcelimited() const; bool is_actlimited() const; - std::vector& act(); - mjtNum& ctrl(); + std::vector& act(const std::string& state_name); + mjtNum& ctrl(const std::string& state_name); private: void Compile(void); // compiler diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 795503da..8f539ee8 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -870,6 +870,8 @@ TEST_F(MujocoTest, PreserveState) { + + @@ -888,6 +890,7 @@ TEST_F(MujocoTest, PreserveState) { + @@ -915,6 +918,14 @@ TEST_F(MujocoTest, PreserveState) { data->ctrl[1] = 2; d_expected->ctrl[0] = 2; + // set mocap + data->mocap_pos[3] = 1; + data->mocap_quat[4] = 0; + data->mocap_quat[5] = 1; + d_expected->mocap_pos[0] = 1; + d_expected->mocap_quat[0] = 0; + d_expected->mocap_quat[1] = 1; + // step models mj_step(model, data); mj_step(m_expected, d_expected); @@ -925,6 +936,11 @@ TEST_F(MujocoTest, PreserveState) { EXPECT_THAT(body, NotNull()); EXPECT_THAT(mjs_detachBody(spec, body), 0); + // detach mocap + mjsBody* mocap_body = mjs_findBody(spec, "mocap_detach"); + EXPECT_THAT(mocap_body, NotNull()); + EXPECT_THAT(mjs_detachBody(spec, mocap_body), 0); + // add body mjsBody* newbody = mjs_addBody(mjs_findBody(spec, "world"), 0); EXPECT_THAT(newbody, NotNull()); @@ -968,6 +984,17 @@ TEST_F(MujocoTest, PreserveState) { EXPECT_EQ(data->act[i], d_expected->act[i]) << i; } + // compare mocap + EXPECT_EQ(model->nmocap, m_expected->nmocap); + for (int i = 0; i < model->nmocap; ++i) { + for (int j = 0; j < 3; ++j) { + EXPECT_EQ(data->mocap_pos[3*i+j], d_expected->mocap_pos[3*i+j]) << i; + } + for (int j = 0; j < 4; ++j) { + EXPECT_EQ(data->mocap_quat[4*i+j], d_expected->mocap_quat[4*i+j]) << i; + } + } + // check that the function is callable with no data mj_deleteData(data); mj_recompile(spec, 0, model, nullptr); @@ -979,5 +1006,68 @@ TEST_F(MujocoTest, PreserveState) { mj_deleteModel(m_expected); } +TEST_F(MujocoTest, AttachMocap) { + std::array er; + mjtNum tol = 0; + std::string field = ""; + + static constexpr char xml[] = R"( + + + + + + + + )"; + + static constexpr char xml_expected[] = R"( + + + + + + + + + + )"; + + mjSpec* spec = mj_parseXMLString(xml, 0, er.data(), er.size()); + EXPECT_THAT(spec, NotNull()) << er.data(); + + mjsBody* body = mjs_findBody(spec, "mocap"); + EXPECT_THAT(body, NotNull()); + + mjsBody* world = mjs_findBody(spec, "world"); + EXPECT_THAT(world, NotNull()); + + mjsFrame* frame = mjs_addFrame(world, NULL); + mjs_attachBody(frame, body, "attached-", "-1"); + + mjsBody* attached_body = mjs_findBody(spec, "attached-mocap-1"); + EXPECT_THAT(attached_body, NotNull()); + attached_body->pos[0] = 3; + attached_body->pos[1] = 3; + attached_body->pos[2] = 3; + attached_body->quat[0] = 0; + attached_body->quat[1] = 0; + attached_body->quat[2] = 1; + attached_body->quat[3] = 0; + + mjModel* model = mj_compile(spec, 0); + EXPECT_THAT(model, NotNull()); + + mjModel* m_expected = LoadModelFromString(xml_expected, er.data(), er.size()); + EXPECT_THAT(m_expected, NotNull()) << er.data(); + EXPECT_LE(CompareModel(model, m_expected, field), tol) + << "Expected and attached models are different!\n" + << "Different field: " << field << '\n'; + + mj_deleteSpec(spec); + mj_deleteModel(model); + mj_deleteModel(m_expected); +} + } // namespace } // namespace mujoco