Check keyframe size before appending it during attach.

Fixes #2365. PiperOrigin-RevId: 718455748 Change-Id: I20c56b7c6fbd868962d19682e72e6294edb6893e
google-deepmind · Jan 22, 2025 · 21d9790 · 21d9790
1 parent aa505a7
commit 21d9790
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
@@ -3473,7 +3473,30 @@ void mjCModel::StoreKeyframes(mjCModel* dest) {
     info.mpos = !key->spec_mpos_.empty();
     info.mquat = !key->spec_mquat_.empty();
     dest->key_pending_.push_back(info);
-    ResizeKeyframe(key, qpos0.data(), body_pos0.data(), body_quat0.data());
+    if (!key->spec_qpos_.empty() && key->spec_qpos_.size() != nq) {
+      throw mjCError(nullptr, "Keyframe '%s' has invalid qpos size, got %d, should be %d",
+                     key->name.c_str(), key->spec_qpos_.size(), nq);
+    }
+    if (!key->spec_qvel_.empty() && key->spec_qvel_.size() != nv) {
+      throw mjCError(nullptr, "Keyframe %s has invalid qvel size, got %d, should be %d",
+                     key->name.c_str(), key->spec_qvel_.size(), nv);
+    }
+    if (!key->spec_act_.empty() && key->spec_act_.size() != na) {
+      throw mjCError(nullptr, "Keyframe %s has invalid act size, got %d, should be %d",
+                     key->name.c_str(), key->spec_act_.size(), na);
+    }
+    if (!key->spec_ctrl_.empty() && key->spec_ctrl_.size() != nu) {
+      throw mjCError(nullptr, "Keyframe %s has invalid ctrl size, got %d, should be %d",
+                     key->name.c_str(), key->spec_ctrl_.size(), nu);
+    }
+    if (!key->spec_mpos_.empty() && key->spec_mpos_.size() != 3*nmocap) {
+      throw mjCError(nullptr, "Keyframe %s has invalid mpos size, got %d, should be %d",
+                     key->name.c_str(), key->spec_mpos_.size(), 3*nmocap);
+    }
+    if (!key->spec_mquat_.empty() && key->spec_mquat_.size() != 4*nmocap) {
+      throw mjCError(nullptr, "Keyframe %s has invalid mquat size, got %d, should be %d",
+                     key->name.c_str(), key->spec_mquat_.size(), 4*nmocap);
+    }
     SaveState(info.name, key->spec_qpos_.data(), key->spec_qvel_.data(),
               key->spec_act_.data(), key->spec_ctrl_.data(),
               key->spec_mpos_.data(), key->spec_mquat_.data());

diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc
@@ -2056,6 +2056,32 @@ TEST_F(MujocoTest, ResizeParentKeyframe) {
   mj_deleteModel(expected);
 }
 
+TEST_F(MujocoTest, KeyframeSizeError) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <replicate count="2" offset="0 1 0">
+        <body name="ball" pos="0 0 1">
+          <geom type="sphere" size="0.1"/>
+          <joint/>
+        </body>
+      </replicate>
+    </worldbody>
+
+    <keyframe>
+      <key name="valid_qpos" qpos="0.5"/>
+      <key name="invalid_qpos" qpos="0.5 0.25"/>
+    </keyframe>
+  </mujoco>
+  )";
+
+  std::array<char, 1000> er;
+  mjSpec* spec = mj_parseXMLString(xml, 0, er.data(), er.size());
+  EXPECT_THAT(spec, IsNull());
+  EXPECT_THAT(er.data(), HasSubstr(
+      "Keyframe 'invalid_qpos' has invalid qpos size, got 2, should be 1"));
+}
+
 TEST_F(MujocoTest, DifferentUnitsAllowed) {
   static constexpr char gchild_xml[] = R"(
   <mujoco model="gchild">