From afd7c73f44e78c262aada0eadb6695c62c32632c Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 14 Aug 2024 04:41:51 -0700 Subject: [PATCH] Attach multiple keyframes. Still missing: support for ctrl, mocap, and time. PiperOrigin-RevId: 662871569 Change-Id: I8a418c4007663e3063d7146bd72db4eba78ddc5d --- doc/changelog.rst | 2 ++ src/user/user_model.cc | 31 +++++++++++++++++----------- src/user/user_model.h | 2 ++ src/user/user_objects.cc | 41 +++++++++++++++++++++++++++++++------- src/user/user_objects.h | 20 ++++++++++++------- test/user/user_api_test.cc | 3 +++ 6 files changed, 73 insertions(+), 26 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index f601660f..e48c492b 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -9,6 +9,8 @@ General ^^^^^^^ 1. Add :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. Version 3.2.2 (Aug 8, 2024) --------------------------- diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 0d25cc59..5ea2124c 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -109,6 +109,9 @@ mjCModel::mjCModel() { // this class allocated the plugins plugin_owner = true; + + // default state name + state_name_ = "state"; } @@ -2862,14 +2865,17 @@ void mjCModel::CopyObjects(mjModel* m) { 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()); + 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()); } for (auto actuator : actuators_) { - if (actuator->actadr_ != -1 && act) { - actuator->act.assign(actuator->actdim_, 0); - mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actdim_); + if (actuator->actadr_ != -1 && actuator->actdim_ != -1 && act) { + actuator->act().assign(actuator->actdim_, 0); + mjuu_copyvec(actuator->act().data(), act + actuator->actadr_, actuator->actdim_); } } } @@ -2893,21 +2899,21 @@ template void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act) { 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()[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()); + 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->actdim_); + if (!actuator->act().empty() && mjuu_defined(actuator->act()[0]) && act) { + mjuu_copyvec(act + actuator->actadr_, actuator->act().data(), actuator->actdim_); } } } @@ -2941,8 +2947,8 @@ void mjCModel::StoreKeyframes() { info.qvel = !key->spec_qvel_.empty(); info.act = !key->spec_act_.empty(); key_pending_.push_back(info); + state_name_ = info.name; 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) { @@ -3512,6 +3518,7 @@ void mjCModel::ResolveKeyframes(const mjModel* m) { 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); + state_name_ = info.name; RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data()); } diff --git a/src/user/user_model.h b/src/user/user_model.h index 22e37135..674ab72a 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -381,5 +381,7 @@ 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 fb0a18fd..0afb39e7 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -1408,8 +1408,8 @@ 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; + joint->qpos_.clear(); + joint->qvel_.clear(); } for (auto body : bodies) { body->ForgetKeyframes(); @@ -1769,8 +1769,8 @@ mjCJoint::mjCJoint(mjCModel* _model, mjCDef* _def) { CopyFromSpec(); // no previous state when a joint is created - qpos[0] = mjNAN; - qvel[0] = mjNAN; + qposadr_ = -1; + dofadr_ = -1; } @@ -1828,6 +1828,24 @@ 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}; + } + return qpos_.at(model->state_name_).data(); +} + + + +mjtNum* mjCJoint::qvel() { + if (qvel_.find(model->state_name_) == qvel_.end()) { + qvel_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0}; + } + return qvel_.at(model->state_name_).data(); +} + + + void mjCJoint::PointToLocal() { spec.element = static_cast(this); spec.name = &name; @@ -5349,7 +5367,8 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) { PointToLocal(); // no previous state when an actuator is created - act.push_back(mjNAN); + actadr_ = -1; + actdim_ = -1; } @@ -5374,8 +5393,7 @@ mjCActuator& mjCActuator::operator=(const mjCActuator& other) { void mjCActuator::ForgetKeyframes() { - act.clear(); - act.push_back(mjNAN); + act_.clear(); } @@ -5386,6 +5404,15 @@ 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); + } + return act_.at(model->state_name_); +} + + + void mjCActuator::PointToLocal() { spec.element = static_cast(this); spec.name = &name; diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 8eed97d9..bd39c7ef 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -16,6 +16,7 @@ #define MUJOCO_SRC_USER_USER_OBJECTS_H_ #include +#include #include #include #include @@ -401,10 +402,10 @@ class mjCJoint_ : public mjCBase { mjCBody* body; // joint's body // variable used for temporarily storing the state of the joint - int qposadr_; // address of dof in data->qpos - int dofadr_; // address of dof in data->qvel - mjtNum qpos[7]; // qpos at the previous step - mjtNum qvel[6]; // qvel at the previous step + int qposadr_; // address of dof in data->qpos + int dofadr_; // address of dof in data->qvel + std::map> qpos_; // qpos at the previous step + std::map> qvel_; // qvel at the previous step // variable-size data std::vector userdata_; @@ -443,6 +444,9 @@ class mjCJoint : public mjCJoint_, private mjsJoint { int nq() const { return nq(spec.type); } int nv() const { return nv(spec.type); } + mjtNum* qpos(); + mjtNum* qvel(); + private: int Compile(void); // compiler; return dofnum void PointToLocal(void); @@ -1390,9 +1394,9 @@ class mjCActuator_ : public mjCBase { int trnid[2]; // id of transmission target // variable used for temporarily storing the state of the actuator - int actadr_; // address of dof in data->act - int actdim_; // number of dofs in data->act - std::vector act; // act at the previous step + int actadr_; // address of dof in data->act + int actdim_; // number of dofs in data->act + std::map> act_; // act at the previous step // variable-size data std::string plugin_name; @@ -1431,6 +1435,8 @@ class mjCActuator : public mjCActuator_, private mjsActuator { bool is_forcelimited() const; bool is_actlimited() const; + std::vector& act(); + private: void Compile(void); // compiler void CopyFromSpec(); diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index e5c581d5..75280243 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -504,6 +504,7 @@ TEST_F(MujocoTest, AttachSame) { + )"; @@ -621,6 +622,7 @@ TEST_F(MujocoTest, AttachDifferent) { + )"; @@ -742,6 +744,7 @@ TEST_F(MujocoTest, AttachFrame) { + )";