Files
ai_mouse/docs/superpowers/specs/2026-07-09-mouse-postprocess-quality-design.md

5.2 KiB
Raw Blame History

Mouse trajectory quality: post-processing rework + Heun sampling

Date: 2026-07-09 Status: approved Scope: inference-side only (src/ai_mouse/). No retraining, no ONNX re-export.

Problem

Generated mouse trajectories look unnatural in two ways (confirmed by diagnostic plots, 4 cases × 6 seeds, bundled weights):

  1. Endpoint artifacts — trajectories hit the target at near-right angles ("vertical wall" of points stacked at the target x), or hook back after overshooting. Start segments show abrupt kinks.
  2. Exaggerated curvature — large dome arcs on straight moves, loops on short moves. Up to 8 direction changes >45° per trace (max 135°).

Diagnosis

Symptom 1 is manufactured by post-processing in _postprocess.py:

  • enforce_forward_monotonic hard-clips forward to [0, 1]. Natural overshoot past the target becomes a stack of points at forward=1 with varying lateral → the vertical wall.
  • snap_endpoints drags the last 6 points toward (1, 0) with quadratic easing. When the raw sample ends off-target, the drag direction fights the trajectory's own direction → hooks.
  • smooth_start multiplies lateral[1] by 1/5 and releases abruptly after point n → start kinks.

Symptom 2 is mostly learned from data (Balabit fixed-window click-anchored segmentation includes mid-gesture starts and composite move+hover gestures) and is out of scope here — deferred to a possible follow-up (gesture re-segmentation + retrain). Coarse 10-step Euler sampling contributes secondary jitter and IS in scope.

Design

1. Post-processing pipeline rework (_postprocess.py, mouse.py)

Current order: snap_endpoints → smooth_start → enforce_forward_monotonic → gaussian_smooth(lateral).

New order (steps run in this sequence):

  1. Soft monotonic (replaces enforce_forward_monotonic):
    • No clip(0, 1).
    • Tolerate small backtracking: enforce forward[i] >= forward[i-1] - 0.02.
    • Allow overshoot past 1.0; soft-compress extremes beyond ~1.08 with tanh so the path never flies far past the target.
  2. Continuous start damping (replaces smooth_start):
    • Smoothstep-ramped lateral damping over the first n points; no abrupt release, no local max() monotonic fix (step 1 owns that).
  3. Smoothinggaussian_smooth applied to both forward and lateral (currently lateral only).
  4. Global residual correction (replaces snap_endpoints, runs last so endpoints stay exact after smoothing):
    • Compute residuals of first/last points vs (0,0)/(1,0).
    • Distribute the correction over the whole curve with smoothstep weights (weight → 1 at the corrected end, → 0 at the opposite end).
    • Endpoints land exactly; approach direction stays natural.

Function signatures, the generate() API, and the exact-endpoint guarantee are preserved.

2. Sampling: Euler → Heun (mouse.py) — REJECTED during implementation

Replace the 10-step first-order Euler loop with 10-step Heun (predictor-corrector): per step, evaluate v at x and at the Euler prediction, advance with the average. NFE 10 → 20; each call is a d_model=128 transformer (~1-2 ms CPU), total latency stays ~40 ms. Seed reproducibility unaffected (randomness is only in the init noise and duration sampling, both unchanged).

Outcome (2026-07-09, implementation): Heun was implemented, measured, and reverted. Per-stage probing showed Heun's raw ODE output contains 40-51 direction changes >90° per trace vs Euler's 2-11; a t-clamped variant was equally bad and Euler-20 gave no meaningful gain. The trained flow field is only self-consistent along its own Euler-discretized paths, so second-order integration injects noise instead of reducing error. The shipped code keeps the original 10-step Euler loop; the new post-processing pipeline alone meets the quality gates (max tail turn 32-58° vs the old pipeline's 53-135°, zero jagged-chain artifacts).

3. Tests and acceptance

  1. Golden regression re-capturetests/unit/data/golden_mouse.npz is re-captured with the new pipeline (expected, intentional behavior change; scroll golden untouched). CHANGELOG entry.
  2. Unit tests (tests/unit/test_postprocess.py) — backtrack tolerance, overshoot compression, exact endpoint hit after global correction, correction weights 0/1 at the ends. The tail-quality guard allows at most ONE >90° reversal (the natural overshoot-and-correct gesture that overshoot support implies); two or more indicate the hook/zigzag artifact class. (Amended 2026-07-09: the original "no turns >90°" wording predated overshoot support and was empirically over-strict — a single 92-135° reversal appears in ~20% of traces and is correct behavior.)
  3. Acceptance — re-run the diagnostic script (same 4 cases × 6 seeds) and compare: turns>45° count drops sharply, no vertical wall in the last 10 points. Final gate: user visually approves the Web UI verify page (restart server; post-processing is Python-side, no ONNX re-export needed).

Out of scope

  • Balabit re-segmentation (velocity-threshold gesture splitting) and retraining — revisit after this lands if curvature is still unsatisfactory.
  • Scroll subsystem — no reported issues.