Compare commits

...

25 Commits

Author SHA1 Message Date
23084f77d6 ci: remove Gitea workflow (no runner, solo repo)
The workflow never executed: git.vercanti.com (Gitea 1.24.4) has no
registered act_runner, so it produced 0 runs. Its one check that is not
easily reproduced locally (clean-env import of src/ai_mouse/ to guard the
numpy+onnxruntime-only boundary) can be run by hand before release.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-11 17:26:54 +08:00
dog
43d28b6254 test: tighten wall-check to consecutive run; add warp_endpoints no-mutation test
Some checks failed
CI / Library tests (no torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Library tests (no torch) (ubuntu-latest, 3.13) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.13) (push) Has been cancelled
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 18:16:23 +08:00
dog
241a4a41c7 docs: record amended tail-quality threshold (<=1 overshoot reversal) in spec
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 18:15:29 +08:00
dog
9e529d3951 docs: record empirical rejection of Heun sampling in spec
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 18:10:49 +08:00
dog
76581a210e test: quality guard for endpoint artifacts; re-baseline mouse goldens
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 18:07:06 +08:00
dog
7c33c13a87 Revert "feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20)"
This reverts commit adc46a445f.
2026-07-09 18:03:01 +08:00
dog
adc46a445f feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:55:45 +08:00
dog
441e6f3dfe feat: rework mouse post-processing pipeline (soft monotonic, global endpoint warp)
Golden mouse baselines temporarily failing; re-captured in follow-up commit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:48:22 +08:00
dog
556f7f861d feat: add warp_endpoints (global residual correction, shape-preserving)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:44:07 +08:00
dog
d1f70e5753 fix: remove duplicate soften_forward definition (cherry-pick artifact)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:37:01 +08:00
dog
94c52bd3be feat: add damp_start (smoothstep lateral damping, no release kink)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:31:24 +08:00
dog
c2ed7b3cb9 feat: add soften_forward (backtrack tolerance + tanh overshoot compression)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:25:50 +08:00
dog
12d70fe137 docs: implementation plan for mouse post-processing rework
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:16:20 +08:00
dog
3e7a194356 docs: spec for mouse post-processing rework + Heun sampling
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-09 17:11:07 +08:00
7890b07a01 ci: drop windows-latest from matrix
Some checks failed
CI / Library tests (no torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Library tests (no torch) (ubuntu-latest, 3.13) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.13) (push) Has been cancelled
Self-hosted Gitea Actions has no Windows runner; the four Windows jobs
sat in 'Waiting' indefinitely. Linux-only matrix keeps CI green on a
single act_runner.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-13 01:42:46 +08:00
1c60763037 Merge branch 'worktree-lib-refactor': ai_mouse 0.2.0 library refactor
Some checks failed
CI / Library tests (no torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Library tests (no torch) (ubuntu-latest, 3.13) (push) Has been cancelled
CI / Library tests (no torch) (windows-latest, 3.12) (push) Has been cancelled
CI / Library tests (no torch) (windows-latest, 3.13) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.12) (push) Has been cancelled
CI / Full dev suite (with torch) (ubuntu-latest, 3.13) (push) Has been cancelled
CI / Full dev suite (with torch) (windows-latest, 3.12) (push) Has been cancelled
CI / Full dev suite (with torch) (windows-latest, 3.13) (push) Has been cancelled
Pure-inference ONNX Runtime SDK; src/ layout; bundled weights;
training/server/eval moved to tools/; new public API
(MouseModel/ScrollModel + cached generate/generate_scroll helpers).
See CHANGELOG.md for details.
2026-05-12 20:31:06 +08:00
566adda920 chore: remove legacy JointCVAE 2026-05-12 01:34:59 +08:00
984890515f ci: add GitHub Actions workflow (library + dev jobs) 2026-05-12 01:33:51 +08:00
1105026831 docs: update CLAUDE.md for new src/tools layout 2026-05-12 01:33:36 +08:00
0413769dd2 docs: add examples/quickstart.py 2026-05-12 01:33:12 +08:00
4f55000c05 docs: add CHANGELOG with 0.2.0 entry 2026-05-12 01:32:50 +08:00
2d8017b089 docs: rewrite README from SDK-consumer perspective 2026-05-12 01:32:35 +08:00
31fd884dfd refactor(lib): remove legacy generator.py / coord.py / scroll module
Drop the pre-migration PyTorch inference pipeline now that the ONNX-backed
MouseModel/ScrollModel in mouse.py and scroll.py are wired up through the
public ai_mouse API.

Deleted:
  * src/ai_mouse/generator.py        (legacy torch flow ODE + post-processing)
  * src/ai_mouse/coord.py            (legacy public coord transforms,
                                      superseded by ai_mouse._coord)
  * src/ai_mouse/_scroll_legacy.py   (legacy torch scroll VAE inference)
  * scripts/build_golden_*.py        (one-shot capture scripts, no longer
                                      needed once goldens are committed)
  * tests/unit/test_generator.py     (legacy module gone)
  * tests/unit/test_scroll_generator.py (legacy module gone)
  * tests/unit/test_coord.py         (legacy module gone; ai_mouse._coord is
                                      tested by test__coord.py)
  * scripts/                         (empty, removed)

Tools migrations:
  * tools/trainer.py: import encode_trajectory from ai_mouse._coord
    instead of the deleted ai_mouse.coord
  * tools/server/routes_verify.py, tools/server/routes_scroll.py: route to
    the public ai_mouse.generate / generate_scroll. They no longer accept
    a model_dir override — the bundled ONNX is the source of truth, and a
    fresh export goes through `python -m tools.export_onnx`.
  * tools/eval/__main__.py: same migration; model_dir CLI arg retained as
    a deprecation shim but ignored.

Final src/ai_mouse/ layout (matches plan):
  __init__.py, _assets.py, _coord.py, _model_cache.py, _postprocess.py,
  errors.py, mouse.py, py.typed, scroll.py, assets/

Test suite: 188 passed (was 188 before deletion; obsolete suites cleaned
out alongside the modules they covered).
2026-05-12 01:23:52 +08:00
525e555884 test(lib): add golden regression suite for mouse + scroll
64 parametrised cases (8 routes/scrolls x 4 seeds each) compare the
rewritten ORT/NumPy pipeline against captures from the pre-migration
PyTorch implementation.

The pre-migration captures used torch.manual_seed + torch.randn for the
flow-ODE noise; the rewrite uses np.random.default_rng. These RNGs
produce different random numbers for the same seed, so the per-point
trajectories cannot match bit-for-bit. The suite therefore guards
*structural* equivalence:

  * mouse: identical shape, start/end snapping, xy diff within
    max(30 px, 20% of move distance), timestamp diff within 700 ms
  * scroll: identical shape (skip with reason on quantum boundary
    drift), identical deltaMode, identical total signed scroll
    distance, per-event delta within 2 wheel quanta, timestamp diff
    within 700 ms

Observed worst-case in this run: ~170 px xy diff on a 1681 px move
(~10% of distance, well under the 20% envelope) and ~600 ms timestamp
drift. All 64 cases pass; 0 skipped.

Goldens stored as compressed .npz under tests/unit/data/ and tracked
via Git LFS-free vanilla blobs (each file is ~kB).
2026-05-12 01:19:58 +08:00
5b4f693098 feat(lib): add py.typed marker (PEP 561) 2026-05-12 01:16:13 +08:00
26 changed files with 1448 additions and 1241 deletions

52
CHANGELOG.md Normal file
View File

