Add control values to keyframes.

- Add basic test for keyframes.
- Add missing documentation for keyframe mocap positions and quaternions.

PiperOrigin-RevId: 459021649
Change-Id: I91cf7ecbddc6262e8c72a868eeb82d627f389fb3

doc/APIreference.rst

@@ -1725,6 +1725,7 @@ mjModel
        mjtNum*   key_act;              // key activation                           (nkey x na)
        mjtNum*   key_mpos;             // key mocap position                       (nkey x 3*nmocap)
        mjtNum*   key_mquat;            // key mocap quaternion                     (nkey x 4*nmocap)
+       mjtNum*   key_ctrl;             // key control                              (nkey x nu)
 
        // names
        int*      name_bodyadr;         // body name pointers                       (nbody x 1)

doc/XMLreference.rst

@@ -1411,7 +1411,9 @@ if present, the parser ignores it. The symbols in the second column of the table
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 |                          |    |    | :at:`name`              | :at:`time`              | :at:`qpos`              | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
-|                          |    |    | :at:`qvel`              | :at:`act`               |                         | |
+|                          |    |    | :at:`qvel`              | :at:`act`               | :at:`ctrl`              | |
+|                          |    |    +-------------------------+-------------------------+-------------------------+ |
+|                          |    |    | :at:`mpos`              | :at:`mquat`             |                         | |
 |                          |    |    +-------------------------+-------------------------+-------------------------+ |
 +--------------------------+----+------------------------------------------------------------------------------------+
 
@@ -5194,3 +5196,9 @@ This element sets the data for one of the keyframes. They are set in the order i
    Vector of joint velocities, copied into mjData.qvel when the simulation state is set to this keyframe.
 :at:`act`: :at-val:`real(mjModel.na), "0 0 ..."`
    Vector of actuator activations, copied into mjData.act when the simulation state is set to this keyframe.
+:at:`ctrl`: :at-val:`real(mjModel.nu), "0 0 ..."`
+   Vector of controls, copied into mjData.ctrl when the simulation state is set to this keyframe.
+:at:`mpos`: :at-val:`real(3*mjModel.nmocap), default = mjModel.body_pos`
+   Vector of mocap body positions, copied into mjData.mocap_pos when the simulation state is set to this keyframe.
+:at:`mquat`: :at-val:`real(4*mjModel.nmocap), default = mjModel.body_quat`
+   Vector of mocap body quaternions, copied into mjData.mocap_quat when the simulation state is set to this keyframe.

include/mujoco/mjmodel.h

@@ -888,6 +888,7 @@ struct mjModel_ {
   mjtNum*   key_act;              // key activation                           (nkey x na)
   mjtNum*   key_mpos;             // key mocap position                       (nkey x 3*nmocap)
   mjtNum*   key_mquat;            // key mocap quaternion                     (nkey x 4*nmocap)
+  mjtNum*   key_ctrl;             // key control                              (nkey x nu)
 
   // names
   int*      name_bodyadr;         // body name pointers                       (nbody x 1)

include/mujoco/mjxmacro.h

@@ -135,6 +135,7 @@
     int nq = m->nq;                         \
     int nv = m->nv;                         \
     int na = m->na;                         \
+    int nu = m->nu;                         \
     int nmocap = m->nmocap;
 
 
@@ -395,6 +396,7 @@
     X( mjtNum,  key_act,               nkey,          MJ_M(na)             ) \
     X( mjtNum,  key_mpos,              nkey,          MJ_M(nmocap)*3       ) \
     X( mjtNum,  key_mquat,             nkey,          MJ_M(nmocap)*4       ) \
+    X( mjtNum,  key_ctrl,              nkey,          MJ_M(nu)             ) \
     X( int,     name_bodyadr,          nbody,         1                    ) \
     X( int,     name_jntadr,           njnt,          1                    ) \
     X( int,     name_geomadr,          ngeom,         1                    ) \

sample/simulate.cc

@@ -1442,6 +1442,7 @@ void uiEvent(mjuiState* state) {
         mju_copy(d->act, m->key_act+i*m->na, m->na);
         mju_copy(d->mocap_pos, m->key_mpos+i*3*m->nmocap, 3*m->nmocap);
         mju_copy(d->mocap_quat, m->key_mquat+i*4*m->nmocap, 4*m->nmocap);
+        mju_copy(d->ctrl, m->key_ctrl+i*m->nu, m->nu);
         mj_forward(m, d);
         profilerupdate();
         sensorupdate();
@@ -1456,6 +1457,7 @@ void uiEvent(mjuiState* state) {
         mju_copy(m->key_act+i*m->na, d->act, m->na);
         mju_copy(m->key_mpos+i*3*m->nmocap, d->mocap_pos, 3*m->nmocap);
         mju_copy(m->key_mquat+i*4*m->nmocap, d->mocap_quat, 4*m->nmocap);
+        mju_copy(m->key_ctrl+i*m->nu, d->ctrl, m->nu);
         break;
       }
     }

src/engine/engine_io.c

@@ -1016,6 +1016,7 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) {
     mju_copy(d->act,  m->key_act+ key*m->na, m->na);
     mju_copy(d->mocap_pos,  m->key_mpos+key*3*m->nmocap, 3*m->nmocap);
     mju_copy(d->mocap_quat, m->key_mquat+key*4*m->nmocap, 4*m->nmocap);
+    mju_copy(d->ctrl, m->key_ctrl+key*m->nu, m->nu);
   }
 }
 

