Merge branch 'main' into dependabot/pip/pillow-10.3.0

Author: nimrod-gileadi

dm_control/_render/__init__.py

@@ -45,6 +45,9 @@ def _import_glfw():
 def _import_osmesa():
   from dm_control._render.pyopengl.osmesa_renderer import OSMesaContext
   return OSMesaContext
+
+
+# Import removed.
 # pylint: enable=g-import-not-at-top
 
 
@@ -59,6 +62,7 @@ def no_renderer(*args, **kwargs):
     (constants.GLFW, _import_glfw),
     (constants.EGL, _import_egl),
     (constants.OSMESA, _import_osmesa),
+    # Option removed.
 )
 
 _NO_RENDERER = (

dm_control/_render/base.py

@@ -35,6 +35,8 @@
 
 from absl import logging
 from dm_control._render import executor
+import numpy as np
+
 
 _CURRENT_CONTEXT_FOR_THREAD = collections.defaultdict(lambda: None)
 _CURRENT_THREAD_FOR_CONTEXT = collections.defaultdict(lambda: None)
@@ -147,6 +149,10 @@ def make_current(self):
           ctx.call(self._platform_make_current)
       yield ctx
 
+  def to_pixels(self, buffer):
+    """Converts the buffer to pixels."""
+    return np.flipud(buffer)
+
   @abc.abstractmethod
   def _platform_init(self, max_width, max_height):
     """Performs an implementation-specific context initialization."""

dm_control/_render/constants.py

@@ -28,5 +28,6 @@
 OSMESA = ('osmesa',)
 GLFW = ('glfw', 'on', 'enable', 'enabled', 'true', '1', '')
 EGL = ('egl',)
+# Constant removed.
 NO_RENDERER = ('off', 'disable', 'disabled', 'false', '0')
 

dm_control/autowrap/header_parsing.py

@@ -78,11 +78,24 @@ def _nested_if_else(if_, pred, else_, endif, match_if_true, match_if_false):
   return ifelse
 
 
+def _nested_ifn_else(ifn_, pred, else_, endif, match_if_true, match_if_false):
+  """Constructs a parser for (possibly nested) if...(else)...endif blocks."""
+  ifnelse = pp.Forward()
+  ifnelse << pp.Group(  # pylint: disable=expression-not-assigned
+      ifn_ +
+      pred("predicate") +
+      pp.ZeroOrMore(match_if_true | ifnelse)("if_false") +
+      pp.Optional(else_ +
+                  pp.ZeroOrMore(match_if_false | ifnelse)("if_true")) +
+      endif)
+  return ifnelse
+
+
 # Some common string patterns to suppress.
 # ------------------------------------------------------------------------------
 (LPAREN, RPAREN, LBRACK, RBRACK, LBRACE, RBRACE, SEMI, COMMA, EQUAL, FSLASH,
  BSLASH) = list(map(pp.Suppress, "()[]{};,=/\\"))
-X = (pp.Keyword("X") | pp.Keyword("XMJV")).suppress()
+X = (pp.Keyword("X") | pp.Keyword("XMJV") | pp.Keyword("XNV")).suppress()
 EOL = pp.LineEnd().suppress()
 
 # Comments, continuation.
@@ -189,7 +202,9 @@ def _nested_if_else(if_, pred, else_, endif, match_if_true, match_if_false):
 UNCOND_DECL = DEF_FLAG | DEF_CONST | TYPE_DECL
 
 # Declarations inside (possibly nested) #if(n)def... #else... #endif... blocks.
-COND_DECL = _nested_if_else(IFDEF, NAME, ELSE, ENDIF, UNCOND_DECL, UNCOND_DECL)
+COND_DECL = _nested_if_else(
+    IFDEF, NAME, ELSE, ENDIF, UNCOND_DECL, UNCOND_DECL
+) | _nested_ifn_else(IFNDEF, NAME, ELSE, ENDIF, UNCOND_DECL, UNCOND_DECL)
 # Note: this doesn't work for '#if defined(FLAG)' blocks
 
 # e.g. "mjtNum gravity[3];              // gravitational acceleration"

dm_control/blender/mujoco_exporter/blender_scene.py

@@ -122,7 +122,7 @@ def is_none(self) -> bool:
 
   @property
   def is_armature(self) -> bool:
-    return (
+    return (  # pytype: disable=bad-return-type
         self.native_obj
         and self.native_obj.type == _ARMATURE
         and not self.native_bone
@@ -138,19 +138,19 @@ def is_bone(self) -> bool:
 
   @property
   def is_mesh(self) -> bool:
-    return self.native_obj and self.native_obj.type == _MESH
+    return self.native_obj and self.native_obj.type == _MESH  # pytype: disable=bad-return-type
 
   @property
   def is_light(self) -> bool:
-    return self.native_obj and self.native_obj.type == _LIGHT
+    return self.native_obj and self.native_obj.type == _LIGHT  # pytype: disable=bad-return-type
 
   @property
   def is_camera(self) -> bool:
-    return self.native_obj and self.native_obj.type == _CAMERA
+    return self.native_obj and self.native_obj.type == _CAMERA  # pytype: disable=bad-return-type
 
   @property
   def is_empty(self) -> bool:
-    return self.native_obj and self.native_obj.type == _EMPTY
+    return self.native_obj and self.native_obj.type == _EMPTY  # pytype: disable=bad-return-type
 
   @property
   def name(self) -> str:

dm_control/blender/mujoco_exporter/mujoco_scene.py

@@ -74,7 +74,7 @@ def light_builder(
   el.setAttribute('pos', vec_to_mjcf(transform.pos))
   el.setAttribute('dir', vec_to_mjcf(transform.rot @ _OZ))
   el.setAttribute('directional', directional)
-  el.setAttribute('castshadow', bool_to_mjcf(light.use_shadow))
+  el.setAttribute('castshadow', bool_to_mjcf(light.use_shadow))  # pytype: disable=wrong-arg-types
   el.setAttribute('diffuse', color_to_mjcf(light.color))
   el.setAttribute('attenuation', attenuation)
   return el

dm_control/composer/entity.py

@@ -480,7 +480,8 @@ def set_pose(self, physics, position=None, quaternion=None):
       if position is not None:
         physics.bind(root_joint).qpos[:3] = position
       if quaternion is not None:
-        physics.bind(root_joint).qpos[3:] = quaternion
+        normalised_quaternion = quaternion / np.linalg.norm(quaternion)
+        physics.bind(root_joint).qpos[3:] = normalised_quaternion
     else:
       attachment_frame = mjcf.get_attachment_frame(self.mjcf_model)
       if attachment_frame is None:
@@ -521,7 +522,7 @@ def shift_pose(self,
     if position is not None:
       new_position = current_position + position
     if quaternion is not None:
-      quaternion = np.array(quaternion, dtype=np.float64, copy=False)
+      quaternion = np.asarray(quaternion, dtype=np.float64)
       new_quaternion = _multiply_quaternions(quaternion, current_quaternion)
       root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
       if root_joint and rotate_velocity:

dm_control/composer/environment.py

@@ -249,8 +249,9 @@ def _recompile_physics_and_update_observables(self):
 
   def _recompile_physics(self):
     """Creates a new Physics using the latest MJCF model from the task."""
-    if getattr(self, '_physics', None):
-      self._physics.free()
+    physics = getattr(self, '_physics', None)
+    if physics:
+      physics.free()
     self._physics = mjcf.Physics.from_mjcf_model(
         self._task.root_entity.mjcf_model)
     self._physics.legacy_step = self._legacy_step
@@ -267,9 +268,10 @@ def physics(self):
     """Returns a `weakref.ProxyType` pointing to the current `mjcf.Physics`.
 
     Note that the underlying `mjcf.Physics` will be destroyed whenever the MJCF
-    model is recompiled. It is therefore unsafe for external objects to hold a
-    reference to `environment.physics`. Attempting to access attributes of a
-    dead `Physics` instance will result in a `ReferenceError`.
+    model is recompiled or environment.close() is called. It is therefore unsafe
+    for external objects to hold a reference to `environment.physics`.
+    Attempting to access attributes of a dead `Physics` instance will result in
+    a `ReferenceError`.
     """
     return self._physics_proxy
 
@@ -462,6 +464,11 @@ def _substep(self, action):
     self._physics.step()
     self._hooks.after_substep(self._physics_proxy, self._random_state)
 
+  def close(self):
+    super().close()
+    self._physics.free()
+    self._physics = None
+
   def action_spec(self):
     """Returns the action specification for this environment."""
     return self._task.action_spec(self._physics_proxy)

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)
@@ -237,22 +241,26 @@ 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(
+        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

@@ -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
@@ -111,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):
@@ -143,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
@@ -160,6 +162,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 +184,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):