From f67e3595329718c80b84a024febbed143d1b4636 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 12 Jun 2023 08:55:58 -0700 Subject: [PATCH] Add `mjtState` enum and related `mj_stateSize`, `mj_getState` and `mj_setState` functions. PiperOrigin-RevId: 539667943 Change-Id: I14c34be8ce4c287e529380257b5c5bffc3ab45ce --- doc/APIreference/APItypes.rst | 12 +++ doc/APIreference/functions.rst | 30 +++++++ doc/APIreference/functions_override.rst | 15 ++++ doc/changelog.rst | 2 + doc/computation.rst | 86 ++++++++++-------- doc/includes/references.h | 26 ++++++ include/mujoco/mjdata.h | 25 ++++++ include/mujoco/mujoco.h | 9 ++ introspect/enums.py | 24 +++++ introspect/functions.py | 78 ++++++++++++++++ python/mujoco/bindings_test.py | 37 ++++++++ python/mujoco/functions.cc | 19 ++++ src/engine/engine_support.c | 115 ++++++++++++++++++++++++ src/engine/engine_support.h | 12 +++ test/engine/engine_support_test.cc | 62 ++++++++++++- unity/Runtime/Bindings/MjBindings.cs | 9 ++ 16 files changed, 520 insertions(+), 41 deletions(-) diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index d3e78000..b390da37 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -346,6 +346,18 @@ Data The enums below are defined in `mjdata.h `_. + +.. _mjtState: + +mjtState +~~~~~~~~ + +State component elements as integer bitflags and several convenient combinations of these flags. Used by +:ref:`mj_getState`, :ref:`mj_setState` and :ref:`mj_stateSize`. + +.. mujoco-include:: mjtState + + .. _mjtWarning: mjtWarning diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 72e234ed..af9289a7 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -151,6 +151,36 @@ These are support functions that need access to :ref:`mjModel` and :ref:`mjData` not need such access. Support functions are called within the simulator but some of them can also be useful for custom computations, and are documented in more detail below. +.. _mj_stateSize: + +mj_stateSize +~~~~~~~~~~~~ + +.. mujoco-include:: mj_stateSize + +Returns the number of :ref:`mjtNum` |-| s required for a given state specification. The bits of the integer ``spec`` +correspond to element fields of :ref:`mjtState`. + +.. _mj_getState: + +mj_getState +~~~~~~~~~~~ + +.. mujoco-include:: mj_getState + +Copy concatenated state components specified by ``spec`` from ``d`` into ``state``. The bits of the integer +``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid. + +.. _mj_setState: + +mj_setState +~~~~~~~~~~~ + +.. mujoco-include:: mj_setState + +Copy concatenated state components specified by ``spec`` from ``state`` into ``d``. The bits of the integer +``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid. + .. _mj_addContact: mj_addContact diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 83dce986..121e2b1b 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -96,6 +96,21 @@ These are support functions that need access to :ref:`mjModel` and :ref:`mjData` not need such access. Support functions are called within the simulator but some of them can also be useful for custom computations, and are documented in more detail below. +.. _mj_stateSize: + +Returns the number of :ref:`mjtNum` |-| s required for a given state specification. The bits of the integer ``spec`` +correspond to element fields of :ref:`mjtState`. + +.. _mj_getState: + +Copy concatenated state components specified by ``spec`` from ``d`` into ``state``. The bits of the integer +``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid. + +.. _mj_setState: + +Copy concatenated state components specified by ``spec`` from ``state`` into ``d``. The bits of the integer +``spec`` correspond to element fields of :ref:`mjtState`. Fails with :ref:`mju_error` if ``spec`` is invalid. + .. _mj_mulJacVec: This function multiplies the constraint Jacobian mjData.efc_J by a vector. Note that the Jacobian can be either dense or diff --git a/doc/changelog.rst b/doc/changelog.rst index 9647a34a..f49e2d18 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -63,6 +63,8 @@ Simulate General ^^^^^^^ +- Added :ref:`mj_getState` and :ref:`mj_setState` for getting and setting the simulation state as a concatenated vector + of floating point numbers. See the :ref:`State` section for details. - Added :ref:`mjContact.solreffriction`, allowing different :ref:`solref` parameters for the normal and frictional axes of contacts when using :ref:`elliptic friction cones`. This attribute is required for elastic frictional collisions, see associated diff --git a/doc/computation.rst b/doc/computation.rst index 13a6f8df..22b87ee5 100644 --- a/doc/computation.rst +++ b/doc/computation.rst @@ -569,41 +569,60 @@ Fast implicit-in-velocity (``implicitfast``) .. _geState: -The **state** -~~~~~~~~~~~~~ +The State +~~~~~~~~~ To complete our description of the general framework we will now discuss the notion of *state*. MuJoCo has a compact, -well-defined internal state which, together with the deterministic computational pipeline, means that operations like -resetting the state and computing dynamics derivatives are also well-defined. The state is entirely encapsulated in the -``mjData`` struct and consists of several components: +well-defined internal state which, together with the :ref:`deterministic computational pipeline`, +means that operations like resetting the state and computing dynamics derivatives are also well-defined. + +The state is entirely encapsulated in the :ref:`mjData` struct and consists of several components. The components are +enumerated in :ref:`mjtState` as bit flags, along with several common combinations, corresponding to the groupings +below. Concatenated state vectors can be conveniently read from and written into :ref:`mjData` using :ref:`mj_getState` +and :ref:`mj_setState`, respectively. .. _gePhysicsState: Physics state ^^^^^^^^^^^^^ -The *physics state* contains all quantities which are time-integrated during stepping. -These are ``mjData.{qpos, qvel, act, time}``: +The *physics state* (:ref:`mjSTATE_PHYSICS`) contains the main quantities which are time-integrated during +stepping. These are ``mjData.{qpos, qvel, act}``: -Mechanical state: ``qpos`` and ``qvel`` - The *mechanical state* of a simulation is given by the generalized position (``mjData.qpos``) and velocity - (``mjData.qvel``) vectors, denoted above as :math:`q` and :math:`v`, respectively. +Position: ``qpos`` + The configuration in generalized coodinates, denoted above as :math:`q`. -Actuator activations: ``act`` +Velocity: ``qvel`` + The generalized velocities, denoted above as :math:`v`. + +Actuator activation: ``act`` ``mjData.act`` contains the internal states of stateful actuators, denoted above as :math:`w`. +.. _geFullPhysics: + +Full physics state +^^^^^^^^^^^^^^^^^^ + +The *full physics state* (:ref:`mjSTATE_FULLPHYSICS`) contains the physics state and two additional +components: + Time: ``time`` - The time of the simulation is given by the scalar ``mjData.time``. Since physics is time-invariant, it is - often excluded from the *physics state*; an exception could be a time-dependent user callback (e.g., an open-loop + The simulation time is given by the scalar ``mjData.time``. Since physics is time-invariant, it is + excluded from the *physics state*; exceptions include time-dependent user callbacks and plugins (e.g., an open-loop controller), in which case time should be included. +Plugin state: ``plugin_state`` + ``mjData.plugin_state`` are states declared by :ref:`engine plugins`. Please see the :ref:`exPluginState` + section for more details. + .. _geInput: User inputs ^^^^^^^^^^^ -These input fields are set by the user and affect the physics simulation, but are untouched by the simulator. All input -fields except for MoCap poses default to 0. -Controls: ``ctrl`` +These input fields (:ref:`mjSTATE_USER`) are set by the user and affect the physics simulation, but are +untouched by the simulator. All input fields except for MoCap poses default to 0. + +Control: ``ctrl`` Controls are defined by the :ref:`actuator` section of the XML. ``mjData.ctrl`` values either produce generalized forces directly (stateless actuators), or affect the actuator activations in ``mjData.act``, which then produce forces. @@ -627,8 +646,8 @@ User data: ``userdata`` .. _geWarmstart: -Warmstart accelerations -^^^^^^^^^^^^^^^^^^^^^^^ +Warmstart acceleration +^^^^^^^^^^^^^^^^^^^^^^ ``qacc_warmstart`` ``mjData.qacc_warmstart`` are accelerations used to warmstart the constraint solver, saved from the previous step. @@ -642,34 +661,27 @@ Warmstart accelerations `__ when time-stepping, quickly leading to divergent trajectories for different warmstarts. -.. _gePlugin: - -Plugin state -^^^^^^^^^^^^ - -``plugin_state`` - ``mjData.plugin_state`` are states declared by :ref:`engine plugins`. Please see the :ref:`exPluginState` - section for more details. - .. _geIntegrationState: Integration state ^^^^^^^^^^^^^^^^^ -The *integration state* is the union of all the above ``mjData`` fields and constitutes the entire set of inputs to -the *forward dynamics*. In the case of *inverse dynamics*, ``mjData.qacc`` is also treated as an input variable. All -other ``mjData`` fields are functions of the integration state. -|br| When saving the integration state in order to reload it elsewhere, it is sensible to avoid saving unused fields -that always remain in their default values. Specifically, ``xfrc_applied`` can be quite large (``6 x nbody``) yet is -often unused. + +The *integration state* (:ref:`mjSTATE_INTEGRATION`) is the union of all the above :ref:`mjData` fields and +constitutes the entire set of inputs to the *forward dynamics*. In the case of *inverse dynamics*, ``mjData.qacc`` is +also treated as an input variable. All other :ref:`mjData` fields are functions of the integration state. + +Note that the full integration state as given by :ref:`mjSTATE_INTEGRATION` is maximalist and includes fields +which are often unused. If a small state size is desired, it might be sensible to avoid saving unused fields. +In particular `xfrc_applied`` can be quite large (``6 x nbody``) yet is often unused. .. _geSimulationState: Simulation state: ``mjData`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -The *simulation state* is the entirety of the ``mjData`` struct and associated memory buffer. This state includes -all derived quantities computed during dynamics computation. Because the ``mjData`` buffers are preallocated for the +The *simulation state* is the entirety of the :ref:`mjData` struct and associated memory buffer. This state includes +all derived quantities computed during dynamics computation. Because the :ref:`mjData` buffers are preallocated for the worst case, it is often significantly faster to recompute derived quantities from the *integration state* rather than -using ``mj_copyData``. +using :ref:`mj_copyData`. .. _Constraint: @@ -1591,7 +1603,7 @@ MuJoCo's simulation pipeline is entirely deterministic and reproducible -- if a saved and reloaded and :ref:`mj_step` called again, the resulting next state will be identical. However, there are some important caveats: -- Save all the required :ref:`integration state` components. In particular :ref:`warmstart +- Save all the required :ref:`integration state` components. In particular :ref:`warmstart accelerations` have only a very small effect on the next state, but should be saved if bit-wise equality is required. - Any numerical difference between states, no matter how small, will become significant upon integration, especially for diff --git a/doc/includes/references.h b/doc/includes/references.h index 986444d3..ad6a3235 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -16,6 +16,29 @@ // Error: C reference not found // NOLINTBEGIN +typedef enum mjtState_ { // state elements + mjSTATE_TIME = 1<<0, // time + mjSTATE_QPOS = 1<<1, // position + mjSTATE_QVEL = 1<<2, // velocity + mjSTATE_ACT = 1<<3, // actuator activation + mjSTATE_WARMSTART = 1<<4, // acceleration used for warmstart + mjSTATE_CTRL = 1<<5, // control + mjSTATE_QFRC_APPLIED = 1<<6, // applied generalized force + mjSTATE_XFRC_APPLIED = 1<<7, // applied Cartesian force/torque + mjSTATE_MOCAP_POS = 1<<8, // positions of mocap bodies + mjSTATE_MOCAP_QUAT = 1<<9, // orientations of mocap bodies + mjSTATE_USERDATA = 1<<10, // user data + mjSTATE_PLUGIN = 1<<11, // plugin state + + mjNSTATE = 12, // number of state elements + + // convenience values for commonly used state specifications + mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT, + mjSTATE_FULLPHYSICS = mjSTATE_PHYSICS | mjSTATE_TIME | mjSTATE_PLUGIN, + mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED | + mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT | mjSTATE_USERDATA, + mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART +} mjtState; typedef enum mjtWarning_ { // warning types mjWARN_INERTIA = 0, // (near) singular inertia matrix mjWARN_CONTACTFULL, // too many contacts in contact list @@ -2176,6 +2199,9 @@ void mj_projectConstraint(const mjModel* m, mjData* d); void mj_referenceConstraint(const mjModel* m, mjData* d); 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 spec); +void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec); +void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec); int mj_addContact(const mjModel* m, mjData* d, const mjContact* con); int mj_isPyramidal(const mjModel* m); int mj_isSparse(const mjModel* m); diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 19962b00..defd4b19 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -23,6 +23,31 @@ //---------------------------------- primitive types (mjt) ----------------------------------------- +typedef enum mjtState_ { // state elements + mjSTATE_TIME = 1<<0, // time + mjSTATE_QPOS = 1<<1, // position + mjSTATE_QVEL = 1<<2, // velocity + mjSTATE_ACT = 1<<3, // actuator activation + mjSTATE_WARMSTART = 1<<4, // acceleration used for warmstart + mjSTATE_CTRL = 1<<5, // control + mjSTATE_QFRC_APPLIED = 1<<6, // applied generalized force + mjSTATE_XFRC_APPLIED = 1<<7, // applied Cartesian force/torque + mjSTATE_MOCAP_POS = 1<<8, // positions of mocap bodies + mjSTATE_MOCAP_QUAT = 1<<9, // orientations of mocap bodies + mjSTATE_USERDATA = 1<<10, // user data + mjSTATE_PLUGIN = 1<<11, // plugin state + + mjNSTATE = 12, // number of state elements + + // convenience values for commonly used state specifications + mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT, + mjSTATE_FULLPHYSICS = mjSTATE_PHYSICS | mjSTATE_TIME | mjSTATE_PLUGIN, + mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED | + mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT | mjSTATE_USERDATA, + mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART +} mjtState; + + typedef enum mjtWarning_ { // warning types mjWARN_INERTIA = 0, // (near) singular inertia matrix mjWARN_CONTACTFULL, // too many contacts in contact list diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 23df98cf..74bc7695 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -357,6 +357,15 @@ MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, //---------------------------------- Support ------------------------------------------------------- +// Return size of state specification. +MJAPI int mj_stateSize(const mjModel* m, unsigned int spec); + +// Get state. +MJAPI void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec); + +// Set state. +MJAPI void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec); + // Add contact to d->contact list; return 0 if success; 1 if buffer full. MJAPI int mj_addContact(const mjModel* m, mjData* d, const mjContact* con); diff --git a/introspect/enums.py b/introspect/enums.py index b9604216..10a4db11 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -373,6 +373,30 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjLRMODE_ALL', 3), ]), )), + ('mjtState', + EnumDecl( + name='mjtState', + declname='enum mjtState_', + values=dict([ + ('mjSTATE_TIME', 1), + ('mjSTATE_QPOS', 2), + ('mjSTATE_QVEL', 4), + ('mjSTATE_ACT', 8), + ('mjSTATE_WARMSTART', 16), + ('mjSTATE_CTRL', 32), + ('mjSTATE_QFRC_APPLIED', 64), + ('mjSTATE_XFRC_APPLIED', 128), + ('mjSTATE_MOCAP_POS', 256), + ('mjSTATE_MOCAP_QUAT', 512), + ('mjSTATE_USERDATA', 1024), + ('mjSTATE_PLUGIN', 2048), + ('mjNSTATE', 12), + ('mjSTATE_PHYSICS', 14), + ('mjSTATE_FULLPHYSICS', 2063), + ('mjSTATE_USER', 2016), + ('mjSTATE_INTEGRATION', 4095), + ]), + )), ('mjtWarning', EnumDecl( name='mjtWarning', diff --git a/introspect/functions.py b/introspect/functions.py index d4a0ba1a..dc54bc80 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -1783,6 +1783,84 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set *cost = s(jar) where jar = Jac*qacc-aref.', # pylint: disable=line-too-long )), + ('mj_stateSize', + FunctionDecl( + name='mj_stateSize', + return_type=ValueType(name='int'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='spec', + type=ValueType(name='unsigned int'), + ), + ), + doc='Return size of state specification.', + )), + ('mj_getState', + FunctionDecl( + name='mj_getState', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData', is_const=True), + ), + ), + FunctionParameterDecl( + name='state', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='spec', + type=ValueType(name='unsigned int'), + ), + ), + doc='Get state.', + )), + ('mj_setState', + FunctionDecl( + name='mj_setState', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + FunctionParameterDecl( + name='state', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='spec', + type=ValueType(name='unsigned int'), + ), + ), + doc='Set state.', + )), ('mj_addContact', FunctionDecl( name='mj_addContact', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index a9c48da3..8813e294 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -667,6 +667,43 @@ class MuJoCoBindingsTest(parameterized.TestCase): with self.assertRaises(TypeError): mujoco.mju_rotVecQuat(vec, quat, res=np.zeros(3, int)) + def test_getsetstate(self): # pylint: disable=invalid-name + mujoco.mj_step(self.model, self.data) + + # Test for invalid state spec + invalid_spec = 2**mujoco.mjtState.mjNSTATE.value + expected_message = ( + f'mj_stateSize: invalid state spec {invalid_spec} >= 2^mjNSTATE' + ) + with self.assertRaisesWithLiteralMatch(mujoco.FatalError, expected_message): + mujoco.mj_stateSize(self.model, invalid_spec) + + spec = mujoco.mjtState.mjSTATE_INTEGRATION + size = mujoco.mj_stateSize(self.model, spec) + + state_bad_size = np.empty(size + 1, np.float64) + expected_message = ('state size should equal mj_stateSize(m, spec)') + with self.assertRaisesWithLiteralMatch(TypeError, expected_message): + mujoco.mj_getState(self.model, self.data, state_bad_size, spec) + + # Get initial state. + state0 = np.empty(size, np.float64) + mujoco.mj_getState(self.model, self.data, state0, spec) + + # Step, get next state. + mujoco.mj_step(self.model, self.data) + state1a = np.empty(size, np.float64) + mujoco.mj_getState(self.model, self.data, state1a, spec) + + # Reset to initial state, step again, get state again. + mujoco.mj_setState(self.model, self.data, state0, spec) + mujoco.mj_step(self.model, self.data) + state1b = np.empty(size, np.float64) + mujoco.mj_getState(self.model, self.data, state1b, spec) + + # Expect next states to be equal. + np.testing.assert_array_equal(state1a, state1b) + def test_mj_jacSite(self): # pylint: disable=invalid-name mujoco.mj_forward(self.model, self.data) site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite') diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index a3b0a16c..f9a0190e 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -269,6 +269,25 @@ PYBIND11_MODULE(_functions, pymodule) { }); // Support + Def(pymodule); + Def( + pymodule, + [](const raw::MjModel* m, const raw::MjData* d, + Eigen::Ref state, unsigned int spec) { + if (state.size() != mj_stateSize(m, spec)) { + throw py::type_error("state size should equal mj_stateSize(m, spec)"); + } + return InterceptMjErrors(::mj_getState)(m, d, state.data(), spec); + }); + Def( + pymodule, + [](const raw::MjModel* m, raw::MjData* d, + const Eigen::Ref state, unsigned int spec) { + if (state.size() != mj_stateSize(m, spec)) { + throw py::type_error("state size should equal mj_stateSize(m, spec)"); + } + return InterceptMjErrors(::mj_setState)(m, d, state.data(), spec); + }); Def(pymodule); Def(pymodule); Def(pymodule); diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 86544881..bc8a2411 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -36,6 +36,8 @@ #endif #endif +//-------------------------- Constants ------------------------------------------------------------- + #define mjVERSION 236 #define mjVERSIONSTRING "2.3.6" @@ -87,6 +89,119 @@ const char* mjTIMERSTRING[mjNTIMER]= { +//-------------------------- get/set state --------------------------------------------------------- + +// return size of a single state element +static inline int mj_stateElemSize(const mjModel* m, mjtState spec) { + switch (spec) { + case mjSTATE_TIME: return 1; + case mjSTATE_QPOS: return m->nq; + case mjSTATE_QVEL: return m->nv; + case mjSTATE_ACT: return m->na; + case mjSTATE_WARMSTART: return m->nv; + case mjSTATE_CTRL: return m->nu; + case mjSTATE_QFRC_APPLIED: return m->nv; + case mjSTATE_XFRC_APPLIED: return 6*m->nbody; + case mjSTATE_MOCAP_POS: return 3*m->nmocap; + case mjSTATE_MOCAP_QUAT: return 4*m->nmocap; + case mjSTATE_USERDATA: return m->nuserdata; + case mjSTATE_PLUGIN: return m->npluginstate; + default: + mju_error("mju_stateElementSize: invalid state element %u", spec); + return 0; + } +} + + + +// return pointer to a single state element +static inline mjtNum* mj_stateElemPtr(const mjModel* m, mjData* d, mjtState spec) { + switch (spec) { + case mjSTATE_TIME: return &d->time; + case mjSTATE_QPOS: return d->qpos; + case mjSTATE_QVEL: return d->qvel; + case mjSTATE_ACT: return d->act; + case mjSTATE_WARMSTART: return d->qacc_warmstart; + case mjSTATE_CTRL: return d->ctrl; + case mjSTATE_QFRC_APPLIED: return d->qfrc_applied; + case mjSTATE_XFRC_APPLIED: return d->xfrc_applied; + case mjSTATE_MOCAP_POS: return d->mocap_pos; + case mjSTATE_MOCAP_QUAT: return d->mocap_quat; + case mjSTATE_USERDATA: return d->userdata; + case mjSTATE_PLUGIN: return d->plugin_state; + default: + mju_error("mju_stateElemPtr: invalid state element %u", spec); + return NULL; + } +} + + + +static inline const mjtNum* mj_stateElemConstPtr(const mjModel* m, const mjData* d, mjtState spec) { + return mj_stateElemPtr(m, (mjData*) d, spec); // discard const qualifier from d +} + + + +// get size of state specification +int mj_stateSize(const mjModel* m, unsigned int spec) { + if (spec >= (1<= 2^mjNSTATE", spec); + } + + int size = 0; + for (int i=0; i < mjNSTATE; i++) { + mjtState element = 1<= (1<= 2^mjNSTATE", spec); + } + + int adr = 0; + for (int i=0; i < mjNSTATE; i++) { + mjtState element = 1<= (1<= 2^mjNSTATE", spec); + } + + int adr = 0; + for (int i=0; i < mjNSTATE; i++) { + mjtState element = 1<nq); mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv); - // all we need for Jacobians are kinematics and CoM-related quantitites + // all we need for Jacobians are kinematics and CoM-related quantities mj_kinematics(model, data); mj_comPos(model, data); @@ -93,7 +93,7 @@ TEST_F(JacobianTest, SubtreeJac) { // compare analytic Jacobian to finite-difference approximation static const mjtNum eps = 1e-6; - for (int i=0; iqpos, reset nudge mju_copy(data->qpos, qpos, model->nq); nudge[i] = 1; @@ -105,7 +105,7 @@ TEST_F(JacobianTest, SubtreeJac) { mj_comPos(model, data); // compare finite-differenced and analytic Jacobian - for (int j=0; j<3; j++) { + for (int j=0; j < 3; j++) { mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps; EXPECT_THAT(jac_subtree[nv*j+i], DoubleNear(findiff, eps)); } @@ -127,7 +127,7 @@ TEST_F(JacobianTest, SubtreeJacNoInternalAcc) { mjData* data = mj_makeData(model); mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv); - // all we need for Jacobians are kinematics and CoM-related quantitites + // all we need for Jacobians are kinematics and CoM-related quantities mj_kinematics(model, data); mj_comPos(model, data); @@ -351,5 +351,59 @@ TEST_F(SupportTest, DifferentiatePosSubQuat) { mj_deleteModel(model); } +static const char* const kDefaultModel = "testdata/model.xml"; + +TEST_F(SupportTest, GetSetStateStepEqual) { + 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); + + // take one step + mj_step(model, data); + + int spec = mjSTATE_INTEGRATION; + int size = mj_stateSize(model, spec); + + // save the initial state and step + std::vector state0a(size); + mj_getState(model, data, state0a.data(), spec); + + // get the initial state, expect equality + std::vector state0b(size); + mj_getState(model, data, state0b.data(), spec); + EXPECT_EQ(state0a, state0b); + + // take one step + mj_step(model, data); + + // save the resulting state + std::vector state1a(size); + mj_getState(model, data, state1a.data(), spec); + + // expect the state to be different after stepping + EXPECT_THAT(state0a, testing::Ne(state1a)); + + // reset to the saved state, step again, get the resulting state + mj_setState(model, data, state0a.data(), spec); + mj_step(model, data); + std::vector state1b(size); + mj_getState(model, data, state1b.data(), spec); + + // expect the state to be the same after re-stepping + EXPECT_EQ(state1a, state1b); + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 8c1d6160..f94103bc 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -3191,6 +3191,15 @@ public static unsafe extern void mj_referenceConstraint(mjModel_* m, mjData_* d) [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_constraintUpdate(mjModel_* m, mjData_* d, double* jar, double* cost, int flg_coneHessian); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern int mj_stateSize(mjModel_* m, uint spec); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_getState(mjModel_* m, mjData_* d, double* state, uint spec); + +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_setState(mjModel_* m, mjData_* d, double* state, uint spec); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern int mj_addContact(mjModel_* m, mjData_* d, mjContact_* con);