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

PiperOrigin-RevId: 539667943
Change-Id: I14c34be8ce4c287e529380257b5c5bffc3ab45ce

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