src/engine/engine_print.c

@@ -592,7 +592,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
         k = 2;
       }
 
-    // print if nozero
+    // print if nonzero
     if (k==2) {
       fprintf(fp, "key_qvel%d   ", i);
       for (int j=0; j<m->nv; j++) {
@@ -665,6 +665,24 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
       fprintf(fp, "\n");
     }
 
+    // check ctrl for nonzero
+    for (int j=0; j<m->nu; j++) {
+      if (m->key_ctrl[i*m->nu + j]) {
+        k = 6;
+        break;
+      }
+    }
+
+    // print if nonzero
+    if (k==6) {
+      fprintf(fp, "key_ctrl%d   ", i);
+      for (int j=0; j<m->nu; j++) {
+        fprintf(fp, float_format, m->key_ctrl[i*m->nu + j]);
+      }
+      fprintf(fp, "\n");
+    }
+
+
     // new line if any data was written
     if (k) {
       fprintf(fp, "\n");

src/user/user_model.cc

@@ -1935,6 +1935,8 @@ void mjCModel::CopyObjects(mjModel* m) {
     for (j=0; j<nmocap; j++) {
       mju_normalize4(m->key_mquat+i*4*nmocap+4*j);
     }
+
+    copyvec(m->key_ctrl+i*nu, keys[i]->ctrl.data(), nu);
   }
 
   // save qpos0 in user model (to recognize changed key_qpos in write)
@@ -2878,6 +2880,9 @@ bool mjCModel::CopyBack(const mjModel* m) {
       copyvec(pk->mpos.data(), m->key_mpos + i*3*nmocap, 3*nmocap);
       copyvec(pk->mquat.data(), m->key_mquat + i*4*nmocap, 4*nmocap);
     }
+    if (nu) {
+      copyvec(pk->ctrl.data(), m->key_ctrl + i*nu, nu);
+    }
   }
 
   return true;

src/user/user_objects.cc

@@ -3980,6 +3980,7 @@ mjCKey::mjCKey(mjCModel* _model) {
   act.clear();
   mpos.clear();
   mquat.clear();
+  ctrl.clear();
 }
 
 
@@ -3991,6 +3992,7 @@ mjCKey::~mjCKey() {
   act.clear();
   mpos.clear();
   mquat.clear();
+  ctrl.clear();
 }
 
 
@@ -4063,4 +4065,15 @@ void mjCKey::Compile(const mjModel* m) {
   } else if (mquat.size()!=4*m->nmocap) {
     throw mjCError(this, "key %d: invalid mquat size", 0, id);
   }
+
+  // ctrl: allocate or check size
+  if (ctrl.empty()) {
+    ctrl.resize(m->nu);
+    for (i=0; i<m->nu; i++) {
+      ctrl[i] = 0;
+    }
+  } else if (ctrl.size()!=m->nu) {
+    throw mjCError(this, "key %d: invalid ctrl size", 0, id);
+  }
+
 }

