4.8 KiB
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):
- 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.
- 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_monotonichard-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_endpointsdrags 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_startmultiplieslateral[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):
- 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.
- No
- 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).
- Smoothstep-ramped lateral damping over the first n points; no
abrupt release, no local
- Smoothing —
gaussian_smoothapplied to both forward and lateral (currently lateral only). - 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
- Golden regression re-capture —
tests/unit/data/golden_mouse.npzis re-captured with the new pipeline (expected, intentional behavior change; scroll golden untouched). CHANGELOG entry. - Unit tests (
tests/unit/test_postprocess.py) — backtrack tolerance, overshoot compression, exact endpoint hit after global correction, correction weights 0/1 at the ends, no turns >90° on smoothed output. - 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.