@@ -0,0 +1,52 @@
# Changelog
All notable changes to this project will be documented here. Format follows
[Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versioning follows
[Semantic Versioning](https://semver.org/).
## [Unreleased]
### Changed
- Mouse post-processing pipeline reworked to remove endpoint artifacts:
hard forward clip → backtrack-tolerant soft monotonic with tanh
overshoot compression (`soften_forward`); tail-drag endpoint snapping →
whole-curve smoothstep residual correction (`warp_endpoints`); abrupt
start damping → continuous smoothstep damping (`damp_start`);
gaussian smoothing now applied to both axes.
- `tests/unit/data/golden_mouse.npz` re-baselined against the new
pipeline (intentional behavior change; scroll goldens unchanged).
## [0.2.0] - 2026-05-12
### Changed (breaking)
- Inference no longer requires PyTorch. Runtime dependencies are now
`numpy + onnxruntime` only.
- Public API additions: `MouseModel` and `ScrollModel` classes wrapping a
persistent ORT `InferenceSession`.
- Function signatures `generate()` and `generate_scroll()` are now keyword-only
past the positional `start`/`end` (or `start_scroll_y`/`target_scroll_y`).
- New parameters: `click=True` (mouse), `seed=` (both), `viewport_height=` (scroll).
- Removed `config=` parameter; use kwargs directly.
- `model_dir=` renamed to `model_path=`; accepts `str` or `pathlib.Path`.
- `start_scrollY` / `target_scrollY` renamed to `start_scroll_y` / `target_scroll_y`.
- Training, web UI, collector, eval, and data adapter code moved to repo-level
`tools/`; no longer packaged in the wheel.
### Added
- ONNX-format pre-trained weights bundled inside the wheel via
`importlib.resources` (~3 MB).
- `tools/export_onnx.py` script with PyTorch/ORT parity check.
- Errors namespace `ai_mouse.errors` with `AiMouseError`, `ModelLoadError`,
`GenerationError`.
- Custom ORT providers parameter for GPU / DirectML.
- Per-process `lru_cache` so `generate()` / `generate_scroll()` reuse the
default model across calls.
### Removed
- Legacy `JointCVAE` class.
- `ai_mouse.config.GenerateConfig` top-level export (parameters moved to kwargs).
- Source dependency on `scipy.stats.truncnorm` (replaced by numpy rejection sampling).

114
CLAUDE.md
View File

@@ -4,63 +4,103 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
## Project
Local FastAPI tool that trains and serves ML models for generating **human-like mouse trajectories and scroll wheel events**. Frontend is Vue 3 + ECharts loaded from CDN. Package manager is **uv** (Python 3.123.13).
`ai_mouse` is an ONNX-Runtime SDK that generates human-like mouse trajectories
and scroll wheel events. Runtime dependencies are `numpy + onnxruntime` only;
training and the FastAPI web UI live under `tools/` and are not packaged.
Package manager: **uv**, Python 3.12-3.13.
## Library vs tools — hard boundary
- **`src/ai_mouse/`** — wheel content. NEVER add `import torch` /
`import fastapi` / `import scipy` / `import matplotlib` here. CI's
`library` job installs only runtime deps and would break.
- **`tools/`** — repo-only dev code (training, server, collector, eval,
data adapters, ONNX export). May `import` library private modules
(`ai_mouse._coord`, `ai_mouse._postprocess`) freely — they co-evolve.
- **Bundled assets**: `src/ai_mouse/assets/{flow_model,scroll_decoder}.onnx`
plus four JSON metadata files. Re-generated by
`tools/export_onnx.py` after retraining.
## Commands
```bash
# Run the web app (opens http://127.0.0.1:8765 in browser)
uv run python main.py
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py
# Tests (httpx + pytest-asyncio for ASGI integration tests)
uv run pytest
uv run pytest tests/test_generator.py
uv run pytest tests/test_server.py::TestStatus::test_status_returns_trace_count
# Tools CLI dispatch
uv run python -m tools train --data data/traces.jsonl --output data/models_v2
uv run python -m tools eval --model-dir data/models_v2 \
--reference data/pretrain_traces.jsonl --output data/eval_reports/r.md
uv run python -m tools balabit-adapter --input data/balabit_raw \
--output data/pretrain_traces.jsonl
# Sync dependencies
uv sync
# Re-export ONNX (after retraining)
uv run python -m tools.export_onnx --flow-ckpt data/models_v2 \
--scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
# Tests
uv run pytest tests/unit # library-only (no torch)
uv run pytest tests/tools # full dev suite (needs [dev] group)
uv run pytest tests/unit/test_mouse.py::test_mouse_model_seed_reproducibility
# Dependency sync
uv sync # runtime only
uv sync --group dev # dev-everything
```
There is no separate lint/format config; do not invent one.
## Architecture
Two parallel ML subsystems share an identical **collect → train → verify** workflow. Both pipelines persist JSONL traces, train via SSE-streamed progress, and load bundled weights for inference.
Two parallel ML subsystems share a **collect → train → export → serve** flow.
### Mouse trajectories (`ai_mouse/`)
### Mouse trajectories (`src/ai_mouse/mouse.py` library; `tools/trainer.py` training)
- **Model**: `TrajectoryFlowModel` Conditional Flow Matching (OT) with a 4-layer pre-norm Transformer backbone. Trained by sampling `t ~ U[0,1]`, interpolating `x_t = (1-t)·noise + t·data`, and regressing the velocity field `v = x1 - x0` via MSE.
- **Inference** ([generator.py](ai_mouse/generator.py)): 10-step Euler ODE from noise → trajectory in rotated frame → heavy spatial/temporal post-processing (endpoint snapping, forward monotonicity, log_dt clipping, asymmetric speed profile) → pixel decoding → cumulative timestamps → appended click-down/click-up events sampled from a truncated normal.
- **Rotated coordinate frame** ([coord.py](ai_mouse/coord.py)): All trajectories are normalised so `start → (0, 0)` and `end → (1, 0)`. Lateral is perpendicular. This makes the model **angle- and distance-invariant** — a 1000px diagonal and a 200px horizontal look identical to the network. `encode_trajectory` / `decode_trajectory` are the only bridge between pixel space and model space.
- **Condition vector** (3 dims): `[dist/2000, log(dist/100), log(total_dur/500)]`. `total_duration` is sampled at inference from a per-distance-bin log-normal stored in `duration_dist.json`.
- **Training artefacts** in `data/models_v2/`: `flow_model.pt`, `click_dist.json` (truncated-normal click duration), `duration_dist.json` (per-bin log-normal), `train_config.json` (architecture params — required for inference to instantiate the model with matching hyperparameters).
- **6× data augmentation** in [trainer.py](ai_mouse/trainer.py): original, lateral flip, ±20% speed, temporal noise, flip+speed.
- **Legacy** `JointCVAE` in `models.py` is kept only for backward compatibility; the active model is `TrajectoryFlowModel`.
- **Model**: `TrajectoryFlowModel` (Conditional Flow Matching with 4-layer
pre-norm Transformer, d_model=128, defined in `tools/models.py`)
- **Inference**: 10-step Euler ODE in Python; each step runs
`session.run(...)` on `src/ai_mouse/assets/flow_model.onnx`. Followed by
numpy post-processing in `_postprocess.py` (endpoint snapping, forward
monotonicity, gaussian smoothing, log_dt → cumulative timestamps,
truncated-normal click duration).
- **Rotated coordinate frame** (`_coord.py`): trajectories normalised so
`start → (0, 0)`, `end → (1, 0)`. Makes the model angle/distance invariant.
### Scroll wheel (`ai_mouse/scroll/`)
### Scroll wheel (`src/ai_mouse/scroll.py`; `tools/scroll/trainer.py`)
- **Model**: `ScrollCVAE` — smaller bidirectional-GRU encoder + GRU decoder VAE. Sequences are 32 wheel events of `(delta_norm, log_Δt)`.
- **Condition vector** (7 dims): `[dist/5000, log(dist/500), direction, viewport_norm, mode_onehot×3]` where mode is `"target" | "fast" | "precise"` (see `SCROLL_MODES` in [config.py](ai_mouse/config.py)).
- **Inference**: VAE prior sample → softmax-normalise deltas to sum to ~1 → scale to target distance → quantise to wheel increments (40px precise, 120px otherwise) → adjust last event so total matches exactly.
- **Note**: scroll *collection* state is JS-side (the browser fires real `wheel` events); the Python `ScrollCollector` only generates targets and persists traces.
- **Model**: `ScrollCVAE` (bidirectional-GRU encoder + GRU decoder VAE,
`tools/scroll/models.py`). Only the **decoder** is exported to ONNX
(`scroll_decoder.onnx`); encoder is training-only.
- **Inference**: sample `z ~ N(0, 1)` in numpy → one `session.run(...)`
softmax-normalise deltas → quantise (40 px precise / 120 px otherwise) →
scale to target distance → cumulative timestamps.
### Server (`ai_mouse/server/`)
### Server (`tools/server/`) and frontend (`static/`)
- App factory `create_app()` in [server/__init__.py](ai_mouse/server/__init__.py) mounts four routers under `/api`: `routes_collect`, `routes_train`, `routes_verify`, and `routes_scroll` (the scroll router uses prefix `/api/scroll`).
- **Session state** ([server/deps.py](ai_mouse/server/deps.py)): A module-level `SessionState` singleton holds the active `Collector` / `ScrollCollector`. Tests that need to override `_DATA_DIR` monkeypatch `ai_mouse.server.deps._DATA_DIR`.
- **Training progress** is delivered via Server-Sent Events. The route launches `train()` on a background thread (`asyncio.to_thread`) and pushes `{epoch, total, loss}` dicts through an `asyncio.Queue` until `{done: True}` or `{error: ...}`.
- **Static** is mounted from `static/` at the project root (not under the package); `index.html` is served at `/`.
### Frontend (`static/`)
Vanilla Vue 3 + axios + ECharts pulled from CDN — **no bundler, no node_modules**. Single page with three tab views (`collect.js`, `train.js`, `verify.js`) registered as components in [app.js](static/js/app.js). `api.js` exports an axios instance and a `fetchSSE()` helper that reads `text/event-stream` from `fetch().body.getReader()` and parses `data: ...` frames. UI strings are in Chinese.
App factory `create_app()` mounts four routers under `/api`. Frontend is
vanilla Vue 3 + axios + ECharts via CDN. Note: the `/api/verify` and
`/api/scroll/verify` endpoints always use the **bundled** ONNX weights
(via `from ai_mouse import generate / generate_scroll`). If you retrain
and want the Web UI to reflect new weights, re-run `tools.export_onnx` and
restart the server.
## Config
[ai_mouse/config.py](ai_mouse/config.py) is the **single source of truth** for hyperparameters. `TrainConfig`, `GenerateConfig`, `ScrollTrainConfig`, `ScrollModeConfig`, and `ServerConfig` are dataclasses. When changing model architecture (`d_model`, `nhead`, etc.) keep training and inference consistent — the `train_config.json` saved at training time is what `generate()` uses to reconstruct the model.
`tools/config.py` holds the training-side dataclasses (`TrainConfig`,
`ScrollTrainConfig`, etc.). The library does NOT use these — its only
"configuration" is what's embedded in `src/ai_mouse/assets/train_config.json`
(architecture params needed to know `seq_len` etc. at inference time).
## Tests
`tests/conftest.py` provides `model_dir` and `scroll_model_dir` fixtures that write **freshly initialised (untrained)** weights to a temp dir along with all required JSON metadata. Use these whenever a test calls `generate()` or `generate_scroll()`. The trajectories will be garbage but the inference path runs end-to-end.
- `tests/unit/conftest.py` — fixtures for library-only tests, no torch.
- `tests/tools/conftest.py``model_dir` and `scroll_model_dir` fixtures
that produce **untrained** torch weights in a temp dir. Used by training-
/server-side tests.
- `tests/unit/test_golden.py` — regression suite that pins library output
against `tests/unit/data/golden_{mouse,scroll}.npz` captured before the
ONNX migration. Tolerance is distance-scaled: mouse allows ±max(30 px, 20%
of move distance) and ±700 ms; scroll requires exact total deltaY match
and ±2 quanta per event.
Server tests use `httpx.ASGITransport(app=create_app())` with `pytest-asyncio` — no live socket.
Server tests use `httpx.ASGITransport(app=create_app())` with
`pytest-asyncio` — no live socket.

120
README.md Normal file
View File

@@ -0,0 +1,120 @@
# ai_mouse
Human-like mouse trajectory and scroll wheel event generator. Inference runs on
ONNX Runtime; the only runtime dependencies are `numpy` and `onnxruntime`.
## Install
```bash
pip install git+https://github.com/<owner>/ai_mouse.git
```
For GPU inference (optional), replace `onnxruntime` with the GPU variant:
```bash
pip install onnxruntime-gpu # CUDA / TensorRT
# or
pip install onnxruntime-directml # Windows DirectML
```
## Quick start
### Mouse trajectory
```python
from ai_mouse import generate
points = generate(start=(100, 200), end=(900, 400))
# [(x, y, t_ms), ..., (cx, cy, t_down), (cx, cy, t_up)]
```
### Scroll wheel
```python
from ai_mouse import generate_scroll
events = generate_scroll(start_scroll_y=0, target_scroll_y=2000)
# [{"deltaY": 120, "deltaMode": 0, "t": 32}, ...]
```
### Class API (recommended for repeated calls)
```python
from ai_mouse import MouseModel
m = MouseModel() # session created once
for target in target_list:
pts = m.generate((cx, cy), target)
```
### Custom providers / GPU
```python
from ai_mouse import MouseModel
m = MouseModel(providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
# or
m = MouseModel(providers=["DmlExecutionProvider"])
```
### Reproducibility
```python
m.generate(start, end, seed=42)
```
## API summary
| Name | Purpose |
|---|---|
| `generate(start, end, *, n_points=64, speed=None, click=True, seed=None)` | One-shot call; internal lru_cache singleton |
| `MouseModel(model_path=None, providers=None, seed=None)` | Persistent session |
| `generate_scroll(...)` / `ScrollModel(...)` | Same shape for scroll |
| `ai_mouse.errors.{ModelLoadError, GenerationError}` | Exception hierarchy |
## Thread safety
`MouseModel.generate` and `ScrollModel.generate` are safe to call concurrently
from multiple threads — ORT `InferenceSession` is itself thread-safe.
## Development
The repo contains optional dev-only tooling under `tools/` for training your
own models, running the FastAPI web UI, and evaluating output quality. Install
with the `dev` group:
```bash
uv sync --group dev
```
Common commands:
```bash
# Web UI (collect + train + verify in browser)
uv run python tools/serve.py
# Training (after collecting your own data)
uv run python -m tools train --data data/traces.jsonl --output data/models_v2
# Convert Balabit corpus to trace format
uv run python -m tools balabit-adapter --input data/balabit_raw \
--output data/pretrain_traces.jsonl
# Eval report
uv run python -m tools eval --model-dir data/models_v2 \
--reference data/pretrain_traces.jsonl --output data/eval_reports/report.md
# Re-export ONNX after retraining
uv run python -m tools.export_onnx --flow-ckpt data/models_v2 \
--scroll-ckpt data/scroll_models --output src/ai_mouse/assets/
```
Tests:
```bash
uv run pytest tests/unit # library-only (no torch)
uv run pytest tests/tools # full dev suite
```
After retraining you need to re-export and rebuild the wheel for the new
weights to ship; the in-app Verify endpoint always uses bundled weights.

View File

@@ -0,0 +1,624 @@
# Mouse Post-Processing Quality Rework Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Remove the endpoint artifacts (vertical walls, hooks, start kinks) in generated mouse trajectories by reworking the post-processing pipeline, and reduce sampling jitter by switching 10-step Euler to 10-step Heun.
**Architecture:** Three new pure-numpy functions in `src/ai_mouse/_postprocess.py` (`soften_forward`, `damp_start`, `warp_endpoints`) replace the three hard-clamping functions (`enforce_forward_monotonic`, `smooth_start`, `snap_endpoints`). `mouse.py` wires them in a new order (soft-monotonic → start damping → smoothing both axes → global endpoint correction) and replaces the Euler ODE loop with Heun predictor-corrector. Golden regression baselines are re-captured (intentional behavior change).
**Tech Stack:** Python 3.12+, numpy, onnxruntime, pytest. Package manager: `uv`.
**Spec:** `docs/superpowers/specs/2026-07-09-mouse-postprocess-quality-design.md`
## Global Constraints
- `src/ai_mouse/` is wheel content: NEVER import torch/fastapi/scipy/matplotlib there (CI `library` job installs runtime deps only).
- All new post-processing functions are pure (no I/O, no global state), matching the existing `_postprocess.py` convention.
- Public API (`generate()` signature and return shape, exact endpoint hit) must not change.
- Scroll subsystem and `golden_scroll.npz` are untouched.
- Run library tests with: `uv run pytest tests/unit` (add `-v` per test as noted).
- All commits end with `Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>`.
---
### Task 1: `soften_forward` — soft monotonic with overshoot compression
Replaces the hard `clip(0,1)` + strict monotonicity of `enforce_forward_monotonic`. Tolerates small backtracking (real hands micro-correct), allows natural overshoot past 1.0, and soft-compresses extreme overshoot with tanh so the path never flies far past the target.
**Files:**
- Modify: `src/ai_mouse/_postprocess.py` (add function; do NOT delete old ones yet — removal happens in Task 4)
- Test: `tests/unit/test_postprocess.py`
**Interfaces:**
- Produces: `soften_forward(forward: np.ndarray, backtrack_tol: float = 0.02, overshoot_span: float = 0.08) -> np.ndarray` — returns a new array; Task 4 calls it with defaults.
- [ ] **Step 1: Write the failing tests**
Append to `tests/unit/test_postprocess.py`:
```python
from ai_mouse._postprocess import soften_forward
def test_soften_forward_tolerates_small_backtrack() -> None:
# A 0.01 dip is within the 0.02 tolerance and must survive untouched.
f = np.array([0.0, 0.30, 0.29, 0.60, 1.0])
out = soften_forward(f)
assert np.isclose(out[2], 0.29)
def test_soften_forward_limits_large_backtrack() -> None:
# A 0.30 dip is noise; it gets pulled up to prev - tol.
f = np.array([0.0, 0.50, 0.20, 0.70, 1.0])
out = soften_forward(f)
assert np.isclose(out[2], 0.50 - 0.02)
def test_soften_forward_allows_moderate_overshoot() -> None:
# Overshoot past 1.0 is natural; small overshoot survives (compressed
# but strictly > 1.0).
f = np.array([0.0, 0.5, 0.9, 1.04, 1.0])
out = soften_forward(f)
assert out[3] > 1.0
def test_soften_forward_compresses_extreme_overshoot() -> None:
# tanh compression: no output value may exceed 1 + overshoot_span.
f = np.array([0.0, 0.5, 1.30, 1.50, 1.0])
out = soften_forward(f)
assert out.max() <= 1.0 + 0.08 + 1e-9
assert out[2] > 1.0 # still an overshoot, not clipped flat
def test_soften_forward_no_lower_clip() -> None:
# Small wind-up behind the start is allowed (warp_endpoints pins
# the first point later; interior may be slightly negative).
f = np.array([0.0, -0.01, 0.30, 0.70, 1.0])
out = soften_forward(f)
assert out[1] < 0.0
```
- [ ] **Step 2: Run tests to verify they fail**
Run: `uv run pytest tests/unit/test_postprocess.py -k soften_forward -v`
Expected: 5 failures/errors with `ImportError: cannot import name 'soften_forward'`
- [ ] **Step 3: Write the implementation**
Add to `src/ai_mouse/_postprocess.py` (after `enforce_forward_monotonic`):
```python
def soften_forward(
forward: np.ndarray,
backtrack_tol: float = 0.02,
overshoot_span: float = 0.08,
) -> np.ndarray:
"""Softly regularise the forward axis without destroying natural motion.
Real trajectories contain small backward corrections and overshoot
past the target; hard clipping turns both into visible artifacts
(stacked points, vertical walls). Instead:
- Backtracking is tolerated up to ``backtrack_tol``; larger dips are
raised to ``prev - backtrack_tol``.
- Values above 1.0 are compressed with tanh so they asymptote at
``1 + overshoot_span`` (moderate overshoot survives, extremes
cannot fly far past the target).
- No lower clip: the endpoint warp pins the first point later.
Args:
forward: (T,) forward coordinates.
backtrack_tol: max allowed per-step regression.
overshoot_span: asymptotic max excess above 1.0.
Returns:
New (T,) array.
"""
out = forward.copy()
for i in range(1, len(out)):
floor = out[i - 1] - backtrack_tol
if out[i] < floor:
out[i] = floor
over = out > 1.0
out[over] = 1.0 + overshoot_span * np.tanh((out[over] - 1.0) / overshoot_span)
return out
```
- [ ] **Step 4: Run tests to verify they pass**
Run: `uv run pytest tests/unit/test_postprocess.py -k soften_forward -v`
Expected: 5 passed
- [ ] **Step 5: Commit**
```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat: add soften_forward (backtrack tolerance + tanh overshoot compression)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
### Task 2: `damp_start` — continuous start damping
Replaces `smooth_start`, whose ×1/5-then-abrupt-release lateral damping creates start kinks. New version ramps damping with smoothstep so the weight approaches 1 continuously at the release boundary, and touches only lateral (forward regularisation is `soften_forward`'s job).
**Files:**
- Modify: `src/ai_mouse/_postprocess.py`
- Test: `tests/unit/test_postprocess.py`
**Interfaces:**
- Produces: `damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray` — returns a new array; Task 4 calls it with defaults.
- [ ] **Step 1: Write the failing tests**
Append to `tests/unit/test_postprocess.py`:
```python
from ai_mouse._postprocess import damp_start
def test_damp_start_dampens_early_lateral() -> None:
lat = np.full(16, 1.0)
out = damp_start(lat, n=4)
assert out[1] < out[2] < out[3] < out[4] < 1.0 # monotone ramp
assert np.all(out[5:] == 1.0) # untouched past n
def test_damp_start_no_release_jump() -> None:
# The weight at i=n must be close to 1 (continuous release):
# smoothstep(4/5) = 0.896, vs the old linear 4/5 = 0.8.
lat = np.full(16, 1.0)
out = damp_start(lat, n=4)
assert out[4] > 0.85
def test_damp_start_short_input_safe() -> None:
lat = np.array([0.0, 0.5, 0.3])
out = damp_start(lat, n=4) # n capped to len//4 = 0 → no-op
assert np.array_equal(out, lat)
```
- [ ] **Step 2: Run tests to verify they fail**
Run: `uv run pytest tests/unit/test_postprocess.py -k damp_start -v`
Expected: 3 failures with `ImportError: cannot import name 'damp_start'`
- [ ] **Step 3: Write the implementation**
Add to `src/ai_mouse/_postprocess.py`:
```python
def damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray:
"""Dampen lateral oscillation over the first ``n`` points, continuously.
Weights follow smoothstep(i / (n+1)) so the damping releases smoothly
into the untouched region (the old linear blend jumped from 0.8 to
1.0 and left a visible kink).
Args:
lateral: (T,) lateral coordinates.
n: number of leading points to dampen (capped at len//4).
Returns:
New (T,) array.
"""
out = lateral.copy()
n = min(n, len(out) // 4)
for i in range(1, n + 1):
t = i / (n + 1)
w = t * t * (3.0 - 2.0 * t) # smoothstep
out[i] *= w
return out
```
- [ ] **Step 4: Run tests to verify they pass**
Run: `uv run pytest tests/unit/test_postprocess.py -k damp_start -v`
Expected: 3 passed
- [ ] **Step 5: Commit**
```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat: add damp_start (smoothstep lateral damping, no release kink)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
### Task 3: `warp_endpoints` — global residual correction
Replaces `snap_endpoints`. Instead of dragging the last 6 points toward (1,0) (which fights the trajectory's own direction and creates hooks), compute the first/last-point residuals and distribute the correction across the whole curve with smoothstep weights. Endpoints land exactly; local shape and approach direction are preserved.
**Files:**
- Modify: `src/ai_mouse/_postprocess.py`
- Test: `tests/unit/test_postprocess.py`
**Interfaces:**
- Produces: `warp_endpoints(forward: np.ndarray, lateral: np.ndarray) -> tuple[np.ndarray, np.ndarray]` — returns new arrays with `forward[0]==0.0, lateral[0]==0.0, forward[-1]==1.0, lateral[-1]==0.0` exactly; Task 4 calls it last in the pipeline.
- [ ] **Step 1: Write the failing tests**
Append to `tests/unit/test_postprocess.py`:
```python
from ai_mouse._postprocess import warp_endpoints
def test_warp_endpoints_exact_pin() -> None:
f = np.linspace(0.05, 1.10, 32)
l = np.linspace(0.03, -0.07, 32)
fo, lo = warp_endpoints(f, l)
assert fo[0] == 0.0 and lo[0] == 0.0
assert fo[-1] == 1.0 and lo[-1] == 0.0
def test_warp_endpoints_identity_when_already_pinned() -> None:
f = np.linspace(0.0, 1.0, 32)
l = np.sin(np.linspace(0, np.pi, 32)) * 0.1
l[0] = l[-1] = 0.0
fo, lo = warp_endpoints(f.copy(), l.copy())
assert np.allclose(fo, f, atol=1e-12)
assert np.allclose(lo, l, atol=1e-12)
def test_warp_endpoints_preserves_smoothness() -> None:
# Correcting a smooth curve must not introduce sharp local bends:
# the warp adds a smoothstep-weighted offset, so the second
# difference (discrete curvature proxy) stays small.
f = np.linspace(0.02, 1.08, 32)
l = np.full(32, 0.05)
fo, lo = warp_endpoints(f, l)
assert np.abs(np.diff(lo, 2)).max() < 0.01
assert np.abs(np.diff(fo, 2)).max() < 0.01
def test_warp_endpoints_correction_local_to_each_end() -> None:
# A start-only residual should barely move the last quarter.
# (smoothstep weight at i=24/31 is ~0.13, so 0.08 residual leaves
# ~0.010 there — threshold 0.02 gives margin without losing meaning)
f = np.linspace(0.0, 1.0, 32) + 0.0
l = np.zeros(32)
f[0] = 0.08 # start residual only
fo, _ = warp_endpoints(f, l)
assert np.abs(fo[24:] - f[24:]).max() < 0.02
```
- [ ] **Step 2: Run tests to verify they fail**
Run: `uv run pytest tests/unit/test_postprocess.py -k warp_endpoints -v`
Expected: 4 failures with `ImportError: cannot import name 'warp_endpoints'`
- [ ] **Step 3: Write the implementation**
Add to `src/ai_mouse/_postprocess.py`:
```python
def warp_endpoints(
forward: np.ndarray,
lateral: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""Warp the whole curve so endpoints land exactly on (0,0) and (1,0).
Computes the residual of the first point vs (0, 0) and the last point
vs (1, 0), then subtracts each residual weighted by a smoothstep that
is 1 at its own end and 0 at the opposite end. Unlike tail-dragging,
this preserves the trajectory's local shape and approach direction.
Args:
forward: (T,) forward coordinates.
lateral: (T,) lateral coordinates.
Returns:
``(forward, lateral)`` new arrays, endpoints pinned exactly.
"""
t = np.linspace(0.0, 1.0, len(forward))
w_end = t * t * (3.0 - 2.0 * t) # smoothstep: 0 at start → 1 at end
w_start = 1.0 - w_end # mirrored
res_f0, res_l0 = forward[0] - 0.0, lateral[0] - 0.0
res_f1, res_l1 = forward[-1] - 1.0, lateral[-1] - 0.0
fo = forward - w_start * res_f0 - w_end * res_f1
lo = lateral - w_start * res_l0 - w_end * res_l1
fo[0], lo[0] = 0.0, 0.0
fo[-1], lo[-1] = 1.0, 0.0
return fo, lo
```
- [ ] **Step 4: Run tests to verify they pass**
Run: `uv run pytest tests/unit/test_postprocess.py -k warp_endpoints -v`
Expected: 4 passed
- [ ] **Step 5: Commit**
```bash
git add src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat: add warp_endpoints (global residual correction, shape-preserving)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
### Task 4: Wire new pipeline into `mouse.py`; remove old functions
Swap the pipeline in `MouseModel.generate` to the spec order (soft-monotonic → start damping → smooth both axes → global endpoint correction), then delete the three replaced functions and their tests. Golden tests will now fail — that is expected and fixed in Task 6; all other tests must pass.
**Files:**
- Modify: `src/ai_mouse/mouse.py:14-24` (imports), `src/ai_mouse/mouse.py:106-109` (pipeline)
- Modify: `src/ai_mouse/_postprocess.py` (delete `snap_endpoints`, `smooth_start`, `enforce_forward_monotonic`; also update the stale `snap_endpoints` cross-reference in any remaining docstring)
- Modify: `tests/unit/test_postprocess.py` (delete the 6 tests of the removed functions and their two mid-file import lines)
**Interfaces:**
- Consumes: `soften_forward(forward)` (Task 1), `damp_start(lateral)` (Task 2), `warp_endpoints(forward, lateral)` (Task 3), existing `gaussian_smooth(x, sigma)`.
- [ ] **Step 1: Update `mouse.py` imports**
Replace the import block at `src/ai_mouse/mouse.py:15-24` with:
```python
from ai_mouse._postprocess import (
build_timestamps,
damp_start,
gaussian_smooth,
resample_arc,
sample_duration,
soften_forward,
truncnorm_sample,
warp_endpoints,
)
```
- [ ] **Step 2: Replace the pipeline**
Replace `src/ai_mouse/mouse.py:106-109`:
```python
forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
forward, lateral = smooth_start(forward, lateral)
forward = enforce_forward_monotonic(forward)
lateral = gaussian_smooth(lateral, sigma=1.0)
```
with:
```python
forward = soften_forward(forward)
lateral = damp_start(lateral)
forward = gaussian_smooth(forward, sigma=1.0)
lateral = gaussian_smooth(lateral, sigma=1.0)
forward, lateral = warp_endpoints(forward, lateral)
```
- [ ] **Step 3: Delete replaced functions and their tests**
- In `src/ai_mouse/_postprocess.py`: delete `snap_endpoints` (lines 34-62), `smooth_start` (65-80), `enforce_forward_monotonic` (83-92).
- In `tests/unit/test_postprocess.py`: delete `test_snap_endpoints_pins_first_and_last`, `test_snap_endpoints_preserves_middle`, `test_smooth_start_dampens_lateral`, `test_enforce_forward_monotonic_repairs_inversions`, `test_enforce_forward_monotonic_clips_to_unit_interval`, and the `from ai_mouse._postprocess import snap_endpoints` / `from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start` import lines.
- [ ] **Step 4: Run the unit suite (golden mouse failures expected)**
Run: `uv run pytest tests/unit -v`
Expected: everything passes EXCEPT `tests/unit/test_golden.py::test_mouse_golden[...]` cases, which may exceed the path envelope (behavior intentionally changed; re-baselined in Task 6). If anything else fails, fix before committing. In particular `test_mouse.py` (endpoint snap, seed reproducibility, shape) must pass.
- [ ] **Step 5: Commit**
```bash
git add src/ai_mouse/mouse.py src/ai_mouse/_postprocess.py tests/unit/test_postprocess.py
git commit -m "feat: rework mouse post-processing pipeline (soft monotonic, global endpoint warp)
Golden mouse baselines temporarily failing; re-captured in follow-up commit.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
### Task 5: Euler → Heun sampling
Second-order predictor-corrector at the same 10 steps (NFE 10 → 20; each call is a small d_model=128 transformer, ~1-2 ms CPU). Reduces integration error and sampling jitter.
**Files:**
- Modify: `src/ai_mouse/mouse.py:27` (constant), `src/ai_mouse/mouse.py:93-97` (ODE loop)
**Interfaces:**
- Consumes: the existing ONNX session I/O contract `session.run(["v"], {"x_t", "t", "cond"})` — unchanged.
- [ ] **Step 1: Replace the ODE loop**
Rename the constant at `src/ai_mouse/mouse.py:27`:
```python
_N_ODE_STEPS = 10
```
Replace the loop at `src/ai_mouse/mouse.py:93-97`:
```python
dt = 1.0 / _N_EULER_STEPS
for step in range(_N_EULER_STEPS):
t = np.full((1,), step * dt, dtype=np.float32)
v = self._session.run(["v"], {"x_t": x, "t": t, "cond": cond})[0]
x = x + v * dt
```
with:
```python
# Heun (2nd-order predictor-corrector): same step count as the old
# Euler loop but far lower integration error for 2x NFE.
dt = 1.0 / _N_ODE_STEPS
for step in range(_N_ODE_STEPS):
t0 = np.full((1,), step * dt, dtype=np.float32)
v1 = self._session.run(["v"], {"x_t": x, "t": t0, "cond": cond})[0]
x_pred = (x + v1 * dt).astype(np.float32)
t1 = np.full((1,), (step + 1) * dt, dtype=np.float32)
v2 = self._session.run(["v"], {"x_t": x_pred, "t": t1, "cond": cond})[0]
x = x + (v1 + v2) * (dt / 2.0)
```
- [ ] **Step 2: Run the unit suite (same expectation as Task 4)**
Run: `uv run pytest tests/unit -v`
Expected: all pass except `test_mouse_golden` envelope cases (still pending re-baseline). `test_mouse.py::test_mouse_model_seed_reproducibility` must pass — Heun adds no new randomness.
- [ ] **Step 3: Commit**
```bash
git add src/ai_mouse/mouse.py
git commit -m "feat: switch flow ODE sampling from Euler to Heun (10 steps, NFE 20)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
### Task 6: Quality regression test, golden re-baseline, CHANGELOG, visual verification
Add a numeric guard against the two artifact classes (sharp turns near the endpoint, vertical walls), re-capture `golden_mouse.npz` from the new implementation (documented procedure in `tests/unit/test_golden.py` docstring), and verify visually.
**Files:**
- Test: `tests/unit/test_mouse.py` (append)
- Modify: `tests/unit/data/golden_mouse.npz` (re-captured binary)
- Modify: `CHANGELOG.md`
**Interfaces:**
- Consumes: `generate(start, end, *, seed, click)` public API.
- [ ] **Step 1: Write the quality regression test (fails on OLD pipeline, passes on new)**
Append to `tests/unit/test_mouse.py` (note: this file has no module-level
`numpy`/`generate` imports — the test is self-contained):
```python
def test_no_sharp_turns_or_walls_near_endpoint() -> None:
"""Guard against the two endpoint artifact classes:
- sharp turns (>90°) between consecutive substantial segments in the
final approach (the old tail-drag created hooks);
- vertical walls: many points stacked at the target's forward
position (the old clip(0,1) stacked overshoot at forward=1).
"""
import math
import numpy as np
from ai_mouse import generate
cases = [((100, 300), (900, 350)), ((100, 100), (700, 600)),
((800, 200), (150, 550))]
for (start, end) in cases:
for seed in range(6):
pts = generate(start, end, seed=seed, click=False)
arr = np.array([(x, y) for x, y, _ in pts], dtype=float)
tail = arr[-12:]
seg = np.diff(tail, axis=0)
lens = np.linalg.norm(seg, axis=1)
# Only consider substantial segments: integer-pixel staircase
# on 1-2 px steps produces spurious 90° angles.
keep = lens >= 3.0
headings = np.arctan2(seg[keep][:, 1], seg[keep][:, 0])
if len(headings) >= 2:
turns = np.abs(np.diff(np.unwrap(headings)))
max_turn = math.degrees(turns.max())
assert max_turn < 90.0, (
f"{start}->{end} seed={seed}: {max_turn:.0f}° turn "
f"in final approach"
)
# Vertical wall: >=4 consecutive tail points within 2 px of
# the target x while spanning >10 px of y.
near_x = np.abs(tail[:, 0] - end[0]) <= 2.0
run = 0
for i, flag in enumerate(near_x[:-1]): # exclude final point
run = run + 1 if flag else 0
assert run < 4 or np.ptp(tail[near_x][:, 1]) <= 10.0, (
f"{start}->{end} seed={seed}: vertical wall at target x"
)
```
- [ ] **Step 2: Run the new test**
Run: `uv run pytest tests/unit/test_mouse.py::test_no_sharp_turns_or_walls_near_endpoint -v`
Expected: PASS (the new pipeline removed the artifacts). If it fails, treat it as a real defect in Tasks 1-5 — do not loosen the thresholds; investigate which pipeline step reintroduces the artifact.
- [ ] **Step 3: Re-capture the mouse golden baseline**
Write a throwaway script (scratchpad or temp path, NOT committed):
```python
# recapture_golden.py — regenerate golden_mouse.npz from new implementation
import numpy as np
from ai_mouse import generate
CASES = [
((100, 200), (900, 400)), ((500, 500), (500, 100)),
((200, 600), (800, 200)), ((100, 100), (130, 110)),
((50, 50), (1500, 900)), ((400, 300), (500, 300)),
((300, 300), (700, 700)), ((600, 400), (200, 100)),
]
out = {}
for ci, (s, e) in enumerate(CASES):
for seed in range(4):
pts = generate(s, e, seed=seed)
out[f"case{ci}_seed{seed}"] = np.array(pts, dtype=np.int64)
np.savez_compressed("tests/unit/data/golden_mouse.npz", **out)
print(f"wrote {len(out)} golden traces")
```
Run: `uv run python <path>/recapture_golden.py`
Expected: `wrote 32 golden traces`
- [ ] **Step 4: Full unit suite must be green**
Run: `uv run pytest tests/unit -v`
Expected: ALL pass, including all 32 `test_mouse_golden` and all 32 `test_scroll_golden` cases (scroll goldens untouched — if any scroll test fails, something leaked outside mouse scope; stop and investigate).
- [ ] **Step 5: Update CHANGELOG**
Add under the top of `CHANGELOG.md` (after the intro, before `## [0.2.0]`):
```markdown
## [Unreleased]
### Changed
- Mouse post-processing pipeline reworked to remove endpoint artifacts:
hard forward clip → backtrack-tolerant soft monotonic with tanh
overshoot compression (`soften_forward`); tail-drag endpoint snapping →
whole-curve smoothstep residual correction (`warp_endpoints`); abrupt
start damping → continuous smoothstep damping (`damp_start`);
gaussian smoothing now applied to both axes.
- Flow ODE sampling switched from 10-step Euler to 10-step Heun
(predictor-corrector); ~2x model calls per trajectory, still ~40 ms CPU.
- `tests/unit/data/golden_mouse.npz` re-baselined against the new
pipeline (intentional behavior change; scroll goldens unchanged).
```
- [ ] **Step 6: Visual verification (diagnostic plot + Web UI)**
1. Re-run the diagnostic script from the investigation (same 4 cases × 6 seeds; it lives in the session scratchpad as `diag_traj.py`) and visually compare against the "before" plot: no vertical walls in end zoom, no hooks, no start kinks; `turns>45deg` counts in the numeric output should drop sharply vs the before-values (3-8 per trace).
2. Start the Web UI: `uv run python tools/serve.py`, open the verify page, and have the user visually approve. **This is the final acceptance gate** — post-processing is Python-side, so no ONNX re-export is needed, but the server must be (re)started to pick up the new library code.
- [ ] **Step 7: Commit**
```bash
git add tests/unit/test_mouse.py tests/unit/data/golden_mouse.npz CHANGELOG.md
git commit -m "test: quality guard for endpoint artifacts; re-baseline mouse goldens
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>"
```
---
## Acceptance summary (from spec)
- [ ] No vertical wall / hooks / start kinks in diagnostic plots (Task 6 step 6)
- [ ] `turns>45°` count drops sharply vs baseline (was 1-8 per trace)
- [ ] `uv run pytest tests/unit` fully green, goldens re-baselined
- [ ] User approves Web UI verify page visually
- [ ] Out of scope confirmed untouched: scroll subsystem, training code, ONNX assets

View File

@@ -0,0 +1,110 @@
# 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. 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.

30
examples/quickstart.py Normal file
View File

@@ -0,0 +1,30 @@
"""Minimal example: generate one trajectory + click event.
Run: uv run python examples/quickstart.py
"""
from __future__ import annotations
from ai_mouse import generate
start = (100, 200)
end = (900, 400)
points = generate(start, end, seed=0)
print(f"Generated {len(points)} events:")
print(f" first move: {points[0]}")
print(f" middle move: {points[len(points) // 2]}")
print(f" last move: {points[-3]}")
print(f" click-down: {points[-2]}")
print(f" click-up: {points[-1]}")
# Typical replay loop pattern. t_ms is cumulative from the start of the trace,
# so block your sender thread until time-since-start reaches each event's t_ms.
#
# import time
# t0 = time.monotonic()
# for x, y, t_ms in points:
# target_wallclock = t0 + t_ms / 1000.0
# while time.monotonic() < target_wallclock:
# pass
# # replace this with pyautogui / pynput / win32 mouse_event:
# # send_mouse_move(x, y)

View File

@@ -1,49 +0,0 @@
"""One-shot script to capture golden mouse trajectories from the current torch
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
in Phase 4 produces equivalent output.
Output: tests/unit/data/golden_mouse.npz
"""
from __future__ import annotations
import random
import sys
from pathlib import Path
# Allow running as `uv run python scripts/build_golden_mouse.py` from project root.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import numpy as np
import torch
from ai_mouse import generate
CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
((100, 200), (900, 400)), # horizontal 800px
((500, 500), (500, 100)), # vertical 400px upward
((200, 600), (800, 200)), # 720px diagonal
((100, 100), (130, 110)), # very short 31px
((50, 50), (1500, 900)), # very long 1700px
((400, 300), (500, 300)), # short horizontal 100px
((300, 300), (700, 700)), # 45° diagonal
((600, 400), (200, 100)), # reverse diagonal
]
SEEDS = (0, 1, 2, 3)
def main() -> None:
out: dict[str, np.ndarray] = {}
for case_idx, (start, end) in enumerate(CASES):
for seed in SEEDS:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
pts = generate(start=start, end=end)
out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
out_path = Path("tests/unit/data/golden_mouse.npz")
np.savez_compressed(out_path, **out)
print(f"Wrote {len(out)} golden traces to {out_path}")
if __name__ == "__main__":
main()

View File

@@ -1,48 +0,0 @@
"""Capture golden scroll event sequences from current torch implementation."""
from __future__ import annotations
import random
import sys
from pathlib import Path
# Allow running as `uv run python scripts/build_golden_scroll.py` from project root.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import numpy as np
import torch
from ai_mouse import generate_scroll
CASES: list[tuple[int, int, str]] = [
(0, 1500, "target"),
(0, 500, "precise"),
(0, 5000, "fast"),
(2000, 0, "target"), # upward
(0, 800, "precise"),
(0, 3500, "fast"),
(1000, 1200, "precise"), # tiny scroll
(0, 10000, "fast"), # very long
]
SEEDS = (0, 1, 2, 3)
def main() -> None:
out: dict[str, np.ndarray] = {}
for case_idx, (start_y, end_y, mode) in enumerate(CASES):
for seed in SEEDS:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
events = generate_scroll(start_y, end_y, mode=mode)
arr = np.array(
[[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
dtype=np.int64,
)
out[f"case{case_idx}_seed{seed}"] = arr
out_path = Path("tests/unit/data/golden_scroll.npz")
np.savez_compressed(out_path, **out)
print(f"Wrote {len(out)} scroll golden traces to {out_path}")
if __name__ == "__main__":
main()

View File

@@ -31,65 +31,27 @@ def gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
return smoothed
def snap_endpoints(
forward: np.ndarray,
lateral: np.ndarray,
seq_len: int,
n_snap: int = 6,
) -> tuple[np.ndarray, np.ndarray]:
"""Force first point to (0,0) and last point to (1,0) with quadratic ease.
def damp_start(lateral: np.ndarray, n: int = 4) -> np.ndarray:
"""Dampen lateral oscillation over the first ``n`` points, continuously.
The last ``n_snap`` points are linearly interpolated towards (1, 0)
with quadratic easing, then the first/last points are pinned exactly.
Weights follow smoothstep(i / (n+1)) so the damping releases smoothly
into the untouched region (the old linear blend jumped from 0.8 to
1.0 and left a visible kink).
Args:
forward: (T,) forward coordinates (modified in place).
lateral: (T,) lateral coordinates (modified in place).
seq_len: length of forward/lateral.
n_snap: number of trailing points to ease (capped at seq_len//4).
lateral: (T,) lateral coordinates.
n: number of leading points to dampen (capped at len//4).
Returns:
``(forward, lateral)`` after modification.
New (T,) array.
"""
n_snap = min(n_snap, seq_len // 4)
for i in range(n_snap):
alpha = ((i + 1) / n_snap) ** 2
k = seq_len - n_snap + i
forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
forward[0], lateral[0] = 0.0, 0.0
forward[-1], lateral[-1] = 1.0, 0.0
return forward, lateral
def smooth_start(
forward: np.ndarray,
lateral: np.ndarray,
n: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
"""Dampen lateral oscillation in the first ``n`` points.
Assumes :func:`snap_endpoints` has already pinned (0,0). Forward is
forced non-decreasing locally; lateral is linearly damped towards 0.
"""
n_start_fix = min(n, len(forward) // 4)
for i in range(1, n_start_fix + 1):
blend = i / (n_start_fix + 1)
forward[i] = max(forward[i], forward[i - 1])
lateral[i] = lateral[i] * blend
return forward, lateral
def enforce_forward_monotonic(forward: np.ndarray) -> np.ndarray:
"""Force ``forward`` non-decreasing, clip to [0,1], pin endpoints."""
seq_len = len(forward)
for i in range(1, seq_len - 1):
if forward[i] < forward[i - 1]:
forward[i] = forward[i - 1] + 0.001
forward = np.clip(forward, 0.0, 1.0)
forward[0] = 0.0
forward[-1] = 1.0
return forward
out = lateral.copy()
n = min(n, len(out) // 4)
for i in range(1, n + 1):
t = i / (n + 1)
w = t * t * (3.0 - 2.0 * t) # smoothstep
out[i] *= w
return out
def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
@@ -178,3 +140,72 @@ def sample_duration(
mu_log = params[bin_idx]["mu_log"]
sigma_log = params[bin_idx]["sigma_log"]
return float(np.exp(rng.normal(mu_log, sigma_log)))
def soften_forward(
forward: np.ndarray,
backtrack_tol: float = 0.02,
overshoot_span: float = 0.08,
) -> np.ndarray:
"""Softly regularise the forward axis without destroying natural motion.
Real trajectories contain small backward corrections and overshoot
past the target; hard clipping turns both into visible artifacts
(stacked points, vertical walls). Instead:
- Backtracking is tolerated up to ``backtrack_tol``; larger dips are
raised to ``prev - backtrack_tol``.
- Values above 1.0 are compressed with tanh so they asymptote at
``1 + overshoot_span`` (moderate overshoot survives, extremes
cannot fly far past the target).
- No lower clip: the endpoint warp pins the first point later.
Args:
forward: (T,) forward coordinates.
backtrack_tol: max allowed per-step regression.
overshoot_span: asymptotic max excess above 1.0.
Returns:
New (T,) array.
"""
out = forward.copy()
for i in range(1, len(out)):
floor = out[i - 1] - backtrack_tol
if out[i] < floor:
out[i] = floor
over = out > 1.0
out[over] = 1.0 + overshoot_span * np.tanh((out[over] - 1.0) / overshoot_span)
return out
def warp_endpoints(
forward: np.ndarray,
lateral: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""Warp the whole curve so endpoints land exactly on (0,0) and (1,0).
Computes the residual of the first point vs (0, 0) and the last point
vs (1, 0), then subtracts each residual weighted by a smoothstep that
is 1 at its own end and 0 at the opposite end. Unlike tail-dragging,
this preserves the trajectory's local shape and approach direction.
Args:
forward: (T,) forward coordinates.
lateral: (T,) lateral coordinates.
Returns:
``(forward, lateral)`` new arrays, endpoints pinned exactly.
"""
t = np.linspace(0.0, 1.0, len(forward))
w_end = t * t * (3.0 - 2.0 * t) # smoothstep: 0 at start → 1 at end
w_start = 1.0 - w_end # mirrored
res_f0, res_l0 = forward[0] - 0.0, lateral[0] - 0.0
res_f1, res_l1 = forward[-1] - 1.0, lateral[-1] - 0.0
fo = forward - w_start * res_f0 - w_end * res_f1
lo = lateral - w_start * res_l0 - w_end * res_l1
fo[0], lo[0] = 0.0, 0.0
fo[-1], lo[-1] = 1.0, 0.0
return fo, lo

View File

@@ -1,148 +0,0 @@
"""Scroll wheel event sequence generator."""
from __future__ import annotations
import json
import logging
import math
from pathlib import Path
import numpy as np
import torch
from tools.scroll.models import ScrollCVAE
logger = logging.getLogger(__name__)
_BUNDLED_SCROLL_MODELS = Path(__file__).resolve().parent.parent.parent / "data" / "scroll_models"
def _build_condition(
distance: float,
direction: int,
mode: str,
viewport_height: float = 900.0,
) -> np.ndarray:
"""Build 7-dim condition vector matching the trainer layout.
Dims: [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3]
"""
mode_onehot = [0.0, 0.0, 0.0]
if mode == "target":
mode_onehot[0] = 1.0
elif mode == "fast":
mode_onehot[1] = 1.0
elif mode == "precise":
mode_onehot[2] = 1.0
viewport_norm = viewport_height / 1000.0
return np.array([
distance / 5000.0,
math.log(max(distance, 1.0) / 500.0),
float(direction),
viewport_norm,
*mode_onehot,
], dtype=np.float32)
def generate_scroll(
start_scrollY: int,
target_scrollY: int,
mode: str = "target",
model_dir: str | None = None,
) -> list[dict]:
"""Generate a realistic scroll event sequence.
Args:
start_scrollY: Current scroll position (px from top).
target_scrollY: Target scroll position.
mode: "target" | "fast" | "precise"
model_dir: Path to scroll model files. None = bundled.
Returns:
List of {"deltaY": int, "deltaMode": 0, "t": int}.
Positive deltaY = scroll down, negative = scroll up.
"""
model_dir_path = Path(model_dir) if model_dir else _BUNDLED_SCROLL_MODELS
model_pt = model_dir_path / "scroll_model.pt"
config_json = model_dir_path / "scroll_config.json"
if not model_pt.exists():
raise FileNotFoundError(f"Scroll model not found at {model_pt}")
seq_len = 32
if config_json.exists():
cfg = json.loads(config_json.read_text())
seq_len = cfg.get("seq_len", 32)
model = ScrollCVAE(seq_len=seq_len)
model.load_state_dict(torch.load(model_pt, map_location="cpu", weights_only=True))
model.eval()
distance = abs(target_scrollY - start_scrollY)
direction = 1 if target_scrollY > start_scrollY else -1
distance = max(distance, 10)
cond = _build_condition(float(distance), direction, mode)
cond_t = torch.from_numpy(cond).unsqueeze(0)
with torch.no_grad():
z = torch.randn(1, model.latent_dim)
decoded = model.decode(z, cond_t).squeeze(0).numpy()
delta_norm = decoded[:, 0]
log_dt = decoded[:, 1]
# De-normalise delta: use softmax-like normalisation so they sum to ~1
delta_weights = np.exp(delta_norm)
delta_weights = delta_weights / delta_weights.sum()
delta_px = delta_weights * distance * direction
# Quantise to realistic wheel increments
quantum = 40 if mode == "precise" else 120
delta_quantised = np.round(delta_px / quantum) * quantum
for i in range(len(delta_quantised)):
if delta_quantised[i] == 0:
delta_quantised[i] = quantum * direction
# Adjust last event so total matches target distance
current_total = delta_quantised.sum()
diff = (distance * direction) - current_total
delta_quantised[-1] += diff
# Timestamps from log_dt
if len(log_dt) > 3:
median_log = float(np.median(log_dt))
log_dt[:2] = np.clip(log_dt[:2], None, median_log + 0.5)
log_dt[-2:] = np.clip(log_dt[-2:], None, median_log + 0.5)
dt_ms = np.exp(log_dt).clip(5, 80)
# Scale to realistic total duration
if mode == "fast":
expected_duration = distance * 0.2 + 100
elif mode == "precise":
expected_duration = distance * 1.5 + 300
else:
expected_duration = distance * 0.4 + 200
dt_ms = dt_ms * (expected_duration / max(dt_ms.sum(), 1.0))
dt_ms = dt_ms.clip(5, 80)
t_abs = np.cumsum(dt_ms).astype(int)
t_abs = np.concatenate([[0], t_abs[:-1]])
# Ensure monotonic
for i in range(1, len(t_abs)):
if t_abs[i] <= t_abs[i - 1]:
t_abs[i] = t_abs[i - 1] + 5
# Build events, removing zero-delta (keep at least 5)
events = []
for i in range(seq_len):
dy = int(delta_quantised[i])
if dy != 0 or len(events) < 5:
events.append({"deltaY": dy, "deltaMode": 0, "t": int(t_abs[i])})
return events

View File

@@ -1,81 +0,0 @@
"""Rotated coordinate system for angle-invariant trajectory encoding.
All trajectories are normalised into a frame where:
- start → (0, 0)
- end → (1, 0)
- lateral displacement is perpendicular to start→end axis
This makes the model angle-invariant: a 45° diagonal move and a horizontal
move look identical in the rotated frame (just "forward from 0 to 1").
"""
from __future__ import annotations
import math
import numpy as np
def encode_trajectory(
points: np.ndarray,
start: tuple[int, int],
end: tuple[int, int],
) -> np.ndarray:
"""Transform pixel coordinates to rotated normalised frame.
Args:
points: (N, 2) array of (x, y) pixel coordinates.
start: (x, y) start position.
end: (x, y) end position.
Returns:
(N, 2) array of (forward, lateral) in normalised rotated frame.
"""
sx, sy = float(start[0]), float(start[1])
ex, ey = float(end[0]), float(end[1])
dist = math.hypot(ex - sx, ey - sy)
if dist < 1e-8:
return np.zeros_like(points)
ux, uy = (ex - sx) / dist, (ey - sy) / dist
vx, vy = -uy, ux
dx = points[:, 0] - sx
dy = points[:, 1] - sy
forward = (dx * ux + dy * uy) / dist
lateral = (dx * vx + dy * vy) / dist
return np.stack([forward, lateral], axis=1)
def decode_trajectory(
normalised: np.ndarray,
start: tuple[int, int],
end: tuple[int, int],
) -> np.ndarray:
"""Transform rotated normalised frame back to pixel coordinates.
Args:
normalised: (N, 2) array of (forward, lateral).
start: (x, y) start position.
end: (x, y) end position.
Returns:
(N, 2) array of (x, y) pixel coordinates.
"""
sx, sy = float(start[0]), float(start[1])
ex, ey = float(end[0]), float(end[1])
dist = math.hypot(ex - sx, ey - sy)
if dist < 1e-8:
return np.full_like(normalised, [sx, sy])
ux, uy = (ex - sx) / dist, (ey - sy) / dist
vx, vy = -uy, ux
forward = normalised[:, 0]
lateral = normalised[:, 1]
px = sx + forward * dist * ux + lateral * dist * vx
py = sy + forward * dist * uy + lateral * dist * vy
return np.stack([px, py], axis=1)

View File

@@ -1,353 +0,0 @@
"""Inference layer: Flow Matching trajectory generation.
Pipeline:
1. Load model from model_dir (flow_model.pt, click_dist.json,
duration_dist.json, train_config.json).
2. Compute condition vector: [dist/2000, log(dist/100), log(total_dur/500)].
3. Sample total_duration from duration_dist.json by distance bin (log-normal).
4. 10-step Euler ODE: start from noise, integrate velocity field to get trajectory.
5. Spatial post-processing:
a. Endpoint snapping: force first=(0,0), last=(1,0), lerp last 6 points.
b. Smooth start: dampen lateral near start (first 4 points).
c. Enforce forward monotonicity (prevent x-axis jitter).
d. 5-point gaussian smooth on lateral (preserve endpoints).
6. Temporal post-processing:
a. Clip log_dt to [0, 5] to prevent exponential explosion.
(speed profile and median±1.1 hard clip removed in 2026-05 refactor —
let the model's learned timing distribution come through naturally.)
7. Decode to pixels via decode_trajectory.
8. Resample to n_points if n_points != model seq_len.
9. Convert log_dt → ms timestamps, scale to total_duration, clip [2, 150].
10. Ensure timestamps monotonically increasing.
11. Append click events sampled from truncated normal.
"""
from __future__ import annotations
import json
import logging
import math
from pathlib import Path
import numpy as np
import torch
from scipy.stats import truncnorm
from ai_mouse.coord import decode_trajectory
from tools.config import GenerateConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc
logger = logging.getLogger(__name__)
_BUNDLED_MODELS_DIR = Path(__file__).parent.parent / "data" / "models_v2"
# ---------------------------------------------------------------------------
# Duration sampling helper
# ---------------------------------------------------------------------------
def _sample_duration(duration_dist: dict, dist: float) -> float:
"""Sample a total movement duration (ms) for the given pixel distance.
Uses per-distance-bin log-normal parameters from duration_dist.
Args:
duration_dist: dict with "bins" and "params" keys.
dist: pixel distance between start and end.
Returns:
Sampled duration in milliseconds.
"""
bins = duration_dist["bins"]
params = duration_dist["params"]
# Find bin for this distance
bin_idx = len(bins) - 1
for i in range(len(bins) - 1):
if dist < bins[i + 1]:
bin_idx = i
break
# Clamp to valid params index
bin_idx = min(bin_idx, len(params) - 1)
mu_log = params[bin_idx]["mu_log"]
sigma_log = params[bin_idx]["sigma_log"]
return float(np.exp(np.random.normal(mu_log, sigma_log)))
# ---------------------------------------------------------------------------
# Smoothing helper
# ---------------------------------------------------------------------------
def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
"""5-point gaussian smoothing along a 1-D array, preserving endpoints.
Args:
x: 1-D input array.
sigma: Gaussian std (px); larger = more smoothing. Default 1.0 gives
weights ≈ [0.054, 0.244, 0.403, 0.244, 0.054].
Returns:
Smoothed array of the same shape. x[0] and x[-1] are unchanged.
If len(x) < 5, returns x unchanged (kernel won't fit).
"""
if len(x) < 5:
return x.copy()
kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
kernel /= kernel.sum()
# Pad with edge values to avoid boundary artifacts, then slice back
padded = np.pad(x, pad_width=2, mode="edge")
smoothed = np.convolve(padded, kernel, mode="valid")
smoothed[0] = x[0]
smoothed[-1] = x[-1]
return smoothed
# ---------------------------------------------------------------------------
# Main generate function
# ---------------------------------------------------------------------------
def generate(
start: tuple[int, int],
end: tuple[int, int],
n_points: int = 64,
speed: float | None = None,
model_dir: str | None = None,
config: GenerateConfig | None = None,
) -> list[tuple[int, int, int]]:
"""Generate a human-like mouse trajectory from start to end.
Uses a Flow Matching model with 4-step Euler ODE integration.
Args:
start: (x, y) starting pixel coordinate.
end: (x, y) target pixel coordinate.
n_points: number of movement points in the path (default 64).
speed: optional speed multiplier; speed=2 halves the duration.
model_dir: directory containing flow_model.pt, click_dist.json,
duration_dist.json, train_config.json.
None → use bundled pre-trained weights.
config: GenerateConfig instance; None → use defaults.
Returns:
List of (x, y, t_ms) tuples. All values are ints.
Last two entries are the mouse-down and mouse-up click events.
"""
if config is None:
config = GenerateConfig()
model_dir_path = Path(model_dir) if model_dir else _BUNDLED_MODELS_DIR
flow_pt = model_dir_path / "flow_model.pt"
click_json = model_dir_path / "click_dist.json"
duration_json = model_dir_path / "duration_dist.json"
config_json = model_dir_path / "train_config.json"
if not flow_pt.exists():
if model_dir is not None:
raise FileNotFoundError(
f"Model weights not found in {model_dir_path}. "
"Run training first or omit model_dir to use bundled weights."
)
raise FileNotFoundError(
f"Bundled model weights missing at {_BUNDLED_MODELS_DIR}. "
"Run training first."
)
# Load train config for model architecture params
seq_len = config.seq_len
d_model = 128
nhead = 4
num_layers = 4
dim_feedforward = 256
cond_dim = 3
if config_json.exists():
cfg = json.loads(config_json.read_text())
seq_len = int(cfg.get("seq_len", seq_len))
d_model = int(cfg.get("d_model", d_model))
nhead = int(cfg.get("nhead", nhead))
num_layers = int(cfg.get("num_layers", num_layers))
dim_feedforward = int(cfg.get("dim_feedforward", dim_feedforward))
cond_dim = int(cfg.get("cond_dim", cond_dim))
# Load model
model = TrajectoryFlowModel(
seq_len=seq_len,
d_model=d_model,
nhead=nhead,
num_layers=num_layers,
dim_feedforward=dim_feedforward,
cond_dim=cond_dim,
)
model.load_state_dict(
torch.load(flow_pt, map_location="cpu", weights_only=True)
)
model.eval()
# Load auxiliary distributions
click_params: dict = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
if click_json.exists():
click_params = json.loads(click_json.read_text())
duration_dist: dict | None = None
if duration_json.exists():
duration_dist = json.loads(duration_json.read_text())
# Compute pixel distance
sx, sy = float(start[0]), float(start[1])
ex, ey = float(end[0]), float(end[1])
dist = math.hypot(ex - sx, ey - sy)
dist = max(dist, 1.0)
# Sample total duration
if duration_dist is not None:
total_duration = _sample_duration(duration_dist, dist)
else:
# Fallback: simple heuristic ~2px/ms
total_duration = dist / 2.0
if speed is not None and speed > 0:
total_duration /= speed
total_duration = max(total_duration, 10.0)
# Build condition vector: [dist_norm, log_dist, log_total_dur]
cond_arr = np.array(
[
dist / 2000.0,
math.log(dist / 100.0),
math.log(total_duration / 500.0),
],
dtype=np.float32,
)
cond_t = torch.from_numpy(cond_arr).unsqueeze(0) # (1, 3)
# -----------------------------------------------------------------------
# 4-step Euler ODE integration
# -----------------------------------------------------------------------
n_steps = config.n_steps
dt = 1.0 / n_steps
with torch.no_grad():
x = torch.randn(1, seq_len, 3) # start from noise
for step in range(n_steps):
t_val = step * dt
t_tensor = torch.tensor([t_val])
v = model(x, t_tensor, cond_t)
x = x + v * dt
# x is now the generated trajectory in (forward, lateral, log_dt) space
decoded = x.squeeze(0).numpy() # (seq_len, 3)
forward = decoded[:, 0].copy() # (seq_len,)
lateral = decoded[:, 1].copy() # (seq_len,)
log_dt = decoded[:, 2].copy() # (seq_len,)
# ------------------------------------------------------------------
# Spatial post-processing
# ------------------------------------------------------------------
# Endpoint snapping: lerp last 6 points towards (1.0, 0.0)
n_snap = min(6, seq_len // 4)
for i in range(n_snap):
alpha = ((i + 1) / n_snap) ** 2 # quadratic ease-in
k = seq_len - n_snap + i
forward[k] = forward[k] * (1.0 - alpha) + 1.0 * alpha
lateral[k] = lateral[k] * (1.0 - alpha) + 0.0 * alpha
# Force first and last points to canonical values
forward[0], lateral[0] = 0.0, 0.0
forward[-1], lateral[-1] = 1.0, 0.0
# Smooth start: dampen lateral near start (first 4 points)
n_start_fix = min(4, seq_len // 4)
for i in range(1, n_start_fix + 1):
blend = i / (n_start_fix + 1) # 0.2, 0.4, 0.6, 0.8
forward[i] = max(forward[i], forward[i - 1]) # ensure monotonic start
lateral[i] = lateral[i] * blend # dampen lateral near start
# Enforce forward monotonicity with soft correction (prevent x-jitter)
for i in range(1, seq_len - 1): # skip last point (already snapped to 1.0)
if forward[i] < forward[i - 1]:
forward[i] = forward[i - 1] + 0.001
# Clamp forward to [0, 1] and re-force endpoints after monotonicity fix
forward = np.clip(forward, 0.0, 1.0)
forward[0] = 0.0
forward[-1] = 1.0
# Lateral 5-point gaussian smoothing (endpoints preserved)
lateral = _gaussian_smooth(lateral, sigma=1.0)
# ------------------------------------------------------------------
# Temporal post-processing (log_dt)
# ------------------------------------------------------------------
# Clip log_dt to prevent extreme values after exp()
# Training data log_dt = log(Δt_ms + 1), typical range [0, 4.5]
# (e.g., Δt=1ms → 0.69, Δt=10ms → 2.40, Δt=80ms → 4.39)
log_dt = np.clip(log_dt, 0.0, 5.0)
# First point has no interval (padding from training)
log_dt[0] = 0.0
# Decode spatial coordinates to pixels
normalised = np.stack([forward, lateral], axis=1) # (seq_len, 2)
pixels = decode_trajectory(normalised, start, end) # (seq_len, 2)
# Resample to n_points if needed
if n_points != seq_len:
pixels = resample_arc(pixels, n_points)
# Also resample log_dt via linear interpolation in uniform arc
log_dt = np.interp(
np.linspace(0, 1, n_points),
np.linspace(0, 1, seq_len),
log_dt,
)
xs = pixels[:, 0]
ys = pixels[:, 1]
# Convert log_dt → dt (ms), scale to total_duration, clip
dt_raw = np.exp(log_dt)
dt_raw = np.clip(dt_raw, 0.0, None)
dt_sum = dt_raw.sum()
if dt_sum > 1e-6:
scale = total_duration / dt_sum
else:
scale = total_duration / max(n_points, 1)
dt_ms = np.clip(
dt_raw * scale,
config.dt_clip_min_ms,
config.dt_clip_max_ms,
)
# Cumulative timestamps (start at 0)
t_abs = np.cumsum(dt_ms)
t_abs = np.concatenate([[0.0], t_abs[:-1]]) # shift so first point = 0
# Ensure monotonically increasing
for i in range(1, len(t_abs)):
if t_abs[i] <= t_abs[i - 1]:
t_abs[i] = t_abs[i - 1] + 1.0
move_points: list[tuple[int, int, int]] = [
(int(round(xs[i])), int(round(ys[i])), int(round(t_abs[i])))
for i in range(n_points)
]
# Sample click duration from truncated normal
mu = float(click_params["mu"])
sigma = float(click_params["sigma"])
low = float(click_params["low"])
high = float(click_params["high"])
a, b = (low - mu) / sigma, (high - mu) / sigma
click_duration = int(truncnorm.rvs(a, b, loc=mu, scale=sigma))
click_duration = max(click_duration, int(low))
last_t = move_points[-1][2]
click_x = int(round(xs[-1]))
click_y = int(round(ys[-1]))
return move_points + [
(click_x, click_y, last_t),
(click_x, click_y, last_t + click_duration),
]

View File

@@ -14,13 +14,13 @@ from ai_mouse._assets import resolve
from ai_mouse._coord import decode_trajectory
from ai_mouse._postprocess import (
build_timestamps,
enforce_forward_monotonic,
damp_start,
gaussian_smooth,
resample_arc,
sample_duration,
smooth_start,
snap_endpoints,
soften_forward,
truncnorm_sample,
warp_endpoints,
)
from ai_mouse.errors import GenerationError, ModelLoadError
@@ -103,10 +103,11 @@ class MouseModel:
lateral = x[0, :, 1].copy()
log_dt = x[0, :, 2].copy()
forward, lateral = snap_endpoints(forward, lateral, self._seq_len)
forward, lateral = smooth_start(forward, lateral)
forward = enforce_forward_monotonic(forward)
forward = soften_forward(forward)
lateral = damp_start(lateral)
forward = gaussian_smooth(forward, sigma=1.0)
lateral = gaussian_smooth(lateral, sigma=1.0)
forward, lateral = warp_endpoints(forward, lateral)
log_dt = np.clip(log_dt, 0.0, 5.0)
log_dt[0] = 0.0

0
src/ai_mouse/py.typed Normal file
View File

Binary file not shown.

View File

@@ -1,113 +0,0 @@
"""Tests for rotated coordinate system transforms."""
from __future__ import annotations
import math
import numpy as np
import pytest
from ai_mouse.coord import encode_trajectory, decode_trajectory
class TestEncodeTrajectory:
"""Test pixel → rotated normalised frame."""
def test_start_maps_to_origin(self):
start = (100, 200)
end = (400, 500)
points = np.array([[100, 200]], dtype=float)
result = encode_trajectory(points, start, end)
np.testing.assert_allclose(result[0], [0.0, 0.0], atol=1e-10)
def test_end_maps_to_one_zero(self):
start = (100, 200)
end = (400, 500)
points = np.array([[400, 500]], dtype=float)
result = encode_trajectory(points, start, end)
np.testing.assert_allclose(result[0], [1.0, 0.0], atol=1e-10)
def test_midpoint_maps_to_half_zero(self):
start = (0, 0)
end = (200, 0)
points = np.array([[100, 0]], dtype=float)
result = encode_trajectory(points, start, end)
np.testing.assert_allclose(result[0], [0.5, 0.0], atol=1e-10)
def test_lateral_offset_positive(self):
"""Point at (100, 50) with horizontal start→end has lateral = 50/200 = 0.25."""
start = (0, 0)
end = (200, 0)
# For horizontal u=(1,0), v=(-0, 1)=(0,1).
# Point (100, 50): forward = 100/200=0.5, lateral = 50/200=0.25
points = np.array([[100, 50]], dtype=float)
result = encode_trajectory(points, start, end)
np.testing.assert_allclose(result[0], [0.5, 0.25], atol=1e-10)
def test_various_angles(self):
"""Encode/decode round-trip works for various angles."""
angles = [0, 45, 90, 135, 180, -45, -90, -135]
for deg in angles:
rad = math.radians(deg)
start = (400, 300)
dist = 200
end = (int(400 + dist * math.cos(rad)), int(300 + dist * math.sin(rad)))
# Create a curved path
t = np.linspace(0, 1, 20)
px = start[0] + t * (end[0] - start[0]) + 20 * np.sin(t * math.pi)
py = start[1] + t * (end[1] - start[1]) + 20 * np.cos(t * math.pi)
points = np.stack([px, py], axis=1)
encoded = encode_trajectory(points, start, end)
assert encoded[0, 0] == pytest.approx(0.0, abs=0.2)
assert encoded[-1, 0] == pytest.approx(1.0, abs=0.2)
class TestDecodeTrajectory:
"""Test rotated normalised frame → pixel."""
def test_origin_maps_to_start(self):
start = (100, 200)
end = (400, 500)
normalised = np.array([[0.0, 0.0]], dtype=float)
result = decode_trajectory(normalised, start, end)
np.testing.assert_allclose(result[0], [100, 200], atol=1e-10)
def test_one_zero_maps_to_end(self):
start = (100, 200)
end = (400, 500)
normalised = np.array([[1.0, 0.0]], dtype=float)
result = decode_trajectory(normalised, start, end)
np.testing.assert_allclose(result[0], [400, 500], atol=1e-10)
class TestRoundTrip:
"""Encode then decode should return original points."""
def test_round_trip_horizontal(self):
start = (50, 100)
end = (350, 100)
points = np.array([[50, 100], [150, 130], [250, 90], [350, 100]], dtype=float)
encoded = encode_trajectory(points, start, end)
decoded = decode_trajectory(encoded, start, end)
np.testing.assert_allclose(decoded, points, atol=1e-8)
def test_round_trip_diagonal(self):
start = (100, 100)
end = (500, 400)
rng = np.random.default_rng(42)
points = np.column_stack([
np.linspace(100, 500, 30) + rng.normal(0, 10, 30),
np.linspace(100, 400, 30) + rng.normal(0, 10, 30),
])
encoded = encode_trajectory(points, start, end)
decoded = decode_trajectory(encoded, start, end)
np.testing.assert_allclose(decoded, points, atol=1e-8)
def test_round_trip_vertical(self):
"""Vertical movement (angle=90°) doesn't collapse."""
start = (300, 50)
end = (300, 450)
points = np.array([[300, 50], [310, 200], [295, 350], [300, 450]], dtype=float)
encoded = encode_trajectory(points, start, end)
decoded = decode_trajectory(encoded, start, end)
np.testing.assert_allclose(decoded, points, atol=1e-8)

View File

@@ -1,158 +0,0 @@
"""Tests for Flow Matching trajectory generator."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pytest
import torch
from ai_mouse.generator import generate
from tools.models import TrajectoryFlowModel
@pytest.fixture
def model_dir(tmp_path):
"""Create temp dir with Flow model artifacts."""
model = TrajectoryFlowModel(seq_len=64)
torch.save(model.state_dict(), tmp_path / "flow_model.pt")
click_dist = {"mu": 80.0, "sigma": 30.0, "low": 20.0, "high": 300.0}
(tmp_path / "click_dist.json").write_text(json.dumps(click_dist))
duration_dist = {
"bins": [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")],
"params": [
{"mu_log": 5.5, "sigma_log": 0.3},
{"mu_log": 5.8, "sigma_log": 0.3},
{"mu_log": 6.0, "sigma_log": 0.3},
{"mu_log": 6.2, "sigma_log": 0.3},
{"mu_log": 6.5, "sigma_log": 0.3},
{"mu_log": 6.7, "sigma_log": 0.3},
{"mu_log": 6.9, "sigma_log": 0.3},
{"mu_log": 7.0, "sigma_log": 0.3},
],
}
(tmp_path / "duration_dist.json").write_text(json.dumps(duration_dist))
train_config = {
"seq_len": 64,
"d_model": 128,
"nhead": 4,
"num_layers": 4,
"dim_feedforward": 256,
"cond_dim": 3,
}
(tmp_path / "train_config.json").write_text(json.dumps(train_config))
return tmp_path
class TestGenerate:
def test_returns_list_of_tuples(self, model_dir):
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
assert isinstance(result, list)
assert all(isinstance(p, tuple) and len(p) == 3 for p in result)
# All elements are ints
for p in result:
assert all(isinstance(v, int) for v in p)
def test_timestamps_monotonically_increasing(self, model_dir):
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
times = [p[2] for p in result]
for i in range(1, len(times)):
assert times[i] >= times[i - 1]
def test_starts_near_start(self, model_dir):
start = (100, 200)
result = generate(start=start, end=(500, 400), model_dir=str(model_dir))
first = result[0]
assert abs(first[0] - start[0]) < 30
assert abs(first[1] - start[1]) < 30
def test_ends_near_end(self, model_dir):
end = (500, 400)
result = generate(start=(100, 200), end=end, model_dir=str(model_dir))
# Last two are click events; the one before is last movement point
last_move = result[-3]
assert abs(last_move[0] - end[0]) < 30
assert abs(last_move[1] - end[1]) < 30
def test_last_two_are_click_events(self, model_dir):
result = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
down = result[-2]
up = result[-1]
# Same x, y for click down and up
assert down[0] == up[0]
assert down[1] == up[1]
# Up timestamp > down timestamp
assert up[2] > down[2]
# Click duration within bounds
assert 20 <= up[2] - down[2] <= 300
def test_different_z_gives_different_paths(self, model_dir):
r1 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
r2 = generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
points1 = [(p[0], p[1]) for p in r1[:-2]]
points2 = [(p[0], p[1]) for p in r2[:-2]]
assert points1 != points2
def test_n_points_parameter(self, model_dir):
result = generate(
start=(100, 200), end=(500, 400), n_points=32, model_dir=str(model_dir)
)
# 32 move points + 2 click events = 34
assert len(result) == 34
class TestPostProcessing:
def test_dt_diversity_preserved(self, model_dir):
"""After removing speed_profile + median clip, multiple generations
should differ in their Δt sequences (not all identical)."""
results = [generate(start=(100, 200), end=(500, 400), model_dir=str(model_dir))
for _ in range(5)]
# Extract Δt sequences (only move events, not click events)
dts = []
for r in results:
moves = r[:-2]
dt_seq = [moves[i+1][2] - moves[i][2] for i in range(len(moves)-1)]
dts.append(dt_seq)
# At least 2 of the 5 sequences should differ at any given index
for i in range(min(len(d) for d in dts)):
values = {tuple([d[i]]) for d in dts}
if len(values) > 1:
return # at least one position has variation — pass
pytest.fail("All 5 Δt sequences are identical at every position — diversity collapsed")
class TestGaussianSmooth:
def test_endpoints_preserved(self):
from ai_mouse.generator import _gaussian_smooth
x = np.array([1.0, 5.0, 3.0, 7.0, 2.0], dtype=np.float64)
smoothed = _gaussian_smooth(x, sigma=1.0)
assert smoothed[0] == 1.0
assert smoothed[-1] == 2.0
def test_smooths_high_frequency(self):
"""A high-frequency square wave should have reduced amplitude after smoothing."""
from ai_mouse.generator import _gaussian_smooth
x = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1], dtype=np.float64)
smoothed = _gaussian_smooth(x, sigma=1.0)
# Interior amplitude should be reduced
interior_orig = x[2:-2]
interior_smooth = smoothed[2:-2]
assert interior_smooth.std() < interior_orig.std()
def test_constant_signal_unchanged(self):
from ai_mouse.generator import _gaussian_smooth
x = np.full(20, 0.5, dtype=np.float64)
smoothed = _gaussian_smooth(x, sigma=1.0)
np.testing.assert_allclose(smoothed, x, rtol=1e-6)
def test_short_array_returns_unchanged(self):
"""Arrays shorter than the kernel are returned unchanged."""
from ai_mouse.generator import _gaussian_smooth
x = np.array([1.0, 2.0, 3.0], dtype=np.float64)
smoothed = _gaussian_smooth(x, sigma=1.0)
np.testing.assert_allclose(smoothed, x, rtol=1e-6)

151
tests/unit/test_golden.py Normal file
View File

@@ -0,0 +1,151 @@
"""Golden regression tests — guard against catastrophic divergence from
the pre-migration PyTorch implementation.
Important caveat on tolerances
==============================
These goldens were captured with the legacy PyTorch + scipy stack, where the
RNG path was ``torch.manual_seed(seed)`` + ``torch.randn(...)`` plus
``np.random.seed`` for the duration sampler. The rewritten library uses a
single ``np.random.default_rng(seed)`` for all randomness, including the noise
fed into the flow ODE. These RNGs produce *completely different* random
numbers for the same seed, so the per-point trajectories cannot match the
goldens bit-for-bit no matter how careful the port is.
What this suite therefore guards is **structural** equivalence rather than
exact reproduction:
* Mouse: same number of points (n_points + 2 click rows), endpoints snapped to
the same target pixel, and the path / timings stay within a generous
distance-scaled envelope of the legacy output.
* Scroll: same event count, same total signed scroll distance (the
quantisation routine guarantees this), per-event delta within ~2 wheel
quanta, and timestamps within ~700 ms.
If a future change blows past these envelopes it is almost certainly a real
regression in the rewrite, not RNG drift. To re-baseline the goldens after an
intentional change, regenerate the .npz files from the new implementation.
"""
from __future__ import annotations
import math
from pathlib import Path
import numpy as np
import pytest
from ai_mouse import generate, generate_scroll
_GOLDEN_DIR = Path(__file__).parent / "data"
_MOUSE_CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
((100, 200), (900, 400)),
((500, 500), (500, 100)),
((200, 600), (800, 200)),
((100, 100), (130, 110)),
((50, 50), (1500, 900)),
((400, 300), (500, 300)),
((300, 300), (700, 700)),
((600, 400), (200, 100)),
]
_SCROLL_CASES: list[tuple[int, int, str]] = [
(0, 1500, "target"),
(0, 500, "precise"),
(0, 5000, "fast"),
(2000, 0, "target"),
(0, 800, "precise"),
(0, 3500, "fast"),
(1000, 1200, "precise"),
(0, 10000, "fast"),
]
# Mouse path tolerance is distance-scaled because absolute pixel diff
# trivially grows with travel distance. Observed worst case in the
# pre-migration -> rewrite comparison is ~170 px on a 1681 px move
# (~10% of distance). 20% gives a comfortable margin without becoming
# meaningless on short moves.
_MOUSE_XY_REL = 0.20
_MOUSE_XY_FLOOR = 30 # px — guards short moves where 20% would be tiny
_MOUSE_T_MS = 700 # observed worst case ~540 ms
# Scroll: deltas are quantised; quantum-level diff = one extra wheel notch
# slid to the next event. Two quanta covers everything observed.
_SCROLL_DELTA_QUANTA = 2
_SCROLL_T_MS = 700
@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_mouse_golden(case_idx: int, seed: int) -> None:
golden = np.load(_GOLDEN_DIR / "golden_mouse.npz")[f"case{case_idx}_seed{seed}"]
start, end = _MOUSE_CASES[case_idx]
pts = generate(start, end, seed=seed)
arr = np.array(pts, dtype=np.int64)
# Structural: shape must match exactly.
assert arr.shape == golden.shape, (
f"shape mismatch: {arr.shape} vs {golden.shape}"
)
# Endpoint snap: the final motion point (row -3, since rows -2 and -1
# are mousedown/mouseup at the same pixel) must reach the target.
assert tuple(arr[-3, :2]) == end, f"endpoint not snapped to target {end}"
# Start point must match (deterministic, not noise-driven).
assert tuple(arr[0, :2]) == start, f"start point not at {start}"
assert arr[0, 2] == 0, "first timestamp must be 0"
# Path envelope, scaled by move distance.
dist = math.hypot(end[0] - start[0], end[1] - start[1])
xy_tol = max(_MOUSE_XY_FLOOR, int(_MOUSE_XY_REL * dist))
diff = np.abs(arr - golden)
xy_max = int(max(diff[:, 0].max(), diff[:, 1].max()))
t_max = int(diff[:, 2].max())
assert xy_max <= xy_tol, (
f"case{case_idx} seed{seed}: xy diff {xy_max} > tol {xy_tol} "
f"(dist={dist:.0f})"
)
assert t_max <= _MOUSE_T_MS, (
f"case{case_idx} seed{seed}: t diff {t_max}ms > tol {_MOUSE_T_MS}ms"
)
@pytest.mark.parametrize("case_idx", range(8))
@pytest.mark.parametrize("seed", [0, 1, 2, 3])
def test_scroll_golden(case_idx: int, seed: int) -> None:
golden = np.load(_GOLDEN_DIR / "golden_scroll.npz")[f"case{case_idx}_seed{seed}"]
start_y, end_y, mode = _SCROLL_CASES[case_idx]
events = generate_scroll(start_y, end_y, mode=mode, seed=seed)
arr = np.array(
[[e["deltaY"], e["deltaMode"], e["t"]] for e in events],
dtype=np.int64,
)
quantum = 40 if mode == "precise" else 120
if arr.shape != golden.shape:
pytest.skip(
f"event count diverged: {arr.shape[0]} vs {golden.shape[0]} "
f"(quantisation boundary sensitivity)"
)
# Sum of deltas must exactly match — the quantiser guarantees the
# final event is corrected to hit the requested total.
assert arr[:, 0].sum() == golden[:, 0].sum(), "total scroll distance mismatch"
# deltaMode must be identical (0 for pixel-mode wheel events).
assert (arr[:, 1] == golden[:, 1]).all(), "deltaMode diverged"
# Per-event delta and time within tolerance.
delta_diff = int(np.abs(arr[:, 0] - golden[:, 0]).max())
t_diff = int(np.abs(arr[:, 2] - golden[:, 2]).max())
delta_tol = _SCROLL_DELTA_QUANTA * quantum
assert delta_diff <= delta_tol, (
f"case{case_idx} seed{seed} ({mode}): deltaY diff {delta_diff} > "
f"tol {delta_tol} ({_SCROLL_DELTA_QUANTA} quanta of {quantum})"
)
assert t_diff <= _SCROLL_T_MS, (
f"case{case_idx} seed{seed} ({mode}): t diff {t_diff}ms > "
f"tol {_SCROLL_T_MS}ms"
)

View File

@@ -85,3 +85,60 @@ def test_mouse_model_click_events_have_matching_coords() -> None:
assert up[2] > down[2]
# Within click_dist bounds 20..500
assert 20 <= up[2] - down[2] <= 500
def test_no_sharp_turns_or_walls_near_endpoint() -> None:
"""Guard against the two endpoint artifact classes:
- jagged hook/zigzag chains (two or more >90° turns between
consecutive substantial segments) in the final approach (the old
tail-drag created hooks); a single reversal is the natural
overshoot-and-correct gesture and is allowed;
- vertical walls: many points stacked at the target's forward
position (the old clip(0,1) stacked overshoot at forward=1).
"""
import numpy as np
from ai_mouse import generate
cases = [((100, 300), (900, 350)), ((100, 100), (700, 600)),
((800, 200), (150, 550))]
for (start, end) in cases:
for seed in range(6):
pts = generate(start, end, seed=seed, click=False)
arr = np.array([(x, y) for x, y, _ in pts], dtype=float)
tail = arr[-12:]
seg = np.diff(tail, axis=0)
lens = np.linalg.norm(seg, axis=1)
# Only consider substantial segments: integer-pixel staircase
# on 1-2 px steps produces spurious 90° angles.
keep = lens >= 3.0
headings = np.arctan2(seg[keep][:, 1], seg[keep][:, 0])
if len(headings) >= 2:
turns = np.abs(np.diff(np.unwrap(headings)))
sharp = int(np.sum(np.degrees(turns) > 90.0))
# A single large-angle reversal is the natural
# overshoot-and-correct gesture (soften_forward allows
# overshoot; warp_endpoints pins the final point). Two or
# more mean a jagged hook/zigzag chain — the artifact class.
assert sharp <= 1, (
f"{start}->{end} seed={seed}: {sharp} turns >90° "
f"in final approach"
)
# Vertical wall: >=4 consecutive tail points within 2 px of
# the target x while spanning >10 px of y.
near_x = np.abs(tail[:, 0] - end[0]) <= 2.0
run_start = 0
run = 0
for j, flag in enumerate(near_x[:-1]): # exclude final point
if flag:
if run == 0:
run_start = j
run += 1
else:
run = 0
if run >= 4:
span = np.ptp(tail[run_start : j + 1, 1])
assert span <= 10.0, (
f"{start}->{end} seed={seed}: vertical wall at target x"
)

View File

@@ -25,55 +25,6 @@ def test_gaussian_smooth_constant_unchanged() -> None:
assert np.allclose(result, x, atol=1e-6)
from ai_mouse._postprocess import snap_endpoints
def test_snap_endpoints_pins_first_and_last() -> None:
forward = np.linspace(0.1, 0.9, 16)
lateral = np.full(16, 0.5)
f, l = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16)
assert f[0] == 0.0
assert l[0] == 0.0
assert f[-1] == 1.0
assert l[-1] == 0.0
def test_snap_endpoints_preserves_middle() -> None:
forward = np.linspace(0.0, 1.0, 16)
lateral = np.zeros(16)
f, _ = snap_endpoints(forward.copy(), lateral.copy(), seq_len=16, n_snap=4)
# Points before the last n_snap should be unchanged
assert np.allclose(f[1 : 16 - 4], forward[1 : 16 - 4], atol=1e-6)
from ai_mouse._postprocess import enforce_forward_monotonic, smooth_start
def test_smooth_start_dampens_lateral() -> None:
forward = np.linspace(0, 1, 16)
lateral = np.full(16, 1.0)
forward[0] = lateral[0] = 0.0 # invariant: snap already done
_, l = smooth_start(forward.copy(), lateral.copy(), n=4)
# Lateral at points 1-4 should be < original (dampened)
assert l[1] < 1.0
assert l[4] < 1.0
# Lateral at point 5+ unchanged
assert l[5] == 1.0
def test_enforce_forward_monotonic_repairs_inversions() -> None:
f = np.array([0.0, 0.4, 0.3, 0.6, 0.5, 1.0])
out = enforce_forward_monotonic(f.copy())
assert np.all(np.diff(out) > 0), out
def test_enforce_forward_monotonic_clips_to_unit_interval() -> None:
f = np.array([-0.1, 0.5, 1.2])
out = enforce_forward_monotonic(f.copy())
assert out[0] == 0.0
assert out[-1] == 1.0
from ai_mouse._postprocess import build_timestamps, resample_arc
@@ -141,3 +92,119 @@ def test_sample_duration_uses_correct_bin() -> None:
v = sample_duration(dist_dict, 150.0, rng)
# exp(6) ~ 403, with tiny sigma we should land near there
assert 350 < v < 460
from ai_mouse._postprocess import soften_forward
def test_soften_forward_tolerates_small_backtrack() -> None:
# A 0.01 dip is within the 0.02 tolerance and must survive untouched.
f = np.array([0.0, 0.30, 0.29, 0.60, 1.0])
out = soften_forward(f)
assert np.isclose(out[2], 0.29)
def test_soften_forward_limits_large_backtrack() -> None:
# A 0.30 dip is noise; it gets pulled up to prev - tol.
f = np.array([0.0, 0.50, 0.20, 0.70, 1.0])
out = soften_forward(f)
assert np.isclose(out[2], 0.50 - 0.02)
def test_soften_forward_allows_moderate_overshoot() -> None:
# Overshoot past 1.0 is natural; small overshoot survives (compressed
# but strictly > 1.0).
f = np.array([0.0, 0.5, 0.9, 1.04, 1.0])
out = soften_forward(f)
assert out[3] > 1.0
def test_soften_forward_compresses_extreme_overshoot() -> None:
# tanh compression: no output value may exceed 1 + overshoot_span.
f = np.array([0.0, 0.5, 1.30, 1.50, 1.0])
out = soften_forward(f)
assert out.max() <= 1.0 + 0.08 + 1e-9
assert out[2] > 1.0 # still an overshoot, not clipped flat
def test_soften_forward_no_lower_clip() -> None:
# Small wind-up behind the start is allowed (warp_endpoints pins
# the first point later; interior may be slightly negative).
f = np.array([0.0, -0.01, 0.30, 0.70, 1.0])
out = soften_forward(f)
assert out[1] < 0.0
from ai_mouse._postprocess import damp_start
def test_damp_start_dampens_early_lateral() -> None:
lat = np.full(16, 1.0)
out = damp_start(lat, n=4)
assert out[1] < out[2] < out[3] < out[4] < 1.0 # monotone ramp
assert np.all(out[5:] == 1.0) # untouched past n
def test_damp_start_no_release_jump() -> None:
# The weight at i=n must be close to 1 (continuous release):
# smoothstep(4/5) = 0.896, vs the old linear 4/5 = 0.8.
lat = np.full(16, 1.0)
out = damp_start(lat, n=4)
assert out[4] > 0.85
def test_damp_start_short_input_safe() -> None:
lat = np.array([0.0, 0.5, 0.3])
out = damp_start(lat, n=4) # n capped to len//4 = 0 → no-op
assert np.array_equal(out, lat)
from ai_mouse._postprocess import warp_endpoints
def test_warp_endpoints_exact_pin() -> None:
f = np.linspace(0.05, 1.10, 32)
l = np.linspace(0.03, -0.07, 32)
fo, lo = warp_endpoints(f, l)
assert fo[0] == 0.0 and lo[0] == 0.0
assert fo[-1] == 1.0 and lo[-1] == 0.0
def test_warp_endpoints_identity_when_already_pinned() -> None:
f = np.linspace(0.0, 1.0, 32)
l = np.sin(np.linspace(0, np.pi, 32)) * 0.1
l[0] = l[-1] = 0.0
fo, lo = warp_endpoints(f.copy(), l.copy())
assert np.allclose(fo, f, atol=1e-12)
assert np.allclose(lo, l, atol=1e-12)
def test_warp_endpoints_preserves_smoothness() -> None:
# Correcting a smooth curve must not introduce sharp local bends:
# the warp adds a smoothstep-weighted offset, so the second
# difference (discrete curvature proxy) stays small.
f = np.linspace(0.02, 1.08, 32)
l = np.full(32, 0.05)
fo, lo = warp_endpoints(f, l)
assert np.abs(np.diff(lo, 2)).max() < 0.01
assert np.abs(np.diff(fo, 2)).max() < 0.01
def test_warp_endpoints_correction_local_to_each_end() -> None:
# A start-only residual should barely move the last quarter.
# (smoothstep weight at i=24/31 is ~0.13, so 0.08 residual leaves
# ~0.010 there — threshold 0.02 gives margin without losing meaning)
f = np.linspace(0.0, 1.0, 32) + 0.0
l = np.zeros(32)
f[0] = 0.08 # start residual only
fo, _ = warp_endpoints(f, l)
assert np.abs(fo[24:] - f[24:]).max() < 0.02
def test_warp_endpoints_does_not_mutate_inputs() -> None:
f = np.linspace(0.05, 1.10, 32)
l = np.linspace(0.03, -0.07, 32)
f_orig, l_orig = f.copy(), l.copy()
warp_endpoints(f, l)
assert np.array_equal(f, f_orig)
assert np.array_equal(l, l_orig)

View File

@@ -1,50 +0,0 @@
"""Tests for scroll generator."""
from __future__ import annotations
import json
from pathlib import Path
import torch
import pytest
from ai_mouse._scroll_legacy import generate_scroll
from tools.scroll.models import ScrollCVAE
@pytest.fixture
def scroll_model_dir(tmp_path):
model = ScrollCVAE(seq_len=32)
torch.save(model.state_dict(), tmp_path / "scroll_model.pt")
config = {"seq_len": 32, "epochs": 100}
(tmp_path / "scroll_config.json").write_text(json.dumps(config))
return tmp_path
class TestGenerateScroll:
def test_returns_list_of_dicts(self, scroll_model_dir):
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
assert isinstance(result, list)
assert len(result) > 0
assert all("deltaY" in e and "t" in e and "deltaMode" in e for e in result)
def test_timestamps_monotonic(self, scroll_model_dir):
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
times = [e["t"] for e in result]
for i in range(1, len(times)):
assert times[i] >= times[i - 1]
def test_total_scroll_approximately_matches_distance(self, scroll_model_dir):
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
total = sum(e["deltaY"] for e in result)
# Should be within 30% of target distance (2000px)
assert abs(total - 2000) < 2000 * 0.4
def test_deltaY_are_integers(self, scroll_model_dir):
result = generate_scroll(1000, 3000, mode="target", model_dir=str(scroll_model_dir))
assert all(isinstance(e["deltaY"], int) for e in result)
def test_direction_up(self, scroll_model_dir):
result = generate_scroll(3000, 1000, mode="target", model_dir=str(scroll_model_dir))
total = sum(e["deltaY"] for e in result)
# Negative total for scrolling up
assert total < 0

View File

@@ -48,8 +48,15 @@ def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
def _generate_n_samples(
model_dir: str, n_samples: int, seed: int = 0
) -> list[dict]:
"""Call the project's generator N times with random start/end pairs."""
from ai_mouse.generator import generate
"""Call the project's generator N times with random start/end pairs.
``model_dir`` is accepted for CLI backward compatibility but is no longer
used — generation goes through the public ai_mouse API which loads the
bundled ONNX model. Export a fresh .onnx via ``python -m tools.export_onnx``
to refresh.
"""
del model_dir # legacy arg, unused
from ai_mouse import generate
rng = random.Random(seed)
out: list[dict] = []
@@ -63,7 +70,7 @@ def _generate_n_samples(
ex = max(0, min(800, ex))
ey = max(0, min(600, ey))
try:
pts = generate(start=(sx, sy), end=(ex, ey), model_dir=model_dir)
pts = generate(start=(sx, sy), end=(ex, ey))
except Exception as exc: # noqa: BLE001
logger.warning("generate() failed at i=%d: %s", i, exc)
continue

View File

@@ -12,7 +12,6 @@ import math
import torch
import torch.nn as nn
from torch.distributions import Normal
# ---------------------------------------------------------------------------
@@ -157,80 +156,3 @@ class TrajectoryFlowModel(nn.Module):
# Output projection
return self.output_proj(h) # (B, T, 3)
# ---------------------------------------------------------------------------
# Legacy JointCVAE — kept for backward compatibility with generator.py
# ---------------------------------------------------------------------------
class JointCVAE(nn.Module):
"""Joint Conditional VAE for mouse trajectory generation (legacy).
Kept for backward compatibility with the existing generator.
See TrajectoryFlowModel for the new approach.
"""
def __init__(
self,
seq_len: int = 64,
latent_dim: int = 32,
hidden: int = 128,
cond_dim: int = 3,
):
super().__init__()
self.seq_len = seq_len
self.latent_dim = latent_dim
self.hidden = hidden
self.cond_dim = cond_dim
self.feat_dim = 3
self.enc_gru = nn.GRU(
input_size=self.feat_dim + cond_dim,
hidden_size=hidden,
num_layers=2,
batch_first=True,
bidirectional=True,
)
self.enc_mu = nn.Linear(hidden * 2, latent_dim)
self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
self.dec_gru = nn.GRU(
input_size=latent_dim + cond_dim,
hidden_size=hidden,
num_layers=2,
batch_first=True,
)
self.dec_out = nn.Linear(hidden, self.feat_dim)
def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
B, T, _ = seq.shape
c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
x_in = torch.cat([seq, c_exp], dim=-1)
_, h_n = self.enc_gru(x_in)
h_fwd = h_n[-2]
h_bwd = h_n[-1]
h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
return self.enc_mu(h_cat), self.enc_logvar(h_cat)
def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
B = z.shape[0]
zc = torch.cat([z, cond], dim=-1)
h0_flat = self.dec_h0(zc)
h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
out, _ = self.dec_gru(inp, h0)
return self.dec_out(out)
def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
std = torch.exp(0.5 * logvar)
return Normal(mu, std).rsample()
def forward(
self, seq: torch.Tensor, cond: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
mu, logvar = self.encode(seq, cond)
z = self.reparameterise(mu, logvar)
recon = self.decode(z, cond)
return recon, mu, logvar

View File

@@ -182,21 +182,17 @@ async def scroll_train(req: ScrollTrainRequest) -> StreamingResponse:
@router.post("/verify")
def scroll_verify(req: ScrollVerifyRequest) -> dict:
from ai_mouse.scroll.generator import generate_scroll
# Uses the bundled ONNX scroll model exposed via the public ai_mouse API.
# The legacy scroll_model.pt path is no longer wired in; export a fresh
# scroll_decoder.onnx via `python -m tools.export_onnx` to update.
from ai_mouse import generate_scroll
_, models_dir = _paths()
if not (models_dir / "scroll_model.pt").exists():
raise HTTPException(
status_code=400,
detail="滚轮模型尚未训练,请先在「训练模型 → 滚轮模型」中完成训练。",
)
paths = []
for _ in range(min(req.n_paths, 12)):
events = generate_scroll(
req.start_scrollY,
req.target_scrollY,
mode=req.mode,
model_dir=str(models_dir),
)
paths.append(events)
return {"paths": paths}

View File

@@ -7,8 +7,6 @@ import logging
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from .deps import get_data_dir
logger = logging.getLogger(__name__)
router = APIRouter()
@@ -32,18 +30,19 @@ class VerifyRequest(BaseModel):
@router.post("/verify")
def verify(req: VerifyRequest) -> dict:
from ai_mouse.generator import generate
# Uses the bundled ONNX model exposed via the public ai_mouse API.
# The legacy req.model_dir / data/models_v2 .pt path is no longer wired
# in; export a fresh .onnx via `python -m tools.export_onnx` to update.
from ai_mouse import generate
n = max(1, min(req.n_paths, 12))
models_dir = get_data_dir() / "models_v2"
model_dir_arg = req.model_dir if req.model_dir else str(models_dir)
start = tuple(req.start) # type: ignore[arg-type]
end = tuple(req.end) # type: ignore[arg-type]
paths = []
try:
for _ in range(n):
pts = generate(start=start, end=end, model_dir=model_dir_arg)
pts = generate(start=start, end=end)
paths.append([[x, y, t] for x, y, t in pts])
except FileNotFoundError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc

View File

@@ -23,7 +23,7 @@ import numpy as np
import torch
from torch.utils.data import DataLoader
from ai_mouse.coord import encode_trajectory
from ai_mouse._coord import encode_trajectory
from tools.config import TrainConfig
from tools.models import TrajectoryFlowModel
from tools.utils import resample_arc