From b76f477ee067c9a918e4904612d10f82f24774fd Mon Sep 17 00:00:00 2001 From: Shawn Maybush Date: Sat, 17 Oct 2026 01:11:36 -0400 Subject: [PATCH 1/3] Add interactive body dragging with spring perturbations and mocap support --- src/mjviser/__init__.py | 1 + src/mjviser/scene.py | 247 +++++++++++++++++++++++++++++++++++++- src/mjviser/viewer.py | 228 ++++++++++++++++++++++++++++++++++- tests/test_scene.py | 81 ++++++++++++- tests/test_viewer_drag.py | 204 +++++++++++++++++++++++++++++++ 5 files changed, 757 insertions(+), 4 deletions(-) create mode 100644 tests/test_viewer_drag.py diff --git a/src/mjviser/__init__.py b/src/mjviser/__init__.py index 80af911..5d8a8a6 100644 --- a/src/mjviser/__init__.py +++ b/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 diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 1b02976..86deec7 100644 --- a/src/mjviser/scene.py +++ b/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: @@ -146,6 +178,9 @@ 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[ + tuple[DragButton, KeyModifier | None, Callable[[BodyDragEvent], Any]] + ] = [] # Visualization settings. self.env_idx = 0 @@ -265,11 +300,71 @@ 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``. + """ + entry = (button, modifier, callback) + with self._update_lock: + if entry in self._body_drag_callbacks: + return + self._body_drag_callbacks.append(entry) + self._bind_drag_entry_to_all_handles(button, modifier, callback) + + 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_callbacks(self) -> None: + """Remove all registered body drag callbacks from the scene.""" + with self._update_lock: + self._body_drag_callbacks.clear() + for mg in self._mesh_groups: + mg.handle.remove_drag_callback() + for handle in self._fixed_geom_handles.values(): + handle.remove_drag_callback() + for handle, _ in self._hull_dynamic_handles: + handle.remove_drag_callback() + def rebuild_visual_handles(self) -> None: """Rebuild handles whose geometry or appearance is baked at creation time.""" with self._update_lock: @@ -1109,6 +1204,7 @@ def _add_fixed_geometry(self) -> None: visible=visible, ) self._fixed_geom_handles[(body_id, group_id, sub_idx)] = handle + self._bind_drag_callbacks_to_fixed_geoms() @staticmethod def _geom_subgroup_fingerprint( @@ -1187,6 +1283,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_to_mesh_groups() def _add_fixed_sites(self) -> None: """Add fixed site geometry to the scene as static nodes.""" @@ -1313,6 +1410,7 @@ def _build_hull_handles(self) -> None: lod="off", ) self._hull_dynamic_handles.append((handle, body_id)) + self._bind_drag_callbacks_to_hull_handles() def _clear_hull_handles(self) -> None: """Remove convex hull handles and clear their caches.""" @@ -1324,6 +1422,153 @@ def _clear_hull_handles(self) -> None: handle.remove() self._hull_dynamic_handles.clear() + def _bind_drag_entry_to_all_handles( + self, + button: DragButton, + modifier: KeyModifier | None, + callback: Callable[[BodyDragEvent], Any], + ) -> None: + """Bind a single (button, modifier, callback) entry to all current handles.""" + for mg in self._mesh_groups: + mg.handle.on_drag(button, modifier=modifier)( + self._make_mesh_group_drag_handler(mg, callback) + ) + for (body_id, _, _), handle in self._fixed_geom_handles.items(): + handle.on_drag(button, modifier=modifier)( + self._make_fixed_drag_handler(body_id, callback) + ) + for handle, body_id in self._hull_dynamic_handles: + handle.on_drag(button, modifier=modifier)( + self._make_hull_drag_handler(body_id, callback) + ) + + def _bind_drag_callbacks_to_mesh_groups(self) -> None: + """Attach all registered body drag callbacks to current mesh group handles.""" + for button, modifier, callback in self._body_drag_callbacks: + for mg in self._mesh_groups: + mg.handle.on_drag(button, modifier=modifier)( + self._make_mesh_group_drag_handler(mg, callback) + ) + + def _bind_drag_callbacks_to_fixed_geoms(self) -> None: + """Attach all registered body drag callbacks to current fixed geom handles.""" + for button, modifier, callback in self._body_drag_callbacks: + for (body_id, _, _), handle in self._fixed_geom_handles.items(): + handle.on_drag(button, modifier=modifier)( + self._make_fixed_drag_handler(body_id, callback) + ) + + def _bind_drag_callbacks_to_hull_handles(self) -> None: + """Attach all registered body drag callbacks to current dynamic hull handles.""" + for button, modifier, callback in self._body_drag_callbacks: + for handle, body_id in self._hull_dynamic_handles: + handle.on_drag(button, modifier=modifier)( + self._make_hull_drag_handler(body_id, callback) + ) + + 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: + 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(): diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index b35fb73..3c729e1 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -6,15 +6,20 @@ import time from collections.abc import Callable from threading import Lock +from typing import Literal import mujoco import numpy as np import viser +import viser.transforms as vtf -from .scene import ViserMujocoScene +from .conversions import is_fixed_body +from .scene import BodyDragEvent, ViserMujocoScene _SPEEDS = [1 / 8, 1 / 4, 1 / 2, 1.0, 2.0, 4.0, 8.0] _FRAME_TIME = 1.0 / 60.0 # 60 Hz render target. +_TRANSLATE_LINE_COLOR = (255, 120, 60) +_ROTATE_LINE_COLOR = (120, 220, 120) def _format_speed(multiplier: float) -> str: @@ -77,6 +82,26 @@ def __init__( self._lock = Lock() self.scene.set_refresh_handler(self._refresh_scene_from_gui) + # Perturbation / drag state. + self.perturbation_stiffness: float = 50.0 + self.perturbation_torque_scale: float = 2.0 + self._drag_body_id: int | None = None + self._drag_mocap_id: int | None = None + self._drag_mode: Literal["translate", "rotate"] | None = None + self._drag_grab_local: np.ndarray | None = None + self._drag_target_world: np.ndarray | None = None + self._drag_start_world: np.ndarray | None = None + self._drag_mocap_offset: np.ndarray | None = None + self._drag_mocap_init_quat: np.ndarray | None = None + self._perturb_line_handle: viser.LineSegmentsHandle | None = None + + self.scene.add_body_drag_callback( + self._on_body_drag, button="left", modifier="cmd/ctrl" + ) + self.scene.add_body_drag_callback( + self._on_body_drag, button="left", modifier="cmd/ctrl+shift" + ) + # Speed. self._speed_idx = _SPEEDS.index(1.0) self._joint_sliders: list = [] # [(slider, qpos_adr), ...] @@ -116,6 +141,7 @@ def _render(self) -> None: self._render_fn(self.scene) else: self.scene.update_from_mjdata(self.data) + self._update_perturb_visual() def _refresh_scene_from_gui(self) -> None: """Refresh cached visualization state under the simulation lock.""" @@ -140,6 +166,7 @@ def _sync_sliders(self) -> None: def _reset(self) -> None: """Reset the simulation.""" + self._clear_drag_state() if self._reset_fn is not None: self._reset_fn(self.model, self.data) else: @@ -215,6 +242,7 @@ def _step_physics(self, dt: float) -> bool: deadline = time.perf_counter() + _FRAME_TIME hit_deadline = False while self._budget >= step_dt: + self._apply_drag_perturbation() self._step_fn(self.model, self.data) self._budget -= step_dt self._step_count += 1 @@ -228,6 +256,168 @@ def _step_physics(self, dt: float) -> bool: self._budget = min(self._budget, step_dt) return True + def _on_body_drag(self, event: BodyDragEvent) -> None: + """Handle body drag events from ``ViserMujocoScene``.""" + mocap_id = event.mocap_id + if mocap_id is None: + weld_id = int(self.model.body_weldid[event.body_id]) + if weld_id > 0 and self.model.body_mocapid[weld_id] >= 0: + mocap_id = int(self.model.body_mocapid[weld_id]) + + if mocap_id is None and is_fixed_body(self.model, event.body_id): + return + + mode: Literal["translate", "rotate"] = ( + "rotate" if event.modifier == "cmd/ctrl+shift" else "translate" + ) + + with self._lock: + if event.phase == "start": + self._clear_drag_state() + if np.all(self.data.xmat == 0): + mujoco.mj_kinematics(self.model, self.data) + + start_w = np.array(event.start_position, dtype=np.float64) + self._drag_body_id = event.body_id + self._drag_mocap_id = mocap_id + self._drag_mode = mode + self._drag_start_world = start_w.copy() + self._drag_target_world = start_w.copy() + + if mocap_id is not None: + self._drag_mocap_offset = self.data.mocap_pos[mocap_id].copy() - start_w + self._drag_mocap_init_quat = self.data.mocap_quat[mocap_id].copy() + self._drag_grab_local = None + else: + xpos = self.data.xpos[event.body_id] + xmat = self.data.xmat[event.body_id].reshape(3, 3) + self._drag_grab_local = xmat.T @ (start_w - xpos) + self._apply_drag_perturbation() + self._update_perturb_visual() + + elif event.phase == "update": + if self._drag_body_id != event.body_id: + return + end_w = np.array(event.end_position, dtype=np.float64) + self._drag_target_world = end_w + + if self._drag_mocap_id is not None: + mid = self._drag_mocap_id + if self._drag_mode == "translate" and self._drag_mocap_offset is not None: + self.data.mocap_pos[mid] = end_w + self._drag_mocap_offset + elif ( + self._drag_mode == "rotate" + and self._drag_start_world is not None + and self._drag_mocap_init_quat is not None + ): + drag_vec = end_w - self._drag_start_world + r_delta = vtf.SO3.exp(self.perturbation_torque_scale * drag_vec) + r_init = vtf.SO3(wxyz=self._drag_mocap_init_quat) + self.data.mocap_quat[mid] = np.asarray((r_delta @ r_init).wxyz) + + if self._paused: + mujoco.mj_forward(self.model, self.data) + self._render() + else: + self._dirty = True + else: + self._apply_drag_perturbation() + self._update_perturb_visual() + self._dirty = True + + elif event.phase == "end": + self._clear_drag_state() + self._dirty = True + + def _apply_drag_perturbation(self) -> None: + """Compute and write active drag force/torque to ``data.xfrc_applied``.""" + body_id = self._drag_body_id + if ( + body_id is None + or self._drag_mocap_id is not None + or self._drag_grab_local is None + or self._drag_target_world is None + ): + return + + xpos = self.data.xpos[body_id] + xmat = self.data.xmat[body_id].reshape(3, 3) + grab_world = xpos + xmat @ self._drag_grab_local + com = self.data.xipos[body_id] + lever = grab_world - com + + if self._drag_mode == "translate": + force = self.perturbation_stiffness * (self._drag_target_world - grab_world) + torque = np.cross(lever, force) + elif self._drag_mode == "rotate" and self._drag_start_world is not None: + force = self.perturbation_stiffness * (self._drag_start_world - grab_world) + drag_vec = self._drag_target_world - self._drag_start_world + ext_torque = self.perturbation_torque_scale * drag_vec + torque = np.cross(lever, force) + ext_torque + else: + return + + self.data.xfrc_applied[body_id, :3] = force + self.data.xfrc_applied[body_id, 3:] = torque + + def _update_perturb_visual(self) -> None: + """Update or hide the 3D perturbation indicator line.""" + body_id = self._drag_body_id + if ( + body_id is None + or self._drag_mocap_id is not None + or self._drag_grab_local is None + or self._drag_target_world is None + ): + if self._perturb_line_handle is not None: + self._perturb_line_handle.visible = False + return + + xpos = self.data.xpos[body_id] + xmat = self.data.xmat[body_id].reshape(3, 3) + grab_world = xpos + xmat @ self._drag_grab_local + + if self._drag_mode == "rotate" and self._drag_start_world is not None: + p_start = self._drag_start_world + color = _ROTATE_LINE_COLOR + else: + p_start = grab_world + color = _TRANSLATE_LINE_COLOR + + offset = self.scene.scene_offset + points = np.array( + [[p_start + offset, self._drag_target_world + offset]], dtype=np.float32 + ) + colors = np.full((1, 2, 3), color, dtype=np.uint8) + + if self._perturb_line_handle is None: + self._perturb_line_handle = self._server.scene.add_line_segments( + "/perturbation_force", + points=points, + colors=colors, + line_width=3.0, + visible=True, + ) + else: + self._perturb_line_handle.points = points + self._perturb_line_handle.colors = colors + self._perturb_line_handle.visible = True + + def _clear_drag_state(self) -> None: + """Clear active drag perturbation state and zero applied forces.""" + if self._drag_body_id is not None and self._drag_mocap_id is None: + self.data.xfrc_applied[self._drag_body_id] = 0.0 + self._drag_body_id = None + self._drag_mocap_id = None + self._drag_mode = None + self._drag_grab_local = None + self._drag_target_world = None + self._drag_start_world = None + self._drag_mocap_offset = None + self._drag_mocap_init_quat = None + if self._perturb_line_handle is not None: + self._perturb_line_handle.visible = False + def _update_stats(self) -> None: if self._paused: return @@ -289,6 +479,7 @@ def _(_) -> None: def _(_) -> None: if self._paused: with self._lock: + self._apply_drag_perturbation() self._step_fn(self.model, self.data) self._step_count += 1 self._render() @@ -338,6 +529,7 @@ def _(event) -> None: def _load_keyframe() -> None: idx = key_names.index(key_dropdown.value) with self._lock: + self._clear_drag_state() mujoco.mj_resetDataKeyframe(self.model, self.data, idx) mujoco.mj_forward(self.model, self.data) if self._reset_fn is not None: @@ -373,10 +565,42 @@ def _(_) -> None: self._setup_joint_sliders() self._setup_actuator_sliders() - # Physics tab (disable/enable flags). + # Physics tab (perturbation + disable/enable flags). with tabs.add_tab("Physics", icon=viser.Icon.ATOM): + self._setup_perturbation_gui() self._setup_physics_flags() + def _setup_perturbation_gui(self) -> None: + """Add perturbation controls and gesture instructions.""" + with self._server.gui.add_folder("Perturbation"): + self._server.gui.add_markdown( + "**Cmd/Ctrl + Drag**: Pull & tumble body \n" + "**Cmd/Ctrl + Shift + Drag**: Rotate around drag vector" + ) + stiffness_slider = self._server.gui.add_slider( + "Stiffness", + min=1.0, + max=500.0, + step=1.0, + initial_value=self.perturbation_stiffness, + ) + + @stiffness_slider.on_update + def _(_) -> None: + self.perturbation_stiffness = float(stiffness_slider.value) + + torque_slider = self._server.gui.add_slider( + "Torque Scale", + min=0.1, + max=50.0, + step=0.1, + initial_value=self.perturbation_torque_scale, + ) + + @torque_slider.on_update + def _(_) -> None: + self.perturbation_torque_scale = float(torque_slider.value) + _MAX_SLIDERS: int = 200 def _setup_joint_sliders(self) -> None: diff --git a/tests/test_scene.py b/tests/test_scene.py index c3c7ea1..144b044 100644 --- a/tests/test_scene.py +++ b/tests/test_scene.py @@ -1,8 +1,10 @@ +import asyncio + import mujoco import numpy as np import viser -from mjviser import ViserMujocoScene +from mjviser import BodyDragEvent, ViserMujocoScene from mjviser.conversions import is_fixed_body _MESH_BODY_XML = """ @@ -417,3 +419,80 @@ def test_autoconnect_capsules_match_mujoco_connector_sizes(): server.stop() except RuntimeError: pass + + +def test_body_drag_callback_registration_and_dispatch(): + xml = """ + + + + + + + + + + + + + + + + """ + server = viser.ViserServer(port=0) + model = mujoco.MjModel.from_xml_string(xml) + scene = ViserMujocoScene(server, model, num_envs=2) + + received: list[BodyDragEvent] = [] + + @scene.on_body_drag("left", modifier="cmd/ctrl") + def _on_drag(event: BodyDragEvent) -> None: + received.append(event) + + assert len(scene._mesh_groups) == 2 + for mg in scene._mesh_groups: + assert mg.handle._has_any_drag_callbacks() + + # Find the batched group containing box_a and box_b (N=2 bodies, 2 envs => 4 instances). + box_mg = next(mg for mg in scene._mesh_groups if mg.mocap_ids is None) + handler = scene._make_mesh_group_drag_handler(box_mg, _on_drag) + + # Instance 3 -> env_idx = 3 // 2 = 1, body_idx = 3 % 2 = 1 (box_b) + fake_event = viser.SceneNodeDragEvent( + client=None, # type: ignore[arg-type] + client_id=0, + target=box_mg.handle, + phase="start", + instance_index=3, + start_position=(1.0, 2.0, 3.0), + start_screen_pos=(0.5, 0.5), + end_position=(1.5, 2.5, 3.5), + end_screen_pos=(0.6, 0.6), + button="left", + modifier="cmd/ctrl", + ) + asyncio.run(handler(fake_event)) + + assert len(received) == 1 + ev = received[0] + assert ev.body_name == "box_b" + assert ev.env_idx == 1 + assert ev.mocap_id is None + assert ev.phase == "start" + assert ev.start_position == (1.0, 2.0, 3.0) + assert ev.end_position == (1.5, 2.5, 3.5) + + # Test rebuild_visual_handles preserves drag callbacks. + scene.rebuild_visual_handles() + for mg in scene._mesh_groups: + assert mg.handle._has_any_drag_callbacks() + + # Test remove_body_drag_callbacks clears callbacks. + scene.remove_body_drag_callbacks() + for mg in scene._mesh_groups: + assert not mg.handle._has_any_drag_callbacks() + + try: + server.stop() + except RuntimeError: + pass diff --git a/tests/test_viewer_drag.py b/tests/test_viewer_drag.py new file mode 100644 index 0000000..87cc0aa --- /dev/null +++ b/tests/test_viewer_drag.py @@ -0,0 +1,204 @@ +"""Tests for body drag interactions in Viewer.""" + +import mujoco +import numpy as np +import pytest +import viser + +from mjviser import BodyDragEvent, Viewer + +_DRAG_TEST_XML = """ + + +""" + + +@pytest.fixture +def viewer(): + server = viser.ViserServer(port=0) + model = mujoco.MjModel.from_xml_string(_DRAG_TEST_XML) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + v = Viewer(model, data, server=server) + yield v + try: + server.stop() + except RuntimeError: + pass + + +def _make_event( + phase: str, + body_id: int, + body_name: str, + start_pos: tuple[float, float, float], + end_pos: tuple[float, float, float], + mocap_id: int | None = None, + modifier: str | None = "cmd/ctrl", +) -> BodyDragEvent: + return BodyDragEvent( + client=None, # type: ignore[arg-type] + phase=phase, # type: ignore[arg-type] + body_id=body_id, + body_name=body_name, + env_idx=0, + mocap_id=mocap_id, + start_position=start_pos, + end_position=end_pos, + start_screen_pos=(0.5, 0.5), + end_screen_pos=(0.6, 0.6), + button="left", + modifier=modifier, # type: ignore[arg-type] + ) + + +def test_viewer_translate_drag_applies_force_and_torque(viewer: Viewer) -> None: + box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") + viewer.perturbation_stiffness = 100.0 + + # Grab off-center at (0.5, 0.0, 1.0) on the box centered at (0, 0, 1). + start_ev = _make_event( + "start", box_id, "box", (0.5, 0.0, 1.0), (0.5, 0.0, 1.0), modifier="cmd/ctrl" + ) + viewer._on_body_drag(start_ev) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id], np.zeros(6)) + + # Drag target to (0.5, 1.0, 1.0) -> delta = (0, 1, 0), F = (0, 100, 0) + # Lever arm r = (0.5, 0, 0), torque = r x F = (0, 0, 50) + update_ev = _make_event( + "update", box_id, "box", (0.5, 0.0, 1.0), (0.5, 1.0, 1.0), modifier="cmd/ctrl" + ) + viewer._on_body_drag(update_ev) + + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 100.0, 0.0]) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, 3:], [0.0, 0.0, 50.0]) + assert viewer._perturb_line_handle is not None + assert viewer._perturb_line_handle.visible + + # End drag resets xfrc_applied and hides visual line. + end_ev = _make_event( + "end", box_id, "box", (0.5, 0.0, 1.0), (0.5, 1.0, 1.0), modifier="cmd/ctrl" + ) + viewer._on_body_drag(end_ev) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id], np.zeros(6)) + assert not viewer._perturb_line_handle.visible + + +def test_viewer_rotate_drag_applies_pin_and_torque(viewer: Viewer) -> None: + box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") + viewer.perturbation_stiffness = 80.0 + viewer.perturbation_torque_scale = 5.0 + + start_ev = _make_event( + "start", + box_id, + "box", + (0.0, 0.0, 1.0), + (0.0, 0.0, 1.0), + modifier="cmd/ctrl+shift", + ) + viewer._on_body_drag(start_ev) + + update_ev = _make_event( + "update", + box_id, + "box", + (0.0, 0.0, 1.0), + (0.0, 0.0, 3.0), + modifier="cmd/ctrl+shift", + ) + viewer._on_body_drag(update_ev) + + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 0.0, 0.0]) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, 3:], [0.0, 0.0, 10.0]) + + +def test_viewer_mocap_body_drag_updates_mocap_pos_and_quat(viewer: Viewer) -> None: + mocap_body_id = mujoco.mj_name2id( + viewer.model, mujoco.mjtObj.mjOBJ_BODY, "mocap_target" + ) + mocap_id = int(viewer.model.body_mocapid[mocap_body_id]) + + # Translate mocap body. + viewer._on_body_drag( + _make_event( + "start", + mocap_body_id, + "mocap_target", + (2.0, 0.0, 1.0), + (2.0, 0.0, 1.0), + mocap_id=mocap_id, + modifier="cmd/ctrl", + ) + ) + viewer._on_body_drag( + _make_event( + "update", + mocap_body_id, + "mocap_target", + (2.0, 0.0, 1.0), + (3.5, -1.0, 2.0), + mocap_id=mocap_id, + modifier="cmd/ctrl", + ) + ) + np.testing.assert_allclose(viewer.data.mocap_pos[mocap_id], [3.5, -1.0, 2.0]) + + viewer._on_body_drag( + _make_event( + "end", + mocap_body_id, + "mocap_target", + (2.0, 0.0, 1.0), + (3.5, -1.0, 2.0), + mocap_id=mocap_id, + modifier="cmd/ctrl", + ) + ) + + # Rotate mocap body. + init_quat = viewer.data.mocap_quat[mocap_id].copy() + viewer._on_body_drag( + _make_event( + "start", + mocap_body_id, + "mocap_target", + (3.5, -1.0, 2.0), + (3.5, -1.0, 2.0), + mocap_id=mocap_id, + modifier="cmd/ctrl+shift", + ) + ) + viewer._on_body_drag( + _make_event( + "update", + mocap_body_id, + "mocap_target", + (3.5, -1.0, 2.0), + (3.5, -1.0, 2.5), + mocap_id=mocap_id, + modifier="cmd/ctrl+shift", + ) + ) + assert not np.allclose(viewer.data.mocap_quat[mocap_id], init_quat) + + +def test_viewer_ignores_fixed_body_drag(viewer: Viewer) -> None: + wall_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "fixed_wall") + viewer._on_body_drag( + _make_event("start", wall_id, "fixed_wall", (0.0, -1.0, 0.0), (0.0, -1.0, 0.0)) + ) + assert viewer._drag_body_id is None From 5cd865e1f6b2dc02db36d245779601f48746dcd8 Mon Sep 17 00:00:00 2001 From: Shawn Maybush Date: Sat, 17 Oct 2026 01:51:28 -0400 Subject: [PATCH 2/3] Address review feedback: use MjvPerturb, fix render lock race, and add selective callback removal --- src/mjviser/scene.py | 128 ++++++++++++---------- src/mjviser/viewer.py | 223 +++++++++++++++++++------------------- tests/test_scene.py | 71 +++++++++++- tests/test_viewer_drag.py | 47 +++++++- 4 files changed, 292 insertions(+), 177 deletions(-) diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 86deec7..703ada8 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -66,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 @@ -178,9 +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[ - tuple[DragButton, KeyModifier | None, Callable[[BodyDragEvent], Any]] - ] = [] + self._body_drag_callbacks: list[_BodyDragRegistration] = [] # Visualization settings. self.env_idx = 0 @@ -324,12 +334,15 @@ def add_body_drag_callback( button: Mouse button that triggers the drag (``"left"``, ``"middle"``, ``"right"``). modifier: Key modifier combination required at drag start, or ``None``. """ - entry = (button, modifier, callback) with self._update_lock: - if entry in self._body_drag_callbacks: - return - self._body_drag_callbacks.append(entry) - self._bind_drag_entry_to_all_handles(button, modifier, callback) + 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, @@ -354,16 +367,27 @@ def _decorator( 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: + remaining: list[_BodyDragRegistration] = [] + for reg in self._body_drag_callbacks: + if reg.callback == callback: + for handle, wrapper in reg.bound_wrappers: + if not handle._impl.removed: + 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: + for reg in self._body_drag_callbacks: + for handle, wrapper in reg.bound_wrappers: + if not handle._impl.removed: + handle.remove_drag_callback(wrapper) self._body_drag_callbacks.clear() - for mg in self._mesh_groups: - mg.handle.remove_drag_callback() - for handle in self._fixed_geom_handles.values(): - handle.remove_drag_callback() - for handle, _ in self._hull_dynamic_handles: - handle.remove_drag_callback() def rebuild_visual_handles(self) -> None: """Rebuild handles whose geometry or appearance is baked at creation time.""" @@ -1204,7 +1228,7 @@ def _add_fixed_geometry(self) -> None: visible=visible, ) self._fixed_geom_handles[(body_id, group_id, sub_idx)] = handle - self._bind_drag_callbacks_to_fixed_geoms() + self._bind_drag_callbacks(kinds=("fixed",)) @staticmethod def _geom_subgroup_fingerprint( @@ -1283,7 +1307,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_to_mesh_groups() + self._bind_drag_callbacks(kinds=("mesh",)) def _add_fixed_sites(self) -> None: """Add fixed site geometry to the scene as static nodes.""" @@ -1410,7 +1434,7 @@ def _build_hull_handles(self) -> None: lod="off", ) self._hull_dynamic_handles.append((handle, body_id)) - self._bind_drag_callbacks_to_hull_handles() + self._bind_drag_callbacks(kinds=("hull",)) def _clear_hull_handles(self) -> None: """Remove convex hull handles and clear their caches.""" @@ -1422,49 +1446,32 @@ def _clear_hull_handles(self) -> None: handle.remove() self._hull_dynamic_handles.clear() - def _bind_drag_entry_to_all_handles( + def _bind_drag_callbacks( self, - button: DragButton, - modifier: KeyModifier | None, - callback: Callable[[BodyDragEvent], Any], + kinds: tuple[Literal["mesh", "fixed", "hull"], ...] = ("mesh", "fixed", "hull"), + registrations: list[_BodyDragRegistration] | None = None, ) -> None: - """Bind a single (button, modifier, callback) entry to all current handles.""" - for mg in self._mesh_groups: - mg.handle.on_drag(button, modifier=modifier)( - self._make_mesh_group_drag_handler(mg, callback) - ) - for (body_id, _, _), handle in self._fixed_geom_handles.items(): - handle.on_drag(button, modifier=modifier)( - self._make_fixed_drag_handler(body_id, callback) - ) - for handle, body_id in self._hull_dynamic_handles: - handle.on_drag(button, modifier=modifier)( - self._make_hull_drag_handler(body_id, callback) - ) - - def _bind_drag_callbacks_to_mesh_groups(self) -> None: - """Attach all registered body drag callbacks to current mesh group handles.""" - for button, modifier, callback in self._body_drag_callbacks: - for mg in self._mesh_groups: - mg.handle.on_drag(button, modifier=modifier)( - self._make_mesh_group_drag_handler(mg, callback) - ) - - def _bind_drag_callbacks_to_fixed_geoms(self) -> None: - """Attach all registered body drag callbacks to current fixed geom handles.""" - for button, modifier, callback in self._body_drag_callbacks: - for (body_id, _, _), handle in self._fixed_geom_handles.items(): - handle.on_drag(button, modifier=modifier)( - self._make_fixed_drag_handler(body_id, callback) - ) - - def _bind_drag_callbacks_to_hull_handles(self) -> None: - """Attach all registered body drag callbacks to current dynamic hull handles.""" - for button, modifier, callback in self._body_drag_callbacks: - for handle, body_id in self._hull_dynamic_handles: - handle.on_drag(button, modifier=modifier)( - self._make_hull_drag_handler(body_id, callback) - ) + """Bind registered body drag callbacks to the specified handle collections.""" + 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 not h._impl.removed + ] + 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, @@ -1541,6 +1548,11 @@ async def _dispatch_body_drag( 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]), diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index 3c729e1..2fe790d 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -6,7 +6,6 @@ import time from collections.abc import Callable from threading import Lock -from typing import Literal import mujoco import numpy as np @@ -82,17 +81,14 @@ def __init__( self._lock = Lock() self.scene.set_refresh_handler(self._refresh_scene_from_gui) - # Perturbation / drag state. - self.perturbation_stiffness: float = 50.0 - self.perturbation_torque_scale: float = 2.0 - self._drag_body_id: int | None = None - self._drag_mocap_id: int | None = None - self._drag_mode: Literal["translate", "rotate"] | None = None - self._drag_grab_local: np.ndarray | None = None - self._drag_target_world: np.ndarray | None = None + # Perturbation / drag state (backed by MuJoCo's MjvPerturb). + self.perturbation_stiffness: float = 100.0 + self.perturbation_torque_scale: float = 1.0 + self._pert = mujoco.MjvPerturb() self._drag_start_world: np.ndarray | None = None - self._drag_mocap_offset: np.ndarray | None = None - self._drag_mocap_init_quat: np.ndarray | None = None + self._drag_end_world: np.ndarray | None = None + self._drag_init_xpos: np.ndarray | None = None + self._drag_init_xquat: np.ndarray | None = None self._perturb_line_handle: viser.LineSegmentsHandle | None = None self.scene.add_body_drag_callback( @@ -127,6 +123,12 @@ def __init__( self._fps = 0.0 self._sps = 0.0 + @property + def _drag_body_id(self) -> int | None: + if self._pert.active != 0 and self._pert.select > 0: + return int(self._pert.select) + return None + @property def speed(self) -> float: return _SPEEDS[self._speed_idx] @@ -222,8 +224,9 @@ def _tick(self) -> None: self._time_until_next_render = 0.0 if self._dirty or self.scene.needs_update: - self._dirty = False - self._render() + with self._lock: + self._dirty = False + self._render() self._stats_frames += 1 self._update_stats() @@ -257,117 +260,115 @@ def _step_physics(self, dt: float) -> bool: return True def _on_body_drag(self, event: BodyDragEvent) -> None: - """Handle body drag events from ``ViserMujocoScene``.""" - mocap_id = event.mocap_id - if mocap_id is None: - weld_id = int(self.model.body_weldid[event.body_id]) - if weld_id > 0 and self.model.body_mocapid[weld_id] >= 0: - mocap_id = int(self.model.body_mocapid[weld_id]) - - if mocap_id is None and is_fixed_body(self.model, event.body_id): - return + """Handle body drag events from ``ViserMujocoScene``. - mode: Literal["translate", "rotate"] = ( - "rotate" if event.modifier == "cmd/ctrl+shift" else "translate" - ) + Runs synchronously on Viser's event loop (invoked by ``_dispatch_body_drag``) + so ``start``/``update``/``end`` phases are strictly ordered, acquiring + ``self._lock`` briefly to update ``self._pert`` and ``self.data``. + """ + body_id = event.body_id + weld_id = int(self.model.body_weldid[body_id]) + if weld_id > 0 and self.model.body_mocapid[weld_id] >= 0: + body_id = weld_id + + if is_fixed_body(self.model, body_id): + return with self._lock: if event.phase == "start": self._clear_drag_state() if np.all(self.data.xmat == 0): - mujoco.mj_kinematics(self.model, self.data) + mujoco.mj_forward(self.model, self.data) start_w = np.array(event.start_position, dtype=np.float64) - self._drag_body_id = event.body_id - self._drag_mocap_id = mocap_id - self._drag_mode = mode - self._drag_start_world = start_w.copy() - self._drag_target_world = start_w.copy() - - if mocap_id is not None: - self._drag_mocap_offset = self.data.mocap_pos[mocap_id].copy() - start_w - self._drag_mocap_init_quat = self.data.mocap_quat[mocap_id].copy() - self._drag_grab_local = None + xpos = self.data.xpos[body_id].copy() + xquat = self.data.xquat[body_id].copy() + xmat = self.data.xmat[body_id].reshape(3, 3) + + self._drag_start_world = start_w + self._drag_end_world = start_w.copy() + self._drag_init_xpos = xpos + self._drag_init_xquat = xquat + + self._pert.select = body_id + self._pert.localpos[:] = xmat.T @ (start_w - xpos) + self._pert.scale = 1.0 + self._pert.localmass = float(self.model.body_subtreemass[body_id]) * ( + self.perturbation_stiffness / 100.0 + ) + self._pert.refpos[:] = xpos + self._pert.refquat[:] = xquat + self._pert.refselpos[:] = start_w + if event.modifier == "cmd/ctrl+shift": + self._pert.active = int( + mujoco.mjtPertBit.mjPERT_TRANSLATE | mujoco.mjtPertBit.mjPERT_ROTATE + ) else: - xpos = self.data.xpos[event.body_id] - xmat = self.data.xmat[event.body_id].reshape(3, 3) - self._drag_grab_local = xmat.T @ (start_w - xpos) - self._apply_drag_perturbation() - self._update_perturb_visual() + self._pert.active = int(mujoco.mjtPertBit.mjPERT_TRANSLATE) + + self._apply_drag_perturbation() + self._update_perturb_visual() + self._dirty = True elif event.phase == "update": - if self._drag_body_id != event.body_id: + if ( + self._pert.active == 0 + or self._pert.select != body_id + or self._drag_start_world is None + or self._drag_init_xpos is None + or self._drag_init_xquat is None + ): return end_w = np.array(event.end_position, dtype=np.float64) - self._drag_target_world = end_w - - if self._drag_mocap_id is not None: - mid = self._drag_mocap_id - if self._drag_mode == "translate" and self._drag_mocap_offset is not None: - self.data.mocap_pos[mid] = end_w + self._drag_mocap_offset - elif ( - self._drag_mode == "rotate" - and self._drag_start_world is not None - and self._drag_mocap_init_quat is not None - ): - drag_vec = end_w - self._drag_start_world - r_delta = vtf.SO3.exp(self.perturbation_torque_scale * drag_vec) - r_init = vtf.SO3(wxyz=self._drag_mocap_init_quat) - self.data.mocap_quat[mid] = np.asarray((r_delta @ r_init).wxyz) - - if self._paused: - mujoco.mj_forward(self.model, self.data) - self._render() - else: - self._dirty = True + self._drag_end_world = end_w + drag_vec = end_w - self._drag_start_world + + if self._pert.active & int(mujoco.mjtPertBit.mjPERT_ROTATE): + self._pert.refpos[:] = self._drag_init_xpos + self._pert.refselpos[:] = self._drag_start_world + r_delta = vtf.SO3.exp(self.perturbation_torque_scale * drag_vec) + r_init = vtf.SO3(wxyz=self._drag_init_xquat) + self._pert.refquat[:] = np.asarray((r_delta @ r_init).wxyz) else: - self._apply_drag_perturbation() - self._update_perturb_visual() - self._dirty = True + self._pert.refpos[:] = self._drag_init_xpos + drag_vec + self._pert.refselpos[:] = end_w + + self._apply_drag_perturbation() + self._update_perturb_visual() + self._dirty = True elif event.phase == "end": self._clear_drag_state() self._dirty = True def _apply_drag_perturbation(self) -> None: - """Compute and write active drag force/torque to ``data.xfrc_applied``.""" - body_id = self._drag_body_id - if ( - body_id is None - or self._drag_mocap_id is not None - or self._drag_grab_local is None - or self._drag_target_world is None - ): + """Apply active ``MjvPerturb`` pose and force to ``self.data``.""" + if self._pert.active == 0 or self._pert.select <= 0: return - - xpos = self.data.xpos[body_id] - xmat = self.data.xmat[body_id].reshape(3, 3) - grab_world = xpos + xmat @ self._drag_grab_local - com = self.data.xipos[body_id] - lever = grab_world - com - - if self._drag_mode == "translate": - force = self.perturbation_stiffness * (self._drag_target_world - grab_world) - torque = np.cross(lever, force) - elif self._drag_mode == "rotate" and self._drag_start_world is not None: - force = self.perturbation_stiffness * (self._drag_start_world - grab_world) - drag_vec = self._drag_target_world - self._drag_start_world - ext_torque = self.perturbation_torque_scale * drag_vec - torque = np.cross(lever, force) + ext_torque + body_id = int(self._pert.select) + self._pert.localmass = float(self.model.body_subtreemass[body_id]) * ( + self.perturbation_stiffness / 100.0 + ) + mujoco.mjv_applyPerturbPose(self.model, self.data, self._pert, int(self._paused)) + if not self._paused: + mujoco.mjv_applyPerturbForce(self.model, self.data, self._pert) else: - return - - self.data.xfrc_applied[body_id, :3] = force - self.data.xfrc_applied[body_id, 3:] = torque + mujoco.mj_forward(self.model, self.data) def _update_perturb_visual(self) -> None: """Update or hide the 3D perturbation indicator line.""" - body_id = self._drag_body_id + body_id = int(self._pert.select) + active = int(self._pert.active) + start_world = self._drag_start_world + end_world = self._drag_end_world + localpos = np.array(self._pert.localpos, dtype=np.float64) + if ( - body_id is None - or self._drag_mocap_id is not None - or self._drag_grab_local is None - or self._drag_target_world is None + active == 0 + or body_id <= 0 + or self.model.body_mocapid[body_id] >= 0 + or start_world is None + or end_world is None ): if self._perturb_line_handle is not None: self._perturb_line_handle.visible = False @@ -375,19 +376,17 @@ def _update_perturb_visual(self) -> None: xpos = self.data.xpos[body_id] xmat = self.data.xmat[body_id].reshape(3, 3) - grab_world = xpos + xmat @ self._drag_grab_local + grab_world = xpos + xmat @ localpos - if self._drag_mode == "rotate" and self._drag_start_world is not None: - p_start = self._drag_start_world + if active & int(mujoco.mjtPertBit.mjPERT_ROTATE): + p_start = start_world color = _ROTATE_LINE_COLOR else: p_start = grab_world color = _TRANSLATE_LINE_COLOR offset = self.scene.scene_offset - points = np.array( - [[p_start + offset, self._drag_target_world + offset]], dtype=np.float32 - ) + points = np.array([[p_start + offset, end_world + offset]], dtype=np.float32) colors = np.full((1, 2, 3), color, dtype=np.uint8) if self._perturb_line_handle is None: @@ -405,16 +404,14 @@ def _update_perturb_visual(self) -> None: def _clear_drag_state(self) -> None: """Clear active drag perturbation state and zero applied forces.""" - if self._drag_body_id is not None and self._drag_mocap_id is None: - self.data.xfrc_applied[self._drag_body_id] = 0.0 - self._drag_body_id = None - self._drag_mocap_id = None - self._drag_mode = None - self._drag_grab_local = None - self._drag_target_world = None + if self._pert.select > 0 and self.model.body_mocapid[self._pert.select] < 0: + self.data.xfrc_applied[self._pert.select] = 0.0 + self._pert.active = 0 + self._pert.select = 0 self._drag_start_world = None - self._drag_mocap_offset = None - self._drag_mocap_init_quat = None + self._drag_end_world = None + self._drag_init_xpos = None + self._drag_init_xquat = None if self._perturb_line_handle is not None: self._perturb_line_handle.visible = False diff --git a/tests/test_scene.py b/tests/test_scene.py index 144b044..6b9be14 100644 --- a/tests/test_scene.py +++ b/tests/test_scene.py @@ -482,12 +482,41 @@ def _on_drag(event: BodyDragEvent) -> None: assert ev.start_position == (1.0, 2.0, 3.0) assert ev.end_position == (1.5, 2.5, 3.5) - # Test rebuild_visual_handles preserves drag callbacks. + # Test show_only_selected routes instance_index to the selected _last_env_idx. + received.clear() + scene.show_only_selected = True + scene._last_env_idx = 1 + fake_event_slice = viser.SceneNodeDragEvent( + client=None, # type: ignore[arg-type] + client_id=0, + target=box_mg.handle, + phase="start", + instance_index=0, # body_idx 0 (box_a) in selected env 1 + start_position=(0.0, 0.0, 1.0), + start_screen_pos=(0.5, 0.5), + end_position=(0.0, 0.0, 1.0), + end_screen_pos=(0.5, 0.5), + button="left", + modifier="cmd/ctrl", + ) + asyncio.run(handler(fake_event_slice)) + assert len(received) == 1 + assert received[0].body_name == "box_a" + assert received[0].env_idx == 1 + + # Test selective remove_body_drag_callback preserves other callbacks. + other_received: list[BodyDragEvent] = [] + scene.add_body_drag_callback(other_received.append, button="left", modifier="alt") + scene.remove_body_drag_callback(_on_drag) + for mg in scene._mesh_groups: + assert mg.handle._has_any_drag_callbacks() + + # Test rebuild_visual_handles preserves remaining drag callbacks. scene.rebuild_visual_handles() for mg in scene._mesh_groups: assert mg.handle._has_any_drag_callbacks() - # Test remove_body_drag_callbacks clears callbacks. + # Test remove_body_drag_callbacks clears all callbacks. scene.remove_body_drag_callbacks() for mg in scene._mesh_groups: assert not mg.handle._has_any_drag_callbacks() @@ -496,3 +525,41 @@ def _on_drag(event: BodyDragEvent) -> None: server.stop() except RuntimeError: pass + + +def test_convex_hull_drag_callback_dispatch(): + server = viser.ViserServer(port=0) + model = mujoco.MjModel.from_xml_string(_MESH_BODY_XML) + scene = ViserMujocoScene(server, model, num_envs=2) + + received: list[BodyDragEvent] = [] + scene.add_body_drag_callback(received.append) + + assert len(scene._hull_dynamic_handles) == 1 + hull_handle, body_id = scene._hull_dynamic_handles[0] + assert hull_handle._has_any_drag_callbacks() + + handler = scene._make_hull_drag_handler(body_id, received.append) + fake_event = viser.SceneNodeDragEvent( + client=None, # type: ignore[arg-type] + client_id=0, + target=hull_handle, + phase="start", + instance_index=1, + start_position=(0.0, 0.0, 1.0), + start_screen_pos=(0.5, 0.5), + end_position=(0.5, 0.0, 1.0), + end_screen_pos=(0.6, 0.5), + button="left", + modifier="cmd/ctrl", + ) + asyncio.run(handler(fake_event)) + assert len(received) == 1 + assert received[0].body_id == body_id + assert received[0].body_name == "mesh_body" + assert received[0].env_idx == 1 + + try: + server.stop() + except RuntimeError: + pass diff --git a/tests/test_viewer_drag.py b/tests/test_viewer_drag.py index 87cc0aa..6f7079c 100644 --- a/tests/test_viewer_drag.py +++ b/tests/test_viewer_drag.py @@ -99,8 +99,8 @@ def test_viewer_translate_drag_applies_force_and_torque(viewer: Viewer) -> None: def test_viewer_rotate_drag_applies_pin_and_torque(viewer: Viewer) -> None: box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") - viewer.perturbation_stiffness = 80.0 - viewer.perturbation_torque_scale = 5.0 + viewer.perturbation_stiffness = 100.0 + viewer.perturbation_torque_scale = 0.5 start_ev = _make_event( "start", @@ -117,13 +117,52 @@ def test_viewer_rotate_drag_applies_pin_and_torque(viewer: Viewer) -> None: box_id, "box", (0.0, 0.0, 1.0), - (0.0, 0.0, 3.0), + (0.0, 0.0, 2.0), modifier="cmd/ctrl+shift", ) viewer._on_body_drag(update_ev) + # Unit cube (half-size 0.5, mass 1.0) has inertia 1/6; mjv_applyPerturbForce + # applies 250 * I * 0.5 = 41.6667 N*m around +Z while pinning linear force at 0. np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 0.0, 0.0]) - np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, 3:], [0.0, 0.0, 10.0]) + np.testing.assert_allclose( + viewer.data.xfrc_applied[box_id, 3:], [0.0, 0.0, 250.0 / 6.0], rtol=1e-4 + ) + + +def test_viewer_damping_and_reset_mid_drag(viewer: Viewer) -> None: + box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") + viewer.perturbation_stiffness = 100.0 + + # Give the box +1 m/s velocity along Y; mjv_applyPerturbForce damps the 100 N pull to 90 N. + viewer.data.qvel[1] = 1.0 + mujoco.mj_forward(viewer.model, viewer.data) + + viewer._on_body_drag( + _make_event("start", box_id, "box", (0.0, 0.0, 1.0), (0.0, 0.0, 1.0)) + ) + viewer._on_body_drag( + _make_event("update", box_id, "box", (0.0, 0.0, 1.0), (0.0, 1.0, 1.0)) + ) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 90.0, 0.0]) + + # Reset mid-drag clears xfrc_applied and active perturbation state. + viewer._reset() + assert viewer._drag_body_id is None + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id], np.zeros(6)) + + +def test_viewer_paused_free_body_drag_updates_qpos(viewer: Viewer) -> None: + box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") + viewer._paused = True + + viewer._on_body_drag( + _make_event("start", box_id, "box", (0.0, 0.0, 1.0), (0.0, 0.0, 1.0)) + ) + viewer._on_body_drag( + _make_event("update", box_id, "box", (0.0, 0.0, 1.0), (1.0, 2.0, 3.0)) + ) + np.testing.assert_allclose(viewer.data.qpos[:3], [1.0, 2.0, 3.0]) def test_viewer_mocap_body_drag_updates_mocap_pos_and_quat(viewer: Viewer) -> None: From 26d4df5c4be32c0278b8efb3aa91b6516cd5f4f4 Mon Sep 17 00:00:00 2001 From: Shawn Maybush Date: Sat, 17 Oct 2026 02:06:43 -0400 Subject: [PATCH 3/3] Normalize perturbation stiffness by model.vis.map.stiffness and avoid Viser private internals Normalize MjvPerturb.localmass by model.vis.map.stiffness so perturbation_stiffness stays in N/(m*kg) across models with custom map stiffness, apply active drag forces when single-stepping while paused, hide the perturbation force line when paused, and track live drag handle IDs in ViserMujocoScene instead of reading handle._impl.removed. --- src/mjviser/scene.py | 18 +++++++++++++----- src/mjviser/viewer.py | 40 +++++++++++++++++++++++++-------------- tests/test_viewer_drag.py | 26 +++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 19 deletions(-) diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 703ada8..d7634cc 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -370,11 +370,12 @@ def _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 not handle._impl.removed: + if id(handle) in live_ids: handle.remove_drag_callback(wrapper) else: remaining.append(reg) @@ -383,9 +384,10 @@ def remove_body_drag_callback(self, callback: Callable[[BodyDragEvent], Any]) -> 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 not handle._impl.removed: + if id(handle) in live_ids: handle.remove_drag_callback(wrapper) self._body_drag_callbacks.clear() @@ -1446,17 +1448,23 @@ 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 not h._impl.removed - ] + 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) diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index 2fe790d..9d9b1c5 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -82,7 +82,10 @@ def __init__( self.scene.set_refresh_handler(self._refresh_scene_from_gui) # Perturbation / drag state (backed by MuJoCo's MjvPerturb). - self.perturbation_stiffness: float = 100.0 + default_stiffness = float(model.vis.map.stiffness) + self.perturbation_stiffness: float = ( + default_stiffness if default_stiffness > 0.0 else 100.0 + ) self.perturbation_torque_scale: float = 1.0 self._pert = mujoco.MjvPerturb() self._drag_start_world: np.ndarray | None = None @@ -293,9 +296,7 @@ def _on_body_drag(self, event: BodyDragEvent) -> None: self._pert.select = body_id self._pert.localpos[:] = xmat.T @ (start_w - xpos) self._pert.scale = 1.0 - self._pert.localmass = float(self.model.body_subtreemass[body_id]) * ( - self.perturbation_stiffness / 100.0 - ) + self._pert.localmass = self._compute_localmass(body_id) self._pert.refpos[:] = xpos self._pert.refquat[:] = xquat self._pert.refselpos[:] = start_w @@ -341,16 +342,24 @@ def _on_body_drag(self, event: BodyDragEvent) -> None: self._clear_drag_state() self._dirty = True - def _apply_drag_perturbation(self) -> None: + def _compute_localmass(self, body_id: int) -> float: + """Compute ``MjvPerturb.localmass`` so ``perturbation_stiffness`` is in N/(m*kg).""" + map_stiffness = float(self.model.vis.map.stiffness) + if map_stiffness <= 0.0: + map_stiffness = 100.0 + return float(self.model.body_subtreemass[body_id]) * ( + self.perturbation_stiffness / map_stiffness + ) + + def _apply_drag_perturbation(self, *, apply_force: bool | None = None) -> None: """Apply active ``MjvPerturb`` pose and force to ``self.data``.""" if self._pert.active == 0 or self._pert.select <= 0: return body_id = int(self._pert.select) - self._pert.localmass = float(self.model.body_subtreemass[body_id]) * ( - self.perturbation_stiffness / 100.0 - ) - mujoco.mjv_applyPerturbPose(self.model, self.data, self._pert, int(self._paused)) - if not self._paused: + self._pert.localmass = self._compute_localmass(body_id) + use_force = (not self._paused) if apply_force is None else apply_force + mujoco.mjv_applyPerturbPose(self.model, self.data, self._pert, int(not use_force)) + if use_force: mujoco.mjv_applyPerturbForce(self.model, self.data, self._pert) else: mujoco.mj_forward(self.model, self.data) @@ -359,12 +368,14 @@ def _update_perturb_visual(self) -> None: """Update or hide the 3D perturbation indicator line.""" body_id = int(self._pert.select) active = int(self._pert.active) + paused = self._paused start_world = self._drag_start_world end_world = self._drag_end_world localpos = np.array(self._pert.localpos, dtype=np.float64) if ( - active == 0 + paused + or active == 0 or body_id <= 0 or self.model.body_mocapid[body_id] >= 0 or start_world is None @@ -476,7 +487,7 @@ def _(_) -> None: def _(_) -> None: if self._paused: with self._lock: - self._apply_drag_perturbation() + self._apply_drag_perturbation(apply_force=True) self._step_fn(self.model, self.data) self._step_count += 1 self._render() @@ -572,12 +583,13 @@ def _setup_perturbation_gui(self) -> None: with self._server.gui.add_folder("Perturbation"): self._server.gui.add_markdown( "**Cmd/Ctrl + Drag**: Pull & tumble body \n" - "**Cmd/Ctrl + Shift + Drag**: Rotate around drag vector" + "**Cmd/Ctrl + Shift + Drag**: Rotate around drag vector \n" + "*While paused, dragging repositions free and mocap bodies directly.*" ) stiffness_slider = self._server.gui.add_slider( "Stiffness", min=1.0, - max=500.0, + max=max(500.0, self.perturbation_stiffness), step=1.0, initial_value=self.perturbation_stiffness, ) diff --git a/tests/test_viewer_drag.py b/tests/test_viewer_drag.py index 6f7079c..a2a8f0d 100644 --- a/tests/test_viewer_drag.py +++ b/tests/test_viewer_drag.py @@ -241,3 +241,29 @@ def test_viewer_ignores_fixed_body_drag(viewer: Viewer) -> None: _make_event("start", wall_id, "fixed_wall", (0.0, -1.0, 0.0), (0.0, -1.0, 0.0)) ) assert viewer._drag_body_id is None + + +def test_viewer_custom_map_stiffness_and_paused_step(viewer: Viewer) -> None: + box_id = mujoco.mj_name2id(viewer.model, mujoco.mjtObj.mjOBJ_BODY, "box") + # Even if the model specifies , perturbation_stiffness=100 + # still produces 100 N/m per kg of subtree mass. + viewer.model.vis.map.stiffness = 400.0 + viewer.perturbation_stiffness = 100.0 + + viewer._on_body_drag( + _make_event("start", box_id, "box", (0.0, 0.0, 1.0), (0.0, 0.0, 1.0)) + ) + viewer._on_body_drag( + _make_event("update", box_id, "box", (0.0, 0.0, 1.0), (0.0, 1.0, 1.0)) + ) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 100.0, 0.0]) + + # While paused, the perturbation force line is hidden, and single-step applies force. + viewer._paused = True + viewer._update_perturb_visual() + assert viewer._perturb_line_handle is not None + assert not viewer._perturb_line_handle.visible + + viewer.data.xfrc_applied[box_id] = 0.0 + viewer._apply_drag_perturbation(apply_force=True) + np.testing.assert_allclose(viewer.data.xfrc_applied[box_id, :3], [0.0, 100.0, 0.0])