src/user/user_objects.h

@@ -951,6 +951,7 @@ class mjCKey : public mjCBase {
   std::vector<double> act;        // act
   std::vector<double> mpos;       // mocap pos
   std::vector<double> mquat;      // mocap quat
+  std::vector<double> ctrl;       // ctrl
 
  private:
   mjCKey(mjCModel*);              // constructor

src/xml/xml_native_reader.cc

@@ -348,7 +348,7 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = {
 
     {"keyframe", "*", "0"},
     {"<"},
-        {"key", "*", "7", "name", "time", "qpos", "qvel", "act", "mpos", "mquat"},
+        {"key", "*", "8", "name", "time", "qpos", "qvel", "act", "mpos", "mquat", "ctrl"},
     {">"},
 {">"}
 };
@@ -2788,7 +2788,7 @@ void mjXReader::Keyframe(XMLElement* section) {
       mjuu_copyvec(pk->qvel.data(), data, n);
     }
 
-    // read qvel
+    // read act
     n = ReadAttr(elem, "act", 1000, data, text, false, false);
     if (n) {
       pk->act.resize(n);
@@ -2809,6 +2809,13 @@ void mjXReader::Keyframe(XMLElement* section) {
       mjuu_copyvec(pk->mquat.data(), data, n);
     }
 
+    // read ctrl
+    n = ReadAttr(elem, "ctrl", 1000, data, text, false, false);
+    if (n) {
+      pk->ctrl.resize(n);
+      mjuu_copyvec(pk->ctrl.data(), data, n);
+    }
+
     // advance to next element
     elem = elem->NextSiblingElement();
   }

src/xml/xml_native_writer.cc

@@ -1662,6 +1662,15 @@ void mjXWriter::Keyframe(XMLElement* root) {
       }
     }
 
