Add mj_copyState

PiperOrigin-RevId: 836588419
Change-Id: I7609e121dc0ac697d4d015d4244bdd5962650def

doc/APIreference/functions.rst

@@ -291,6 +291,15 @@ is not a subset of the bits set in ``srcsig``.
 Copy concatenated state components specified by ``sig`` from  ``state`` into ``d``. The bits of the integer
 ``sig`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``sig`` is invalid.
 
+.. _mj_copyState:
+
+`mj_copyState <#mj_copyState>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_copyState
+
+Copy state from src to dst.
+
 .. _mj_setKeyframe:
 
 `mj_setKeyframe <#mj_setKeyframe>`__

doc/changelog.rst

@@ -20,6 +20,7 @@ General
 - Increase Windows stack size to 16MB to enable models with deep nested body hierarchies.
 - Added a new :ref:`mj_extractState` function that allows a subset of a state that was previously returned by
   :ref:`mj_getState` to be extracted without having to be written back into ``mjData`` first.
+- Added a new :ref:`mj_copyState` function that copies state components from one ``mjData`` to another.
 - Tendon paths can now be queried from Python via ``MjsTendon.path``, the returned object
   is iterable and indexing it will give the ``MjsWrap`` at the given index in the path.
 - ``MjsWrap`` now exposes:

doc/includes/references.h

@@ -3181,6 +3181,7 @@ void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int
 void mj_extractState(const mjModel* m, const mjtNum* src, unsigned int srcsig,
                      mjtNum* dst, unsigned int dstsig);
 void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int sig);
+void mj_copyState(const mjModel* m, const mjData* src, mjData* dst, unsigned int sig);
 void mj_setKeyframe(mjModel* m, const mjData* d, int k);
 int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
 int mj_isPyramidal(const mjModel* m);

include/mujoco/mujoco.h

@@ -477,6 +477,9 @@ MJAPI void mj_extractState(const mjModel* m, const mjtNum* src, unsigned int src
 // Set state.
 MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int sig);
 
+// Copy state from src to dst.
+MJAPI void mj_copyState(const mjModel* m, const mjData* src, mjData* dst, unsigned int sig);
+
 // Copy current state to the k-th model keyframe.
 MJAPI void mj_setKeyframe(mjModel* m, const mjData* d, int k);
 

python/mujoco/bindings_test.py

@@ -806,6 +806,13 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     # Expect next states to be equal.
     np.testing.assert_array_equal(state1a, state1b)
 
+    # Test mj_copyState
+    data2 = mujoco.MjData(self.model)
+    mujoco.mj_copyState(self.model, self.data, data2, sig)
+    state1c = np.empty(size, np.float64)
+    mujoco.mj_getState(self.model, data2, state1c, sig)
+    np.testing.assert_array_equal(state1a, state1c)
+
   def test_mj_setKeyframe(self):  # pylint: disable=invalid-name
     mujoco.mj_step(self.model, self.data)
 

python/mujoco/functions.cc

@@ -350,6 +350,7 @@ PYBIND11_MODULE(_functions, pymodule) {
         }
         return InterceptMjErrors(::mj_setState)(m, d, state.data(), sig);
       });
+  Def<traits::mj_copyState>(pymodule);
   Def<traits::mj_setKeyframe>(pymodule);
   Def<traits::mj_addContact>(pymodule);
   Def<traits::mj_isPyramidal>(pymodule);

python/mujoco/introspect/functions.py

@@ -2520,6 +2520,36 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Set state.',
      )),
