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..10ca260 --- /dev/null +++ b/src/mjviser/interaction.py @@ -0,0 +1,471 @@ +"""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 + +from threading import Lock + +import mujoco +import numpy as np +import viser + +_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-perturb on the viewer's MjData.""" + + def __init__( + self, + server: viser.ViserServer, + mj_model: mujoco.MjModel, + ) -> 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._lock = Lock() + + self.selected_body_id: int | 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 + + # 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 + + 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 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: 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._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, + body_ids: np.ndarray, + ) -> None: + """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(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]) + + @handle.on_click + def _(event: viser.SceneNodePointerEvent) -> None: # type: ignore[type-arg] + bid = _body_id(event.instance_index) + if bid is None or bid == 0: + return + with self._lock: + self.selected_body_id = bid + 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._reset_drag_state() + 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, + body_xpos: np.ndarray, + body_xmat: np.ndarray, + env_idx: int, + scene_offset: np.ndarray, + ) -> None: + """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 (viser - mujoco). + """ + 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() + + 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 + + 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) + 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) + 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 ( + 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] + ): + 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 + 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 self._body_xmat 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) + 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() + 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 46b76bf..13b6ce4 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 @@ -150,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] @@ -183,6 +184,12 @@ 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._geom_group_checkboxes: dict[int, viser.GuiCheckboxHandle] = {} + 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 @@ -203,6 +210,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 +277,11 @@ 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.""" + 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: @@ -277,6 +290,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() @@ -300,6 +319,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() @@ -444,6 +464,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) @@ -461,6 +490,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 +528,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 +745,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: @@ -727,6 +774,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: @@ -802,6 +850,92 @@ def create_visualization_gui( self.create_groups_gui() return tabs + def _apply_camera_selection(self) -> None: + """Switch camera based on dropdown value.""" + if self._camera_dropdown is None: + 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) -> 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)), + ("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)) + + # 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): + geom_cmd = server.gui.add_command( + f"Toggle Geom Group {i}", + hotkey=str(i), # type: ignore[arg-type] + ) + geom_cmd.on_trigger( + lambda _, _idx=i: self._toggle_group(self._geom_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 self._camera_dropdown is not None: + self._camera_dropdown.value = "Free" + else: + self.camera_tracking_enabled = True + def update_from_arrays( self, body_xpos: np.ndarray, @@ -943,6 +1077,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..03329f1 100644 --- a/src/mjviser/viewer.py +++ b/src/mjviser/viewer.py @@ -184,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 @@ -215,6 +220,7 @@ def _step_physics(self, dt: float) -> bool: deadline = time.perf_counter() + _FRAME_TIME hit_deadline = False while self._budget >= step_dt: + 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 @@ -258,53 +264,77 @@ 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 + # 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._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 +343,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 +385,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 +409,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() + _MAX_SLIDERS: int = 200 def _setup_joint_sliders(self) -> None: diff --git a/tests/test_interaction.py b/tests/test_interaction.py new file mode 100644 index 0000000..2344a38 --- /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) + 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) 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]]