Add interactive body dragging with MjvPerturb forces and mocap support

src/mjviser/__init__.py

@@ -1,4 +1,5 @@
 """Web-based MuJoCo viewer powered by viser."""
 
+from mjviser.scene import BodyDragEvent as BodyDragEvent
 from mjviser.scene import ViserMujocoScene as ViserMujocoScene
 from mjviser.viewer import Viewer as Viewer

src/mjviser/scene.py

@@ -3,9 +3,11 @@
 from __future__ import annotations
 
 import dataclasses
+import inspect
 from collections import defaultdict
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from threading import RLock
+from typing import Any, Literal
 
 import mujoco
 import numpy as np
@@ -23,6 +25,36 @@
   merge_sites,
 )
 
+DragButton = Literal["left", "middle", "right"]
+DragPhase = Literal["start", "update", "end"]
+KeyModifier = Literal[
+  "cmd/ctrl",
+  "alt",
+  "shift",
+  "cmd/ctrl+alt",
+  "cmd/ctrl+shift",
+  "alt+shift",
+  "cmd/ctrl+alt+shift",
+]
+
+
+@dataclasses.dataclass(frozen=True)
+class BodyDragEvent:
+  """Event passed to body drag callbacks registered on ``ViserMujocoScene``."""
+
+  client: viser.ClientHandle
+  phase: DragPhase
+  body_id: int
+  body_name: str
+  env_idx: int
+  mocap_id: int | None
+  start_position: tuple[float, float, float]
+  end_position: tuple[float, float, float]
+  start_screen_pos: tuple[float, float]
+  end_screen_pos: tuple[float, float]
+  button: DragButton
+  modifier: KeyModifier | None
+
 
 @dataclasses.dataclass
 class _MeshGroup:
@@ -34,6 +66,18 @@ class _MeshGroup:
   mocap_ids: np.ndarray | None  # (N,) int32 mocap IDs, or None if non-mocap
 
 
+@dataclasses.dataclass
+class _BodyDragRegistration:
+  """Registered body drag callback and its per-handle Viser wrapper functions."""
+
+  button: DragButton
+  modifier: KeyModifier | None
+  callback: Callable[[BodyDragEvent], Any]
+  bound_wrappers: list[tuple[Any, Callable[..., Any]]] = dataclasses.field(
+    default_factory=list
+  )
+
+
 # Viser visualization defaults.
 _DEFAULT_FOV_DEGREES = 60
 _DEFAULT_FOV_MIN = 20
@@ -146,6 +190,7 @@ def __init__(
     self._show_convex_hull: bool = False
     self._hull_hide_meshes: bool = False
     self._autoconnect_hide_meshes: bool = False
+    self._body_drag_callbacks: list[_BodyDragRegistration] = []
 
     # Visualization settings.
     self.env_idx = 0
@@ -265,11 +310,87 @@ def show_convex_hull(self, value: bool) -> None:
         handle.visible = value
       self._sync_visibilities()
 
+  @property
+  def scene_offset(self) -> np.ndarray:
+    """Current world-to-scene translation offset applied for camera tracking."""
+    with self._update_lock:
+      return self._scene_offset.copy()
+
   def set_refresh_handler(self, handler: Callable[[], None] | None) -> None:
     """Install a callback used to refresh cached visualization state."""
     with self._update_lock:
       self._refresh_handler = handler
 
+  def add_body_drag_callback(
+    self,
+    callback: Callable[[BodyDragEvent], Any],
+    button: DragButton = "left",
+    modifier: KeyModifier | None = "cmd/ctrl",
+  ) -> None:
+    """Register a callback to receive ``BodyDragEvent`` updates when a body is dragged.
+
+    Args:
+      callback: Function or coroutine function taking a ``BodyDragEvent``.
+      button: Mouse button that triggers the drag (``"left"``, ``"middle"``, ``"right"``).
+      modifier: Key modifier combination required at drag start, or ``None``.
+    """
+    with self._update_lock:
+      for reg in self._body_drag_callbacks:
+        if (
+          reg.button == button and reg.modifier == modifier and reg.callback == callback
+        ):
+          return
+      reg = _BodyDragRegistration(button=button, modifier=modifier, callback=callback)
+      self._body_drag_callbacks.append(reg)
+      self._bind_drag_callbacks(registrations=[reg])
+
+  def on_body_drag(
+    self,
+    button: DragButton | Callable[[BodyDragEvent], Any] = "left",
+    *,
+    modifier: KeyModifier | None = "cmd/ctrl",
+  ) -> Any:
+    """Decorator to register a body drag callback on the scene.
+
+    Can be used as ``@scene.on_body_drag`` or
+    ``@scene.on_body_drag("left", modifier="cmd/ctrl")``.
+    """
+    if callable(button):
+      self.add_body_drag_callback(button, button="left", modifier=modifier)
+      return button
+
+    def _decorator(
+      cb: Callable[[BodyDragEvent], Any],
+    ) -> Callable[[BodyDragEvent], Any]:
+      self.add_body_drag_callback(cb, button=button, modifier=modifier)
+      return cb
+
+    return _decorator
+
+  def remove_body_drag_callback(self, callback: Callable[[BodyDragEvent], Any]) -> None:
+    """Remove a specific body drag callback without disturbing other registrations."""
+    with self._update_lock:
+      live_ids = self._live_drag_handle_ids()
+      remaining: list[_BodyDragRegistration] = []
+      for reg in self._body_drag_callbacks:
+        if reg.callback == callback:
+          for handle, wrapper in reg.bound_wrappers:
+            if id(handle) in live_ids:
+              handle.remove_drag_callback(wrapper)
+        else:
+          remaining.append(reg)
+      self._body_drag_callbacks = remaining
+
+  def remove_body_drag_callbacks(self) -> None:
+    """Remove all registered body drag callbacks from the scene."""
+    with self._update_lock:
+      live_ids = self._live_drag_handle_ids()
+      for reg in self._body_drag_callbacks:
+        for handle, wrapper in reg.bound_wrappers:
+          if id(handle) in live_ids:
+            handle.remove_drag_callback(wrapper)
+      self._body_drag_callbacks.clear()
+
   def rebuild_visual_handles(self) -> None:
     """Rebuild handles whose geometry or appearance is baked at creation time."""
     with self._update_lock:
@@ -1109,6 +1230,7 @@ def _add_fixed_geometry(self) -> None:
           visible=visible,
         )
         self._fixed_geom_handles[(body_id, group_id, sub_idx)] = handle
