From e8ba9d654b1ee92475bd0db0c7364bed3900bff6 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 6 May 2026 22:01:55 -0700 Subject: [PATCH 1/6] Add keyboard shortcuts, body selection, and drag perturbation --- src/mjviser/__init__.py | 1 + src/mjviser/interaction.py | 283 +++++++++++++++++++++++++++++++++++++ src/mjviser/scene.py | 99 +++++++++++++ src/mjviser/viewer.py | 124 ++++++++++++---- 4 files changed, 477 insertions(+), 30 deletions(-) create mode 100644 src/mjviser/interaction.py diff --git a/src/mjviser/__init__.py b/src/mjviser/__init__.py index 80af911..6f4da2e 100644 --- a/src/mjviser/__init__.py +++ b/src/mjviser/__init__.py @@ -1,4 +1,5 @@ """Web-based MuJoCo viewer powered by viser.""" +from mjviser.interaction import PerturbationHandler as PerturbationHandler from mjviser.scene import ViserMujocoScene as ViserMujocoScene from mjviser.viewer import Viewer as Viewer diff --git a/src/mjviser/interaction.py b/src/mjviser/interaction.py new file mode 100644 index 0000000..9c3db84 --- /dev/null +++ b/src/mjviser/interaction.py @@ -0,0 +1,283 @@ +"""Interactive perturbation: click-to-select and drag-to-apply-force.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from threading import Lock + +import mujoco +import numpy as np +import viser + + +@dataclass +class PerturbationState: + """Current perturbation force to be applied by the simulation loop.""" + + body_id: int + force: np.ndarray + torque: np.ndarray + point: np.ndarray + + +class PerturbationHandler: + """Manages body selection and drag-to-apply-force perturbation. + + Uses MuJoCo's critically damped spring model: the force is proportional + to stiffness * localmass, with critical damping from velocity feedback. + """ + + def __init__( + self, + server: viser.ViserServer, + mj_model: mujoco.MjModel, + mj_data: mujoco.MjData, + ) -> None: + self._server = server + self._model = mj_model + self._data = mj_data + self._lock = Lock() + + self.selected_body_id: int | None = None + self._drag_body_id: int | None = None + self._drag_grab_local: np.ndarray | None = None + self._drag_target: np.ndarray | None = None + self._drag_mode: str | None = None + + self._scene_offset = np.zeros(3) + self._body_xpos: np.ndarray | None = None + self._body_xmat: np.ndarray | None = None + + self._info_text: viser.GuiTextHandle | None = None + + self._prev_grab_world: np.ndarray | None = None + self._prev_time: float | None = None + + def setup_gui(self) -> None: + """Add selection info display to the GUI.""" + with self._server.gui.add_folder("Selection"): + self._info_text = self._server.gui.add_text( + "Body", initial_value="(none)", disabled=True + ) + + def register_click_handler(self) -> None: + """Register scene-level click handler for body selection via ray cast.""" + + @self._server.scene.on_click() + def _(event: viser.SceneClickEvent) -> None: + origin = np.array(event.ray_origin) - self._scene_offset + direction = np.array(event.ray_direction) + geomid = np.zeros(1, dtype=np.int32) + dist = mujoco.mj_ray( + self._model, + self._data, + origin, + direction, + None, + 1, + -1, + geomid, + ) + if dist < 0: + with self._lock: + self.selected_body_id = None + if self._info_text is not None: + self._info_text.value = "(none)" + return + + body_id = int(self._model.geom_bodyid[geomid[0]]) + with self._lock: + self.selected_body_id = body_id + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, body_id) + if self._info_text is not None: + self._info_text.value = name or f"body_{body_id}" + + def register_drag_handlers( + self, + handle: viser.BatchedGlbHandle, + body_ids: np.ndarray, + ) -> None: + """Attach drag and click handlers to a batched mesh handle.""" + n_bodies = len(body_ids) + + def _body_id_from_event( + event: viser.SceneNodeDragEvent, # type: ignore[type-arg] + ) -> int | None: + idx = event.instance_index + if idx is None: + return None + return int(body_ids[idx % n_bodies]) + + @handle.on_click + def _(event: viser.SceneNodePointerEvent) -> None: # type: ignore[type-arg] + idx = event.instance_index + if idx is None: + return + bid = int(body_ids[idx % n_bodies]) + if bid == 0: + return + with self._lock: + self.selected_body_id = bid + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) + if self._info_text is not None: + self._info_text.value = name or f"body_{bid}" + + @handle.on_drag_start("left", modifier="cmd/ctrl") + async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + bid = _body_id_from_event(event) + if bid is None or bid == 0: + return + grab_world = np.array(event.start_position) - self._scene_offset + grab_local = self._world_to_body_local(bid, grab_world) + with self._lock: + self._drag_body_id = bid + self._drag_grab_local = grab_local + self._drag_target = grab_world + self._drag_mode = "translate" + self.selected_body_id = bid + self._prev_grab_world = grab_world.copy() + self._prev_time = time.perf_counter() + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) + if self._info_text is not None: + self._info_text.value = name or f"body_{bid}" + + @handle.on_drag_update("left", modifier="cmd/ctrl") + async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + with self._lock: + if self._drag_body_id is not None: + self._drag_target = np.array(event.end_position) - self._scene_offset + + @handle.on_drag_end("left", modifier="cmd/ctrl") + async def _(_: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + with self._lock: + self._drag_body_id = None + self._drag_grab_local = None + self._drag_target = None + self._drag_mode = None + self._prev_grab_world = None + self._prev_time = None + + @handle.on_drag_start("left", modifier="cmd/ctrl+shift") + async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + bid = _body_id_from_event(event) + if bid is None or bid == 0: + return + grab_world = np.array(event.start_position) - self._scene_offset + grab_local = self._world_to_body_local(bid, grab_world) + with self._lock: + self._drag_body_id = bid + self._drag_grab_local = grab_local + self._drag_target = grab_world + self._drag_mode = "rotate" + self.selected_body_id = bid + self._prev_grab_world = grab_world.copy() + self._prev_time = time.perf_counter() + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) + if self._info_text is not None: + self._info_text.value = name or f"body_{bid}" + + @handle.on_drag_update("left", modifier="cmd/ctrl+shift") + async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + with self._lock: + if self._drag_body_id is not None: + self._drag_target = np.array(event.end_position) - self._scene_offset + + @handle.on_drag_end("left", modifier="cmd/ctrl+shift") + async def _(_: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + with self._lock: + self._drag_body_id = None + self._drag_grab_local = None + self._drag_target = None + self._drag_mode = None + self._prev_grab_world = None + self._prev_time = None + + def update_state( + self, + body_xpos: np.ndarray, + body_xmat: np.ndarray, + env_idx: int, + scene_offset: np.ndarray, + ) -> None: + """Update cached body state from the latest simulation frame. + + Args: + body_xpos: Shape ``(num_envs, nbody, 3)``. + body_xmat: Shape ``(num_envs, nbody, 3, 3)``. + env_idx: Active environment index. + scene_offset: Current camera tracking offset. + """ + self._body_xpos = body_xpos[env_idx] + self._body_xmat = body_xmat[env_idx] + self._scene_offset = scene_offset + + def get_perturbation(self) -> PerturbationState | None: + """Return the current perturbation force using MuJoCo's spring model.""" + with self._lock: + if ( + self._drag_body_id is None + or self._drag_grab_local is None + or self._drag_target is None + or self._body_xpos is None + or self._body_xmat is None + ): + return None + + bid = self._drag_body_id + xpos = self._body_xpos[bid] + xmat = self._body_xmat[bid] + + grab_world = xpos + xmat @ self._drag_grab_local + diff = grab_world - self._drag_target + + # Compute localmass from body_invweight0 (translational inverse weight). + invweight_trn = float(self._model.body_invweight0[bid, 0]) + localmass = 1.0 if invweight_trn == 0 else 1.0 / max(invweight_trn, 1e-10) + + # Estimate velocity of the grab point via finite difference. + now = time.perf_counter() + vel = np.zeros(3) + if self._prev_grab_world is not None and self._prev_time is not None: + dt = now - self._prev_time + if dt > 1e-6: + vel = (grab_world - self._prev_grab_world) / dt + self._prev_grab_world = grab_world.copy() + self._prev_time = now + + if self._drag_mode == "translate": + stiffness = float(self._model.vis.map.stiffness) + # Spring force: -k * m * displacement. + force = -stiffness * localmass * diff + # Critical damping: -2*sqrt(k)*m * velocity. + force -= np.sqrt(stiffness) * localmass * vel + + # Torque from moment arm. + moment_arm = grab_world - xpos + torque = np.cross(moment_arm, force) + else: + stiffnessrot = float(self._model.vis.map.stiffnessrot) + # Rotational inertia from body_invweight0. + invweight_rot = float(self._model.body_invweight0[bid, 1]) + inertia = 1.0 if invweight_rot == 0 else 1.0 / max(invweight_rot, 1e-10) + # Torque proportional to displacement (simplified from quat error). + torque = -stiffnessrot * inertia * diff + torque -= np.sqrt(stiffnessrot) * inertia * vel + # Smaller translational force in rotate mode. + stiffness = float(self._model.vis.map.stiffness) + force = -stiffness * localmass * diff * 0.1 + + return PerturbationState( + body_id=bid, + force=force, + torque=torque, + point=grab_world, + ) + + def _world_to_body_local(self, body_id: int, world_pos: np.ndarray) -> np.ndarray: + """Convert a world position to body-local coordinates.""" + if self._body_xpos is None or self._body_xmat is None: + return world_pos + xpos = self._body_xpos[body_id] + xmat = self._body_xmat[body_id] + return xmat.T @ (world_pos - xpos) diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 46b76bf..9c1e817 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -22,6 +22,7 @@ merge_geoms_hull, merge_sites, ) +from .interaction import PerturbationHandler @dataclasses.dataclass @@ -183,6 +184,9 @@ def __init__( self.needs_update = False self.paused = False + self._vis_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} + self._simple_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} + self.perturbation = PerturbationHandler(server, mj_model, self.mj_data) # Cached visualization state for re-rendering when settings change. self._tracked_body_id: int | None = None @@ -203,6 +207,7 @@ def __init__( mujoco.mj_kinematics(mj_model, self.mj_data) self._add_fixed_geometry() self._create_mesh_handles_by_group() + self._register_perturbation_handlers() self._add_fixed_sites() self._create_site_handles_by_group() self._compute_hull_body_meshes() @@ -269,6 +274,12 @@ def show_convex_hull(self, value: bool) -> None: handle.visible = value self._sync_visibilities() + def _register_perturbation_handlers(self) -> None: + """Attach drag handlers to mesh groups for interactive perturbation.""" + self.perturbation.register_click_handler() + for mg in self._mesh_groups: + self.perturbation.register_drag_handlers(mg.handle, mg.body_ids) + def set_refresh_handler(self, handler: Callable[[], None] | None) -> None: """Install a callback used to refresh cached visualization state.""" with self._update_lock: @@ -300,6 +311,7 @@ def rebuild_visual_handles(self) -> None: mujoco.mj_kinematics(self.mj_model, self.mj_data) self._add_fixed_geometry() self._create_mesh_handles_by_group() + self._register_perturbation_handlers() self._add_fixed_sites() self._create_site_handles_by_group() self._compute_hull_body_meshes() @@ -461,6 +473,22 @@ def _(_) -> None: for client in self.server.get_clients().values(): client.camera.fov = np.radians(slider_fov.value) + if self.mj_model.ncam > 0: + cam_names = ["Free"] + for i in range(self.mj_model.ncam): + name = mj_id2name(self.mj_model, mjtObj.mjOBJ_CAMERA, i) + cam_names.append(name or f"camera_{i}") + self._camera_dropdown = self.server.gui.add_dropdown( + "Camera", + options=cam_names, + initial_value="Free", + ) + self._camera_offset = _camera_offset + + @self._camera_dropdown.on_update + def _(_) -> None: + self._apply_camera_selection() + @self.server.on_client_connect def _(client: viser.ClientHandle) -> None: client.camera.fov = np.radians(slider_fov.value) @@ -483,6 +511,7 @@ def _vis_flag_cb(flag_idx: int, initial: bool = False) -> None: """Add a vis-flag checkbox at the current GUI level.""" _opt.flags[flag_idx] = int(initial) cb = self.server.gui.add_checkbox("Enabled", initial_value=initial) + self._vis_flag_checkboxes[flag_idx] = cb @cb.on_update def _(event, _idx=flag_idx) -> None: @@ -699,6 +728,7 @@ def _mutate() -> None: for label, flag_idx, initial in _simple_flags: cb = self.server.gui.add_checkbox(label, initial_value=initial) + self._simple_flag_checkboxes[flag_idx] = cb @cb.on_update def _(event, _idx=flag_idx) -> None: @@ -802,6 +832,70 @@ def create_visualization_gui( self.create_groups_gui() return tabs + def _apply_camera_selection(self) -> None: + """Switch camera based on dropdown value.""" + if not hasattr(self, "_camera_dropdown"): + return + name = self._camera_dropdown.value + if name == "Free": + self.camera_tracking_enabled = True + for client in self.server.get_clients().values(): + client.camera.position = self._camera_offset + client.camera.look_at = np.zeros(3) + return + + idx = self._camera_dropdown.options.index(name) - 1 + mujoco.mj_forward(self.mj_model, self.mj_data) + pos = self.mj_data.cam_xpos[idx].copy() + mat = self.mj_data.cam_xmat[idx].reshape(3, 3) + look_dir = -mat[:, 2] + look_at = pos + look_dir * 3.0 + self.camera_tracking_enabled = False + for client in self.server.get_clients().values(): + client.camera.position = pos + client.camera.look_at = look_at + + def _toggle_vis_flag(self, flag_idx: int) -> None: + """Toggle a visualization flag and sync the corresponding checkbox.""" + new_val = not bool(self._mjv_option.flags[flag_idx]) + + def _mutate() -> None: + self._mjv_option.flags[flag_idx] = int(new_val) + + cb = self._vis_flag_checkboxes.get(flag_idx) or self._simple_flag_checkboxes.get( + flag_idx + ) + if cb is not None: + cb.value = new_val + else: + self._apply_visualization_change(_mutate) + + def register_keybindings(self, server: viser.ViserServer) -> None: + """Register single-letter keyboard shortcuts for visualization toggles.""" + _flag_keys: list[tuple[str, str, int]] = [ + ("Toggle Contacts", "C", int(mujoco.mjtVisFlag.mjVIS_CONTACTPOINT)), + ("Toggle Forces", "F", int(mujoco.mjtVisFlag.mjVIS_CONTACTFORCE)), + ("Toggle Joints", "J", int(mujoco.mjtVisFlag.mjVIS_JOINT)), + ("Toggle Inertia", "I", int(mujoco.mjtVisFlag.mjVIS_INERTIA)), + ("Toggle Tendons", "V", int(mujoco.mjtVisFlag.mjVIS_TENDON)), + ("Toggle COM", "M", int(mujoco.mjtVisFlag.mjVIS_COM)), + ("Toggle Actuators", "U", int(mujoco.mjtVisFlag.mjVIS_ACTUATOR)), + ("Toggle Constraints", "T", int(mujoco.mjtVisFlag.mjVIS_CONSTRAINT)), + ] + for label, key, flag_idx in _flag_keys: + cmd = server.gui.add_command(label, hotkey=key) # type: ignore[arg-type] + cmd.on_trigger(lambda _, _idx=flag_idx: self._toggle_vis_flag(_idx)) + + cmd = server.gui.add_command("Free Camera", hotkey="escape") + cmd.on_trigger(lambda _: self._set_free_camera()) + + def _set_free_camera(self) -> None: + """Switch to free camera mode.""" + if hasattr(self, "_camera_dropdown"): + self._camera_dropdown.value = "Free" + else: + self.camera_tracking_enabled = True + def update_from_arrays( self, body_xpos: np.ndarray, @@ -943,6 +1037,11 @@ def _update_visualization_locked( if mj_data is not None: self._last_mj_data = mj_data + body_xmat_3x3 = body_xmat + if body_xmat_3x3.ndim == 3 and body_xmat_3x3.shape[-1] == 9: + body_xmat_3x3 = body_xmat_3x3.reshape(*body_xmat_3x3.shape[:-1], 3, 3) + self.perturbation.update_state(body_xpos, body_xmat_3x3, env_idx, scene_offset) + self.fixed_bodies_frame.position = scene_offset slice_single = self.show_only_selected and self.num_envs > 1 with self.server.atomic(): diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index b35fb73..8162d66 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -200,6 +200,23 @@ def _tick(self) -> None: self._stats_frames += 1 self._update_stats() + def _apply_perturbation(self) -> None: + """Apply any pending interactive perturbation force.""" + pert = self.scene.perturbation.get_perturbation() + if pert is not None: + self.data.qfrc_applied[:] = 0.0 + mujoco.mj_applyFT( + self.model, + self.data, + pert.force, + pert.torque, + pert.point, + pert.body_id, + self.data.qfrc_applied, + ) + else: + self.data.qfrc_applied[:] = 0.0 + def _step_physics(self, dt: float) -> bool: """Run physics steps for this frame's sim-time budget. @@ -215,6 +232,7 @@ def _step_physics(self, dt: float) -> bool: deadline = time.perf_counter() + _FRAME_TIME hit_deadline = False while self._budget >= step_dt: + self._apply_perturbation() self._step_fn(self.model, self.data) self._budget -= step_dt self._step_count += 1 @@ -258,53 +276,74 @@ def _update_status_display(self) -> None: """ + def _toggle_pause(self) -> None: + self._paused = not self._paused + if not self._paused: + self._budget = 0.0 + self._last_tick = time.perf_counter() + self._pause_btn.label = "Pause" if not self._paused else "Play" + self._pause_btn.icon = ( + viser.Icon.PLAYER_PAUSE if not self._paused else viser.Icon.PLAYER_PLAY + ) + for sl, _ in self._joint_sliders: + sl.disabled = not self._paused + self._update_status_display() + + def _single_step(self) -> None: + if self._paused: + with self._lock: + self._step_fn(self.model, self.data) + self._step_count += 1 + self._render() + self._update_status_display() + + def _do_reset(self) -> None: + with self._lock: + self._reset() + self._step_count = 0 + self._budget = 0.0 + self._render() + self._update_status_display() + self._sync_sliders() + + def _speed_up(self) -> None: + self._speed_idx = min(len(_SPEEDS) - 1, self._speed_idx + 1) + self._update_status_display() + + def _speed_down(self) -> None: + self._speed_idx = max(0, self._speed_idx - 1) + self._update_status_display() + + def _reset_speed(self) -> None: + self._speed_idx = _SPEEDS.index(1.0) + self._update_status_display() + def _setup_gui(self) -> None: tabs = self._server.gui.add_tab_group() with tabs.add_tab("Controls", icon=viser.Icon.SETTINGS): self._status_html = self._server.gui.add_html("") - pause_btn = self._server.gui.add_button( + self._pause_btn = self._server.gui.add_button( "Pause" if not self._paused else "Play", icon=(viser.Icon.PLAYER_PAUSE if not self._paused else viser.Icon.PLAYER_PLAY), ) - @pause_btn.on_click + @self._pause_btn.on_click def _(_) -> None: - self._paused = not self._paused - if not self._paused: - self._budget = 0.0 - self._last_tick = time.perf_counter() - pause_btn.label = "Pause" if not self._paused else "Play" - pause_btn.icon = ( - viser.Icon.PLAYER_PAUSE if not self._paused else viser.Icon.PLAYER_PLAY - ) - for sl, _ in self._joint_sliders: - sl.disabled = not self._paused - self._update_status_display() + self._toggle_pause() step_btn = self._server.gui.add_button("Step", icon=viser.Icon.PLAYER_TRACK_NEXT) @step_btn.on_click def _(_) -> None: - if self._paused: - with self._lock: - self._step_fn(self.model, self.data) - self._step_count += 1 - self._render() - self._update_status_display() + self._single_step() reset_btn = self._server.gui.add_button("Reset") @reset_btn.on_click def _(_) -> None: - with self._lock: - self._reset() - self._step_count = 0 - self._budget = 0.0 - self._render() - self._update_status_display() - self._sync_sliders() + self._do_reset() speed_btns = self._server.gui.add_button_group( "Speed", options=["Slower", "1x", "Faster"] @@ -313,12 +352,11 @@ def _(_) -> None: @speed_btns.on_click def _(event) -> None: if event.target.value == "Slower": - self._speed_idx = max(0, self._speed_idx - 1) + self._speed_down() elif event.target.value == "Faster": - self._speed_idx = min(len(_SPEEDS) - 1, self._speed_idx + 1) + self._speed_up() else: - self._speed_idx = _SPEEDS.index(1.0) - self._update_status_display() + self._reset_speed() # Keyframe selector (only if model has keyframes). if self.model.nkey > 0: @@ -356,6 +394,9 @@ def _(_) -> None: def _(_) -> None: _load_keyframe() + # Interactive perturbation info. + self.scene.perturbation.setup_gui() + # Scene controls (camera, environment). with self._server.gui.add_folder("Scene"): self.scene.create_scene_gui() @@ -377,6 +418,29 @@ def _(_) -> None: with tabs.add_tab("Physics", icon=viser.Icon.ATOM): self._setup_physics_flags() + self._register_keybindings() + + def _register_keybindings(self) -> None: + """Register keyboard shortcuts for simulation control.""" + s = self._server + + cmd = s.gui.add_command("Toggle Pause", hotkey="space") + cmd.on_trigger(lambda _: self._toggle_pause()) + + cmd = s.gui.add_command("Step Forward", hotkey="N") + cmd.on_trigger(lambda _: self._single_step()) + + cmd = s.gui.add_command("Reset", hotkey="backspace") + cmd.on_trigger(lambda _: self._do_reset()) + + cmd = s.gui.add_command("Speed Down") + cmd.on_trigger(lambda _: self._speed_down()) + + cmd = s.gui.add_command("Speed Up") + cmd.on_trigger(lambda _: self._speed_up()) + + self.scene.register_keybindings(s) + _MAX_SLIDERS: int = 200 def _setup_joint_sliders(self) -> None: From 9368b068cab602c726adfa0b23e2f1f724595796 Mon Sep 17 00:00:00 2001 From: brentyi Date: Sun, 17 May 2026 22:03:42 -0700 Subject: [PATCH 2/6] bump Viser, `on_drag()`, fix minor frame issue --- src/mjviser/interaction.py | 241 +++++++++++++++++++------------------ src/mjviser/scene.py | 7 +- src/mjviser/viewer.py | 50 +++++--- 3 files changed, 168 insertions(+), 130 deletions(-) diff --git a/src/mjviser/interaction.py b/src/mjviser/interaction.py index 9c3db84..fe3294d 100644 --- a/src/mjviser/interaction.py +++ b/src/mjviser/interaction.py @@ -51,6 +51,16 @@ def __init__( self._info_text: viser.GuiTextHandle | None = None + # Grab-point velocity, computed once per render frame in + # ``update_state`` and reused across all physics substeps that + # frame. Computing it inside ``get_perturbation`` (which runs once + # per substep) was both wasteful and noisy: the cached body pose + # only refreshes between render frames, so substeps 2..N within + # a frame would see zero motion divided by a sub-millisecond + # wall-clock dt -- damping collapsed to zero for most substeps and + # spiked on the first one. Render-frame cadence matches the rate + # at which the grab point actually moves on screen. + self._grab_vel: np.ndarray = np.zeros(3) self._prev_grab_world: np.ndarray | None = None self._prev_time: float | None = None @@ -61,37 +71,21 @@ def setup_gui(self) -> None: "Body", initial_value="(none)", disabled=True ) - def register_click_handler(self) -> None: - """Register scene-level click handler for body selection via ray cast.""" - - @self._server.scene.on_click() - def _(event: viser.SceneClickEvent) -> None: - origin = np.array(event.ray_origin) - self._scene_offset - direction = np.array(event.ray_direction) - geomid = np.zeros(1, dtype=np.int32) - dist = mujoco.mj_ray( - self._model, - self._data, - origin, - direction, - None, - 1, - -1, - geomid, - ) - if dist < 0: - with self._lock: - self.selected_body_id = None - if self._info_text is not None: - self._info_text.value = "(none)" - return - - body_id = int(self._model.geom_bodyid[geomid[0]]) - with self._lock: - self.selected_body_id = body_id - name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, body_id) - if self._info_text is not None: - self._info_text.value = name or f"body_{body_id}" + def clear(self) -> None: + """Drop selection and any in-flight drag. Call on reset or after + the model is rebuilt -- otherwise a cached body id can outlive the + body it pointed to (out-of-range index, silently retargets a + different body) and the GUI label keeps showing a stale name.""" + with self._lock: + self.selected_body_id = None + self._drag_body_id = None + self._drag_grab_local = None + self._drag_target = None + self._drag_mode = None + self._prev_grab_world = None + self._prev_time = None + if self._info_text is not None: + self._info_text.value = "(none)" def register_drag_handlers( self, @@ -99,6 +93,14 @@ def register_drag_handlers( body_ids: np.ndarray, ) -> None: """Attach drag and click handlers to a batched mesh handle.""" + # Skip groups whose only body is the world body (id 0): perturbing + # it is meaningless, and registering any handler makes the mesh + # "interactive" in viser -- which flips the cursor to "pointer" + # whenever the user hovers it. The world body holds the ground + # plane, which spans the visible canvas, so the pointer cursor + # would show across the entire scene. + if not bool(np.any(body_ids != 0)): + return n_bodies = len(body_ids) def _body_id_from_event( @@ -123,75 +125,48 @@ def _(event: viser.SceneNodePointerEvent) -> None: # type: ignore[type-arg] if self._info_text is not None: self._info_text.value = name or f"body_{bid}" - @handle.on_drag_start("left", modifier="cmd/ctrl") - async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - bid = _body_id_from_event(event) - if bid is None or bid == 0: - return - grab_world = np.array(event.start_position) - self._scene_offset - grab_local = self._world_to_body_local(bid, grab_world) - with self._lock: - self._drag_body_id = bid - self._drag_grab_local = grab_local - self._drag_target = grab_world - self._drag_mode = "translate" - self.selected_body_id = bid - self._prev_grab_world = grab_world.copy() - self._prev_time = time.perf_counter() - name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) - if self._info_text is not None: - self._info_text.value = name or f"body_{bid}" - - @handle.on_drag_update("left", modifier="cmd/ctrl") - async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - with self._lock: - if self._drag_body_id is not None: - self._drag_target = np.array(event.end_position) - self._scene_offset - - @handle.on_drag_end("left", modifier="cmd/ctrl") - async def _(_: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - with self._lock: - self._drag_body_id = None - self._drag_grab_local = None - self._drag_target = None - self._drag_mode = None - self._prev_grab_world = None - self._prev_time = None - - @handle.on_drag_start("left", modifier="cmd/ctrl+shift") - async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - bid = _body_id_from_event(event) - if bid is None or bid == 0: - return - grab_world = np.array(event.start_position) - self._scene_offset - grab_local = self._world_to_body_local(bid, grab_world) - with self._lock: - self._drag_body_id = bid - self._drag_grab_local = grab_local - self._drag_target = grab_world - self._drag_mode = "rotate" - self.selected_body_id = bid - self._prev_grab_world = grab_world.copy() - self._prev_time = time.perf_counter() - name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) - if self._info_text is not None: - self._info_text.value = name or f"body_{bid}" - - @handle.on_drag_update("left", modifier="cmd/ctrl+shift") - async def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - with self._lock: - if self._drag_body_id is not None: - self._drag_target = np.array(event.end_position) - self._scene_offset - - @handle.on_drag_end("left", modifier="cmd/ctrl+shift") - async def _(_: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - with self._lock: - self._drag_body_id = None - self._drag_grab_local = None - self._drag_target = None - self._drag_mode = None - self._prev_grab_world = None - self._prev_time = None + def _make_drag_handler(mode: str): + async def _handler(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + if event.phase == "start": + bid = _body_id_from_event(event) + if bid is None or bid == 0: + return + # Target and grab position are tracked in viser-world frame + # (no scene_offset subtraction) so that camera tracking, + # which shifts scene_offset every frame, doesn't drift the + # target between throttled mouse-move events. Body-local + # offset still needs mujoco-frame coords for the conversion. + target_viser = np.array(event.start_position) + grab_world_mj = target_viser - self._scene_offset + grab_local = self._world_to_body_local(bid, grab_world_mj) + with self._lock: + self._drag_body_id = bid + self._drag_grab_local = grab_local + self._drag_target = target_viser + self._drag_mode = mode + self.selected_body_id = bid + self._prev_grab_world = target_viser.copy() + self._prev_time = time.perf_counter() + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) + if self._info_text is not None: + self._info_text.value = name or f"body_{bid}" + elif event.phase == "update": + with self._lock: + if self._drag_body_id is not None: + self._drag_target = np.array(event.end_position) + else: # "end" + with self._lock: + self._drag_body_id = None + self._drag_grab_local = None + self._drag_target = None + self._drag_mode = None + self._prev_grab_world = None + self._prev_time = None + + return _handler + + handle.on_drag("left", modifier="cmd/ctrl")(_make_drag_handler("translate")) + handle.on_drag("left", modifier="cmd/ctrl+shift")(_make_drag_handler("rotate")) def update_state( self, @@ -208,10 +183,35 @@ def update_state( env_idx: Active environment index. scene_offset: Current camera tracking offset. """ - self._body_xpos = body_xpos[env_idx] - self._body_xmat = body_xmat[env_idx] + xpos_env = body_xpos[env_idx] + xmat_env = body_xmat[env_idx] + self._body_xpos = xpos_env + self._body_xmat = xmat_env self._scene_offset = scene_offset + # Refresh the grab-point velocity once per render frame against + # the new pose. Substeps within the next frame all consume this + # cached value -- see ``_grab_vel`` field comment. + bid = self._drag_body_id + if bid is None or self._drag_grab_local is None or bid >= xpos_env.shape[0]: + self._grab_vel = np.zeros(3) + self._prev_grab_world = None + self._prev_time = None + return + # Velocity is measured in viser-world frame -- the same frame the + # cursor (and therefore _drag_target) lives in -- so damping + # actually opposes motion *relative to the target*. With camera + # tracking on, the dragged body sits near the viser origin and + # mujoco-frame velocity would diverge from screen-frame velocity. + grab_world = xpos_env[bid] + xmat_env[bid] @ self._drag_grab_local + scene_offset + now = time.perf_counter() + if self._prev_grab_world is not None and self._prev_time is not None: + dt = now - self._prev_time + if dt > 1e-6: + self._grab_vel = (grab_world - self._prev_grab_world) / dt + self._prev_grab_world = grab_world.copy() + self._prev_time = now + def get_perturbation(self) -> PerturbationState | None: """Return the current perturbation force using MuJoCo's spring model.""" with self._lock: @@ -225,25 +225,36 @@ def get_perturbation(self) -> PerturbationState | None: return None bid = self._drag_body_id + # Bounds-check against the *current* model: if the model was + # rebuilt mid-drag, the cached body id may now reference a + # different body or be out of range entirely. Drop the drag + # silently rather than indexing past the end. + if bid >= self._body_xpos.shape[0] or bid >= self._model.nbody: + self._drag_body_id = None + self._drag_grab_local = None + self._drag_target = None + self._drag_mode = None + self._prev_grab_world = None + self._prev_time = None + return None xpos = self._body_xpos[bid] xmat = self._body_xmat[bid] - grab_world = xpos + xmat @ self._drag_grab_local + grab_world_mj = xpos + xmat @ self._drag_grab_local + # Displacement is computed in viser-world frame (cursor frame). + # Force/torque vectors are unchanged by the pure-translation + # offset, so we apply them in mujoco frame as-is; only the + # application point passed to mj_applyFT needs mujoco coords. + grab_world = grab_world_mj + self._scene_offset diff = grab_world - self._drag_target # Compute localmass from body_invweight0 (translational inverse weight). invweight_trn = float(self._model.body_invweight0[bid, 0]) localmass = 1.0 if invweight_trn == 0 else 1.0 / max(invweight_trn, 1e-10) - # Estimate velocity of the grab point via finite difference. - now = time.perf_counter() - vel = np.zeros(3) - if self._prev_grab_world is not None and self._prev_time is not None: - dt = now - self._prev_time - if dt > 1e-6: - vel = (grab_world - self._prev_grab_world) / dt - self._prev_grab_world = grab_world.copy() - self._prev_time = now + # Velocity is refreshed per render frame in ``update_state``; + # we just read the cached value here. + vel = self._grab_vel if self._drag_mode == "translate": stiffness = float(self._model.vis.map.stiffness) @@ -252,8 +263,8 @@ def get_perturbation(self) -> PerturbationState | None: # Critical damping: -2*sqrt(k)*m * velocity. force -= np.sqrt(stiffness) * localmass * vel - # Torque from moment arm. - moment_arm = grab_world - xpos + # Torque from moment arm (mujoco-frame coords). + moment_arm = grab_world_mj - xpos torque = np.cross(moment_arm, force) else: stiffnessrot = float(self._model.vis.map.stiffnessrot) @@ -271,7 +282,7 @@ def get_perturbation(self) -> PerturbationState | None: body_id=bid, force=force, torque=torque, - point=grab_world, + point=grab_world_mj, ) def _world_to_body_local(self, body_id: int, world_pos: np.ndarray) -> np.ndarray: diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 9c1e817..9ac8969 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -276,7 +276,6 @@ def show_convex_hull(self, value: bool) -> None: def _register_perturbation_handlers(self) -> None: """Attach drag handlers to mesh groups for interactive perturbation.""" - self.perturbation.register_click_handler() for mg in self._mesh_groups: self.perturbation.register_drag_handlers(mg.handle, mg.body_ids) @@ -288,6 +287,12 @@ def set_refresh_handler(self, handler: Callable[[], None] | None) -> None: def rebuild_visual_handles(self) -> None: """Rebuild handles whose geometry or appearance is baked at creation time.""" with self._update_lock: + # Drop selection / drag state before mesh handles are torn down: + # body ids may not survive a rebuild, and an in-flight drag + # holding a stale id would silently retarget a different body + # (or index out of range) on the next physics step. + self.perturbation.clear() + for group in self._mesh_groups: group.handle.remove() self._mesh_groups.clear() diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index 8162d66..5a274e9 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -87,6 +87,12 @@ def __init__( self._dirty = True self._step_count = 0 + # Snapshot of ``qfrc_applied`` taken just before the perturbation + # adds its force, so we can undo our addition before the next step + # without clobbering user-supplied generalized forces. ``None`` + # means no perturbation force is currently in flight. + self._pert_saved_qfrc: np.ndarray | None = None + # Timing. self._budget = 0.0 self._last_tick = 0.0 @@ -201,21 +207,32 @@ def _tick(self) -> None: self._update_stats() def _apply_perturbation(self) -> None: - """Apply any pending interactive perturbation force.""" + """Apply any pending interactive perturbation force. + + ``mj_applyFT`` *adds* to ``qfrc_applied`` rather than overwriting + it, so a force applied last substep would compound this substep + if we just called it again. We undo the previous addition by + restoring a pre-perturbation snapshot, then re-snapshot and re-add + on top -- which preserves any user-supplied ``qfrc_applied`` + (a documented MuJoCo extension point) instead of zeroing it out. + """ + if self._pert_saved_qfrc is not None: + self.data.qfrc_applied[:] = self._pert_saved_qfrc + self._pert_saved_qfrc = None + pert = self.scene.perturbation.get_perturbation() - if pert is not None: - self.data.qfrc_applied[:] = 0.0 - mujoco.mj_applyFT( - self.model, - self.data, - pert.force, - pert.torque, - pert.point, - pert.body_id, - self.data.qfrc_applied, - ) - else: - self.data.qfrc_applied[:] = 0.0 + if pert is None: + return + self._pert_saved_qfrc = self.data.qfrc_applied.copy() + mujoco.mj_applyFT( + self.model, + self.data, + pert.force, + pert.torque, + pert.point, + pert.body_id, + self.data.qfrc_applied, + ) def _step_physics(self, dt: float) -> bool: """Run physics steps for this frame's sim-time budget. @@ -302,6 +319,11 @@ def _do_reset(self) -> None: self._reset() self._step_count = 0 self._budget = 0.0 + # Drop any in-flight drag and the saved-qfrc snapshot: the + # snapshot was taken against the pre-reset state and would + # otherwise be restored on top of the freshly-reset data. + self.scene.perturbation.clear() + self._pert_saved_qfrc = None self._render() self._update_status_display() self._sync_sliders() From b2e0e23e867dc26ae66244c7d932b481acc0e7ac Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 10 Jun 2026 14:28:24 -0700 Subject: [PATCH 3/6] Rework perturbation on MuJoCo's engine; add rotate and camera fixes Drive a mujoco.MjvPerturb with mjv_applyPerturbForce / mjv_applyPerturbPose so the spring force, critical damping (via mj_objectVelocity), and pose math match MuJoCo's simulate and studio. Input is web-native via viser's absolute drag target: cmd/ctrl+left-drag pulls the grab point toward the cursor (off-center grabs rotate through the moment arm) with a connector line, cmd/ctrl+right-drag rotates toward a screen-built orientation shown by a ghost box, and while paused both move free-joint and mocap bodies kinematically with mj_forward so contacts and forces refresh live. Also: number keys toggle geom groups and shift+number toggles site groups, camera tracking defaults off, and untracking shifts the camera so the view no longer teleports. --- src/mjviser/interaction.py | 599 ++++++++++++++++++++++++------------- src/mjviser/scene.py | 43 ++- src/mjviser/viewer.py | 47 +-- tests/test_interaction.py | 206 +++++++++++++ 4 files changed, 644 insertions(+), 251 deletions(-) create mode 100644 tests/test_interaction.py diff --git a/src/mjviser/interaction.py b/src/mjviser/interaction.py index fe3294d..8145937 100644 --- a/src/mjviser/interaction.py +++ b/src/mjviser/interaction.py @@ -1,32 +1,38 @@ -"""Interactive perturbation: click-to-select and drag-to-apply-force.""" +"""Interactive perturbation: click to select, cmd/ctrl+drag to apply a force. + +Backed by MuJoCo's own perturbation engine (``mjvPerturb`` + +``mjv_applyPerturbForce`` / ``mjv_applyPerturbPose``), so the spring force, +critical damping, and pose math match MuJoCo's simulate and studio exactly -- +including velocity damping computed analytically from ``mj_objectVelocity``. + +The only web-specific part is the input. MuJoCo's native apps accumulate the +reference point from relative mouse deltas scaled by a depth factor, because +they only have 2D cursor motion. viser drag events instead report absolute +world-space coordinates (the grab point on the body, and the cursor projected +onto a camera-aligned plane), so we set the reference selection point directly. + +cmd/ctrl + left-drag applies a force at the grab point (off-center grabs also +rotate the body through the moment arm). cmd/ctrl + right-drag rotates the body +toward an orientation built from the screen drag. While paused, both drags move +free-joint / mocap bodies kinematically instead of applying forces. +""" from __future__ import annotations -import time -from dataclasses import dataclass from threading import Lock import mujoco import numpy as np import viser - -@dataclass -class PerturbationState: - """Current perturbation force to be applied by the simulation loop.""" - - body_id: int - force: np.ndarray - torque: np.ndarray - point: np.ndarray +_CONNECTOR_NAME = "/perturb/connector" +_CONNECTOR_COLOR = (255, 90, 90) +_GHOST_NAME = "/perturb/ghost" +_GHOST_COLOR = (255, 200, 60) class PerturbationHandler: - """Manages body selection and drag-to-apply-force perturbation. - - Uses MuJoCo's critically damped spring model: the force is proportional - to stiffness * localmass, with critical damping from velocity feedback. - """ + """Manages body selection and drag-to-perturb on the viewer's MjData.""" def __init__( self, @@ -40,50 +46,72 @@ def __init__( self._lock = Lock() self.selected_body_id: int | None = None - self._drag_body_id: int | None = None - self._drag_grab_local: np.ndarray | None = None - self._drag_target: np.ndarray | None = None - self._drag_mode: str | None = None + # MuJoCo's perturbation state. We populate select / localpos / localmass / + # refselpos / refpos / refquat ourselves and let MuJoCo do the physics. + self._pert = mujoco.MjvPerturb() + self._pert.active = 0 + + self._drag_body_id: int | None = None + self._grab_local: np.ndarray | None = None # grab point in body-local coords + self._target_viser: np.ndarray | None = None # cursor target, viser frame + + # Rotate drag (cmd + right-drag): a target orientation built from the + # camera-relative screen drag, applied as a spring torque toward it. + self._drag_rotate: bool = False + self._rotate_delta_quat: np.ndarray | None = None # screen-driven rotation + self._rotate_initial_quat: np.ndarray | None = None # body orientation at start + self._rotate_anchor: np.ndarray | None = None # drag-start screen position + # Ghost-box (MuJoCo-style orientation indicator) state, captured at start. + self._rotate_aabb_center: np.ndarray | None = None # body-frame box center + self._rotate_aabb_half: np.ndarray | None = None # body-frame box half-extents + self._rotate_world_quat: np.ndarray | None = None # body world orientation + + # Body pose for the active env, cached each render frame. Used to convert + # the drag-start grab point into body-local coords and to draw the + # connector; the force/pose physics reads the live MjData passed to apply(). self._scene_offset = np.zeros(3) self._body_xpos: np.ndarray | None = None self._body_xmat: np.ndarray | None = None - self._info_text: viser.GuiTextHandle | None = None + # Body whose xfrc_applied we last wrote, so we can zero exactly that + # entry when the drag ends without disturbing other bodies' applied forces. + self._applied_body_id: int | None = None - # Grab-point velocity, computed once per render frame in - # ``update_state`` and reused across all physics substeps that - # frame. Computing it inside ``get_perturbation`` (which runs once - # per substep) was both wasteful and noisy: the cached body pose - # only refreshes between render frames, so substeps 2..N within - # a frame would see zero motion divided by a sub-millisecond - # wall-clock dt -- damping collapsed to zero for most substeps and - # spiked on the first one. Render-frame cadence matches the rate - # at which the grab point actually moves on screen. - self._grab_vel: np.ndarray = np.zeros(3) - self._prev_grab_world: np.ndarray | None = None - self._prev_time: float | None = None + self._info_text: viser.GuiTextHandle | None = None + self._connector: viser.LineSegmentsHandle | None = None + self._ghost: viser.BoxHandle | None = None def setup_gui(self) -> None: - """Add selection info display to the GUI.""" + """Add a selection-info display to the GUI.""" with self._server.gui.add_folder("Selection"): self._info_text = self._server.gui.add_text( "Body", initial_value="(none)", disabled=True ) def clear(self) -> None: - """Drop selection and any in-flight drag. Call on reset or after - the model is rebuilt -- otherwise a cached body id can outlive the - body it pointed to (out-of-range index, silently retargets a - different body) and the GUI label keeps showing a stale name.""" + """Drop selection and any in-flight drag. + + Call on reset or after the model is rebuilt: a cached body id can + otherwise outlive the body it pointed to (out-of-range index, or a + silently retargeted body) and the GUI label would keep a stale name. + The dragged body's ``xfrc_applied`` is zeroed by the next ``apply``. + """ with self._lock: self.selected_body_id = None self._drag_body_id = None - self._drag_grab_local = None - self._drag_target = None - self._drag_mode = None - self._prev_grab_world = None - self._prev_time = None + self._grab_local = None + self._target_viser = None + self._drag_rotate = False + self._rotate_delta_quat = None + self._rotate_initial_quat = None + self._rotate_anchor = None + self._rotate_aabb_center = None + self._rotate_aabb_half = None + self._rotate_world_quat = None + self._pert.active = 0 + self._hide_connector() + self._hide_ghost() if self._info_text is not None: self._info_text.value = "(none)" @@ -92,81 +120,163 @@ def register_drag_handlers( handle: viser.BatchedGlbHandle, body_ids: np.ndarray, ) -> None: - """Attach drag and click handlers to a batched mesh handle.""" - # Skip groups whose only body is the world body (id 0): perturbing - # it is meaningless, and registering any handler makes the mesh - # "interactive" in viser -- which flips the cursor to "pointer" - # whenever the user hovers it. The world body holds the ground - # plane, which spans the visible canvas, so the pointer cursor - # would show across the entire scene. + """Attach click (select) and cmd/ctrl+drag (perturb) handlers.""" + # Skip groups whose only body is the world body (id 0): perturbing it is + # meaningless, and registering any handler makes the mesh interactive in + # viser, flipping the cursor to "pointer" across the whole ground plane. if not bool(np.any(body_ids != 0)): return n_bodies = len(body_ids) - def _body_id_from_event( - event: viser.SceneNodeDragEvent, # type: ignore[type-arg] - ) -> int | None: - idx = event.instance_index + def _body_id(idx: int | None) -> int | None: if idx is None: return None return int(body_ids[idx % n_bodies]) @handle.on_click def _(event: viser.SceneNodePointerEvent) -> None: # type: ignore[type-arg] - idx = event.instance_index - if idx is None: - return - bid = int(body_ids[idx % n_bodies]) - if bid == 0: + bid = _body_id(event.instance_index) + if bid is None or bid == 0: return with self._lock: self.selected_body_id = bid - name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) - if self._info_text is not None: - self._info_text.value = name or f"body_{bid}" - - def _make_drag_handler(mode: str): - async def _handler(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] - if event.phase == "start": - bid = _body_id_from_event(event) - if bid is None or bid == 0: - return - # Target and grab position are tracked in viser-world frame - # (no scene_offset subtraction) so that camera tracking, - # which shifts scene_offset every frame, doesn't drift the - # target between throttled mouse-move events. Body-local - # offset still needs mujoco-frame coords for the conversion. - target_viser = np.array(event.start_position) - grab_world_mj = target_viser - self._scene_offset - grab_local = self._world_to_body_local(bid, grab_world_mj) - with self._lock: - self._drag_body_id = bid - self._drag_grab_local = grab_local - self._drag_target = target_viser - self._drag_mode = mode - self.selected_body_id = bid - self._prev_grab_world = target_viser.copy() - self._prev_time = time.perf_counter() - name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) - if self._info_text is not None: - self._info_text.value = name or f"body_{bid}" - elif event.phase == "update": - with self._lock: - if self._drag_body_id is not None: - self._drag_target = np.array(event.end_position) - else: # "end" - with self._lock: - self._drag_body_id = None - self._drag_grab_local = None - self._drag_target = None - self._drag_mode = None - self._prev_grab_world = None - self._prev_time = None - - return _handler - - handle.on_drag("left", modifier="cmd/ctrl")(_make_drag_handler("translate")) - handle.on_drag("left", modifier="cmd/ctrl+shift")(_make_drag_handler("rotate")) + self._set_info(bid) + + # cmd/ctrl + left-drag moves the body in the camera-facing plane. Depth + # (toward/away from the camera) isn't a separate gesture: viser freezes + # the drag modifier at drag-start and has no scroll-during-drag, so there + # is no fluid way to switch planes mid-drag. Orbit the view and drag again + # to reach a different depth. + @handle.on_drag("left", modifier="cmd/ctrl") + def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + if event.phase == "start": + self._on_drag_start(_body_id(event.instance_index), event.start_position) + elif event.phase == "update": + with self._lock: + if self._drag_body_id is not None and not self._drag_rotate: + self._target_viser = np.array(event.end_position) + else: # "end" + self._end_drag() + + # cmd/ctrl + right-drag rotates the body: the screen drag builds a target + # orientation (yaw from horizontal motion, pitch from vertical), and the + # body is sprung toward it -- MuJoCo's rotate perturbation. + @handle.on_drag("right", modifier="cmd/ctrl") + def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] + if event.phase == "start": + self._on_rotate_start(_body_id(event.instance_index), event.start_screen_pos) + elif event.phase == "update": + dq = self._rotation_from_drag(event) + with self._lock: + if self._drag_body_id is not None and self._drag_rotate: + self._rotate_delta_quat = dq + else: # "end" + self._end_drag() + + def _end_drag(self) -> None: + with self._lock: + self._drag_body_id = None + self._grab_local = None + self._target_viser = None + self._drag_rotate = False + self._rotate_delta_quat = None + self._rotate_initial_quat = None + self._rotate_anchor = None + self._rotate_aabb_center = None + self._rotate_aabb_half = None + self._rotate_world_quat = None + self._pert.active = 0 + self._hide_connector() + self._hide_ghost() + + def _on_drag_start( + self, bid: int | None, start_position: tuple[float, float, float] + ) -> None: + if bid is None or bid == 0: + return + with self._lock: + if ( + self._body_xpos is None + or self._body_xmat is None + or bid >= self._body_xpos.shape[0] + ): + return + # viser reports world coords in the (camera-tracking) viser frame; + # subtract the scene offset to get MuJoCo-frame coords. + grab_world = np.array(start_position) - self._scene_offset + xpos = self._body_xpos[bid] + xmat = self._body_xmat[bid] + self._grab_local = xmat.T @ (grab_world - xpos) + self._drag_body_id = bid + self._drag_rotate = False + self._target_viser = np.array(start_position) + self.selected_body_id = bid + self._set_info(bid) + + def _on_rotate_start( + self, bid: int | None, start_screen_pos: tuple[float, float] + ) -> None: + if bid is None or bid == 0: + return + with self._lock: + self._drag_body_id = bid + self._drag_rotate = True + self._rotate_anchor = np.array(start_screen_pos, dtype=float) + self._rotate_delta_quat = np.array([1.0, 0.0, 0.0, 0.0]) + self._rotate_initial_quat = None # captured on first apply + self._rotate_aabb_center, self._rotate_aabb_half = self._body_aabb(bid) + if self._body_xmat is not None and bid < self._body_xmat.shape[0]: + quat = np.empty(4) + mujoco.mju_mat2Quat(quat, self._body_xmat[bid].reshape(9)) + self._rotate_world_quat = quat + self.selected_body_id = bid + self._set_info(bid) + + def _body_aabb(self, bid: int) -> tuple[np.ndarray, np.ndarray]: + """Axis-aligned bounding box (center, half-extent) of a body's geoms, + in the body frame, from each geom's bounding sphere.""" + lo = np.full(3, np.inf) + hi = np.full(3, -np.inf) + found = False + for g in range(self._model.ngeom): + if int(self._model.geom_bodyid[g]) != bid: + continue + pos = np.asarray(self._model.geom_pos[g], dtype=float) + r = float(self._model.geom_rbound[g]) + if r <= 0.0: + r = float(np.max(self._model.geom_size[g])) or 0.05 + lo = np.minimum(lo, pos - r) + hi = np.maximum(hi, pos + r) + found = True + if not found: + return np.zeros(3), np.full(3, 0.05) + return (lo + hi) / 2.0, (hi - lo) / 2.0 + + def _rotation_from_drag( + self, + event: viser.SceneNodeDragEvent, # type: ignore[type-arg] + ) -> np.ndarray: + """Build a delta rotation quaternion from the camera-relative screen drag.""" + identity = np.array([1.0, 0.0, 0.0, 0.0]) + if self._rotate_anchor is None: + return identity + delta = np.array(event.end_screen_pos, dtype=float) - self._rotate_anchor + cam = event.client.camera + pos = np.asarray(cam.position, dtype=float) + forward = np.asarray(cam.look_at, dtype=float) - pos + forward /= max(np.linalg.norm(forward), 1e-9) + right = np.cross(forward, np.asarray(cam.up_direction, dtype=float)) + right /= max(np.linalg.norm(right), 1e-9) + up = np.cross(right, forward) + # Horizontal drag yaws about the camera up axis; vertical drag (OpenCV y is + # down) pitches about the camera right axis. A full-screen drag ~= 180 deg. + rotvec = up * (delta[0] * np.pi) + right * (delta[1] * np.pi) + angle = float(np.linalg.norm(rotvec)) + if angle < 1e-9: + return identity + quat = np.empty(4) + mujoco.mju_axisAngle2Quat(quat, rotvec / angle, angle) + return quat def update_state( self, @@ -175,120 +285,187 @@ def update_state( env_idx: int, scene_offset: np.ndarray, ) -> None: - """Update cached body state from the latest simulation frame. + """Cache body pose and scene offset from the latest render frame. Args: body_xpos: Shape ``(num_envs, nbody, 3)``. body_xmat: Shape ``(num_envs, nbody, 3, 3)``. env_idx: Active environment index. - scene_offset: Current camera tracking offset. + scene_offset: Current camera-tracking offset (viser - mujoco). """ - xpos_env = body_xpos[env_idx] - xmat_env = body_xmat[env_idx] - self._body_xpos = xpos_env - self._body_xmat = xmat_env + self._body_xpos = body_xpos[env_idx] + self._body_xmat = body_xmat[env_idx] self._scene_offset = scene_offset + self._update_connector() + self._update_rotate_ghost() - # Refresh the grab-point velocity once per render frame against - # the new pose. Substeps within the next frame all consume this - # cached value -- see ``_grab_vel`` field comment. - bid = self._drag_body_id - if bid is None or self._drag_grab_local is None or bid >= xpos_env.shape[0]: - self._grab_vel = np.zeros(3) - self._prev_grab_world = None - self._prev_time = None - return - # Velocity is measured in viser-world frame -- the same frame the - # cursor (and therefore _drag_target) lives in -- so damping - # actually opposes motion *relative to the target*. With camera - # tracking on, the dragged body sits near the viser origin and - # mujoco-frame velocity would diverge from screen-frame velocity. - grab_world = xpos_env[bid] + xmat_env[bid] @ self._drag_grab_local + scene_offset - now = time.perf_counter() - if self._prev_grab_world is not None and self._prev_time is not None: - dt = now - self._prev_time - if dt > 1e-6: - self._grab_vel = (grab_world - self._prev_grab_world) / dt - self._prev_grab_world = grab_world.copy() - self._prev_time = now - - def get_perturbation(self) -> PerturbationState | None: - """Return the current perturbation force using MuJoCo's spring model.""" + def apply(self, model: mujoco.MjModel, data: mujoco.MjData, paused: bool) -> bool: + """Apply the active drag to ``data`` using MuJoCo's perturb engine. + + When running, writes a critically-damped spring force to + ``xfrc_applied`` (consumed by the next ``mj_step``). When paused, + repositions free-joint / mocap bodies so the grab point follows the + cursor. Returns True if it moved the body's pose (paused drag), so the + caller can trigger a re-render. + """ + with self._lock: + bid = self._drag_body_id + if bid is None or bid <= 0 or bid >= model.nbody: + self._clear_applied(data) + return False + self._pert.select = bid + + if self._drag_rotate: + return self._apply_rotate(model, data, paused, bid) + + if self._grab_local is None or self._target_viser is None: + self._clear_applied(data) + return False + + xmat = data.xmat[bid].reshape(3, 3) + selpos = data.xpos[bid] + xmat @ self._grab_local + target = self._target_viser - self._scene_offset + + self._pert.localpos = self._grab_local + + if paused: + # Reposition: translate the body so the grab point reaches the + # target, keeping orientation. mjv_applyPerturbPose moves free + # joints / mocap bodies and is a no-op for anything else. + self._clear_applied(data) + xiquat = np.empty(4) + mujoco.mju_mulQuat(xiquat, data.xquat[bid], model.body_iquat[bid]) + self._pert.active = int(mujoco.mjtPertBit.mjPERT_TRANSLATE) + self._pert.refpos = data.xipos[bid] + (target - selpos) + self._pert.refquat = xiquat + mujoco.mjv_applyPerturbPose(model, data, self._pert, 1) + # mj_forward (not just mj_kinematics) so collisions, contacts, and + # constraint forces refresh for the dragged pose -- matches MuJoCo + # simulate, which runs mj_forward every frame while paused. + mujoco.mj_forward(model, data) + return True + + # Running: spring force toward the cursor. localmass is MuJoCo's + # precomputed effective translational mass for the body. + invweight = float(model.body_invweight0[bid, 0]) + self._pert.localmass = 1.0 / invweight if invweight > 0 else 1.0 + self._pert.refselpos = target + self._pert.active = int(mujoco.mjtPertBit.mjPERT_TRANSLATE) + mujoco.mjv_applyPerturbForce(model, data, self._pert) + self._applied_body_id = bid + return False + + def _apply_rotate( + self, model: mujoco.MjModel, data: mujoco.MjData, paused: bool, bid: int + ) -> bool: + """Spring the body toward the drag-built target orientation. + + The caller holds ``self._lock``. + """ + if self._rotate_delta_quat is None: + self._clear_applied(data) + return False + # Capture the body's orientation once, at the first apply of this rotate. + if self._rotate_initial_quat is None: + initial = np.empty(4) + mujoco.mju_mulQuat(initial, data.xquat[bid], model.body_iquat[bid]) + self._rotate_initial_quat = initial + refquat = np.empty(4) + mujoco.mju_mulQuat(refquat, self._rotate_delta_quat, self._rotate_initial_quat) + mujoco.mju_normalize4(refquat) + self._pert.refpos = data.xipos[bid].copy() + self._pert.refquat = refquat + self._pert.active = int(mujoco.mjtPertBit.mjPERT_ROTATE) + if paused: + self._clear_applied(data) + mujoco.mjv_applyPerturbPose(model, data, self._pert, 1) + mujoco.mj_forward(model, data) + return True + mujoco.mjv_applyPerturbForce(model, data, self._pert) + self._applied_body_id = bid + return False + + def _clear_applied(self, data: mujoco.MjData) -> None: + """Zero the xfrc_applied entry we last wrote, if any.""" + bid = self._applied_body_id + if bid is not None and bid < data.xfrc_applied.shape[0]: + data.xfrc_applied[bid] = 0.0 + self._applied_body_id = None + + def _set_info(self, bid: int) -> None: + if self._info_text is not None: + name = mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, bid) + self._info_text.value = name or f"body_{bid}" + + def _update_connector(self) -> None: + """Draw or refresh the grab-point-to-cursor connector line.""" + selpos = target = None with self._lock: + bid = self._drag_body_id if ( - self._drag_body_id is None - or self._drag_grab_local is None - or self._drag_target is None - or self._body_xpos is None - or self._body_xmat is None + bid is not None + and self._grab_local is not None + and self._target_viser is not None + and self._body_xpos is not None + and self._body_xmat is not None + and bid < self._body_xpos.shape[0] ): - return None - + xmat = self._body_xmat[bid] + selpos = self._body_xpos[bid] + xmat @ self._grab_local + self._scene_offset + target = self._target_viser.copy() + if selpos is None or target is None: + self._hide_connector() + return + points = np.array([[selpos, target]], dtype=np.float32) + self._connector = self._server.scene.add_line_segments( + _CONNECTOR_NAME, points=points, colors=_CONNECTOR_COLOR, line_width=3.0 + ) + + def _hide_connector(self) -> None: + if self._connector is not None: + self._connector.remove() + self._connector = None + + def _update_rotate_ghost(self) -> None: + """Draw a wireframe box at the body's target orientation while rotating.""" + center = half = quat = None + with self._lock: bid = self._drag_body_id - # Bounds-check against the *current* model: if the model was - # rebuilt mid-drag, the cached body id may now reference a - # different body or be out of range entirely. Drop the drag - # silently rather than indexing past the end. - if bid >= self._body_xpos.shape[0] or bid >= self._model.nbody: - self._drag_body_id = None - self._drag_grab_local = None - self._drag_target = None - self._drag_mode = None - self._prev_grab_world = None - self._prev_time = None - return None - xpos = self._body_xpos[bid] - xmat = self._body_xmat[bid] - - grab_world_mj = xpos + xmat @ self._drag_grab_local - # Displacement is computed in viser-world frame (cursor frame). - # Force/torque vectors are unchanged by the pure-translation - # offset, so we apply them in mujoco frame as-is; only the - # application point passed to mj_applyFT needs mujoco coords. - grab_world = grab_world_mj + self._scene_offset - diff = grab_world - self._drag_target - - # Compute localmass from body_invweight0 (translational inverse weight). - invweight_trn = float(self._model.body_invweight0[bid, 0]) - localmass = 1.0 if invweight_trn == 0 else 1.0 / max(invweight_trn, 1e-10) - - # Velocity is refreshed per render frame in ``update_state``; - # we just read the cached value here. - vel = self._grab_vel - - if self._drag_mode == "translate": - stiffness = float(self._model.vis.map.stiffness) - # Spring force: -k * m * displacement. - force = -stiffness * localmass * diff - # Critical damping: -2*sqrt(k)*m * velocity. - force -= np.sqrt(stiffness) * localmass * vel - - # Torque from moment arm (mujoco-frame coords). - moment_arm = grab_world_mj - xpos - torque = np.cross(moment_arm, force) - else: - stiffnessrot = float(self._model.vis.map.stiffnessrot) - # Rotational inertia from body_invweight0. - invweight_rot = float(self._model.body_invweight0[bid, 1]) - inertia = 1.0 if invweight_rot == 0 else 1.0 / max(invweight_rot, 1e-10) - # Torque proportional to displacement (simplified from quat error). - torque = -stiffnessrot * inertia * diff - torque -= np.sqrt(stiffnessrot) * inertia * vel - # Smaller translational force in rotate mode. - stiffness = float(self._model.vis.map.stiffness) - force = -stiffness * localmass * diff * 0.1 - - return PerturbationState( - body_id=bid, - force=force, - torque=torque, - point=grab_world_mj, - ) - - def _world_to_body_local(self, body_id: int, world_pos: np.ndarray) -> np.ndarray: - """Convert a world position to body-local coordinates.""" - if self._body_xpos is None or self._body_xmat is None: - return world_pos - xpos = self._body_xpos[body_id] - xmat = self._body_xmat[body_id] - return xmat.T @ (world_pos - xpos) + if ( + self._drag_rotate + and bid is not None + and self._rotate_delta_quat is not None + and self._rotate_world_quat is not None + and self._rotate_aabb_center is not None + and self._rotate_aabb_half is not None + and self._body_xpos is not None + and bid < self._body_xpos.shape[0] + ): + # Target world orientation = screen rotation applied to the body's + # orientation at drag start. + quat = np.empty(4) + mujoco.mju_mulQuat(quat, self._rotate_delta_quat, self._rotate_world_quat) + rot = np.empty(9) + mujoco.mju_quat2Mat(rot, quat) + center = ( + self._body_xpos[bid] + + rot.reshape(3, 3) @ self._rotate_aabb_center + + self._scene_offset + ) + half = self._rotate_aabb_half.copy() + if center is None or half is None or quat is None: + self._hide_ghost() + return + self._ghost = self._server.scene.add_box( + _GHOST_NAME, + color=_GHOST_COLOR, + dimensions=2.0 * half, + wireframe=True, + wxyz=quat, + position=center, + ) + + def _hide_ghost(self) -> None: + if self._ghost is not None: + self._ghost.remove() + self._ghost = None diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 9ac8969..c602fcd 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -151,7 +151,7 @@ def __init__( # Visualization settings. self.env_idx = 0 - self.camera_tracking_enabled = True + self.camera_tracking_enabled = False self.show_only_selected = False self.geom_groups_visible = [True, True, True, False, False, False] self.site_groups_visible = [True, True, True, False, False, False] @@ -186,6 +186,8 @@ def __init__( self.paused = False self._vis_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} self._simple_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} + self._geom_group_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} + self._site_group_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} self.perturbation = PerturbationHandler(server, mj_model, self.mj_data) # Cached visualization state for re-rendering when settings change. @@ -461,6 +463,15 @@ def _mutate() -> None: for client in self.server.get_clients().values(): client.camera.position = _camera_offset client.camera.look_at = np.zeros(3) + else: + # Untrack without teleporting: the scene offset is about to drop + # to zero (bodies snap back to world coords), so shift the camera + # by the same amount to keep the view exactly where it is now. + # viser's position setter translates look_at along with the camera, + # so setting position alone preserves the orientation. + delta = -self._scene_offset + for client in self.server.get_clients().values(): + client.camera.position = np.asarray(client.camera.position) + delta self._apply_visualization_change(_mutate) @@ -762,6 +773,7 @@ def create_groups_gui(self) -> None: initial_value=self.geom_groups_visible[i], hint=f"Show/hide geometry in group {i}", ) + self._geom_group_checkboxes[i] = cb @cb.on_update def _(event, group_idx=i) -> None: @@ -778,6 +790,7 @@ def _mutate() -> None: initial_value=self.site_groups_visible[i], hint=f"Show/hide sites in group {i}", ) + self._site_group_checkboxes[i] = cb @cb.on_update def _(event, group_idx=i) -> None: @@ -891,9 +904,37 @@ def register_keybindings(self, server: viser.ViserServer) -> None: cmd = server.gui.add_command(label, hotkey=key) # type: ignore[arg-type] cmd.on_trigger(lambda _, _idx=flag_idx: self._toggle_vis_flag(_idx)) + # Number keys toggle geom groups; shift+number toggles site groups + # (mirrors MuJoCo simulate). + for i in range(6): + digit = str(i) + geom_cmd = server.gui.add_command( + f"Toggle Geom Group {i}", + hotkey=digit, # type: ignore[arg-type] + ) + geom_cmd.on_trigger( + lambda _, _idx=i: self._toggle_group(self._geom_group_checkboxes, _idx) + ) + site_cmd = server.gui.add_command( + f"Toggle Site Group {i}", + hotkey=digit, # type: ignore[arg-type] + modifier="shift", + ) + site_cmd.on_trigger( + lambda _, _idx=i: self._toggle_group(self._site_group_checkboxes, _idx) + ) + cmd = server.gui.add_command("Free Camera", hotkey="escape") cmd.on_trigger(lambda _: self._set_free_camera()) + def _toggle_group( + self, checkboxes: dict[int, viser.GuiCheckboxHandle], group_idx: int + ) -> None: + """Flip a geom/site group checkbox; its on_update applies the change.""" + cb = checkboxes.get(group_idx) + if cb is not None: + cb.value = not cb.value + def _set_free_camera(self) -> None: """Switch to free camera mode.""" if hasattr(self, "_camera_dropdown"): diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index 5a274e9..ccb4b16 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -87,12 +87,6 @@ def __init__( self._dirty = True self._step_count = 0 - # Snapshot of ``qfrc_applied`` taken just before the perturbation - # adds its force, so we can undo our addition before the next step - # without clobbering user-supplied generalized forces. ``None`` - # means no perturbation force is currently in flight. - self._pert_saved_qfrc: np.ndarray | None = None - # Timing. self._budget = 0.0 self._last_tick = 0.0 @@ -190,6 +184,11 @@ def _tick(self) -> None: with self._lock: if self._step_physics(dt): self._dirty = True + else: + # Paused: a drag repositions free-joint / mocap bodies kinematically. + with self._lock: + if self.scene.perturbation.apply(self.model, self.data, paused=True): + self._dirty = True # Render at fixed frame rate, but only when state changed. self._time_until_next_render -= dt @@ -206,34 +205,6 @@ def _tick(self) -> None: self._stats_frames += 1 self._update_stats() - def _apply_perturbation(self) -> None: - """Apply any pending interactive perturbation force. - - ``mj_applyFT`` *adds* to ``qfrc_applied`` rather than overwriting - it, so a force applied last substep would compound this substep - if we just called it again. We undo the previous addition by - restoring a pre-perturbation snapshot, then re-snapshot and re-add - on top -- which preserves any user-supplied ``qfrc_applied`` - (a documented MuJoCo extension point) instead of zeroing it out. - """ - if self._pert_saved_qfrc is not None: - self.data.qfrc_applied[:] = self._pert_saved_qfrc - self._pert_saved_qfrc = None - - pert = self.scene.perturbation.get_perturbation() - if pert is None: - return - self._pert_saved_qfrc = self.data.qfrc_applied.copy() - mujoco.mj_applyFT( - self.model, - self.data, - pert.force, - pert.torque, - pert.point, - pert.body_id, - self.data.qfrc_applied, - ) - def _step_physics(self, dt: float) -> bool: """Run physics steps for this frame's sim-time budget. @@ -249,7 +220,7 @@ def _step_physics(self, dt: float) -> bool: deadline = time.perf_counter() + _FRAME_TIME hit_deadline = False while self._budget >= step_dt: - self._apply_perturbation() + self.scene.perturbation.apply(self.model, self.data, paused=False) self._step_fn(self.model, self.data) self._budget -= step_dt self._step_count += 1 @@ -319,11 +290,9 @@ def _do_reset(self) -> None: self._reset() self._step_count = 0 self._budget = 0.0 - # Drop any in-flight drag and the saved-qfrc snapshot: the - # snapshot was taken against the pre-reset state and would - # otherwise be restored on top of the freshly-reset data. + # Drop any in-flight drag so a stale body id can't outlive the + # reset. mj_resetData already zeroed xfrc_applied. self.scene.perturbation.clear() - self._pert_saved_qfrc = None self._render() self._update_status_display() self._sync_sliders() diff --git a/tests/test_interaction.py b/tests/test_interaction.py new file mode 100644 index 0000000..9f997fa --- /dev/null +++ b/tests/test_interaction.py @@ -0,0 +1,206 @@ +"""Tests for interactive drag perturbation. + +These drive the handler's state the way the viser drag callbacks would +(grab point + cursor target in world coords) and check that the MuJoCo +perturbation engine is wired up correctly: force direction and magnitude, +off-center torque, critical damping, force clearing, paused repositioning, +the viser/mujoco frame offset, and model-rebuild bounds safety. +""" + +import mujoco +import numpy as np +import pytest +import viser + +from mjviser.interaction import PerturbationHandler + +_FREE_BODY_XML = """ + + +""" + +_BID = 1 # the free body + + +def _stop(server: viser.ViserServer) -> None: + try: + server.stop() + except RuntimeError: + pass + + +@pytest.fixture +def env(): + model = mujoco.MjModel.from_xml_string(_FREE_BODY_XML) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + server = viser.ViserServer(port=0) + handler = PerturbationHandler(server, model, data) + yield handler, model, data + _stop(server) + + +def _start_drag( + handler: PerturbationHandler, + model: mujoco.MjModel, + data: mujoco.MjData, + bid: int, + grab_world: np.ndarray, + target_world: np.ndarray, + scene_offset: np.ndarray | None = None, +) -> None: + """Mimic a viser drag: cache pose, grab a point, then set the target.""" + if scene_offset is None: + scene_offset = np.zeros(3) + body_xpos = data.xpos[None].copy() + body_xmat = data.xmat.reshape(model.nbody, 3, 3)[None].copy() + handler.update_state(body_xpos, body_xmat, 0, scene_offset) + handler._on_drag_start(bid, tuple(grab_world)) + handler._target_viser = np.array(target_world, dtype=float) + + +def test_no_drag_applies_no_force(env): + handler, model, data = env + assert handler.apply(model, data, paused=False) is False + assert np.all(data.xfrc_applied == 0.0) + + +def test_force_points_toward_target(env): + handler, model, data = env + com = data.xipos[_BID].copy() + _start_drag(handler, model, data, _BID, com, com + [0.3, 0.0, 0.0]) + handler.apply(model, data, paused=False) + force = data.xfrc_applied[_BID, :3] + # stiffness(100) * localmass(=mass 2) * displacement(0.3) = 60, +x only. + np.testing.assert_allclose(force, [60.0, 0.0, 0.0], atol=1e-6) + np.testing.assert_allclose(data.xfrc_applied[_BID, 3:], 0.0, atol=1e-6) + + +def test_drag_moves_body_toward_target(env): + handler, model, data = env + com = data.xipos[_BID].copy() + target = com + [0.5, 0.0, 0.0] + _start_drag(handler, model, data, _BID, com, target) + for _ in range(200): + handler.apply(model, data, paused=False) + mujoco.mj_step(model, data) + assert data.xpos[_BID][0] > 0.3 # pulled most of the way to +0.5 + + +def test_offcenter_grab_produces_torque(env): + handler, model, data = env + grab = data.xipos[_BID] + [0.0, 0.0, 0.1] # top face, above COM + _start_drag(handler, model, data, _BID, grab, grab + [0.3, 0.0, 0.0]) + handler.apply(model, data, paused=False) + torque = data.xfrc_applied[_BID, 3:] + # +x force applied 0.1 above COM -> torque about +y. + assert torque[1] > 1.0 + assert abs(torque[0]) < 1e-6 and abs(torque[2]) < 1e-6 + + +def test_damping_reduces_force_when_moving_toward_target(env): + handler, model, data = env + com = data.xipos[_BID].copy() + _start_drag(handler, model, data, _BID, com, com + [0.3, 0.0, 0.0]) + handler.apply(model, data, paused=False) + rest_force = data.xfrc_applied[_BID, 0] + # Now give the body velocity toward the target; damping should cut force. + data.qvel[0] = 2.0 + mujoco.mj_forward(model, data) + handler.apply(model, data, paused=False) + assert data.xfrc_applied[_BID, 0] < rest_force + + +def test_drag_end_clears_applied_force(env): + handler, model, data = env + com = data.xipos[_BID].copy() + _start_drag(handler, model, data, _BID, com, com + [0.3, 0.0, 0.0]) + handler.apply(model, data, paused=False) + assert np.any(data.xfrc_applied[_BID] != 0.0) + # End the drag and apply again: the body's force must be zeroed. + handler._drag_body_id = None + handler.apply(model, data, paused=False) + assert np.all(data.xfrc_applied[_BID] == 0.0) + + +def test_clear_resets_state_and_force(env): + handler, model, data = env + com = data.xipos[_BID].copy() + _start_drag(handler, model, data, _BID, com, com + [0.3, 0.0, 0.0]) + handler.apply(model, data, paused=False) + handler.clear() + assert handler.selected_body_id is None + assert handler._drag_body_id is None + handler.apply(model, data, paused=False) + assert np.all(data.xfrc_applied[_BID] == 0.0) + + +def test_paused_repositions_free_body(env): + handler, model, data = env + com = data.xipos[_BID].copy() + target = com + [0.3, 0.0, 0.0] + _start_drag(handler, model, data, _BID, com, target) + assert handler.apply(model, data, paused=True) is True + np.testing.assert_allclose(data.qpos[:3], [0.3, 0.0, 0.0], atol=1e-6) + assert np.all(data.xfrc_applied == 0.0) # pose drag applies no force + + +def test_scene_offset_is_subtracted(env): + handler, model, data = env + com = data.xipos[_BID].copy() + offset = np.array([10.0, 0.0, 0.0]) + # Grab and target are reported in viser frame (= mujoco + offset). + _start_drag( + handler, model, data, _BID, com + offset, com + offset + [0.3, 0.0, 0.0], offset + ) + handler.apply(model, data, paused=False) + # Despite the large offset, the force is the same as without it. + np.testing.assert_allclose(data.xfrc_applied[_BID, :3], [60.0, 0.0, 0.0], atol=1e-6) + + +def _start_rotate(handler, bid, delta_quat): + """Mimic a rotate drag: latch the body and set the target rotation.""" + handler._drag_body_id = bid + handler._drag_rotate = True + handler._rotate_initial_quat = None + handler._rotate_delta_quat = np.asarray(delta_quat, dtype=float) + + +def test_rotate_drag_applies_torque(env): + handler, model, data = env + q = np.empty(4) + mujoco.mju_axisAngle2Quat(q, np.array([0.0, 0.0, 1.0]), np.pi / 2) + _start_rotate(handler, _BID, q) + handler.apply(model, data, paused=False) + torque = data.xfrc_applied[_BID, 3:] + # Spring torque points along +z (toward the target yaw), no net force. + assert torque[2] > 1.0 + np.testing.assert_allclose(data.xfrc_applied[_BID, :3], 0.0, atol=1e-6) + + +def test_rotate_drag_paused_reorients_free_body(env): + handler, model, data = env + q = np.empty(4) + mujoco.mju_axisAngle2Quat(q, np.array([0.0, 0.0, 1.0]), np.pi / 2) + _start_rotate(handler, _BID, q) + assert handler.apply(model, data, paused=True) is True + np.testing.assert_allclose( + data.qpos[3:7], [np.cos(np.pi / 4), 0.0, 0.0, np.sin(np.pi / 4)], atol=1e-6 + ) + + +def test_stale_body_id_after_rebuild_is_safe(env): + handler, model, data = env + com = data.xipos[_BID].copy() + _start_drag(handler, model, data, _BID, com, com + [0.3, 0.0, 0.0]) + # Simulate a model rebuild that dropped bodies: id now out of range. + handler._drag_body_id = model.nbody + 5 + assert handler.apply(model, data, paused=False) is False + assert np.all(data.xfrc_applied == 0.0) From d73aed2d42a7b77fe1bbb3d4ea43b6a0123a4833 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 10 Jun 2026 14:28:24 -0700 Subject: [PATCH 4/6] Bump idna and urllib3 to clear Dependabot alerts --- uv.lock | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/uv.lock b/uv.lock index 0ccb1eb..54cd677 100644 --- a/uv.lock +++ b/uv.lock @@ -237,11 +237,11 @@ wheels = [ [[package]] name = "idna" -version = "3.11" +version = "3.18" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/0703ccc57f3a7233505399edb88de3cbd678da106337b9fcde432b65ed60/idna-3.11.tar.gz", hash = "sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902", size = 194582, upload-time = "2025-10-12T14:55:20.501Z" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, + { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, ] [[package]] @@ -905,11 +905,11 @@ wheels = [ [[package]] name = "urllib3" -version = "2.6.3" +version = "2.7.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] [[package]] From b29df69b28aa913df27650e8286cccd0bc737f17 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 10 Jun 2026 14:51:59 -0700 Subject: [PATCH 5/6] Pivot rotate ghost box about the body's center of mass --- src/mjviser/interaction.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/mjviser/interaction.py b/src/mjviser/interaction.py index 8145937..fa1d817 100644 --- a/src/mjviser/interaction.py +++ b/src/mjviser/interaction.py @@ -439,6 +439,7 @@ def _update_rotate_ghost(self) -> None: and self._rotate_aabb_center is not None and self._rotate_aabb_half is not None and self._body_xpos is not None + and self._body_xmat is not None and bid < self._body_xpos.shape[0] ): # Target world orientation = screen rotation applied to the body's @@ -447,11 +448,13 @@ def _update_rotate_ghost(self) -> None: mujoco.mju_mulQuat(quat, self._rotate_delta_quat, self._rotate_world_quat) rot = np.empty(9) mujoco.mju_quat2Mat(rot, quat) - center = ( - self._body_xpos[bid] - + rot.reshape(3, 3) @ self._rotate_aabb_center - + self._scene_offset - ) + rot = rot.reshape(3, 3) + # Pivot the ghost about the body's center of mass -- the actual + # rotation pivot -- not the frame origin, so it stays anchored to the + # body instead of orbiting away as the target orientation sweeps. + ipos = np.asarray(self._model.body_ipos[bid], dtype=float) + com = self._body_xpos[bid] + self._body_xmat[bid] @ ipos + center = com + rot @ (self._rotate_aabb_center - ipos) + self._scene_offset half = self._rotate_aabb_half.copy() if center is None or half is None or quat is None: self._hide_ghost() From e2f41ddb919132426fa5ec3f6c99fdd094205bd2 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 10 Jun 2026 15:55:05 -0700 Subject: [PATCH 6/6] Clean up perturbation review: dead state, dup resets, broken site hotkeys Remove the handler's unused MjData reference (it always operates on the data passed in per call, and the stored copy was the scene's, not the viewer's stepped data). Factor the duplicated drag-state reset out of clear() and _end_drag(), and stop computing the grab-point selpos on every running-mode apply() since only the paused branch uses it. In the scene, initialize the camera dropdown attribute in __init__ instead of guarding with hasattr, and drop the redundant server argument from register_keybindings (the scene already holds it). Remove the shift+digit site-group hotkeys: viser matches hotkeys on event.key, where shift+1 is "!" not "1", so they never fired. The geom-group digit keys stay; sites remain toggleable via the Groups GUI checkboxes. The now-unused site-group checkbox dict goes with them. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mjviser/interaction.py | 51 ++++++++++++++++++-------------------- src/mjviser/scene.py | 30 +++++++++------------- src/mjviser/viewer.py | 2 +- tests/test_interaction.py | 2 +- 4 files changed, 38 insertions(+), 47 deletions(-) diff --git a/src/mjviser/interaction.py b/src/mjviser/interaction.py index fa1d817..10ca260 100644 --- a/src/mjviser/interaction.py +++ b/src/mjviser/interaction.py @@ -38,11 +38,12 @@ def __init__( self, server: viser.ViserServer, mj_model: mujoco.MjModel, - mj_data: mujoco.MjData, ) -> None: + # The live MjData flows in per call (apply/update_state); the handler only + # keeps the static model, since the data it acts on is the viewer's, not + # the scene's. self._server = server self._model = mj_model - self._data = mj_data self._lock = Lock() self.selected_body_id: int | None = None @@ -99,22 +100,26 @@ def clear(self) -> None: """ with self._lock: self.selected_body_id = None - self._drag_body_id = None - self._grab_local = None - self._target_viser = None - self._drag_rotate = False - self._rotate_delta_quat = None - self._rotate_initial_quat = None - self._rotate_anchor = None - self._rotate_aabb_center = None - self._rotate_aabb_half = None - self._rotate_world_quat = None - self._pert.active = 0 + self._reset_drag_state() self._hide_connector() self._hide_ghost() if self._info_text is not None: self._info_text.value = "(none)" + def _reset_drag_state(self) -> None: + """Clear all in-flight drag state. The caller holds ``self._lock``.""" + self._drag_body_id = None + self._grab_local = None + self._target_viser = None + self._drag_rotate = False + self._rotate_delta_quat = None + self._rotate_initial_quat = None + self._rotate_anchor = None + self._rotate_aabb_center = None + self._rotate_aabb_half = None + self._rotate_world_quat = None + self._pert.active = 0 + def register_drag_handlers( self, handle: viser.BatchedGlbHandle, @@ -129,6 +134,9 @@ def register_drag_handlers( n_bodies = len(body_ids) def _body_id(idx: int | None) -> int | None: + # The batched handle lays instances out as (num_envs, n_bodies) + # flattened, so any env's instance maps back to its body via modulo. + # The drag always acts on the active env (the pose we cache). if idx is None: return None return int(body_ids[idx % n_bodies]) @@ -175,17 +183,7 @@ def _(event: viser.SceneNodeDragEvent) -> None: # type: ignore[type-arg] def _end_drag(self) -> None: with self._lock: - self._drag_body_id = None - self._grab_local = None - self._target_viser = None - self._drag_rotate = False - self._rotate_delta_quat = None - self._rotate_initial_quat = None - self._rotate_anchor = None - self._rotate_aabb_center = None - self._rotate_aabb_half = None - self._rotate_world_quat = None - self._pert.active = 0 + self._reset_drag_state() self._hide_connector() self._hide_ghost() @@ -322,10 +320,7 @@ def apply(self, model: mujoco.MjModel, data: mujoco.MjData, paused: bool) -> boo self._clear_applied(data) return False - xmat = data.xmat[bid].reshape(3, 3) - selpos = data.xpos[bid] + xmat @ self._grab_local target = self._target_viser - self._scene_offset - self._pert.localpos = self._grab_local if paused: @@ -333,6 +328,8 @@ def apply(self, model: mujoco.MjModel, data: mujoco.MjData, paused: bool) -> boo # target, keeping orientation. mjv_applyPerturbPose moves free # joints / mocap bodies and is a no-op for anything else. self._clear_applied(data) + xmat = data.xmat[bid].reshape(3, 3) + selpos = data.xpos[bid] + xmat @ self._grab_local xiquat = np.empty(4) mujoco.mju_mulQuat(xiquat, data.xquat[bid], model.body_iquat[bid]) self._pert.active = int(mujoco.mjtPertBit.mjPERT_TRANSLATE) diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index c602fcd..13b6ce4 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -187,8 +187,9 @@ def __init__( self._vis_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} self._simple_flag_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} self._geom_group_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} - self._site_group_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} - self.perturbation = PerturbationHandler(server, mj_model, self.mj_data) + self._camera_dropdown: viser.GuiDropdownHandle[str] | None = None + self._camera_offset = np.zeros(3) + self.perturbation = PerturbationHandler(server, mj_model) # Cached visualization state for re-rendering when settings change. self._tracked_body_id: int | None = None @@ -790,7 +791,6 @@ def _mutate() -> None: initial_value=self.site_groups_visible[i], hint=f"Show/hide sites in group {i}", ) - self._site_group_checkboxes[i] = cb @cb.on_update def _(event, group_idx=i) -> None: @@ -852,7 +852,7 @@ def create_visualization_gui( def _apply_camera_selection(self) -> None: """Switch camera based on dropdown value.""" - if not hasattr(self, "_camera_dropdown"): + if self._camera_dropdown is None: return name = self._camera_dropdown.value if name == "Free": @@ -888,8 +888,9 @@ def _mutate() -> None: else: self._apply_visualization_change(_mutate) - def register_keybindings(self, server: viser.ViserServer) -> None: + def register_keybindings(self) -> None: """Register single-letter keyboard shortcuts for visualization toggles.""" + server = self.server _flag_keys: list[tuple[str, str, int]] = [ ("Toggle Contacts", "C", int(mujoco.mjtVisFlag.mjVIS_CONTACTPOINT)), ("Toggle Forces", "F", int(mujoco.mjtVisFlag.mjVIS_CONTACTFORCE)), @@ -904,25 +905,18 @@ def register_keybindings(self, server: viser.ViserServer) -> None: cmd = server.gui.add_command(label, hotkey=key) # type: ignore[arg-type] cmd.on_trigger(lambda _, _idx=flag_idx: self._toggle_vis_flag(_idx)) - # Number keys toggle geom groups; shift+number toggles site groups - # (mirrors MuJoCo simulate). + # Number keys toggle geom groups. (MuJoCo simulate also maps shift+number + # to site groups, but viser routes hotkeys through event.key, where + # shift+digit is a symbol, not a digit -- so that binding can't be + # expressed here. Use the Groups GUI checkboxes for sites.) for i in range(6): - digit = str(i) geom_cmd = server.gui.add_command( f"Toggle Geom Group {i}", - hotkey=digit, # type: ignore[arg-type] + hotkey=str(i), # type: ignore[arg-type] ) geom_cmd.on_trigger( lambda _, _idx=i: self._toggle_group(self._geom_group_checkboxes, _idx) ) - site_cmd = server.gui.add_command( - f"Toggle Site Group {i}", - hotkey=digit, # type: ignore[arg-type] - modifier="shift", - ) - site_cmd.on_trigger( - lambda _, _idx=i: self._toggle_group(self._site_group_checkboxes, _idx) - ) cmd = server.gui.add_command("Free Camera", hotkey="escape") cmd.on_trigger(lambda _: self._set_free_camera()) @@ -937,7 +931,7 @@ def _toggle_group( def _set_free_camera(self) -> None: """Switch to free camera mode.""" - if hasattr(self, "_camera_dropdown"): + if self._camera_dropdown is not None: self._camera_dropdown.value = "Free" else: self.camera_tracking_enabled = True diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py index ccb4b16..03329f1 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -430,7 +430,7 @@ def _register_keybindings(self) -> None: cmd = s.gui.add_command("Speed Up") cmd.on_trigger(lambda _: self._speed_up()) - self.scene.register_keybindings(s) + self.scene.register_keybindings() _MAX_SLIDERS: int = 200 diff --git a/tests/test_interaction.py b/tests/test_interaction.py index 9f997fa..2344a38 100644 --- a/tests/test_interaction.py +++ b/tests/test_interaction.py @@ -42,7 +42,7 @@ def env(): data = mujoco.MjData(model) mujoco.mj_forward(model, data) server = viser.ViserServer(port=0) - handler = PerturbationHandler(server, model, data) + handler = PerturbationHandler(server, model) yield handler, model, data _stop(server)