+    // check ctrl and write
+    for (int j=0; j<model->nu; j++) {
+      if (pk->ctrl[j]!=0) {
+        WriteAttr(elem, "ctrl", model->nu, pk->ctrl.data());
+        change = true;
+        break;
+      }
+    }
+
     // remove elem if empty
     if (!change) {
       section->DeleteChild(elem);

test/user/testdata/keyframe.xml

@@ -0,0 +1,29 @@
+<mujoco>
+
+  <worldbody>
+    <geom type="plane" size="1 1 .01"/>
+    <light pos="0 0 2"/>
+    <body pos="0 -.1 .3">
+      <joint name="slide" type="slide" axis="1 0 0"/>
+      <geom type="box" size=".05 .05 .05"/>
+    </body>
+    <body pos="0 .1 .3" mocap="true">
+      <geom type="box" size=".05 .05 .05"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor name="stateless" joint="slide"/>
+    <general name="stateful" joint="slide" dyntype="integrator"/>
+  </actuator>
+
+  <keyframe>
+    <key time=".1"/>
+    <key qpos=".2"/>
+    <key qvel=".3"/>
+    <key act=".4"/>
+    <key ctrl=".5 .6"/>
+    <key mpos=".1 .2 .3"/>
+    <key mquat="1 1 1 1"/>
+  </keyframe>
+</mujoco>

test/user/user_objects_test.cc

@@ -28,9 +28,8 @@
 namespace mujoco {
 namespace {
 
-static std::vector<mjtNum> GetRow(const mjtNum* array, int ncolumn, int row) {
+std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
-  return std::vector<mjtNum>(array + ncolumn * row,
+  return std::vector<mjtNum>(array, array + n);
-                             array + ncolumn * (row + 1));
 }
 
 using ::testing::ElementsAre;
@@ -38,6 +37,59 @@ using ::testing::HasSubstr;
 using ::testing::IsNull;
 using ::testing::NotNull;
 
+// ------------------------ test keyframes -------------------------------------
+
+static const char* const kKeyframePath = "user/testdata/keyframe.xml";
+
+TEST_F(MujocoTest, KeyFrameTest) {
+  const std::string xml_path = GetTestDataFilePath(kKeyframePath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  ASSERT_THAT(model, NotNull());
+  EXPECT_EQ(model->nkey, 7);
+  EXPECT_EQ(model->key_time[0 * 1], 0.1);
+  EXPECT_EQ(model->key_qpos[1 * model->nq], 0.2);
+  EXPECT_EQ(model->key_qvel[2 * model->nv], 0.3);
+  EXPECT_EQ(model->key_act[3 * model->na], 0.4);
+  EXPECT_THAT(AsVector(model->key_ctrl + 4*model->nu, model->nu),
+              ElementsAre(0.5, 0.6));
+  EXPECT_THAT(AsVector(model->key_mpos + 3*model->nmocap*5, 3),
+              ElementsAre(.1, .2, .3));
+  EXPECT_THAT(AsVector(model->key_mquat + 4*model->nmocap*6, 4),
+              ElementsAre(.5, .5, .5, .5));
+  mj_deleteModel(model);
+}
+
+TEST_F(MujocoTest, ResetDataKeyframeTest) {
+  const std::string xml_path = GetTestDataFilePath(kKeyframePath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  ASSERT_THAT(model, NotNull());
+  mjData* data = mj_makeData(model);
+
+  mj_resetDataKeyframe(model, data, 0);
+  EXPECT_EQ(data->time, 0.1);
+
+  mj_resetDataKeyframe(model, data, 1);
+  EXPECT_EQ(data->qpos[0], 0.2);
+
+  mj_resetDataKeyframe(model, data, 2);
+  EXPECT_EQ(data->qvel[0], 0.3);
+
+  mj_resetDataKeyframe(model, data, 3);
+  EXPECT_EQ(data->act[0], 0.4);
+
+  mj_resetDataKeyframe(model, data, 4);
+  EXPECT_EQ(data->ctrl[0], 0.5);
+  EXPECT_EQ(data->ctrl[1], 0.6);
+
+  mj_resetDataKeyframe(model, data, 5);
+  EXPECT_THAT(AsVector(data->mocap_pos, 3), ElementsAre(.1, .2, .3));
+
+  mj_resetDataKeyframe(model, data, 6);
+  EXPECT_THAT(AsVector(data->mocap_quat, 4), ElementsAre(.5, .5, .5, .5));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
 
 // ------------- test relative frame sensor compilation-------------------------
 
@@ -214,10 +266,10 @@ TEST_F(QuatNorm, QuatNotNormalized) {
   )";
   std::array<char, 1024> error;
   mjModel* m = LoadModelFromString(xml, error.data(), error.size());
-  EXPECT_THAT(GetRow(m->body_quat, 4, 1), ElementsAre(1./5, 2./5, 2./5, 4./5));
+  EXPECT_THAT(AsVector(m->body_quat+4, 4), ElementsAre(1./5, 2./5, 2./5, 4./5));
-  EXPECT_THAT(GetRow(m->geom_quat, 4, 0), ElementsAre(1./5, 2./5, 2./5, 4./5));
+  EXPECT_THAT(AsVector(m->geom_quat, 4), ElementsAre(1./5, 2./5, 2./5, 4./5));
-  EXPECT_THAT(GetRow(m->site_quat, 4, 0), ElementsAre(1./5, 2./5, 2./5, 4./5));
+  EXPECT_THAT(AsVector(m->site_quat, 4), ElementsAre(1./5, 2./5, 2./5, 4./5));
-  EXPECT_THAT(GetRow(m->cam_quat, 4, 0), ElementsAre(1./5, 2./5, 2./5, 4./5));
+  EXPECT_THAT(AsVector(m->cam_quat, 4), ElementsAre(1./5, 2./5, 2./5, 4./5));
   mj_deleteModel(m);
 }
 

unity/Runtime/Bindings/MujocoBindings.cs

@@ -2147,6 +2147,7 @@ public unsafe struct mjModel_ {
   public double* key_act;
   public double* key_mpos;
   public double* key_mquat;
+  public double* key_ctrl;
   public int* name_bodyadr;
   public int* name_jntadr;
   public int* name_geomadr;