Add keyframe support to attach.

Limitations:
- Supports attaching a single keyframe. If multiple keyframes are present in the child, only the first one will be attached.
- Does not support time, ctrl, mpos, mquat.
- If mjs_detachBody is used, all keyframes are removed.
PiperOrigin-RevId: 662517666
Change-Id: If6e3839471e1884afc14e1a75e7ce1957c2f3581

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;
   }
 }

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<T*>& 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 <class T>
+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 <class T>
+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<mjtNum>(const mjtNum* qpos,
+                                          const mjtNum* qvel,
+                                          const mjtNum* act);
+
+template void mjCModel::RestoreState<mjtNum>(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 <class T>
@@ -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; i<m->nq; 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; i<keys_.size(); i++) {
     keys_[i]->Compile(m);
   }

src/user/user_model.h

@@ -34,6 +34,13 @@
 typedef std::map<std::string, int, std::less<> > mjKeyMap;
 typedef std::array<mjKeyMap, mjNOBJECT> 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 <class T> void SaveState(const T* qpos, const T* qvel, const T* act);
+  template <class T> 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<std::string, mjCDef*> 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<mjKeyInfo> key_pending_;  // attached keyframes
 };
 #endif  // MUJOCO_SRC_USER_USER_MODEL_H_

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<mjCKey>(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<mjCKey>(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<mjCKey>(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); }

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<mjtNum> 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
 };
 

test/user/user_api_test.cc

@@ -421,13 +421,18 @@ static constexpr char xml_child[] = R"(
     </tendon>
 
     <actuator>
-      <position name="hinge" joint="hinge"/>
-      <position name="fixed" tendon="fixed"/>
+      <position name="hinge" joint="hinge" timeconst=".01"/>
+      <position name="fixed" tendon="fixed" timeconst=".01"/>
     </actuator>
 
     <contact>
       <exclude body1="body" body2="targetbody"/>
     </contact>
+
+    <keyframe>
+      <key name="two" qpos="2" act="2 2"/>
+      <key name="three" qpos="3" act="3 3"/>
+    </keyframe>
   </mujoco>)";
 
 TEST_F(MujocoTest, AttachSame) {
@@ -484,16 +489,22 @@ TEST_F(MujocoTest, AttachSame) {
     </tendon>
 
     <actuator>
-      <position name="hinge" joint="hinge"/>
-      <position name="fixed" tendon="fixed"/>
-      <position name="attached-hinge-1" joint="attached-hinge-1"/>
-      <position name="attached-fixed-1" tendon="attached-fixed-1"/>
+      <position name="hinge" joint="hinge" timeconst=".01"/>
+      <position name="fixed" tendon="fixed" timeconst=".01"/>
+      <position name="attached-hinge-1" joint="attached-hinge-1" timeconst=".01"/>
+      <position name="attached-fixed-1" tendon="attached-fixed-1" timeconst=".01"/>
     </actuator>
 
     <contact>
       <exclude body1="body" body2="targetbody"/>
       <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
     </contact>
+
+    <keyframe>
+      <key name="two" qpos="2 0" act="2 2 0 0"/>
+      <key name="three" qpos="3 0" act="3 3 0 0"/>
+      <key name="attached-two-1" qpos="0 2" act="0 0 2 2"/>
+    </keyframe>
   </mujoco>)";
 
   // create parent
@@ -600,13 +611,17 @@ TEST_F(MujocoTest, AttachDifferent) {
     </tendon>
 
     <actuator>
-      <position name="attached-hinge-1" joint="attached-hinge-1"/>
-      <position name="attached-fixed-1" tendon="attached-fixed-1"/>
+      <position name="attached-hinge-1" joint="attached-hinge-1" timeconst=".01"/>
+      <position name="attached-fixed-1" tendon="attached-fixed-1" timeconst=".01"/>
     </actuator>
 
     <contact>
       <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
     </contact>
+
+    <keyframe>
+      <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+    </keyframe>
   </mujoco>)";
 
   // model with one free sphere and a frame
@@ -717,13 +732,17 @@ TEST_F(MujocoTest, AttachFrame) {
     </tendon>
 
     <actuator>
-      <position name="attached-hinge-1" joint="attached-hinge-1"/>
-      <position name="attached-fixed-1" tendon="attached-fixed-1"/>
+      <position name="attached-hinge-1" joint="attached-hinge-1" timeconst=".01"/>
+      <position name="attached-fixed-1" tendon="attached-fixed-1" timeconst=".01"/>
     </actuator>
 
     <contact>
       <exclude body1="attached-body-1" body2="attached-targetbody-1"/>
     </contact>
+
+    <keyframe>
+      <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+    </keyframe>
   </mujoco>)";
 
   // model with one free sphere and a frame
@@ -790,6 +809,7 @@ void TestDetachBody(bool compile) {
       <body name="ignore"/>
       <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
     </worldbody>
+
     <sensor>
       <framepos name="ignore" objtype="body" objname="ignore"/>
     </sensor>