From f1d4f8f0a17e61f8f9439fbbd517dba5732ecb6e Mon Sep 17 00:00:00 2001
From: Nimrod Gileadi <nimrod@google.com>
Date: Fri, 16 Feb 2024 05:07:57 -0800
Subject: [PATCH] Test that mj_resetDataKeyframe treats negative keyframe IDs
 as a normal reset.

PiperOrigin-RevId: 607655812
Change-Id: Ia9bb1031468f04e8f9563b9228a46983356dafef
---
 doc/APIreference/functions.rst |  2 +-
 include/mujoco/mujoco.h        |  2 +-
 introspect/functions.py        |  2 +-
 src/engine/engine_io.c         |  2 +-
 src/engine/engine_io.h         |  2 +-
 test/user/user_objects_test.cc | 22 ++++++++++++++++++++++
 6 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index 04d8802a..0abccf7a 100644
--- a/doc/APIreference/functions.rst
+++ b/doc/APIreference/functions.rst
@@ -1280,7 +1280,7 @@ mj_resetDataKeyframe
 
 .. mujoco-include:: mj_resetDataKeyframe
 
-Reset data, set fields from specified keyframe.
+Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 
 .. _mj_markStack:
 
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 78bc9178..ecfaee76 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -182,7 +182,7 @@ MJAPI void mj_resetData(const mjModel* m, mjData* d);
 // Reset data to defaults, fill everything else with debug_value.
 MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
 
-// Reset data, set fields from specified keyframe.
+// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
 // Mark a new frame on the mjData stack.
diff --git a/introspect/functions.py b/introspect/functions.py
index 0e9d2568..9ad7bd08 100644
--- a/introspect/functions.py
+++ b/introspect/functions.py
@@ -675,7 +675,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  type=ValueType(name='int'),
              ),
          ),
-         doc='Reset data, set fields from specified keyframe.',
+         doc='Reset data. If 0 <= key < nkey, set fields from specified keyframe.',  # pylint: disable=line-too-long
      )),
     ('mj_markStack',
      FunctionDecl(
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index bf99d38f..7704e363 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -1701,7 +1701,7 @@ void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value) {
 
 
 
-// reset data, set fields from specified keyframe
+// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) {
   _resetData(m, d, 0);
 
diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h
index 4f8b530c..0bc8aaed 100644
--- a/src/engine/engine_io.h
+++ b/src/engine/engine_io.h
@@ -104,7 +104,7 @@ MJAPI void mj_resetData(const mjModel* m, mjData* d);
 // set data to defaults, fill everything else with debug_value
 MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
 
-// reset data, set fields from specified keyframe
+// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
 // mjData arena allocate
diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc
index d652a54c..f77e52f5 100644
--- a/test/user/user_objects_test.cc
+++ b/test/user/user_objects_test.cc
@@ -433,6 +433,28 @@ TEST_F(KeyframeTest, ResetDataKeyframe) {
   mj_deleteModel(model);
 }
 
+TEST_F(KeyframeTest, ResetDataKeyframeAcceptsNegativeKeyframe) {
+  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);
+  data->time = data->qpos[0] = data->qvel[0] = data->act[0] = data->ctrl[0] =
+      data->mocap_pos[0] = data->mocap_quat[0] = 1337.0;
+
+  mj_resetDataKeyframe(model, data, -1);
+
+  EXPECT_EQ(data->time, 0.0);
+  EXPECT_EQ(data->qpos[0], 0.0);
+  EXPECT_EQ(data->qvel[0], 0.0);
+  EXPECT_EQ(data->act[0], 0.0);
+  EXPECT_EQ(data->ctrl[0], 0.0);
+  EXPECT_EQ(data->mocap_pos[0], 0.0);
+  EXPECT_EQ(data->mocap_quat[0], 1.0);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 TEST_F(KeyframeTest, BadSize) {
   static constexpr char xml[] = R"(
   <mujoco>