Compare commits
14 Commits
1c60763037
...
v0.2.0
| Author | SHA1 | Date | |
|---|---|---|---|
| 43d28b6254 | |||
| 241a4a41c7 | |||
| 9e529d3951 | |||
| 76581a210e | |||
| 7c33c13a87 | |||
| adc46a445f | |||
| 441e6f3dfe | |||
| 556f7f861d | |||
| d1f70e5753 | |||
| 94c52bd3be | |||
| c2ed7b3cb9 | |||
| 12d70fe137 | |||
| 3e7a194356 | |||
| 7890b07a01 |
4
.github/workflows/ci.yml
vendored
4
.github/workflows/ci.yml
vendored
@@ -13,7 +13,7 @@ jobs:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest]
|
||||
os: [ubuntu-latest]
|
||||
python: ["3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -28,7 +28,7 @@ jobs:
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest]
|
||||
os: [ubuntu-latest]
|
||||
python: ["3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
13
CHANGELOG.md
13
CHANGELOG.md
@@ -4,6 +4,19 @@ 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)
|
||||
|
||||
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.
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Binary file not shown.
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user