Add Flex component.

PiperOrigin-RevId: 572830650
Change-Id: Ia47527e0acdb5760173549014e23d30036ee2971

Author: DeepMind

dm_control/mujoco/engine.py

@@ -61,7 +61,7 @@
 
 Contexts = collections.namedtuple('Contexts', ['gl', 'mujoco'])
 Selected = collections.namedtuple(
-    'Selected', ['body', 'geom', 'skin', 'world_position'])
+    'Selected', ['body', 'geom', 'flex', 'skin', 'world_position'])
 NamedIndexStructs = collections.namedtuple(
     'NamedIndexStructs', ['model', 'data'])
 Pose = collections.namedtuple(
@@ -937,12 +937,14 @@ def select(self, cursor_position):
     cursor_x, cursor_y = cursor_position
     pos = np.empty(3, np.double)
     geom_id_arr = np.intc([-1])
+    flex_id_arr = np.intc([-1])
     skin_id_arr = np.intc([-1])
     body_id = mujoco.mjv_select(self._physics.model.ptr, self._physics.data.ptr,
                                 self._scene_option.ptr, aspect_ratio, cursor_x,
                                 cursor_y, self._scene.ptr, pos, geom_id_arr,
-                                skin_id_arr)
+                                flex_id_arr, skin_id_arr)
     [geom_id] = geom_id_arr
+    [flex_id] = flex_id_arr
     [skin_id] = skin_id_arr
 
     # Validate IDs
@@ -954,6 +956,10 @@ def select(self, cursor_position):
       assert 0 <= geom_id < self._physics.model.ngeom
     else:
       geom_id = None
+    if flex_id != -1:
+      assert 0 <= flex_id < self._physics.model.nflex
+    else:
+      flex_id = None
     if skin_id != -1:
       assert 0 <= skin_id < self._physics.model.nskin
     else:
@@ -963,7 +969,12 @@ def select(self, cursor_position):
       pos = None
 
     return Selected(
-        body=body_id, geom=geom_id, skin=skin_id, world_position=pos)
+        body=body_id,
+        geom=geom_id,
+        flex=flex_id,
+        skin=skin_id,
+        world_position=pos,
+    )
 
 
 class MovableCamera(Camera):

dm_control/viewer/renderer.py

@@ -494,6 +494,7 @@ def raycast(self, viewport, screen_pos):
     viewport_pos = viewport.screen_to_inverse_viewport(screen_pos)
     grab_world_pos = np.empty(3, dtype=np.double)
     selected_geom_id_arr = np.intc([-1])
+    selected_flex_id_arr = np.intc([-1])
     selected_skin_id_arr = np.intc([-1])
     selected_body_id = mujoco.mjv_select(
         self._model.ptr,
@@ -505,9 +506,14 @@ def raycast(self, viewport, screen_pos):
         self._scene.ptr,
         grab_world_pos,
         selected_geom_id_arr,
+        selected_flex_id_arr,
         selected_skin_id_arr,
     )
-    del selected_geom_id_arr, selected_skin_id_arr  # Unused.
+    del (
+        selected_geom_id_arr,
+        selected_skin_id_arr,
+        selected_flex_id_arr,
+    )  # Unused.
     if selected_body_id < 0:
       selected_body_id = _INVALID_BODY_INDEX
       grab_world_pos = None

dm_control/viewer/renderer_test.py

@@ -422,12 +422,25 @@ def initialize_camera(self, enable):
 
   def test_raycast_mapping_geom_to_body_id(self):
     def build_mjv_select(mock_body_id, mock_geom_id, mock_position):
+
       def mock_select(
-          m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid, skinid):
-        del m, d, vopt, aspectratio, relx, rely, scn, skinid  # Unused.
+          m,
+          d,
+          vopt,
+          aspectratio,
+          relx,
+          rely,
+          scn,
+          selpnt,
+          geomid,
+          flexid,
+          skinid,
+      ):
+        del m, d, vopt, aspectratio, relx, rely, scn, flexid, skinid  # Unused.
         selpnt[:] = mock_position
         geomid[:] = mock_geom_id
         return mock_body_id
+
       return mock_select
 
     geom_id = 0
@@ -444,9 +457,11 @@ def mock_select(
       np.testing.assert_array_equal(hit_world_pos, world_pos)
 
   def test_raycast_hitting_empty_space(self):
+
     def mock_select(
-        m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid, skinid):
-      del (m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid,
+        m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid, flexid, skinid
+    ):
+      del (m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid, flexid,
            skinid)  # Unused.
       mock_body_id = -1  # Nothing selected.
       return mock_body_id
@@ -459,13 +474,26 @@ def mock_select(
 
   def test_raycast_maps_coordinates_to_viewport_space(self):
     def build_mjv_select(expected_aspect_ratio, expected_viewport_pos):
+
       def mock_select(
-          m, d, vopt, aspectratio, relx, rely, scn, selpnt, geomid, skinid):
-        del m, d, vopt, scn, selpnt, geomid, skinid  # Unused.
+          m,
+          d,
+          vopt,
+          aspectratio,
+          relx,
+          rely,
+          scn,
+          selpnt,
+          geomid,
+          flexid,
+          skinid,
+      ):
+        del m, d, vopt, scn, selpnt, geomid, flexid, skinid  # Unused.
         self.assertEqual(expected_aspect_ratio, aspectratio)
         np.testing.assert_array_equal(expected_viewport_pos, [relx, rely])
         mock_body_id = 0
         return mock_body_id
+
       return mock_select
 
     viewport_pos = [.5, .5]