Compare commits
25 Commits
d39db46170
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 23084f77d6 | |||
| 43d28b6254 | |||
| 241a4a41c7 | |||
| 9e529d3951 | |||
| 76581a210e | |||
| 7c33c13a87 | |||
| adc46a445f | |||
| 441e6f3dfe | |||
| 556f7f861d | |||
| d1f70e5753 | |||
| 94c52bd3be | |||
| c2ed7b3cb9 | |||
| 12d70fe137 | |||
| 3e7a194356 | |||
| 7890b07a01 | |||
| 1c60763037 | |||
| 566adda920 | |||
| 984890515f | |||
| 1105026831 | |||
| 0413769dd2 | |||
| 4f55000c05 | |||
| 2d8017b089 | |||
| 31fd884dfd | |||
| 525e555884 | |||
| 5b4f693098 |
52
CHANGELOG.md
Normal file
52
CHANGELOG.md
Normal 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
114
CLAUDE.md
@@ -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.12–3.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
120
README.md
Normal 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.
|
||||
624
docs/superpowers/plans/2026-07-09-mouse-postprocess-quality.md
Normal file
624
docs/superpowers/plans/2026-07-09-mouse-postprocess-quality.md
Normal 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
|
||||
@@ -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
30
examples/quickstart.py
Normal 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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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),
|
||||
]
|
||||
@@ -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
0
src/ai_mouse/py.typed
Normal file
Binary file not shown.
@@ -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)
|
||||
@@ -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
151
tests/unit/test_golden.py
Normal 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"
|
||||
)
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user