diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 6136d2af..9e2f369e 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -259,6 +259,17 @@ correspond to element fields of :ref:`mjtState`. Copy concatenated state components specified by ``sig`` from ``d`` into ``state``. The bits of the integer ``sig`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``sig`` is invalid. +.. _mj_extractState: + +`mj_extractState <#mj_extractState>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mj_extractState + +Extract into ``dst`` the subset of components specified by ``dstsig`` from a state ``src`` previously obtained via +:ref:`mj_getState` with components specified by ``srcsig``. Fails with :ref:`mju_error` if the bits set in ``dstsig`` +is not a subset of the bits set in ``srcsig``. + .. _mj_setState: `mj_setState <#mj_setState>`__ diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 1cd1bca3..0fae1683 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -195,6 +195,12 @@ correspond to element fields of :ref:`mjtState`. Copy concatenated state components specified by ``sig`` from ``d`` into ``state``. The bits of the integer ``sig`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``sig`` is invalid. +.. _mj_extractState: + +Extract into ``dst`` the subset of components specified by ``dstsig`` from a state ``src`` previously obtained via +:ref:`mj_getState` with components specified by ``srcsig``. Fails with :ref:`mju_error` if the bits set in ``dstsig`` +is not a subset of the bits set in ``srcsig``. + .. _mj_setState: Copy concatenated state components specified by ``sig`` from ``state`` into ``d``. The bits of the integer diff --git a/doc/changelog.rst b/doc/changelog.rst index 27c4c12b..902400de 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,8 @@ General - Raise an error if there are name collisions also during parsing. - 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. Version 3.3.7 (October 13, 2025) ----------------------------------- diff --git a/doc/includes/references.h b/doc/includes/references.h index 464b4751..e4b8671b 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3118,6 +3118,8 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian); int mj_stateSize(const mjModel* m, unsigned int sig); void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int sig); +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_setKeyframe(mjModel* m, const mjData* d, int k); int mj_addContact(const mjModel* m, mjData* d, const mjContact* con); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 42272804..b5c422c2 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -465,6 +465,10 @@ MJAPI int mj_stateSize(const mjModel* m, unsigned int sig); // Get state. MJAPI void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int sig); +// Extract a subset of components from a state previously obtained via mj_getState. +MJAPI void mj_extractState(const mjModel* m, const mjtNum* src, unsigned int srcsig, + mjtNum* dst, unsigned int dstsig); + // Set state. MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int sig); diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 480af1fc..78bf94b3 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -329,10 +329,24 @@ PYBIND11_MODULE(_functions, pymodule) { } return InterceptMjErrors(::mj_getState)(m, d, state.data(), sig); }); + Def( + pymodule, + [](const raw::MjModel* m, + Eigen::Ref src, unsigned int srcsig, + Eigen::Ref dst, unsigned int dstsig) { + if (src.size() != mj_stateSize(m, srcsig)) { + throw py::type_error("src size should equal mj_stateSize(m, srcsig)"); + } + if (dst.size() != mj_stateSize(m, dstsig)) { + throw py::type_error("dst size should equal mj_stateSize(m, dstsig)"); + } + return InterceptMjErrors(::mj_extractState)(m, src.data(), srcsig, + dst.data(), dstsig); + }); Def( pymodule, [](const raw::MjModel* m, raw::MjData* d, - const Eigen::Ref state, unsigned int sig) { + Eigen::Ref state, unsigned int sig) { if (state.size() != mj_stateSize(m, sig)) { throw py::type_error("state size should equal mj_stateSize(m, sig)"); } diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py index 6d33975a..d791436d 100644 --- a/python/mujoco/introspect/functions.py +++ b/python/mujoco/introspect/functions.py @@ -2416,6 +2416,40 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Get state.', )), + ('mj_extractState', + FunctionDecl( + name='mj_extractState', + 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='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='srcsig', + type=ValueType(name='unsigned int'), + ), + FunctionParameterDecl( + name='dst', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='dstsig', + type=ValueType(name='unsigned int'), + ), + ), + doc='Extract a subset of components from a state previously obtained via mj_getState.', # pylint: disable=line-too-long + )), ('mj_setState', FunctionDecl( name='mj_setState', diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 6c07b760..e1cca49e 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -211,6 +211,28 @@ void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int } +// extract a sub-state from a state +void mj_extractState(const mjModel* m, const mjtNum* src, unsigned int srcsig, + mjtNum* dst, unsigned int dstsig) { + if ((srcsig & dstsig) != dstsig) { + mjERROR("dstsig is not a subset of srcsig"); + return; + } + + for (int i=0; i < mjNSTATE; i++) { + mjtState element = 1<= (1< #include #include #include +#include #include #include @@ -742,6 +744,72 @@ TEST_F(SupportTest, GetSetStateStepEqual) { mj_deleteModel(model); } +TEST_F(SupportTest, ExtractState) { + const std::string xml_path = GetTestDataFilePath(kDefaultModel); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + + // make distribution using seed + std::mt19937_64 rng; + rng.seed(3); + std::normal_distribution dist(0, .01); + + // set controls and applied joint forces to random values + for (int i=0; i < model->nu; i++) data->ctrl[i] = dist(rng); + for (int i=0; i < model->nv; i++) data->qfrc_applied[i] = dist(rng); + for (int i=0; i < model->neq; i++) data->eq_active[i] = dist(rng) > 0; + + // take one step + mj_step(model, data); + + // take a state that will be used as src + int srcsig = mjSTATE_TIME | mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_CTRL; + int srcsize = mj_stateSize(model, srcsig); + vector srcstate(srcsize); + mj_getState(model, data, srcstate.data(), srcsig); + + // extract a subset consisting of only a single bit in srcsig + int dstsig1 = mjSTATE_CTRL; + int dstsize1 = mj_stateSize(model, dstsig1); + EXPECT_LT(dstsize1, srcsize); + EXPECT_EQ(dstsize1, model->nu); + vector dststate1(dstsize1); + mj_extractState(model, srcstate.data(), srcsig, dststate1.data(), dstsig1); + EXPECT_EQ(dststate1, AsVector(data->ctrl, model->nu)); + + // extract a subset consisting of multiple non-consecutive bits in srcsig + int dstsig2 = mjSTATE_QPOS | mjSTATE_CTRL; + int dstsize2 = mj_stateSize(model, dstsig2); + EXPECT_LT(dstsize2, srcsize); + EXPECT_EQ(dstsize2, model->nq + model->nu); + vector dststate2(dstsize2); + mj_extractState(model, srcstate.data(), srcsig, dststate2.data(), dstsig2); + EXPECT_EQ(AsVector(dststate2.data(), model->nq), + AsVector(data->qpos, model->nq)); + EXPECT_EQ(AsVector(dststate2.data() + model->nq, model->nu), + AsVector(data->ctrl, model->nu)); + + // test that an error is correctly raised if dstsig is not a subset of srcsig + static int error_count; + static char last_error_msg[128]; + error_count = 0; + last_error_msg[0] = '\0'; + auto* error_handler = +[](const char* msg) { + std::strncpy(last_error_msg, msg, sizeof(last_error_msg)); + ++error_count; + }; + auto* old_mju_user_error = mju_user_error; + mju_user_error = error_handler; + mj_extractState(model, nullptr, srcsig, nullptr, mjSTATE_QFRC_APPLIED); + mju_user_error = old_mju_user_error; + EXPECT_EQ(error_count, 1); + EXPECT_EQ(std::string_view(last_error_msg), + "mj_extractState: dstsig is not a subset of srcsig"); + + mj_deleteData(data); + mj_deleteModel(model); +} + using InertiaTest = MujocoTest; static const char* const kInertiaPath = "engine/testdata/inertia.xml"; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index e9a7aa43..4f464e1a 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -6621,6 +6621,9 @@ public static unsafe extern int mj_stateSize(mjModel_* m, uint sig); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_getState(mjModel_* m, mjData_* d, double* state, uint sig); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_extractState(mjModel_* m, double* src, uint srcsig, double* dst, uint dstsig); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_setState(mjModel_* m, mjData_* d, double* state, uint sig);