+    ('mj_copyState',
+     FunctionDecl(
+         name='mj_copyState',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='src',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='dst',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='sig',
+                 type=ValueType(name='unsigned int'),
+             ),
+         ),
+         doc='Copy state from src to dst.',
+     )),
     ('mj_setKeyframe',
      FunctionDecl(
          name='mj_setKeyframe',

src/engine/engine_support.c

@@ -265,6 +265,36 @@ void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int
 }
 
 
+// copy state from src to dst
+void mj_copyState(const mjModel* m, const mjData* src, mjData* dst, unsigned int sig) {
+  if (sig >= (1<<mjNSTATE)) {
+    mjERROR("invalid state signature %u >= 2^mjNSTATE", sig);
+  }
+
+  for (int i=0; i < mjNSTATE; i++) {
+    mjtState element = 1<<i;
+    if (element & sig) {
+      int size = mj_stateElemSize(m, element);
+
+      // special handling of eq_active (mjtByte)
+      if (element == mjSTATE_EQ_ACTIVE) {
+        int neq = m->neq;
+        for (int j=0; j < neq; j++) {
+          dst->eq_active[j] = src->eq_active[j];
+        }
+      }
+
+      // regular state components (mjtNum)
+      else {
+        mjtNum* dst_ptr = mj_stateElemPtr(m, dst, element);
+        const mjtNum* src_ptr = mj_stateElemConstPtr(m, src, element);
+        mju_copy(dst_ptr, src_ptr, size);
+      }
+    }
+  }
+}
+
+
 // copy current state to the k-th model keyframe
 void mj_setKeyframe(mjModel* m, const mjData* d, int k) {
   // check keyframe index

src/engine/engine_support.h

@@ -48,6 +48,9 @@ MJAPI void mj_extractState(const mjModel* m, const mjtNum* src, unsigned int src
 // set state
 MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int sig);
 
+// copy state from src to dst
+MJAPI void mj_copyState(const mjModel* m, const mjData* src, mjData* dst, unsigned int sig);
+
 // copy current state to the k-th model keyframe
 MJAPI void mj_setKeyframe(mjModel* m, const mjData* d, int k);
 

test/engine/engine_support_test.cc

@@ -37,6 +37,7 @@ using ::testing::ContainsRegex;  // NOLINT
 using ::testing::DoubleNear;
 using ::testing::Eq;
 using ::testing::MatchesRegex;
+using ::testing::Ne;
 using ::testing::NotNull;
 using ::testing::Pointwise;
 
@@ -744,6 +745,51 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
   mj_deleteModel(model);
 }
 
+TEST_F(SupportTest, CopyState) {
+  const std::string xml_path = GetTestDataFilePath(kDefaultModel);
+  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+
+  mjData* src = mj_makeData(m);
+  mjData* dst = mj_makeData(m);
+
+  // init both datas to default
+  mj_resetData(m, src);
+  mj_resetData(m, dst);
+
+  // modify d_src
+  src->time = 1.23;
+  for (int i=0; i < m->nq; ++i) src->qpos[i] = i*0.1;
+  for (int i=0; i < m->nv; ++i) src->qvel[i] = i*0.2;
+  for (int i=0; i < m->na; ++i) src->act[i] = i*0.3;
+  for (int i=0; i < m->nu; ++i) src->ctrl[i] = i*0.4;
+
+  for (int i=0; i < m->neq; ++i) src->eq_active[i] = 1 - m->eq_active0[i];
+
+  // check that states differ
+  EXPECT_NE(src->time, dst->time);
+  EXPECT_THAT(AsVector(src->qpos, m->nq), Ne(AsVector(dst->qpos, m->nq)));
+  EXPECT_THAT(AsVector(src->ctrl, m->nu), Ne(AsVector(dst->ctrl, m->nu)));
+
+  // copy state with signature
+  int signature = mjSTATE_FULLPHYSICS | mjSTATE_EQ_ACTIVE;
+  mj_copyState(m, src, dst, signature);
+
+  // check copied components
+  EXPECT_EQ(dst->time, src->time);
+  EXPECT_EQ(AsVector(dst->qpos, m->nq), AsVector(src->qpos, m->nq));
+  EXPECT_EQ(AsVector(dst->qvel, m->nv), AsVector(src->qvel, m->nv));
+  EXPECT_EQ(AsVector(dst->act, m->na), AsVector(src->act, m->na));
+  EXPECT_EQ(AsVector(dst->eq_active, m->neq), AsVector(src->eq_active, m->neq));
+
+  // check non-copied components (CTRL not in signature)
+  EXPECT_THAT(AsVector(dst->ctrl, m->nu), Ne(AsVector(src->ctrl, m->nu)));
+  EXPECT_EQ(AsVector(dst->ctrl, m->nu), vector<mjtNum>(m->nu, 0.0));
+
+  mj_deleteData(src);
+  mj_deleteData(dst);
+  mj_deleteModel(m);
+}
+
 TEST_F(SupportTest, ExtractState) {
   const std::string xml_path = GetTestDataFilePath(kDefaultModel);
   mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);

unity/Runtime/Bindings/MjBindings.cs

@@ -6660,6 +6660,9 @@ public static unsafe extern void mj_extractState(mjModel_* m, double* src, uint
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_setState(mjModel_* m, mjData_* d, double* state, uint sig);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_copyState(mjModel_* m, mjData_* src, mjData_* dst, uint sig);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_setKeyframe(mjModel_* m, mjData_* d, int k);
 

wasm/codegen/generated/bindings.cc

@@ -10404,6 +10404,10 @@ void mj_setState_wrapper(const MjModel& m, MjData& d, const NumberArray& state,
   mj_setState(m.get(), d.get(), state_.data(), sig);
 }
 
+void mj_copyState_wrapper(const MjModel& m, const MjData& src, MjData& dst, unsigned int sig) {
+  mj_copyState(m.get(), src.get(), dst.get(), sig);
+}
+
 void mj_setKeyframe_wrapper(MjModel& m, const MjData& d, int k) {
   mj_setKeyframe(m.get(), d.get(), k);
 }
@@ -12451,6 +12455,7 @@ EMSCRIPTEN_BINDINGS(mujoco_functions) {
   function("mj_getState", &mj_getState_wrapper);
   function("mj_extractState", &mj_extractState_wrapper);
   function("mj_setState", &mj_setState_wrapper);
+  function("mj_copyState", &mj_copyState_wrapper);
   function("mj_setKeyframe", &mj_setKeyframe_wrapper);
   function("mj_addContact", &mj_addContact_wrapper);
   function("mj_isPyramidal", &mj_isPyramidal_wrapper);