mujoco: draw steps (#1622)

* implement sensor delay

* 3 it is

* fix type issues in NaoInterface

* fix formatting

* fix some more type issues

* --wip--

* arrow interface for plotting arrows

* track with camera

* ruff check

* remove commented code

* remove more commented code

* move ringbuffer

* move ringbuffer test

tools/machine-learning/mujoco/model/nao.xml

@@ -306,31 +306,35 @@
 
                   <site
                     name="force_sensitive_resistors.left.rear_right"
-                    pos="-0.02 0.025 -0.04519"
-                    size="0.04 0.025 0.005"
+                    pos="-0.02 0.03 -0.04519"
+                    size="0.04 0.03 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.left.rear_left"
-                    pos="-0.02 -0.017 -0.04519"
-                    size="0.04 0.017 0.005"
+                    pos="-0.02 -0.023 -0.04519"
+                    size="0.04 0.023 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.left.front_right"
-                    pos="0.06 0.025 -0.04519"
-                    size="0.04 0.025 0.005"
+                    pos="0.065 0.03 -0.04519"
+                    size="0.045 0.03 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.left.front_left"
-                    pos="0.06 -0.017 -0.04519"
-                    size="0.04 0.017 0.005"
+                    pos="0.065 -0.023 -0.04519"
+                    size="0.045 0.023 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                 </body>
               </body>
@@ -484,31 +488,35 @@
 
                   <site
                     name="force_sensitive_resistors.right.rear_right"
-                    pos="-0.02 -0.025 -0.04519"
-                    size="0.04 0.025 0.005"
+                    pos="-0.02 -0.03 -0.04519"
+                    size="0.04 0.03 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.right.rear_left"
-                    pos="-0.02 0.017 -0.04519"
-                    size="0.04 0.017 0.005"
+                    pos="-0.02 0.023 -0.04519"
+                    size="0.04 0.023 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.right.front_right"
-                    pos="0.06 -0.025 -0.04519"
-                    size="0.04 0.025 0.005"
+                    pos="0.065 -0.03 -0.04519"
+                    size="0.045 0.03 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                   <site
                     name="force_sensitive_resistors.right.front_left"
-                    pos="0.06 0.017 -0.04519"
-                    size="0.04 0.017 0.005"
+                    pos="0.065 0.023 -0.04519"
+                    size="0.045 0.023 0.005"
                     type="box"
                     rgba="1 0 0 0.3"
+                    group="3"
                   />
                 </body>
               </body>

tools/machine-learning/mujoco/model/scene.xml

@@ -8,6 +8,7 @@
     <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
     <rgba haze="0.15 0.25 0.35 1" />
     <global azimuth="160" elevation="-20" />
+    <scale contactwidth="0.075" contactheight="0.025" forcewidth="0.025" />
   </visual>
 
   <asset>

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/__init__.py

@@ -1,5 +1,6 @@
+from .context import current_viewer, set_global_viewer
 from .figure import Figure
 from .marker import Marker
 from .viewer import Viewer
 
-__all__ = ["Figure", "Marker", "Viewer"]
+__all__ = ["Figure", "Marker", "Viewer", "current_viewer", "set_global_viewer"]

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/context.py

@@ -0,0 +1,12 @@
+from mujoco_interactive_viewer.viewer import Viewer
+
+_viewer: Viewer | None = None
+
+
+def set_global_viewer(viewer: Viewer) -> None:
+    global _viewer
+    _viewer = viewer
+
+
+def current_viewer() -> Viewer | None:
+    return _viewer

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/figure.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 import mujoco
+import numpy as np
 
 
 @dataclass
@@ -140,7 +141,7 @@ def _line_id(self, line_name: str) -> int:
     def push_data_to_line(
         self,
         line_name: str,
-        line_data: float,
+        line_data: np.floating,
     ) -> None:
         line_id = self._line_id(line_name)
 

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/marker.py

@@ -10,5 +10,5 @@ class Marker:
     kind: mujoco.mjtGeom
     size: NDArray[np.float64]
     position: NDArray[np.float64]
-    material: NDArray[np.float64]
+    rotation_matrix: NDArray[np.float64]
     rgba: NDArray[np.float32]

tools/machine-learning/mujoco/packages/mujoco-interactive-viewer/src/mujoco_interactive_viewer/viewer.py

@@ -5,7 +5,8 @@
 import glfw
 import mujoco
 import numpy as np
-from numpy.typing import NDArray
+from mujoco._structs import MjvGeom
+from numpy.typing import ArrayLike, NDArray
 
 from mujoco_interactive_viewer.figure import Figure
 from mujoco_interactive_viewer.interaction import InteractionState
@@ -15,6 +16,32 @@
 from mujoco_interactive_viewer.visualization import VisualizationState
 
 
+def normalize(v: NDArray) -> NDArray:
+    return v / np.linalg.norm(v)
+
+
+def look_at(origin: ArrayLike, target: ArrayLike) -> NDArray:
+    origin = np.array(origin)
+    target = np.array(target)
+
+    reference_direction = np.array([0.0, 0.0, 1.0])
+    direction = normalize(target - origin)
+
+    axis_of_rotation = normalize(np.cross(reference_direction, direction))
+
+    cos_theta = np.dot(reference_direction, direction)
+    angle = np.arccos(cos_theta)
+    K = np.array(
+        [
+            [0, -axis_of_rotation[2], axis_of_rotation[1]],
+            [axis_of_rotation[2], 0, -axis_of_rotation[0]],
+            [-axis_of_rotation[1], axis_of_rotation[0], 0],
+        ]
+    )
+
+    return np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * np.dot(K, K)
+
+
 class Viewer:
     def __init__(
         self,
@@ -81,30 +108,76 @@ def __init__(
     def add_marker(
         self,
         kind: mujoco.mjtGeom,
-        size: NDArray[np.float64] | list[float],
-        position: NDArray[np.float64] | list[float],
-        rgba: NDArray[np.float32] | list[float],
-        material: NDArray[np.float64] | list[float] | None = None,
+        size: ArrayLike,
+        position: ArrayLike,
+        rgba: ArrayLike,
+        rotation_matrix: ArrayLike | None = None,
+    ) -> None:
+        size = np.asanyarray(size)
+        position = np.asanyarray(position)
+        rgba = np.asanyarray(rgba)
+        if rotation_matrix is None:
+            rotation_matrix = np.eye(3).flatten()
+        else:
+            rotation_matrix = np.asanyarray(rotation_matrix)
+        self._markers.append(
+            Marker(kind, size, position, rotation_matrix, rgba)
+        )
+
+    def mark_sphere(
+        self,
+        position: ArrayLike,
+        radius: np.floating | float,
+        rgba: ArrayLike,
     ) -> None:
-        size = np.array(size)
-        position = np.array(position)
-        if material is None:
-            material = np.eye(3).flatten()
-        material = np.array(material)
-        rgba = np.array(rgba)
-        self._markers.append(Marker(kind, size, position, material, rgba))
+        self.add_marker(
+            mujoco.mjtGeom.mjGEOM_SPHERE,
+            np.array([radius, 0, 0]),
+            position,
+            rgba,
+        )
+
+    def mark_box(
+        self,
+        position: ArrayLike,
+        size: ArrayLike,
+        rgba: ArrayLike,
+    ) -> None:
+        self.add_marker(
+            mujoco.mjtGeom.mjGEOM_BOX,
+            size,
+            position,
+            rgba,
+        )
+
+    def mark_arrow(
+        self,
+        position: ArrayLike,
+        direction: ArrayLike,
+        width: np.floating | float,
+        rgba: ArrayLike,
+    ) -> None:
+        size = np.array([width, width, np.linalg.norm(direction) * 2])
+        rotation_matrix = look_at(np.zeros(3), direction).flatten()
+        self.add_marker(
+            mujoco.mjtGeom.mjGEOM_ARROW,
+            size=size,
+            position=position,
+            rotation_matrix=rotation_matrix,
+            rgba=rgba,
+        )
 
     def _add_marker_to_scene(self, marker: Marker) -> None:
         if self.scene.ngeom >= self.scene.maxgeom:
             raise OutOfGeomsError()
 
-        geom = self.scene.geoms[self.scene.ngeom]
+        geom: MjvGeom = self.scene.geoms[self.scene.ngeom]
         mujoco.mjv_initGeom(
             geom,
             type=marker.kind.value,
             size=marker.size,
             pos=marker.position,
-            mat=marker.material,
+            mat=marker.rotation_matrix,
             rgba=marker.rgba,
         )
         self.scene.ngeom += 1
@@ -515,6 +588,12 @@ def _handle_perturbation(self, mods: int) -> None:
                 )
         self.perturbation.active = perturbation_kind
 
+    def track_with_camera(self, body: int | str) -> None:
+        bodyid = self.model.body(body).id
+        self.camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING.value
+        self.camera.trackbodyid = bodyid
+        self.camera.fixedcamid = -1
+
     def _handle_selection(self, mods: int) -> None:
         class Mode(Enum):
             Select = 1
@@ -529,7 +608,10 @@ class Mode(Enum):
             mode = None
             if self._interaction_state.left_double_click_pressed:
                 mode = Mode.Select
-            elif self._interaction_state.right_double_click_pressed:
+            elif (
+                self._interaction_state.right_double_click_pressed
+                and mods != glfw.MOD_CONTROL
+            ):
                 mode = Mode.LookAt
             elif (
                 self._interaction_state.right_double_click_pressed