Reset physics time after settling physics in PropPlacer.

Before this change, physics.data.time would be nonzero when PropPlacer is used with settle_physics=True, leading to episodes of variable length when a time limit is set.

PiperOrigin-RevId: 689823887
Change-Id: Ie759cb8331906b224b635ed3321f7dd2216dc90b

Author: nimrod-gileadi

dm_control/composer/initializers/prop_initializer.py

@@ -237,19 +237,23 @@ def place_and_settle():
 
         # Step physics and check prop states.
         original_time = physics.data.time
-        props_isolator = utils.JointStaticIsolator(physics, self._prop_joints)
-        prop_joints_mj = physics.bind(self._prop_joints)
-        while physics.data.time - original_time < self._max_settle_physics_time:
-          with props_isolator:
-            physics.step()
-          max_qvel = np.max(np.abs(prop_joints_mj.qvel))
-          max_qacc = np.max(np.abs(prop_joints_mj.qacc))
-          if (max_qvel < self._max_qvel_tol) and (
-              max_qacc < self._max_qacc_tol) and (
-                  physics.data.time - original_time
-                  ) > self._min_settle_physics_time:
-            return True
-        physics.data.time = original_time
+        try:
+          props_isolator = utils.JointStaticIsolator(physics, self._prop_joints)
+          prop_joints_mj = physics.bind(self._prop_joints)
+          while (
+              physics.data.time - original_time < self._max_settle_physics_time
+          ):
+            with props_isolator:
+              physics.step()
+            max_qvel = np.max(np.abs(prop_joints_mj.qvel))
+            max_qacc = np.max(np.abs(prop_joints_mj.qacc))
+            if (max_qvel < self._max_qvel_tol) and (
+                max_qacc < self._max_qacc_tol) and (
+                    physics.data.time - original_time
+                    ) > self._min_settle_physics_time:
+              return True
+        finally:
+          physics.data.time = original_time
 
       if self._raise_exception_on_settle_failure:
         raise RuntimeError(

dm_control/composer/initializers/prop_initializer_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for prop_initializer."""
-
 from absl.testing import absltest
 from absl.testing import parameterized
 from dm_control import composer
@@ -160,6 +158,8 @@ def test_ignore_contacts_with_entities(self):
   def test_settle_physics(self, settle_physics):
     radius = 0.1
     physics, spheres = _make_spheres(num_spheres=2, radius=radius, nconmax=1)
+    physics_start_time = 1337.0
+    physics.data.time = physics_start_time
 
     # Only place the first sphere.
     prop_placer = prop_initializer.PropPlacer(
@@ -180,7 +180,10 @@ def test_settle_physics(self, settle_physics):
     del second_quaternion  # Unused.
 
     # The sphere that we were not placing should not have moved.
-    self.assertEqual(second_position[2], 0.)
+    self.assertEqual(second_position[2], 0.0)
+    self.assertEqual(
+        physics.data.time, physics_start_time, 'Physics time should be reset.'
+    )
 
   @parameterized.parameters([0, 1, 2, 3])
   def test_settle_physics_multiple_attempts(self, max_settle_physics_attempts):