diff --git a/src/user/user_api.cc b/src/user/user_api.cc index e5ff263b..1a03b274 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -101,7 +101,7 @@ void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) { modelC->Compile(vfs, &m); if (d) { modelC->MakeData(m, &d); - modelC->RestoreState(d->qpos, d->qvel, d->act); + modelC->RestoreState(m->qpos0, d->qpos, d->qvel, d->act); d->time = time; } } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index f0d2810e..0d25cc59 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -134,6 +134,9 @@ mjCModel& mjCModel::operator=(const mjCModel& other) { // add everything else *this += other; + // add keyframes + CopyList(keys_, other.keys_); + // create new default tree mjCDef* subtree = new mjCDef(*other.defaults_[0]); *this += *subtree; @@ -168,6 +171,7 @@ void mjCModel::CopyList(std::vector& dest, continue; } // copy the element from the other model to this model + source[i]->ForgetKeyframes(); dest.push_back(candidate); dest.back()->model = this; dest.back()->id = -1; @@ -204,6 +208,32 @@ void mjCModel::ResetTreeLists() { +// save associated state addresses in related elements +void mjCModel::SaveDofOffsets() { + int qposadr = 0; + int dofadr = 0; + int actadr = 0; + + for (auto joint : joints_) { + joint->qposadr_ = qposadr; + joint->dofadr_ = dofadr; + qposadr += joint->nq(); + dofadr += joint->nv(); + } + + for (auto actuator : actuators_) { + if (actuator->spec.actdim > 0) { + actuator->actdim_ = actuator->spec.actdim; + } else { + actuator->actdim_ = (actuator->spec.dyntype != mjDYN_NONE); + } + actuator->actadr_ = actuator->actdim_ ? actadr : -1; + actadr += actuator->actdim_; + } +} + + + mjCModel& mjCModel::operator+=(const mjCModel& other) { // create global lists mjCBody *world = bodies_[0]; @@ -222,7 +252,9 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) { CopyList(hfields_, other.hfields_); CopyList(textures_, other.textures_); CopyList(materials_, other.materials_); - CopyList(keys_, other.keys_); + for (const auto& key : other.key_pending_) { + key_pending_.push_back(key); + } } CopyList(flexes_, other.flexes_); CopyList(pairs_, other.pairs_); @@ -2693,7 +2725,7 @@ void mjCModel::CopyObjects(mjModel* m) { m->actuator_actnum[i] = pac->actdim; m->actuator_actadr[i] = m->actuator_actnum[i] ? adr : -1; pac->actadr_ = m->actuator_actadr[i]; - pac->actnum_ = m->actuator_actnum[i]; + pac->actdim_ = m->actuator_actnum[i]; adr += m->actuator_actnum[i]; m->actuator_group[i] = pac->group; m->actuator_ctrllimited[i] = (mjtByte)pac->is_ctrllimited(); @@ -2827,7 +2859,8 @@ void mjCModel::CopyObjects(mjModel* m) { // save the current state -void mjCModel::SaveState(const mjtNum* qpos, const mjtNum* qvel, const mjtNum* act) { +template +void mjCModel::SaveState(const T* qpos, const T* qvel, const T* act) { for (auto joint : joints_) { if (qpos) mjuu_copyvec(joint->qpos, qpos + joint->qposadr_, joint->nq()); if (qvel) mjuu_copyvec(joint->qvel, qvel + joint->dofadr_, joint->nv()); @@ -2835,8 +2868,8 @@ void mjCModel::SaveState(const mjtNum* qpos, const mjtNum* qvel, const mjtNum* a for (auto actuator : actuators_) { if (actuator->actadr_ != -1 && act) { - actuator->act.assign(actuator->actnum_, 0); - mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actnum_); + actuator->act.assign(actuator->actdim_, 0); + mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actdim_); } } } @@ -2856,25 +2889,69 @@ void mjCModel::MakeData(const mjModel* m, mjData** dest) { // restore the previous state -void mjCModel::RestoreState(mjtNum* qpos, mjtNum* qvel, mjtNum* act) { +template +void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act) { for (auto joint : joints_) { - if (mjuu_defined(joint->qpos[0]) && qpos) { - mjuu_copyvec(qpos + joint->qposadr_, joint->qpos, joint->nq()); + if (qpos) { + if (mjuu_defined(joint->qpos[0])) { + mjuu_copyvec(qpos + joint->qposadr_, joint->qpos, 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()); } } + // restore act for (auto actuator : actuators_) { if (mjuu_defined(actuator->act[0]) && act) { - mjuu_copyvec(act + actuator->actadr_, actuator->act.data(), actuator->actnum_); + mjuu_copyvec(act + actuator->actadr_, actuator->act.data(), actuator->actdim_); } } } +// force explicit instantiations +template void mjCModel::SaveState(const mjtNum* qpos, + const mjtNum* qvel, + const mjtNum* act); + +template void mjCModel::RestoreState(const mjtNum* qpos0, mjtNum* qpos, + mjtNum* qvel, mjtNum* act); + + + +// resolve keyframe references +void mjCModel::StoreKeyframes() { + bool resetlists = false; + if (joints_.empty()) { + MakeLists(bodies_[0]); + resetlists = true; + } + + SaveDofOffsets(); + + for (auto key : keys_) { + mjKeyInfo info; + info.name = prefix + key->name + suffix; + info.qpos = !key->spec_qpos_.empty(); + info.qvel = !key->spec_qvel_.empty(); + info.act = !key->spec_act_.empty(); + key_pending_.push_back(info); + SaveState(key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data()); + break; // (b/350784262) save only the first keyframe for now + } + + if (resetlists) { + ResetTreeLists(); + } +} + + + //------------------------------- FUSE STATIC ------------------------------------------------------ template @@ -3401,6 +3478,49 @@ void mjCModel::CompileMeshes(const mjVFS* vfs) { +// convert pending keyframes info to actual keyframes +void mjCModel::ResolveKeyframes(const mjModel* m) { + if (key_pending_.empty()) { + return; + } + + // resize non-pending keyframes to the new number of dofs + 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); + } + } + + // store dof offsets in joints and actuators + SaveDofOffsets(); + + // copy state stored in joints and actuators to keyframes + for (const auto& info : key_pending_) { + mjCKey* key = (mjCKey*)FindObject(mjOBJ_KEY, info.name); + key->name = info.name; + if (info.qpos) key->spec_qpos_.assign(nq, 0); + if (info.qvel) key->spec_qvel_.assign(nv, 0); + if (info.act) key->spec_act_.assign(na, 0); + RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data()); + } + + // the attached keyframes have been copied into the model + key_pending_.clear(); +} + + + void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // check if nan test works double test = mjNAN; @@ -3447,6 +3567,12 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { SetDefaultNames(textures_); CheckEmptyNames(); + // create pending keyframes + for (const auto& info : key_pending_) { + mjCKey* key = AddKey(); + key->name = info.name; + } + // set object ids, check for repeated names ProcessLists(); @@ -3671,6 +3797,8 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { } // keyframe compilation needs access to nq, nv, na, nmocap, qpos0 + ResolveKeyframes(m); + for (int i=0; iCompile(m); } diff --git a/src/user/user_model.h b/src/user/user_model.h index 9f88bb5b..22e37135 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -34,6 +34,13 @@ typedef std::map > mjKeyMap; typedef std::array mjListKeyMap; +typedef struct mjKeyInfo_ { + std::string name; + bool qpos; + bool qvel; + bool act; +} mjKeyInfo; + class mjCModel_ : public mjsElement { public: // attach namespaces @@ -276,9 +283,14 @@ class mjCModel : public mjCModel_, private mjSpec { std::string_view name = ""); // save/restore the current state - void SaveState(const mjtNum* qpos, const mjtNum* qvel, const mjtNum* act); + template void SaveState(const T* qpos, const T* qvel, const T* act); + template void RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act); + + // clear existing data void MakeData(const mjModel* m, mjData** dest); - void RestoreState(mjtNum* qpos, mjtNum* qvel, mjtNum* act); + + // resolve keyframe references + void StoreKeyframes(); // map from default class name to default class pointer std::unordered_map def_map; @@ -359,8 +371,15 @@ class mjCModel : public mjCModel_, private mjSpec { // reset lists of kinematic tree void ResetTreeLists(); + // save dof offsets in joints and actuators + void SaveDofOffsets(); + + // convert pending keyframes info to actual keyframes + void ResolveKeyframes(const mjModel* m); + mjListKeyMap ids; // map from object names to ids mjCError errInfo; // last error info bool plugin_owner; // this class allocated the plugins + std::vector key_pending_; // attached keyframes }; #endif // MUJOCO_SRC_USER_USER_MODEL_H_ diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 3f0e96ff..fb0a18fd 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -845,6 +845,7 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) { mjCBody* subtree = other.body; other.model->prefix = other.prefix; other.model->suffix = other.suffix; + other.model->StoreKeyframes(); // attach defaults if (other.model != model) { @@ -887,14 +888,12 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) { bodies.back()->frame = subtree->bodies[i]->frame ? frames[fmap[subtree->bodies[i]->frame]] : nullptr; bodies.back()->NameSpace_(other.model, /*propagate=*/ false); + subtree->bodies[i]->ForgetKeyframes(); } // attach referencing elements *model += *other.model; - // (b/350784262) delete keyframes - model->DeleteAll(model->keys_); - // clear namespace and return body other.model->prefix.clear(); other.model->suffix.clear(); @@ -937,6 +936,10 @@ mjCBody& mjCBody::operator-=(const mjCBody& subtree) { } *bodies[i] -= subtree; } + + // (b/350784262) delete keyframes + model->DeleteAll(model->keys_); + return *this; } @@ -1402,6 +1405,19 @@ void mjCBody::ComputeBVH() { +// reset keyframe references for allowing self-attach +void mjCBody::ForgetKeyframes() const { + for (auto joint : joints) { + joint->qpos[0] = mjNAN; + joint->qvel[0] = mjNAN; + } + for (auto body : bodies) { + body->ForgetKeyframes(); + } +} + + + // compiler void mjCBody::Compile(void) { CopyFromSpec(); @@ -1631,7 +1647,14 @@ mjCFrame& mjCFrame::operator=(const mjCFrame& other) { // attach body to frame mjCFrame& mjCFrame::operator+=(const mjCBody& other) { + other.model->prefix = other.prefix; + other.model->suffix = other.suffix; + other.model->StoreKeyframes(); + other.model->prefix = ""; + other.model->suffix = ""; + mjCBody* subtree = new mjCBody(other, model); + other.ForgetKeyframes(); other.model->prefix = subtree->prefix; other.model->suffix = subtree->suffix; subtree->SetFrame(this); @@ -1650,9 +1673,6 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) { // attach referencing elements *model += *other.model; - // (b/350784262) delete keyframes - model->DeleteAll(model->keys_); - // clear suffixes and return other.model->suffix.clear(); other.model->prefix.clear(); @@ -5353,6 +5373,13 @@ mjCActuator& mjCActuator::operator=(const mjCActuator& other) { +void mjCActuator::ForgetKeyframes() { + act.clear(); + act.push_back(mjNAN); +} + + + bool mjCActuator::is_ctrllimited() const { return islimited(ctrllimited, ctrlrange); } bool mjCActuator::is_forcelimited() const { return islimited(forcelimited, forcerange); } bool mjCActuator::is_actlimited() const { return islimited(actlimited, actrange); } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 42dff904..8eed97d9 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -215,6 +215,10 @@ class mjCBase : public mjCBase_ { virtual ~mjCBase() = default; // destructor + // reset keyframe references for allowing self-attach + virtual void ForgetKeyframes() {} + virtual void ForgetKeyframes() const {} + protected: mjCBase(); // constructor mjCBase(const mjCBase& other); // copy constructor @@ -315,6 +319,9 @@ class mjCBody : public mjCBody_, private mjsBody { // get next child of given type mjsElement* NextChild(mjsElement* child, mjtObj type = mjOBJ_UNKNOWN); + // reset keyframe references for allowing self-attach + void ForgetKeyframes() const; + private: mjCBody(const mjCBody& other, mjCModel* _model); // copy constructor mjCBody& operator=(const mjCBody& other); // copy assignment @@ -1384,7 +1391,7 @@ class mjCActuator_ : public mjCBase { // variable used for temporarily storing the state of the actuator int actadr_; // address of dof in data->act - int actnum_; // number of dofs in data->act + int actdim_; // number of dofs in data->act std::vector act; // act at the previous step // variable-size data @@ -1431,6 +1438,9 @@ class mjCActuator : public mjCActuator_, private mjsActuator { void ResolveReferences(const mjCModel* m); void NameSpace(const mjCModel* m); + // reset keyframe references for allowing self-attach + void ForgetKeyframes(); + mjCBase* ptarget; // transmission target }; diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index ae4fa924..e5c581d5 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -421,13 +421,18 @@ static constexpr char xml_child[] = R"( - - + + + + + + + )"; TEST_F(MujocoTest, AttachSame) { @@ -484,16 +489,22 @@ TEST_F(MujocoTest, AttachSame) { - - - - + + + + + + + + + + )"; // create parent @@ -600,13 +611,17 @@ TEST_F(MujocoTest, AttachDifferent) { - - + + + + + + )"; // model with one free sphere and a frame @@ -717,13 +732,17 @@ TEST_F(MujocoTest, AttachFrame) { - - + + + + + + )"; // model with one free sphere and a frame @@ -790,6 +809,7 @@ void TestDetachBody(bool compile) { +