+    self._bind_drag_callbacks(kinds=("fixed",))
 
   @staticmethod
   def _geom_subgroup_fingerprint(
@@ -1187,6 +1309,7 @@ def _create_mesh_handles_by_group(self) -> None:
             [self.mj_model.body_mocapid[b] for b in body_ids], dtype=np.int32
           )
         self._mesh_groups.append(_MeshGroup(handle, body_ids_arr, group_id, mocap_ids))
+    self._bind_drag_callbacks(kinds=("mesh",))
 
   def _add_fixed_sites(self) -> None:
     """Add fixed site geometry to the scene as static nodes."""
@@ -1313,6 +1436,7 @@ def _build_hull_handles(self) -> None:
         lod="off",
       )
       self._hull_dynamic_handles.append((handle, body_id))
+    self._bind_drag_callbacks(kinds=("hull",))
 
   def _clear_hull_handles(self) -> None:
     """Remove convex hull handles and clear their caches."""
@@ -1324,6 +1448,147 @@ def _clear_hull_handles(self) -> None:
       handle.remove()
     self._hull_dynamic_handles.clear()
 
+  def _live_drag_handle_ids(self) -> set[int]:
+    """Return object IDs of currently active scene handles that support body dragging."""
+    live: set[int] = {id(mg.handle) for mg in self._mesh_groups}
+    live.update(id(handle) for handle in self._fixed_geom_handles.values())
+    live.update(id(handle) for handle, _ in self._hull_dynamic_handles)
+    return live
+
+  def _bind_drag_callbacks(
+    self,
+    kinds: tuple[Literal["mesh", "fixed", "hull"], ...] = ("mesh", "fixed", "hull"),
+    registrations: list[_BodyDragRegistration] | None = None,
+  ) -> None:
+    """Bind registered body drag callbacks to the specified handle collections."""
+    live_ids = self._live_drag_handle_ids()
+    regs = registrations if registrations is not None else self._body_drag_callbacks
+    for reg in regs:
+      reg.bound_wrappers = [(h, w) for h, w in reg.bound_wrappers if id(h) in live_ids]
+      if "mesh" in kinds:
+        for mg in self._mesh_groups:
+          wrapper = self._make_mesh_group_drag_handler(mg, reg.callback)
+          mg.handle.on_drag(reg.button, modifier=reg.modifier)(wrapper)
+          reg.bound_wrappers.append((mg.handle, wrapper))
+      if "fixed" in kinds:
+        for (body_id, _, _), handle in self._fixed_geom_handles.items():
+          wrapper = self._make_fixed_drag_handler(body_id, reg.callback)
+          handle.on_drag(reg.button, modifier=reg.modifier)(wrapper)
+          reg.bound_wrappers.append((handle, wrapper))
+      if "hull" in kinds:
+        for handle, body_id in self._hull_dynamic_handles:
+          wrapper = self._make_hull_drag_handler(body_id, reg.callback)
+          handle.on_drag(reg.button, modifier=reg.modifier)(wrapper)
+          reg.bound_wrappers.append((handle, wrapper))
+
+  def _make_mesh_group_drag_handler(
+    self,
+    mg: _MeshGroup,
+    callback: Callable[[BodyDragEvent], Any],
+  ) -> Callable[[viser.SceneNodeDragEvent[Any]], Coroutine[Any, Any, None]]:
+    async def _handler(event: viser.SceneNodeDragEvent[Any]) -> None:
+      idx = event.instance_index
+      if idx is None:
+        return
+      n_bodies = len(mg.body_ids)
+      if n_bodies == 0:
+        return
+      with self._update_lock:
+        slice_single = self.show_only_selected and self.num_envs > 1
+        last_env_idx = self._last_env_idx
+        offset = self._scene_offset.copy()
+      body_idx = idx % n_bodies
+      env_idx = last_env_idx if slice_single else idx // n_bodies
+      if env_idx < 0 or env_idx >= self.num_envs:
+        return
+      body_id = int(mg.body_ids[body_idx])
+      mocap_id = int(mg.mocap_ids[body_idx]) if mg.mocap_ids is not None else None
+      await self._dispatch_body_drag(
+        event, body_id, env_idx, mocap_id, offset, callback
+      )
+
+    return _handler
+
+  def _make_hull_drag_handler(
+    self,
+    body_id: int,
+    callback: Callable[[BodyDragEvent], Any],
+  ) -> Callable[[viser.SceneNodeDragEvent[Any]], Coroutine[Any, Any, None]]:
+    raw_mocap = int(self.mj_model.body_mocapid[body_id])
+    mocap_id = raw_mocap if raw_mocap >= 0 else None
+
+    async def _handler(event: viser.SceneNodeDragEvent[Any]) -> None:
+      idx = event.instance_index
+      if idx is None:
+        return
+      with self._update_lock:
+        slice_single = self.show_only_selected and self.num_envs > 1
+        last_env_idx = self._last_env_idx
+        offset = self._scene_offset.copy()
+      env_idx = last_env_idx if slice_single else idx
+      if env_idx < 0 or env_idx >= self.num_envs:
+        return
+      await self._dispatch_body_drag(
+        event, body_id, env_idx, mocap_id, offset, callback
+      )
+
+    return _handler
+
+  def _make_fixed_drag_handler(
+    self,
+    body_id: int,
+    callback: Callable[[BodyDragEvent], Any],
+  ) -> Callable[[viser.SceneNodeDragEvent[Any]], Coroutine[Any, Any, None]]:
+    async def _handler(event: viser.SceneNodeDragEvent[Any]) -> None:
+      with self._update_lock:
+        env_idx = self._last_env_idx
+        offset = self._scene_offset.copy()
+      await self._dispatch_body_drag(event, body_id, env_idx, None, offset, callback)
+
+    return _handler
+
+  async def _dispatch_body_drag(
+    self,
+    event: viser.SceneNodeDragEvent[Any],
+    body_id: int,
+    env_idx: int,
+    mocap_id: int | None,
+    offset: np.ndarray,
+    callback: Callable[[BodyDragEvent], Any],
+  ) -> None:
+    # Handlers are declared ``async def`` and invoke synchronous callbacks inline on
+    # Viser's event loop on purpose: Viser submits non-async callbacks to a thread pool
+    # where ``update`` / ``end`` phases can race ahead of ``start``. Running inline on
+    # the event loop guarantees strict ``start`` -> ``update`` -> ``end`` ordering, at
+    # the cost of briefly blocking the event loop if a callback waits on a lock.
+    start_pos = (
+      float(event.start_position[0] - offset[0]),
+      float(event.start_position[1] - offset[1]),
+      float(event.start_position[2] - offset[2]),
+    )
+    end_pos = (
+      float(event.end_position[0] - offset[0]),
+      float(event.end_position[1] - offset[1]),
+      float(event.end_position[2] - offset[2]),
+    )
+    body_event = BodyDragEvent(
+      client=event.client,
+      phase=event.phase,
+      body_id=body_id,
+      body_name=get_body_name(self.mj_model, body_id),
+      env_idx=env_idx,
+      mocap_id=mocap_id,
+      start_position=start_pos,
+      end_position=end_pos,
+      start_screen_pos=event.start_screen_pos,
+      end_screen_pos=event.end_screen_pos,
+      button=event.button,
+      modifier=event.modifier,
+    )
+    res = callback(body_event)
+    if inspect.isawaitable(res):
+      await res
+
   def _clear_decor_handles(self) -> None:
     """Remove decor handles and clear their cache."""
     for handle in self._decor_handles.values():