Suppress exceptions due to physics warnings when initializing observation updaters

yuvaltassa · copybara-github · commit 409a0579c7f8 · 2019-12-03T01:04:06.000-08:00
In order to initialize the buffers used by the observation updater we need the shape and dtype of the array returned by each observation callable. We usually get this by calling the observation callable and inspecting the array it returns. For `MJCFFeature` observations, this may involve calling `physics.forward` if the feature depends on a field that requires recalculation due to other changes to the physics state. However, at the point when this happens there is no guarantee that the physics state will be valid, since it has not yet been fully initialized by `initialize_episode`. Because of this, the call to `physics.forward` might raise a `PhysicsError`.

In practice it does not matter whether or not the physics state is valid, since we only care about the shape and dtype of the observation array, not its contents. We therefore suppress `PhysicsError`s originating from this initial call to the observation callable.

PiperOrigin-RevId: 283497708
Change-Id: I0c7b1bf1e968dadfe8101fbcbe072b1687f5a4de
diff --git a/dm_control/composer/observation/fake_physics.py b/dm_control/composer/observation/fake_physics.py
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import contextlib
+
 from dm_control.composer.observation import observable
 from dm_control.rl import control
 import numpy as np
@@ -68,3 +70,7 @@ def reset(self):
 
   def after_reset(self):
     pass
+
+  @contextlib.contextmanager
+  def suppress_physics_errors(self):
+    yield
diff --git a/dm_control/composer/observation/updater.py b/dm_control/composer/observation/updater.py
@@ -46,11 +46,16 @@ def __init__(self, observable, physics, random_state,
 
     obs_spec = self.observable.array_spec
     if obs_spec is None:
-      # We take an observation here to determine the shape and size.
+      # We take an observation to determine the shape and dtype of the array.
       # This occurs outside of an episode and doesn't affect environment
-      # behavior.
-      obs_value = np.array(self.observation_callable())
-      obs_spec = specs.Array(shape=obs_value.shape, dtype=obs_value.dtype)
+      # behavior. At this point the physics state is not guaranteed to be valid,
+      # so we might get a `PhysicsError` if the observation callable calls
+      # `physics.forward`. We suppress such errors since they do not matter as
+      # far as the shape and dtype of the observation are concerned.
+      with physics.suppress_physics_errors():
+        obs_array = self.observation_callable()
+      obs_array = np.asarray(obs_array)
+      obs_spec = specs.Array(shape=obs_array.shape, dtype=obs_array.dtype)
     self.buffer = obs_buffer.Buffer(
         buffer_size=(observable.buffer_size or DEFAULT_BUFFER_SIZE),
         shape=obs_spec.shape, dtype=obs_spec.dtype,
diff --git a/dm_control/composer/observation/updater_test.py b/dm_control/composer/observation/updater_test.py
@@ -65,7 +65,8 @@ def testNestedSpecsAndValues(self, list_or_tuple):
     observables[1]['five'].enabled = True
 
     observation_updater = updater.Updater(observables)
-    observation_updater.reset(physics=None, random_state=None)
+    observation_updater.reset(physics=fake_physics.FakePhysics(),
+                              random_state=None)
 
     def make_spec(obs):
       array = np.array(obs.observation_callable(None, None)())
diff --git a/dm_control/mujoco/engine.py b/dm_control/mujoco/engine.py
@@ -120,8 +120,19 @@ def __init__(self, data):
     Args:
       data: Instance of `wrapper.MjData`.
     """
+    self._warnings_cause_exception = True
     self._reload_from_data(data)
 
+  @contextlib.contextmanager
+  def suppress_physics_errors(self):
+    """Physics warnings will be logged rather than raise exceptions."""
+    prev_state = self._warnings_cause_exception
+    self._warnings_cause_exception = False
+    try:
+      yield
+    finally:
+      self._warnings_cause_exception = prev_state
+
   def enable_profiling(self):
     """Enables Mujoco timing profiling."""
     wrapper.enable_timer(True)
@@ -263,15 +274,28 @@ def forward(self):
 
   @contextlib.contextmanager
   def check_invalid_state(self):
-    """Raises a `base.PhysicsError` if the simulation state is invalid."""
+    """Checks whether the physics state is invalid at exit.
+
+    Yields:
+      None
+
+    Raises:
+      PhysicsError: if the simulation state is invalid at exit, unless this
+        context is nested inside a `suppress_physics_errors` context, in which
+        case a warning will be logged instead.
+    """
     # `np.copyto(dst, src)` is marginally faster than `dst[:] = src`.
     np.copyto(self._warnings_before, self._warnings)
     yield
     np.greater(self._warnings, self._warnings_before, out=self._new_warnings)
     if any(self._new_warnings):
       warning_names = np.compress(self._new_warnings, enums.mjtWarning._fields)
-      raise _control.PhysicsError(
-          _INVALID_PHYSICS_STATE.format(warning_names=', '.join(warning_names)))
+      message = _INVALID_PHYSICS_STATE.format(
+          warning_names=', '.join(warning_names))
+      if self._warnings_cause_exception:
+        raise _control.PhysicsError(message)
+      else:
+        logging.warn(message)
 
   def __getstate__(self):
     return self.data  # All state is assumed to reside within `self.data`.
diff --git a/dm_control/mujoco/engine_test.py b/dm_control/mujoco/engine_test.py
@@ -346,6 +346,30 @@ def testNanControl(self):
         self._physics.data.ctrl[0] = float('nan')
         self._physics.step()
 
+  def testSuppressPhysicsError(self):
+    bad_value = float('nan')
+    message = engine._INVALID_PHYSICS_STATE.format(
+        warning_names='mjWARN_BADCTRL')
+
+    def assert_physics_error():
+      self._physics.data.ctrl[0] = bad_value
+      with self.assertRaisesWithLiteralMatch(control.PhysicsError, message):
+        self._physics.forward()
+
+    def assert_warning():
+      self._physics.data.ctrl[0] = bad_value
+      with mock.patch.object(engine.logging, 'warn') as mock_warn:
+        self._physics.forward()
+      mock_warn.assert_called_once_with(message)
+
+    assert_physics_error()
+    with self._physics.suppress_physics_errors():
+      assert_warning()
+      with self._physics.suppress_physics_errors():
+        assert_warning()
+      assert_warning()
+    assert_physics_error()
+
   @parameterized.named_parameters(
       ('_copy', lambda x: x.copy()),
       ('_pickle_and_unpickle', lambda x: cPickle.loads(cPickle.dumps(x))),
