diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 3f5a28a5..5f33b0f8 100644 --- a/doc/APIreference/functions.rst +++ b/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>`__ diff --git a/doc/changelog.rst b/doc/changelog.rst index 8dcff282..2b2e1671 100644 --- a/doc/changelog.rst +++ b/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: diff --git a/doc/includes/references.h b/doc/includes/references.h index b3e36fe5..e7ce9cba 100644 --- a/doc/includes/references.h +++ b/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); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 4c816e91..c5b50aa4 100644 --- a/include/mujoco/mujoco.h +++ b/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); diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 2b3856ba..b15dc261 100644 --- a/python/mujoco/bindings_test.py +++ b/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) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 9aae704f..f03d28ff 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -350,6 +350,7 @@ PYBIND11_MODULE(_functions, pymodule) { } return InterceptMjErrors(::mj_setState)(m, d, state.data(), sig); }); + Def(pymodule); Def(pymodule); Def(pymodule); Def(pymodule); diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py index f5b63c75..4d1823e5 100644 --- a/python/mujoco/introspect/functions.py +++ b/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', diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 63678380..56361467 100644 --- a/src/engine/engine_support.c +++ b/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<= 2^mjNSTATE", sig); + } + + for (int i=0; i < mjNSTATE; i++) { + mjtState element = 1<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 diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index 283028d1..089b50fa 100644 --- a/src/engine/engine_support.h +++ b/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); diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index 89b2bd8d..036179b9 100644 --- a/test/engine/engine_support_test.cc +++ b/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(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); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index a1c13b76..8ae4e76f 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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); diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index f1580184..9102eb68 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/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);