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

PiperOrigin-RevId: 539667943
Change-Id: I14c34be8ce4c287e529380257b5c5bffc3ab45ce

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