Throw EpisodeInitializationError instead of raw RuntimeError when prop initializer fails.

The more specific error allows callers to catch it if need be. Existing code should continue to work because EpisodeInitializationError is a subclass of RuntimeError.

PiperOrigin-RevId: 718344531
Change-Id: I344dd518dab878b36e9a99cd0f4ea3ca4fffcf2a

Author: nimrod-gileadi

dm_control/composer/initializers/prop_initializer.py

@@ -88,7 +88,7 @@ def __init__(self,
         velocities are less than this threshold.
       max_attempts_per_prop: The maximum number of rejection sampling attempts
         per prop. If a non-colliding pose cannot be found before this limit is
-        reached, a `RuntimeError` will be raised.
+        reached, an `EpisodeInitializationError` will be raised.
       settle_physics: (optional) If True, the physics simulation will be
         advanced for a few steps to allow the prop positions to settle.
       min_settle_physics_time: (optional) When `settle_physics` is True, lower
@@ -170,8 +170,9 @@ def __call__(self, physics, random_state, ignore_contacts_with_entities=None):
         subsequently).
 
     Raises:
-      RuntimeError: If `ignore_collisions == False` and a non-colliding prop
-        pose could not be found within `max_attempts_per_prop`.
+      EpisodeInitializationError: If `ignore_collisions == False` and a
+        non-colliding prop pose could not be found within
+        `max_attempts_per_prop`.
     """
     if ignore_contacts_with_entities is None:
       ignore_contacts_with_entities = []
@@ -222,9 +223,12 @@ def place_props():
             break
 
         if not success:
-          raise RuntimeError(_REJECTION_SAMPLING_FAILED.format(
-              model_name=prop.mjcf_model.model,
-              max_attempts=self._max_attempts_per_prop))
+          raise composer.EpisodeInitializationError(
+              _REJECTION_SAMPLING_FAILED.format(
+                  model_name=prop.mjcf_model.model,
+                  max_attempts=self._max_attempts_per_prop,
+              )
+          )
 
       for prop in ignore_contacts_with_entities:
         self._restore_contact_parameters(physics, prop, cached_contact_params)
@@ -256,7 +260,7 @@ def place_and_settle():
           physics.data.time = original_time
 
       if self._raise_exception_on_settle_failure:
-        raise RuntimeError(
+        raise composer.EpisodeInitializationError(
             _SETTLING_PHYSICS_FAILED.format(
                 max_attempts=self._max_settle_physics_attempts,
                 max_time=self._max_settle_physics_time,

dm_control/composer/initializers/prop_initializer_test.py

@@ -109,7 +109,9 @@ def test_rejection_sampling_failure(self):
     expected_message = prop_initializer._REJECTION_SAMPLING_FAILED.format(
         model_name=spheres[1].mjcf_model.model,  # Props are placed in order.
         max_attempts=max_attempts_per_prop)
-    with self.assertRaisesWithLiteralMatch(RuntimeError, expected_message):
+    with self.assertRaisesWithLiteralMatch(
+        composer.EpisodeInitializationError, expected_message
+    ):
       prop_placer(physics, random_state=np.random.RandomState(0))
 
   def test_ignore_contacts_with_entities(self):
@@ -141,7 +143,9 @@ def test_ignore_contacts_with_entities(self):
     prop_placer_init(physics, random_state=np.random.RandomState(0))
     expected_message = prop_initializer._REJECTION_SAMPLING_FAILED.format(
         model_name=spheres[0].mjcf_model.model, max_attempts=1)
-    with self.assertRaisesWithLiteralMatch(RuntimeError, expected_message):
+    with self.assertRaisesWithLiteralMatch(
+        composer.EpisodeInitializationError, expected_message
+    ):
       prop_placer_seq[0](physics, random_state=np.random.RandomState(0))
 
     # Placing the first sphere should succeed if we ignore contacts involving

dm_control/composer/initializers/tcp_initializer.py

@@ -131,8 +131,8 @@ def __call__(self, physics, random_state):
       random_state: An `np.random.RandomState` instance.
 
     Raises:
-      RuntimeError: If a collision-free pose could not be found within
-        `max_ik_attempts`.
+      composer.EpisodeInitializationError: If a collision-free pose could not be
+      found within `max_ik_attempts`.
     """
     if self._hand is not None:
       target_site = self._hand.tool_center_point
@@ -162,6 +162,9 @@ def __call__(self, physics, random_state):
       # positions and try again with a new target.
       physics.bind(self._arm.joints).qpos = initial_qpos
 
-    raise RuntimeError(_REJECTION_SAMPLING_FAILED.format(
-        max_rejection_samples=self._max_rejection_samples,
-        max_ik_attempts=self._max_ik_attempts))
+    raise composer.EpisodeInitializationError(
+        _REJECTION_SAMPLING_FAILED.format(
+            max_rejection_samples=self._max_rejection_samples,
+            max_ik_attempts=self._max_ik_attempts,
+        )
+    )

dm_control/composer/initializers/tcp_initializer_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for tcp_initializer."""
-
 import functools
 
 from absl.testing import absltest
@@ -105,7 +103,7 @@ def test_exception_if_hand_colliding_with_fixed_body(self):
 
     initializer = make_initializer()
     with self.assertRaisesWithLiteralMatch(
-        RuntimeError,
+        composer.EpisodeInitializationError,
         tcp_initializer._REJECTION_SAMPLING_FAILED.format(
             max_rejection_samples=max_rejection_samples,
             max_ik_attempts=max_ik_attempts)):
@@ -145,7 +143,7 @@ def test_exception_if_self_collision(self, with_hand):
 
     initializer = make_initializer()
     with self.assertRaisesWithLiteralMatch(
-        RuntimeError,
+        composer.EpisodeInitializationError,
         tcp_initializer._REJECTION_SAMPLING_FAILED.format(
             max_rejection_samples=max_rejection_samples,
             max_ik_attempts=max_ik_attempts)):