106 lines
4.8 KiB
Markdown
106 lines
4.8 KiB
Markdown
# 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. **Smoothing** — `gaussian_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-capture** — `tests/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, no turns >90° on
|
||
smoothed output.
|
||
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.
|