Attach multiple keyframes.

Still missing: support for ctrl, mocap, and time.

PiperOrigin-RevId: 662871569
Change-Id: I8a418c4007663e3063d7146bd72db4eba78ddc5d

doc/changelog.rst

@@ -9,6 +9,8 @@ General
 ^^^^^^^
 1. Add :ref:`mjSpec` option for creating a texture from a buffer.
 2. :ref:`shellinertia <body-geom-shellinertia>` is now supported by all geom types.
+3. Add support for :ref:`attaching<meAttachment>` keyframes. Note: this only supports keyframe containing qpos, qvel,
+   and act.
 
 Version 3.2.2 (Aug 8, 2024)
 ---------------------------

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 <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 (joint->qposadr_ == -1 || joint->dofadr_ == -1) {
-    if (qvel) mjuu_copyvec(joint->qvel, qvel + joint->dofadr_, joint->nv());
+      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) {
+    if (actuator->actadr_ != -1 && actuator->actdim_ != -1 && act) {
-      actuator->act.assign(actuator->actdim_, 0);
+      actuator->act().assign(actuator->actdim_, 0);
-      mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actdim_);
+      mjuu_copyvec(actuator->act().data(), act + actuator->actadr_, actuator->actdim_);
     }
   }
 }
@@ -2893,21 +2899,21 @@ template <class T>
 void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act) {
   for (auto joint : joints_) {
     if (qpos) {
-      if (mjuu_defined(joint->qpos[0])) {
+      if (mjuu_defined(joint->qpos()[0])) {
-        mjuu_copyvec(qpos + joint->qposadr_, joint->qpos, joint->nq());
+        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) {
+    if (mjuu_defined(joint->qvel()[0]) && qvel) {
-      mjuu_copyvec(qvel + joint->dofadr_, joint->qvel, joint->nv());
+      mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(), joint->nv());
     }
   }
 
   // restore act
   for (auto actuator : actuators_) {
-    if (mjuu_defined(actuator->act[0]) && act) {
+    if (!actuator->act().empty() && mjuu_defined(actuator->act()[0]) && act) {
-      mjuu_copyvec(act + actuator->actadr_, actuator->act.data(), actuator->actdim_);
+      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());
   }
 

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<mjKeyInfo> key_pending_;  // attached keyframes
+
+  std::string state_name_;
 };
 #endif  // MUJOCO_SRC_USER_USER_MODEL_H_

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->qpos_.clear();
-    joint->qvel[0] = mjNAN;
+    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;
+  qposadr_ = -1;
-  qvel[0] = mjNAN;
+  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<mjsElement*>(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_.clear();
-  act.push_back(mjNAN);
 }
 
 
@@ -5386,6 +5404,15 @@ bool mjCActuator::is_actlimited() const { return islimited(actlimited, actrange)
 
 
 
+std::vector<mjtNum>& mjCActuator::act() {
+  if (act_.find(model->state_name_) == act_.end()) {
+    act_[model->state_name_] = std::vector<mjtNum>(model->nu, mjNAN);
+  }
+  return act_.at(model->state_name_);
+}
+
+
+
 void mjCActuator::PointToLocal() {
   spec.element = static_cast<mjsElement*>(this);
   spec.name = &name;

src/user/user_objects.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_USER_USER_OBJECTS_H_
 
 #include <cstddef>
+#include <array>
 #include <cstdlib>
 #include <functional>
 #include <map>
@@ -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 qposadr_;                                        // address of dof in data->qpos
-  int dofadr_;                     // address of dof in data->qvel
+  int dofadr_;                                         // address of dof in data->qvel
-  mjtNum qpos[7];                  // qpos at the previous step
+  std::map<std::string, std::array<mjtNum, 7>> qpos_;  // qpos at the previous step
-  mjtNum qvel[6];                  // qvel at the previous step
+  std::map<std::string, std::array<mjtNum, 6>> qvel_;  // qvel at the previous step
 
   // variable-size data
   std::vector<double> 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 actadr_;                                      // address of dof in data->act
-  int actdim_;              // number of dofs in data->act
+  int actdim_;                                      // number of dofs in data->act
-  std::vector<mjtNum> act;  // act at the previous step
+  std::map<std::string, std::vector<mjtNum>> 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<mjtNum>& act();
+
  private:
   void Compile(void);                       // compiler
   void CopyFromSpec();

test/user/user_api_test.cc

@@ -504,6 +504,7 @@ TEST_F(MujocoTest, AttachSame) {
       <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"/>
+      <key name="attached-three-1" qpos="0 3" act="0 0 3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -621,6 +622,7 @@ TEST_F(MujocoTest, AttachDifferent) {
 
     <keyframe>
       <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+      <key name="attached-three-1" qpos="0 0 0 1 0 0 0 3" act="3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -742,6 +744,7 @@ TEST_F(MujocoTest, AttachFrame) {
 
     <keyframe>
       <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+      <key name="attached-three-1" qpos="0 0 0 1 0 0 0 3" act="3 3"/>
     </keyframe>
   </mujoco>)";