Add mjtState enum and related mj_stateSize, mj_getState and mj_setState functions.

PiperOrigin-RevId: 539667943
Change-Id: I14c34be8ce4c287e529380257b5c5bffc3ab45ce

doc/APIreference/APItypes.rst

@@ -346,6 +346,18 @@ Data
 The enums below are defined in `mjdata.h <https://github.com/deepmind/mujoco/blob/main/include/mujoco/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

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

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

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<geState>`  section for details.
 - Added :ref:`mjContact.solreffriction<mjContact>`, allowing different :ref:`solref<CSolver>` parameters for the normal
   and frictional axes of contacts when using :ref:`elliptic friction cones<option-cone>`.  This attribute is required
   for elastic frictional collisions, see associated

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
+well-defined internal state which, together with the :ref:`deterministic computational pipeline<piReproducibility>`,
-resetting the state and computing dynamics derivatives are also well-defined. The state is entirely encapsulated in the
+means that operations like resetting the state and computing dynamics derivatives are also well-defined.
-``mjData`` struct and consists of several components:
+
+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.
+The *physics state* (:ref:`mjSTATE_PHYSICS<mjtState>`) contains the main quantities which are time-integrated during
-These are ``mjData.{qpos, qvel, act, time}``:
+stepping. These are ``mjData.{qpos, qvel, act}``:
 
-Mechanical state: ``qpos`` and ``qvel``
+Position: ``qpos``
-  The *mechanical state* of a simulation is given by the generalized position (``mjData.qpos``) and velocity
+  The configuration in generalized coodinates, denoted above as :math:`q`.
-  (``mjData.qvel``) vectors, denoted above as :math:`q` and :math:`v`, respectively.
 
-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<mjtState>`) 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
+  The simulation time 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
+  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<exPlugin>`. 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<mjtState>`) 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<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
   <https://en.wikipedia.org/wiki/Lyapunov_exponent>`__ 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<exPlugin>`. 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
+The *integration state* (:ref:`mjSTATE_INTEGRATION<mjtState>`) is the union of all the above :ref:`mjData` fields and
-other ``mjData`` fields are functions of the integration state.
+constitutes the entire set of inputs to the *forward dynamics*. In the case of *inverse dynamics*, ``mjData.qacc`` is
-|br| When saving the integration state in order to reload it elsewhere, it is sensible to avoid saving unused fields
+also treated as an input variable. All other :ref:`mjData` fields are functions of the integration state.
-that always remain in their default values. Specifically, ``xfrc_applied`` can be quite large (``6 x nbody``) yet is
+
-often unused.
+Note that the full integration state as given by :ref:`mjSTATE_INTEGRATION<mjtState>` 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
+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 ``mjData`` buffers are preallocated for the
+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<IntegrationState>` components. In particular :ref:`warmstart
+- Save all the required :ref:`integration state<geIntegrationState>` components. In particular :ref:`warmstart
   accelerations<geWarmstart>` 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

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);

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

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);
 

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',

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',

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')

python/mujoco/functions.cc

@@ -269,6 +269,25 @@ PYBIND11_MODULE(_functions, pymodule) {
       });
 
   // Support
+  Def<traits::mj_stateSize>(pymodule);
+  Def<traits::mj_getState>(
+      pymodule,
+      [](const raw::MjModel* m, const raw::MjData* d,
+         Eigen::Ref<EigenVectorX> 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<traits::mj_setState>(
+      pymodule,
+      [](const raw::MjModel* m, raw::MjData* d,
+         const Eigen::Ref<EigenVectorX> 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<traits::mj_addContact>(pymodule);
   Def<traits::mj_isPyramidal>(pymodule);
   Def<traits::mj_isSparse>(pymodule);

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<<mjNSTATE)) {
+    mju_error("mj_stateSize: invalid state spec %u >= 2^mjNSTATE", spec);
+  }
+
+  int size = 0;
+  for (int i=0; i < mjNSTATE; i++) {
+    mjtState element = 1<<i;
+    if (element & spec) {
+      size += mj_stateElemSize(m, element);
+    }
+  }
+
+  return size;
+}
+
+
+
+// get state
+void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, unsigned int spec) {
+  if (spec >= (1<<mjNSTATE)) {
+    mju_error("mj_getState: invalid state spec %u >= 2^mjNSTATE", spec);
+  }
+
+  int adr = 0;
+  for (int i=0; i < mjNSTATE; i++) {
+    mjtState element = 1<<i;
+    if (element & spec) {
+      int size = mj_stateElemSize(m, element);
+      const mjtNum* ptr = mj_stateElemConstPtr(m, d, element);
+      mju_copy(state + adr, ptr, size);
+      adr += size;
+    }
+  }
+}
+
+
+
+// set state
+void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, unsigned int spec) {
+  if (spec >= (1<<mjNSTATE)) {
+    mju_error("mj_setState: invalid state spec %u >= 2^mjNSTATE", spec);
+  }
+
+  int adr = 0;
+  for (int i=0; i < mjNSTATE; i++) {
+    mjtState element = 1<<i;
+    if (element & spec) {
+      int size = mj_stateElemSize(m, element);
+      mjtNum* ptr = mj_stateElemPtr(m, d, element);
+      mju_copy(ptr, state + adr, size);
+      adr += size;
+    }
+  }
+}
+
+
+
 //-------------------------- Jacobians -------------------------------------------------------------
 
 // compute 3/6-by-nv Jacobian of global point attached to given body

src/engine/engine_support.h

@@ -29,6 +29,18 @@ MJAPI extern const char* mjENABLESTRING[mjNENABLE];
 MJAPI extern const char* mjTIMERSTRING[mjNTIMER];
 
 
+//-------------------------- get/set state ---------------------------------------------------------
+
+// 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);
+
+
 //-------------------------- Jacobians -------------------------------------------------------------
 
 // compute 3/6-by-nv Jacobian of global point attached to given body

test/engine/engine_support_test.cc

@@ -78,7 +78,7 @@ TEST_F(JacobianTest, SubtreeJac) {
   mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->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; i<nv; i++) {
+  for (int i=0; i < nv; i++) {
     // reset qpos, nudge i-th dof, update data->qpos, 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<double> 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<mjtNum> state0a(size);
+  mj_getState(model, data, state0a.data(), spec);
+
+  // get the initial state, expect equality
+  std::vector<mjtNum> 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<mjtNum> 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<mjtNum> 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

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);