Allow dm_control.Physics.{get,set}_state to accept a sig parameter.

When sig is provided, these functions behave like mj_{get,set}State.

PiperOrigin-RevId: 823171024
Change-Id: I9d42f3f59c6404e6894f3ab384e92ca928fa6b56

Author: saran-t

dm_control/mujoco/engine.py

@@ -232,35 +232,57 @@ def render(
     camera._scene.free()  # pylint: disable=protected-access
     return image
 
-  def get_state(self):
+  def get_state(self, sig=None):
     """Returns the physics state.
 
+    Args:
+      sig: Optional integer, if specified then the returned array corresponds to
+        the state obtained by calling `mj_getState` with `sig`.
+
     Returns:
       NumPy array containing full physics simulation state.
     """
-    return np.concatenate(self._physics_state_items())
+    if sig is None:
+      return np.concatenate(self._physics_state_items())
+    else:
+      retval = np.empty(mujoco.mj_stateSize(self.model.ptr, sig), np.float64)
+      mujoco.mj_getState(self.model.ptr, self.data.ptr, retval, sig)
+      return retval
 
-  def set_state(self, physics_state):
+  def set_state(self, physics_state, sig=None):
     """Sets the physics state.
 
     Args:
       physics_state: NumPy array containing the full physics simulation state.
+      sig: Optional integer, if specified then physics_state is passed directly
+        to `mj_setState` with `sig`.
 
     Raises:
       ValueError: If `physics_state` has invalid size.
     """
-    state_items = self._physics_state_items()
-
-    expected_shape = (sum(item.size for item in state_items),)
-    if expected_shape != physics_state.shape:
-      raise ValueError('Input physics state has shape {}. Expected {}.'.format(
-          physics_state.shape, expected_shape))
-
-    start = 0
-    for state_item in state_items:
-      size = state_item.size
-      np.copyto(state_item, physics_state[start:start + size])
-      start += size
+    if sig is None:
+      state_items = self._physics_state_items()
+
+      expected_shape = (sum(item.size for item in state_items),)
+      if expected_shape != physics_state.shape:
+        raise ValueError(
+            'Input physics state has shape {}. Expected {}.'.format(
+                physics_state.shape, expected_shape
+            )
+        )
+
+      start = 0
+      for state_item in state_items:
+        size = state_item.size
+        np.copyto(state_item, physics_state[start:start + size])
+        start += size
+    else:
+      mujoco.mj_setState(
+          self.model.ptr,
+          self.data.ptr,
+          np.asarray(physics_state, np.float64),
+          sig,
+      )
 
   def copy(self, share_model=False):
     """Creates a copy of this `Physics` instance.

dm_control/mujoco/engine_test.py

@@ -370,7 +370,7 @@ def testNamedViews(self):
     self.assertEqual(0., self._physics.time())
     self.assertEqual(0.01, self._physics.timestep())
 
-  def testSetGetPhysicsState(self):
+  def testSetGetPhysicsStateLegacy(self):
     physics_state = self._physics.get_state()
 
     # qpos, qvel, act
@@ -384,6 +384,27 @@ def testSetGetPhysicsState(self):
     np.testing.assert_allclose(new_physics_state,
                                self._physics.get_state())
 
+  def testSetGetPhysicsState(self):
+    actual_physics_state = self._physics.get_state(
+        sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+    )
+    expected_physics_state = np.zeros(mujoco.mj_stateSize(
+        self._physics.model.ptr, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+    ))
+    mujoco.mj_getState(self._physics.model.ptr, self._physics.data.ptr,
+                       expected_physics_state,
+                       mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
+    np.testing.assert_array_equal(expected_physics_state, actual_physics_state)
+
+    new_physics_state = np.random.random_sample(expected_physics_state.shape)
+    self._physics.set_state(
+        new_physics_state, sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+    )
+    np.testing.assert_array_equal(
+        new_physics_state,
+        self._physics.get_state(sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
+    )
+
   def testSetGetPhysicsStateWithPlugin(self):
     # Model copied from mujoco/test/plugin/elasticity/elasticity_test.cc
     model_with_cable_plugin = """