Add reset-to-keyframe functionality to Physics.reset()

yuvaltassa · alimuldal · commit e066d77a38f2 · 2020-05-04T09:44:49.000+01:00
PiperOrigin-RevId: 309392225
Change-Id: Id3425c8c057a6ac6d29e2bb8c6b31b1b409b2a04
diff --git a/dm_control/mujoco/engine.py b/dm_control/mujoco/engine.py
@@ -83,6 +83,8 @@
 _RENDER_FLAG_OVERRIDES_NOT_SUPPORTED_FOR_DEPTH_OR_SEGMENTATION = (
     '`render_flag_overrides` are not supported for depth or segmentation '
     'rendering.')
+_KEYFRAME_ID_OUT_OF_RANGE = (
+    '`keyframe_id` must be between 0 and {max_valid} inclusive, got: {actual}.')
 
 
 class Physics(_control.Physics):
@@ -267,9 +269,22 @@ def copy(self, share_model=False):
     new_obj._reload_from_data(new_data)  # pylint: disable=protected-access
     return new_obj
 
-  def reset(self):
-    """Resets internal variables of the physics simulation."""
-    mjlib.mj_resetData(self.model.ptr, self.data.ptr)
+  def reset(self, keyframe_id=None):
+    """Resets internal variables of the simulation, possibly to a keyframe.
+
+    Args:
+      keyframe_id: Optional int. If specified, the keyframe (saved state) to
+        which to set the state.
+
+    """
+    if keyframe_id is None:
+      mjlib.mj_resetData(self.model.ptr, self.data.ptr)
+    else:
+      if not 0 <= keyframe_id < self.model.nkey:
+        raise ValueError(_KEYFRAME_ID_OUT_OF_RANGE.format(
+            max_valid=self.model.nkey-1, actual=keyframe_id))
+      mjlib.mj_resetDataKeyframe(self.model.ptr, self.data.ptr, keyframe_id)
+
     # Disable actuation since we don't yet have meaningful control inputs.
     with self.model.disable('actuation'):
       self.forward()
diff --git a/dm_control/mujoco/engine_test.py b/dm_control/mujoco/engine_test.py
@@ -311,6 +311,22 @@ def testNamedIndexing(self):
   def testReload(self):
     self._physics.reload_from_xml_path(MODEL_PATH)
 
+  def testReset(self):
+    self._physics.reset()
+    self.assertEqual(self._physics.data.qpos[1], 0)
+    keyframe_id = 0
+    self._physics.reset(keyframe_id=keyframe_id)
+    self.assertEqual(self._physics.data.qpos[1],
+                     self._physics.model.key_qpos[keyframe_id, 1])
+    out_of_range = [-1, 3]
+    max_valid = self._physics.model.nkey - 1
+    for actual in out_of_range:
+      with self.assertRaisesWithLiteralMatch(
+          ValueError,
+          engine._KEYFRAME_ID_OUT_OF_RANGE.format(
+              max_valid=max_valid, actual=actual)):
+        self._physics.reset(keyframe_id=actual)
+
   def testLoadAndReloadFromStringWithAssets(self):
     physics = engine.Physics.from_xml_string(
         MODEL_WITH_ASSETS, assets=ASSETS)
diff --git a/dm_control/mujoco/testing/assets/cartpole.xml b/dm_control/mujoco/testing/assets/cartpole.xml
@@ -14,6 +14,8 @@ Actuators (name/actuator/parameter):
 <mujoco model='test_cartpole'>
   <compiler inertiafromgeom='true' coordinate='local'/>
 
+  <size nkey="1"/>
+
   <custom>
     <numeric name="control_timestep" data="0.04" />
     <numeric name="three_numbers" data="1.0 2.0 3.0" />
@@ -66,4 +68,8 @@ Actuators (name/actuator/parameter):
     <touch name="collision" site="cart sensor"/>
   </sensor>
 
+  <keyframe>
+    <key name="hanging_down" qpos="0 1.57"/>
+  </keyframe>
+
 </mujoco>
diff --git a/setup.py b/setup.py
@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.309253177',
+    version='0.0.309392225',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
