yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
2d8017b089d58a474705d0780c5f9085bfdf66c0
ai_mouse/docs/superpowers/plans/2026-05-10-balabit-pretrain-refactor.md
Huang Qi 4d414fd180 chore: initialize git repo, add matplotlib dep, extend config 
- Add .gitignore for Python/data/models
- Add matplotlib>=3.8.0 for eval plots
- Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
2026-05-10 12:24:07 +08:00

3152 lines
100 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# Balabit 预训练 + Fine-tune 重构 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:** Use Balabit Mouse Dynamics Challenge dataset to pretrain `TrajectoryFlowModel`, then fine-tune on user's 605 traces. Fix high-frequency lateral jitter and template-y Δt curves by aggressively removing deterministic post-processing. Add a quantitative eval module that produces Markdown reports with kinematic metrics and FFT spectra.
**Architecture:** Two-stage training (Balabit pretrain → 605 fine-tune via `resume_from`). New modules `ai_mouse/data_adapters/balabit.py` and `ai_mouse/eval/`. Existing model architecture (`TrajectoryFlowModel`), scroll subsystem, and frontend are unchanged. Data format remains identical to current `traces.jsonl`.
**Tech Stack:** Python 3.12+, PyTorch (CPU; optional CUDA), NumPy, SciPy, matplotlib (NEW dep for eval plots), uv package manager.
**Spec reference:** [docs/superpowers/specs/2026-05-10-balabit-pretrain-refactor-design.md](../specs/2026-05-10-balabit-pretrain-refactor-design.md)
**Prerequisites:**
1. User must download Balabit Mouse Dynamics Challenge dataset from https://github.com/balabit/Mouse-Dynamics-Challenge (clone or zip download). Plan assumes path is configurable via CLI arg.
2. Project will be initialized as git repo in Task 1 (currently not a git repo).
---
## Task 1: Project Setup (git init, dependencies, config)
**Files:**
- Create: `.gitignore`
- Modify: `pyproject.toml`
- Modify: `ai_mouse/config.py`
- Modify: `tests/conftest.py` (no functional change yet — just to verify import works after config changes)
- [ ] **Step 1: Initialize git repo (requires user OK)**
This step touches user-controlled state. Confirm with user before running.
```bash
cd /d/code/python/side/ai_mouse
git init
```
Expected: `Initialized empty Git repository in .../ai_mouse/.git/`
- [ ] **Step 2: Create .gitignore**
Create `.gitignore` at project root:
```
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
.venv/
.pytest_cache/
# IDE
.idea/
.vscode/
# uv
uv.lock.bak
# Data & models — large binary, do NOT commit
data/traces.jsonl
data/scroll_traces.jsonl
data/pretrain_traces.jsonl
data/models_v2/
data/models_v2_pretrained/
data/scroll_models/
data/eval_reports/
data/balabit_raw/
# OS
.DS_Store
Thumbs.db
# Playwright artifacts
.playwright-mcp/
```
- [ ] **Step 3: Add matplotlib to pyproject.toml**
Modify `pyproject.toml` — add matplotlib to `dependencies`:
```toml
[project]
name = "ai-mouse"
version = "0.1.0"
requires-python = ">=3.12,<3.14"
dependencies = [
"torch>=2.2.0",
"numpy>=1.26.0",
"fastapi>=0.111.0",
"uvicorn>=0.29.0",
"scipy>=1.10.0",
"matplotlib>=3.8.0",
]
[dependency-groups]
dev = ["pytest>=8.0.0", "pytest-asyncio>=0.23.0", "httpx>=0.27.0"]
```
Then sync:
```bash
uv sync
```
Expected: `Resolved N packages` and matplotlib appears in `.venv`.
- [ ] **Step 4: Add new config dataclasses**
Append to `ai_mouse/config.py` (after existing `ServerConfig`):
```python
# ---------------------------------------------------------------------------
# Pretraining (Balabit) configuration
# ---------------------------------------------------------------------------
@dataclass
class PretrainConfig:
"""Hyperparameters for Balabit pretraining stage."""
epochs: int = 200
batch_size: int = 128
lr: float = 3e-4
seq_len: int = 64
@dataclass
class FinetuneConfig:
"""Hyperparameters for fine-tuning on user-collected data."""
epochs: int = 50
batch_size: int = 64
lr: float = 1e-5 # 比预训练小一个数量级，防止灾难性遗忘
seq_len: int = 64
# ---------------------------------------------------------------------------
# Balabit adapter configuration
# ---------------------------------------------------------------------------
@dataclass
class BalabitAdapterConfig:
"""Settings for Balabit CSV → traces.jsonl conversion."""
window_ms: int = 1200 # click 前回溯窗口
min_dist: int = 50 # 最小起终点距离 (px)
min_events: int = 5 # 最小 Move 事件数
max_span_ms: int = 5000 # 最大段时间跨度 (ms)
max_gap_ms: int = 200 # 段内相邻 Move 最大时间差
# ---------------------------------------------------------------------------
# Eval configuration
# ---------------------------------------------------------------------------
@dataclass
class EvalConfig:
"""Settings for evaluation report generation."""
n_samples: int = 1000
fft_freq_band: tuple[float, float] = (4.0, 12.0) # 生理震颤频段 (Hz)
kl_bins: int = 50
```
- [ ] **Step 5: Verify imports**
```bash
uv run python -c "from ai_mouse.config import PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig; print('OK')"
```
Expected: `OK`
- [ ] **Step 6: Run all existing tests to confirm no regressions**
```bash
uv run pytest -x
```
Expected: all tests pass.
- [ ] **Step 7: Initial commit**
```bash
git add .gitignore pyproject.toml uv.lock ai_mouse/ tests/ static/ main.py docs/
git commit -m "chore: initialize git repo, add matplotlib dep, extend config
- Add .gitignore for Python/data/models
- Add matplotlib>=3.8.0 for eval plots
- Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses"
```
Expected: commit succeeds with hash printed.
---
## Task 2: Balabit Adapter — Core Module Skeleton + Types
**Files:**
- Create: `ai_mouse/data_adapters/__init__.py`
- Create: `ai_mouse/data_adapters/balabit.py`
- Create: `tests/test_balabit_adapter.py`
- [ ] **Step 1: Write failing tests for the public API surface**
Create `tests/test_balabit_adapter.py`:
```python
"""Tests for Balabit Mouse Dynamics Challenge data adapter."""
from __future__ import annotations
from pathlib import Path
import pytest
def test_module_exports():
"""The adapter module must export the public functions used by CLI."""
from ai_mouse.data_adapters import balabit
assert hasattr(balabit, "parse_session_csv")
assert hasattr(balabit, "segment_by_clicks")
assert hasattr(balabit, "filter_segments")
assert hasattr(balabit, "process_session")
assert hasattr(balabit, "MouseEvent")
assert hasattr(balabit, "Segment")
def test_mouse_event_dataclass():
"""MouseEvent has expected fields."""
from ai_mouse.data_adapters.balabit import MouseEvent
e = MouseEvent(t_ms=100, button="NoButton", state="Move", x=300, y=400)
assert e.t_ms == 100
assert e.state == "Move"
assert e.x == 300
def test_segment_dataclass():
"""Segment has expected fields."""
from ai_mouse.data_adapters.balabit import MouseEvent, Segment
events = [MouseEvent(t_ms=0, button="NoButton", state="Move", x=10, y=20)]
s = Segment(events=events, click_x=100, click_y=200, click_t_ms=500, session_id="user1_s1")
assert s.events == events
assert s.click_x == 100
assert s.session_id == "user1_s1"
```
- [ ] **Step 2: Run test, verify it fails**
```bash
uv run pytest tests/test_balabit_adapter.py -v
```
Expected: FAIL with `ModuleNotFoundError: No module named 'ai_mouse.data_adapters'`
- [ ] **Step 3: Create package skeleton**
Create `ai_mouse/data_adapters/__init__.py`:
```python
"""Data adapters: convert external datasets to the project's traces.jsonl format."""
```
Create `ai_mouse/data_adapters/balabit.py`:
```python
"""Adapter for the Balabit Mouse Dynamics Challenge dataset.
Source: https://github.com/balabit/Mouse-Dynamics-Challenge
Each session is a CSV file with columns:
record timestamp, client timestamp, button, state, x, y
Where:
state ∈ {Move, Pressed, Released, Drag, Scroll}
button ∈ {NoButton, Left, Right, Wheel}
We extract "click-anchored" trajectory segments: each Pressed event
defines a target, and the W ms of Move events preceding it form one
training trace.
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from pathlib import Path
logger = logging.getLogger(__name__)
@dataclass
class MouseEvent:
"""A single mouse event from a Balabit CSV row."""
t_ms: int # client timestamp in milliseconds (relative to session start)
button: str # "NoButton", "Left", "Right", "Wheel"
state: str # "Move", "Pressed", "Released", "Drag", "Scroll"
x: int
y: int
@dataclass
class Segment:
"""A click-anchored trajectory segment ready to be written to JSONL."""
events: list[MouseEvent] # only Move events, sorted by t_ms ascending
click_x: int # the Pressed event's x coordinate
click_y: int # the Pressed event's y coordinate
click_t_ms: int # the Pressed event's timestamp
session_id: str # e.g. "user7_session_42"
def parse_session_csv(path: Path) -> list[MouseEvent]:
"""Stub — implemented in Task 3."""
raise NotImplementedError
def segment_by_clicks(
events: list[MouseEvent],
window_ms: int,
session_id: str,
) -> list[Segment]:
"""Stub — implemented in Task 4."""
raise NotImplementedError
def filter_segments(
segments: list[Segment],
min_events: int,
min_dist: int,
max_span_ms: int,
max_gap_ms: int,
) -> list[Segment]:
"""Stub — implemented in Task 5."""
raise NotImplementedError
def process_session(
csv_path: Path,
output_jsonl: Path,
config,
) -> int:
"""Stub — implemented in Task 6."""
raise NotImplementedError
```
- [ ] **Step 4: Run tests, verify they pass**
```bash
uv run pytest tests/test_balabit_adapter.py -v
```
Expected: 3 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/data_adapters/ tests/test_balabit_adapter.py
git commit -m "feat(adapter): scaffold balabit data adapter package"
```
---
## Task 3: Balabit Adapter — CSV Parsing
**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `parse_session_csv`)
- Modify: `tests/test_balabit_adapter.py` (add tests)
- [ ] **Step 1: Write failing tests for CSV parsing**
Append to `tests/test_balabit_adapter.py`:
```python
def _write_csv(path: Path, rows: list[str]) -> None:
"""Helper to write a Balabit-format CSV with header."""
header = "record timestamp,client timestamp,button,state,x,y"
path.write_text(header + "\n" + "\n".join(rows) + "\n", encoding="utf-8")
class TestParseSessionCsv:
def test_parses_basic_rows(self, tmp_path):
from ai_mouse.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_1"
_write_csv(csv, [
"1500000000.000,0.000,NoButton,Move,100,200",
"1500000000.050,0.050,NoButton,Move,110,210",
"1500000000.100,0.100,Left,Pressed,120,220",
])
events = parse_session_csv(csv)
assert len(events) == 3
assert events[0].t_ms == 0
assert events[0].state == "Move"
assert events[0].x == 100
assert events[2].t_ms == 100
assert events[2].state == "Pressed"
assert events[2].button == "Left"
def test_handles_float_timestamps(self, tmp_path):
"""Client timestamps are floats in seconds; we convert to int ms."""
from ai_mouse.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_2"
_write_csv(csv, [
"0,1.234,NoButton,Move,50,60",
"0,1.250,NoButton,Move,55,65",
])
events = parse_session_csv(csv)
assert events[0].t_ms == 1234
assert events[1].t_ms == 1250
def test_skips_malformed_rows(self, tmp_path):
"""Rows with bad data are logged and skipped, not raised."""
from ai_mouse.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_3"
_write_csv(csv, [
"0,0.000,NoButton,Move,100,200",
"BROKEN_ROW",
"0,abc,NoButton,Move,100,200", # bad timestamp
"0,0.100,NoButton,Move,150,250",
])
events = parse_session_csv(csv)
assert len(events) == 2
assert events[0].x == 100
assert events[1].x == 150
def test_returns_empty_list_for_empty_file(self, tmp_path):
from ai_mouse.data_adapters.balabit import parse_session_csv
csv = tmp_path / "session_4"
csv.write_text("record timestamp,client timestamp,button,state,x,y\n", encoding="utf-8")
events = parse_session_csv(csv)
assert events == []
```
- [ ] **Step 2: Run tests, verify they fail**
```bash
uv run pytest tests/test_balabit_adapter.py::TestParseSessionCsv -v
```
Expected: FAIL with `NotImplementedError`.
- [ ] **Step 3: Implement parse_session_csv**
Replace the stub in `ai_mouse/data_adapters/balabit.py`:
```python
def parse_session_csv(path: Path) -> list[MouseEvent]:
"""Parse a Balabit session CSV file into MouseEvent objects.
Malformed rows are logged and skipped (not raised).
Client timestamps (seconds, float) are converted to int milliseconds.
Args:
path: Path to a Balabit session CSV file.
Returns:
List of MouseEvent in original order. Empty list if file is empty.
"""
import csv as csv_module
events: list[MouseEvent] = []
with path.open("r", encoding="utf-8", newline="") as f:
reader = csv_module.DictReader(f)
for row_idx, row in enumerate(reader, 2): # 1-based, header is line 1
try:
client_ts = float(row["client timestamp"])
t_ms = int(round(client_ts * 1000))
button = row["button"].strip()
state = row["state"].strip()
x = int(row["x"])
y = int(row["y"])
except (KeyError, ValueError, TypeError) as exc:
logger.debug("Skipping malformed row %d in %s: %s", row_idx, path.name, exc)
continue
events.append(MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y))
return events
```
- [ ] **Step 4: Run tests, verify they pass**
```bash
uv run pytest tests/test_balabit_adapter.py::TestParseSessionCsv -v
```
Expected: 4 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement Balabit CSV parser"
```
---
## Task 4: Balabit Adapter — Click-Anchored Segmentation
**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `segment_by_clicks`)
- Modify: `tests/test_balabit_adapter.py` (add tests)
- [ ] **Step 1: Write failing tests**
Append to `tests/test_balabit_adapter.py`:
```python
class TestSegmentByClicks:
def _make_event(self, t_ms: int, state: str, x: int, y: int, button: str = "NoButton"):
from ai_mouse.data_adapters.balabit import MouseEvent
return MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y)
def test_one_click_one_segment(self):
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Move", 50, 60),
self._make_event(500, "Move", 100, 100),
self._make_event(600, "Pressed", 110, 110, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="test_s1")
assert len(segments) == 1
seg = segments[0]
assert seg.click_x == 110
assert seg.click_y == 110
assert seg.click_t_ms == 600
assert len(seg.events) == 3
assert seg.session_id == "test_s1"
def test_window_excludes_old_events(self):
"""Move events earlier than (click_t - window_ms) are dropped."""
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20), # too old
self._make_event(100, "Move", 20, 30), # too old
self._make_event(900, "Move", 30, 40), # in window
self._make_event(1000, "Pressed", 40, 50, button="Left"),
]
segments = segment_by_clicks(events, window_ms=200, session_id="s")
assert len(segments) == 1
assert len(segments[0].events) == 1
assert segments[0].events[0].t_ms == 900
def test_multiple_clicks_multiple_segments(self):
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Pressed", 50, 50, button="Left"),
self._make_event(200, "Released", 50, 50, button="Left"),
self._make_event(300, "Move", 60, 60),
self._make_event(400, "Move", 70, 70),
self._make_event(500, "Pressed", 80, 80, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 2
assert segments[0].click_x == 50
assert segments[1].click_x == 80
# Second segment's events must not include the first Pressed
for e in segments[1].events:
assert e.state == "Move"
def test_skips_pressed_with_non_left_button(self):
"""Right-clicks and wheel-clicks don't anchor segments (only Left)."""
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Pressed", 50, 50, button="Right"), # ignored
self._make_event(200, "Move", 60, 60),
self._make_event(300, "Pressed", 70, 70, button="Left"), # anchor
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 1
assert segments[0].click_x == 70
def test_no_clicks_returns_empty(self):
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Move", 20, 30),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert segments == []
def test_excludes_drag_events(self):
"""Drag events are not Move; segment should only include Move."""
from ai_mouse.data_adapters.balabit import segment_by_clicks
events = [
self._make_event(0, "Move", 10, 20),
self._make_event(100, "Drag", 30, 40), # not Move
self._make_event(200, "Move", 50, 60),
self._make_event(300, "Pressed", 70, 80, button="Left"),
]
segments = segment_by_clicks(events, window_ms=1200, session_id="s")
assert len(segments) == 1
# Drag event should not appear in seg.events
assert all(e.state == "Move" for e in segments[0].events)
assert len(segments[0].events) == 2
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_balabit_adapter.py::TestSegmentByClicks -v
```
Expected: FAIL with `NotImplementedError`.
- [ ] **Step 3: Implement segment_by_clicks**
Replace the stub in `ai_mouse/data_adapters/balabit.py`:
```python
def segment_by_clicks(
events: list[MouseEvent],
window_ms: int,
session_id: str,
) -> list[Segment]:
"""Extract click-anchored segments from a session.
For each Left-button Pressed event, collect all Move events within
[click_t - window_ms, click_t) into one segment.
Args:
events: Full session events (any state, any order is OK but typically sorted).
window_ms: How far back to look before each click.
session_id: String tag attached to every segment for debugging.
Returns:
List of Segment, one per Left Pressed event that has at least one preceding Move.
"""
segments: list[Segment] = []
for ev in events:
if ev.state != "Pressed" or ev.button != "Left":
continue
click_t = ev.t_ms
window_start = click_t - window_ms
moves = [
m for m in events
if m.state == "Move" and window_start <= m.t_ms < click_t
]
if not moves:
continue
moves.sort(key=lambda m: m.t_ms)
segments.append(Segment(
events=moves,
click_x=ev.x,
click_y=ev.y,
click_t_ms=click_t,
session_id=session_id,
))
return segments
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_balabit_adapter.py::TestSegmentByClicks -v
```
Expected: 6 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement click-anchored segmentation"
```
---
## Task 5: Balabit Adapter — Filter Rules
**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `filter_segments`)
- Modify: `tests/test_balabit_adapter.py`
- [ ] **Step 1: Write failing tests**
Append to `tests/test_balabit_adapter.py`:
```python
class TestFilterSegments:
def _seg(self, events_data: list[tuple[int, int, int]], click=(500, 500), session="s"):
"""events_data: list of (t_ms, x, y) tuples."""
from ai_mouse.data_adapters.balabit import MouseEvent, Segment
events = [MouseEvent(t_ms=t, button="NoButton", state="Move", x=x, y=y)
for (t, x, y) in events_data]
click_t = events[-1].t_ms + 50 if events else 100
return Segment(events=events, click_x=click[0], click_y=click[1],
click_t_ms=click_t, session_id=session)
def test_drops_segment_with_too_few_events(self):
from ai_mouse.data_adapters.balabit import filter_segments
seg = self._seg([(0, 100, 100), (50, 105, 105)]) # 2 events, min=5
result = filter_segments([seg], min_events=5, min_dist=10,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_short_distance(self):
from ai_mouse.data_adapters.balabit import filter_segments
# Start (100,100), end click=(105,100) → dist=5, min_dist=50
events = [(i*10, 100+i, 100) for i in range(10)]
seg = self._seg(events, click=(105, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_too_long_span(self):
from ai_mouse.data_adapters.balabit import filter_segments
# Span 6000ms, max_span=5000
events = [(0, 100, 100), (2000, 200, 200), (4000, 300, 300), (6000, 400, 400),
(6010, 410, 400), (6020, 420, 400)]
seg = self._seg(events, click=(500, 500))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=10000)
assert result == []
def test_drops_segment_with_gap(self):
from ai_mouse.data_adapters.balabit import filter_segments
# Gap of 500ms between events 2 and 3, max_gap=200
events = [(0, 100, 100), (50, 110, 110), (100, 120, 120),
(600, 200, 200), (650, 210, 210), (700, 220, 220)]
seg = self._seg(events, click=(300, 300))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_out_of_range_coords(self):
from ai_mouse.data_adapters.balabit import filter_segments
# x=10000 out of range
events = [(i*10, 10000, 100) for i in range(6)]
seg = self._seg(events, click=(10100, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_drops_segment_with_short_arc_length(self):
"""Total arc < 50 even though endpoints are far → high-frequency jitter only."""
from ai_mouse.data_adapters.balabit import filter_segments
# Tiny back-and-forth with click 100px away — unrealistic, drop it
events = [(i*10, 100 + (i % 2), 100) for i in range(10)]
seg = self._seg(events, click=(200, 100))
# arc length = sum of |Δp| ≈ 9 (alternating ±1) which is < 50
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert result == []
def test_keeps_valid_segment(self):
from ai_mouse.data_adapters.balabit import filter_segments
# Smooth 100→500 px straight line, 10 events, span 500ms
events = [(i*50, 100 + i*40, 100) for i in range(10)]
seg = self._seg(events, click=(500, 100))
result = filter_segments([seg], min_events=5, min_dist=50,
max_span_ms=5000, max_gap_ms=200)
assert len(result) == 1
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_balabit_adapter.py::TestFilterSegments -v
```
Expected: FAIL with `NotImplementedError`.
- [ ] **Step 3: Implement filter_segments**
Replace the stub in `ai_mouse/data_adapters/balabit.py`:
```python
def filter_segments(
segments: list[Segment],
min_events: int,
min_dist: int,
max_span_ms: int,
max_gap_ms: int,
coord_max: int = 5000,
) -> list[Segment]:
"""Drop segments that fail any quality check.
A segment is dropped if any of these are true:
- len(events) < min_events
- Euclidean dist(events[0], (click_x, click_y)) < min_dist
- events[-1].t_ms - events[0].t_ms > max_span_ms
- any adjacent Move pair has dt > max_gap_ms (sampling drop-out)
- any coord (start/end/click) outside [0, coord_max]
- total arc length < min_dist (high-frequency jitter only)
Args:
segments: Candidate segments.
min_events: Minimum number of Move events.
min_dist: Minimum start→click pixel distance AND minimum total arc length.
max_span_ms: Maximum time span of the segment (events[-1] - events[0]).
max_gap_ms: Maximum allowed gap between adjacent Move events.
coord_max: Maximum allowed pixel coordinate value (5000 catches multi-monitor anomalies).
Returns:
Filtered list, original order preserved.
"""
import math
keep: list[Segment] = []
for seg in segments:
if len(seg.events) < min_events:
continue
sx, sy = seg.events[0].x, seg.events[0].y
ex, ey = seg.click_x, seg.click_y
# Coord range check
if any(c < 0 or c > coord_max for c in (sx, sy, ex, ey)):
continue
# Endpoint distance
dist = math.hypot(ex - sx, ey - sy)
if dist < min_dist:
continue
# Time span
span = seg.events[-1].t_ms - seg.events[0].t_ms
if span > max_span_ms:
continue
# Gap check + total arc length
total_arc = 0.0
bad_gap = False
for i in range(1, len(seg.events)):
dt = seg.events[i].t_ms - seg.events[i - 1].t_ms
if dt > max_gap_ms:
bad_gap = True
break
dx = seg.events[i].x - seg.events[i - 1].x
dy = seg.events[i].y - seg.events[i - 1].y
total_arc += math.hypot(dx, dy)
if bad_gap:
continue
if total_arc < min_dist:
continue
keep.append(seg)
return keep
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_balabit_adapter.py::TestFilterSegments -v
```
Expected: 7 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/data_adapters/balabit.py tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement segment quality filters"
```
---
## Task 6: Balabit Adapter — Process Session + CLI
**Files:**
- Modify: `ai_mouse/data_adapters/balabit.py` (implement `process_session`, add `main`)
- Create: `ai_mouse/data_adapters/__main__.py`
- Modify: `tests/test_balabit_adapter.py`
- [ ] **Step 1: Write failing tests for process_session**
Append to `tests/test_balabit_adapter.py`:
```python
class TestProcessSession:
def test_writes_jsonl_in_expected_format(self, tmp_path):
from ai_mouse.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
# Construct a Balabit-format CSV with one valid segment
csv_path = tmp_path / "user_session_42"
rows = []
# 10 Move events going from (100,100) to (500,200) over 500ms
for i in range(10):
t = i * 0.05
x = 100 + i * 40
y = 100 + i * 10
rows.append(f"0,{t:.3f},NoButton,Move,{x},{y}")
rows.append(f"0,0.550,Left,Pressed,510,210")
_write_csv(csv_path, rows)
out = tmp_path / "out.jsonl"
config = BalabitAdapterConfig(
window_ms=1200, min_dist=50, min_events=5,
max_span_ms=5000, max_gap_ms=200,
)
n = process_session(csv_path, out, config)
assert n == 1
assert out.exists()
line = out.read_text(encoding="utf-8").strip()
record = json.loads(line)
assert record["meta"]["start"] == [100, 100]
assert record["meta"]["end"] == [510, 210]
assert record["meta"]["dist"] > 0
assert record["meta"]["source"] == "balabit"
assert record["meta"]["session_id"] == "user_session_42"
# Events: only move events, no down/up
assert all(e["type"] == "move" for e in record["events"])
# Timestamps relative (start at 0)
assert record["events"][0]["t"] == 0
def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
from ai_mouse.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
csv_path = tmp_path / "empty_session"
_write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"]) # no clicks
out = tmp_path / "out.jsonl"
config = BalabitAdapterConfig()
n = process_session(csv_path, out, config)
assert n == 0
# File may not exist or may be empty — both acceptable
if out.exists():
assert out.read_text() == ""
def test_appends_to_existing_jsonl(self, tmp_path):
from ai_mouse.config import BalabitAdapterConfig
from ai_mouse.data_adapters.balabit import process_session
out = tmp_path / "out.jsonl"
out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")
csv_path = tmp_path / "session_x"
rows = [f"0,{i*0.05:.3f},NoButton,Move,{100+i*40},100" for i in range(10)]
rows.append("0,0.550,Left,Pressed,510,100")
_write_csv(csv_path, rows)
config = BalabitAdapterConfig()
process_session(csv_path, out, config)
lines = out.read_text(encoding="utf-8").strip().split("\n")
assert len(lines) == 2 # original + appended
```
`json` import — make sure top of test file already has `import json`. If not, add.
Add this near the top of `tests/test_balabit_adapter.py`:
```python
import json
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_balabit_adapter.py::TestProcessSession -v
```
Expected: FAIL with `NotImplementedError`.
- [ ] **Step 3: Implement process_session**
Replace the `process_session` stub in `ai_mouse/data_adapters/balabit.py`:
```python
def process_session(
csv_path: Path,
output_jsonl: Path,
config,
) -> int:
"""Convert one Balabit session CSV to JSONL traces, append to output.
Args:
csv_path: Path to a Balabit session CSV.
output_jsonl: Output JSONL file (will be appended to).
config: BalabitAdapterConfig with window_ms / min_dist / etc.
Returns:
Number of valid segments written.
"""
import math
session_id = csv_path.stem # e.g. "session_42"
events = parse_session_csv(csv_path)
if not events:
return 0
raw_segments = segment_by_clicks(
events, window_ms=config.window_ms, session_id=session_id
)
valid_segments = filter_segments(
raw_segments,
min_events=config.min_events,
min_dist=config.min_dist,
max_span_ms=config.max_span_ms,
max_gap_ms=config.max_gap_ms,
)
if not valid_segments:
return 0
output_jsonl.parent.mkdir(parents=True, exist_ok=True)
n_written = 0
with output_jsonl.open("a", encoding="utf-8") as f:
for seg in valid_segments:
sx, sy = seg.events[0].x, seg.events[0].y
ex, ey = seg.click_x, seg.click_y
dist = math.hypot(ex - sx, ey - sy)
angle = math.degrees(math.atan2(ey - sy, ex - sx))
t0 = seg.events[0].t_ms
record = {
"meta": {
"start": [sx, sy],
"end": [ex, ey],
"dist": int(round(dist)),
"angle": round(angle, 1),
"source": "balabit",
"session_id": seg.session_id,
},
"events": [
{"type": "move", "x": e.x, "y": e.y, "t": e.t_ms - t0}
for e in seg.events
],
}
f.write(json.dumps(record, ensure_ascii=False) + "\n")
n_written += 1
return n_written
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_balabit_adapter.py::TestProcessSession -v
```
Expected: 3 tests pass.
- [ ] **Step 5: Add CLI main function to balabit.py**
Append to `ai_mouse/data_adapters/balabit.py`:
```python
def main(argv: list[str] | None = None) -> int:
"""CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
import argparse
from ai_mouse.config import BalabitAdapterConfig
parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
parser.add_argument(
"--input", type=Path, required=True,
help="Directory containing Balabit session CSV files (recursive)",
)
parser.add_argument(
"--output", type=Path, default=Path("data/pretrain_traces.jsonl"),
help="Output JSONL path (default: data/pretrain_traces.jsonl)",
)
parser.add_argument("--window-ms", type=int, default=1200)
parser.add_argument("--min-dist", type=int, default=50)
parser.add_argument("--min-events", type=int, default=5)
parser.add_argument("--max-span-ms", type=int, default=5000)
parser.add_argument("--max-gap-ms", type=int, default=200)
parser.add_argument(
"--overwrite", action="store_true",
help="Truncate output file before writing (default: append)",
)
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
if not args.input.is_dir():
logger.error("Input is not a directory: %s", args.input)
return 2
config = BalabitAdapterConfig(
window_ms=args.window_ms,
min_dist=args.min_dist,
min_events=args.min_events,
max_span_ms=args.max_span_ms,
max_gap_ms=args.max_gap_ms,
)
if args.overwrite and args.output.exists():
args.output.unlink()
# Recursively walk the input directory looking for session files.
# Balabit files have no extension; we accept any regular file.
csv_files = sorted([p for p in args.input.rglob("*") if p.is_file() and not p.name.startswith(".")])
if not csv_files:
logger.error("No session files found under %s", args.input)
return 2
total = 0
for i, csv_path in enumerate(csv_files, 1):
try:
n = process_session(csv_path, args.output, config)
except Exception as exc: # noqa: BLE001
logger.warning("Skipping %s due to error: %s", csv_path.name, exc)
continue
total += n
if i % 10 == 0 or i == len(csv_files):
logger.info("Processed %d/%d sessions, %d segments so far", i, len(csv_files), total)
logger.info("Done. Wrote %d segments to %s", total, args.output)
return 0
```
- [ ] **Step 6: Create __main__.py for CLI dispatch**
Create `ai_mouse/data_adapters/__main__.py`:
```python
"""CLI dispatch: `python -m ai_mouse.data_adapters.balabit ...`
Note: This file makes `python -m ai_mouse.data_adapters` invokable but for
clarity prefer the explicit form `python -m ai_mouse.data_adapters.balabit`.
"""
from __future__ import annotations
import sys
from ai_mouse.data_adapters.balabit import main
if __name__ == "__main__":
sys.exit(main())
```
Also add a `__main__` guard at the bottom of `balabit.py` so `python -m ai_mouse.data_adapters.balabit` works directly:
```python
if __name__ == "__main__":
import sys
sys.exit(main())
```
- [ ] **Step 7: Smoke-test the CLI**
```bash
uv run python -m ai_mouse.data_adapters.balabit --help
```
Expected: argparse help text printed.
- [ ] **Step 8: Run all adapter tests**
```bash
uv run pytest tests/test_balabit_adapter.py -v
```
Expected: all tests pass.
- [ ] **Step 9: Commit**
```bash
git add ai_mouse/data_adapters/ tests/test_balabit_adapter.py
git commit -m "feat(adapter): implement process_session and CLI"
```
---
## Task 7: Trainer — TrajectoryDataset (Streaming Augmentation)
**Files:**
- Modify: `ai_mouse/trainer.py` (replace `_augment` usage with `TrajectoryDataset`)
- Modify: `tests/test_trainer.py` (add tests, keep existing ones passing)
- [ ] **Step 1: Write failing test for the new Dataset class**
Append to `tests/test_trainer.py`:
```python
class TestTrajectoryDataset:
def test_dataset_length_with_augmentation(self):
"""Dataset length = N * 6 when augment=True."""
from ai_mouse.trainer import TrajectoryDataset
seq = np.zeros((10, 64, 3), dtype=np.float32)
cond = np.zeros((10, 3), dtype=np.float32)
ds = TrajectoryDataset(seq, cond, augment=True)
assert len(ds) == 60
def test_dataset_length_without_augmentation(self):
from ai_mouse.trainer import TrajectoryDataset
seq = np.zeros((10, 64, 3), dtype=np.float32)
cond = np.zeros((10, 3), dtype=np.float32)
ds = TrajectoryDataset(seq, cond, augment=False)
assert len(ds) == 10
def test_getitem_returns_tensors(self):
from ai_mouse.trainer import TrajectoryDataset
import torch
seq = np.random.randn(5, 64, 3).astype(np.float32)
cond = np.random.randn(5, 3).astype(np.float32)
ds = TrajectoryDataset(seq, cond, augment=True)
s, c = ds[0]
assert isinstance(s, torch.Tensor)
assert isinstance(c, torch.Tensor)
assert s.shape == (64, 3)
assert c.shape == (3,)
def test_aug_id_zero_returns_original(self):
"""Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
from ai_mouse.trainer import TrajectoryDataset
import torch
seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
ds = TrajectoryDataset(seq, cond, augment=True)
s0, c0 = ds[0]
np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5)
np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5)
def test_aug_id_one_flips_lateral(self):
"""Aug id 1 should flip the sign of the lateral channel (index 1)."""
from ai_mouse.trainer import TrajectoryDataset
seq = np.zeros((1, 64, 3), dtype=np.float32)
seq[0, :, 1] = 0.5 # lateral all positive
cond = np.zeros((1, 3), dtype=np.float32)
ds = TrajectoryDataset(seq, cond, augment=True)
# idx=1 → base_idx=0, aug_id=1 → flip
s1, _ = ds[1]
assert (s1[:, 1] < 0).all()
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_trainer.py::TestTrajectoryDataset -v
```
Expected: FAIL with `ImportError` (TrajectoryDataset not defined).
- [ ] **Step 3: Implement TrajectoryDataset class**
Add to `ai_mouse/trainer.py`, just below the `_augment` function (keep `_augment` for now — Task removes it):
```python
class TrajectoryDataset(torch.utils.data.Dataset):
"""Trajectory dataset with on-the-fly 6× augmentation.
Replaces the old eager `_augment(seq, cond)` which expanded the dataset
6× in memory before training. With this class, the original (N, T, 3)
arrays stay as-is and each `__getitem__` call computes one of the 6
augmentation variants on demand.
Augmentation variants (matching legacy `_augment` semantics):
0 — original
1 — lateral flip (lateral → −lateral)
2 — speed ×0.8 (log_dt[1:] += log(1.25), cond[2] += log(1.25))
3 — speed ×1.2 (log_dt[1:] += log(1/1.2), cond[2] += log(1/1.2))
4 — temporal noise (log_dt[1:] += N(0, 0.05))
5 — flip + speed ×0.9 (lateral flip, log_dt[1:] += log(1/0.9), cond[2] += log(1/0.9))
"""
_LOG_1_25 = math.log(1.25)
_LOG_INV_1_2 = math.log(1.0 / 1.2)
_LOG_1_1 = math.log(1.0 / 0.9)
def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
self.seq = seq
self.cond = cond
self.augment = augment
self._n_aug = 6 if augment else 1
def __len__(self) -> int:
return len(self.seq) * self._n_aug
def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
base = idx // self._n_aug
aug_id = idx % self._n_aug
s = self.seq[base].copy()
c = self.cond[base].copy()
if aug_id == 1:
s[:, 1] = -s[:, 1]
elif aug_id == 2:
s[1:, 2] += self._LOG_1_25
c[2] += self._LOG_1_25
elif aug_id == 3:
s[1:, 2] += self._LOG_INV_1_2
c[2] += self._LOG_INV_1_2
elif aug_id == 4:
noise = np.random.normal(0.0, 0.05, size=s[1:, 2].shape).astype(np.float32)
s[1:, 2] += noise
elif aug_id == 5:
s[:, 1] = -s[:, 1]
s[1:, 2] += self._LOG_1_1
c[2] += self._LOG_1_1
return torch.from_numpy(s), torch.from_numpy(c)
```
- [ ] **Step 4: Run new tests, verify pass**
```bash
uv run pytest tests/test_trainer.py::TestTrajectoryDataset -v
```
Expected: 5 tests pass.
- [ ] **Step 5: Switch `train()` to use TrajectoryDataset**
In `ai_mouse/trainer.py`, find the existing block (around line 332):
```python
# ---- Augment ----
seq_np, cond_np = _augment(seq_np, cond_np)
logger.info("After augmentation: %d samples", len(seq_np))
seq_t = torch.from_numpy(seq_np) # (N, seq_len, 3)
cond_t = torch.from_numpy(cond_np) # (N, 3)
```
Replace with:
```python
# ---- Build streaming dataset (on-the-fly 6× augmentation) ----
if config.augment:
logger.info("Using on-the-fly 6× augmentation, base samples: %d", len(seq_np))
ds = TrajectoryDataset(seq_np, cond_np, augment=config.augment)
logger.info("Effective dataset size: %d", len(ds))
```
Then find the existing DataLoader construction (around line 353):
```python
ds = TensorDataset(seq_t, cond_t)
loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
```
Replace with:
```python
loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
```
(The new `ds` already exists from the previous edit — just remove the TensorDataset construction.)
Also remove the now-unused import line `from torch.utils.data import DataLoader, TensorDataset` and replace with:
```python
from torch.utils.data import DataLoader
```
- [ ] **Step 6: Verify legacy augmentation tests still pass**
The existing test `TestAugment::test_augmentation_multiplies_data` calls `_augment` directly, which we kept. Run all trainer tests:
```bash
uv run pytest tests/test_trainer.py -v
```
Expected: all tests pass (including TestAugment, TestLoadAndPrepare, TestTrain, TestTrajectoryDataset).
- [ ] **Step 7: Commit**
```bash
git add ai_mouse/trainer.py tests/test_trainer.py
git commit -m "feat(trainer): replace eager _augment with streaming TrajectoryDataset"
```
---
## Task 8: Trainer — Resume from Checkpoint
**Files:**
- Modify: `ai_mouse/trainer.py` (add `resume_from` parameter to `train()`)
- Modify: `tests/test_trainer.py`
- [ ] **Step 1: Write failing test for resume_from**
Append to `tests/test_trainer.py`:
```python
class TestResumeFrom:
def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):
"""train() with resume_from should load weights from given checkpoint dir."""
import torch
from ai_mouse.trainer import train
from ai_mouse.models import TrajectoryFlowModel
# First, train an initial model and save it
ckpt_dir = tmp_path / "pretrain"
train(
data_path=synthetic_traces_file,
output_dir=ckpt_dir,
epochs=2,
batch_size=8,
seq_len=64,
)
assert (ckpt_dir / "flow_model.pt").exists()
# Read its weights to compare later
m_pretrain = TrajectoryFlowModel(seq_len=64)
m_pretrain.load_state_dict(torch.load(ckpt_dir / "flow_model.pt", weights_only=True))
first_param_pre = next(m_pretrain.parameters()).clone()
# Now train with resume_from for 0 epochs — weights should still be loaded
out_dir = tmp_path / "finetune"
train(
data_path=synthetic_traces_file,
output_dir=out_dir,
epochs=1,
batch_size=8,
seq_len=64,
resume_from=ckpt_dir,
)
m_after = TrajectoryFlowModel(seq_len=64)
m_after.load_state_dict(torch.load(out_dir / "flow_model.pt", weights_only=True))
first_param_after = next(m_after.parameters())
# After 1 epoch, weights should be close to pre-train, not random init
# (random init would be O(1) magnitude apart; 1 epoch on small data shifts O(0.1))
diff = (first_param_pre - first_param_after).abs().mean().item()
assert diff < 0.5, f"Resume_from weights diverged too much: {diff}"
def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path):
from ai_mouse.trainer import train
with pytest.raises(FileNotFoundError):
train(
data_path=synthetic_traces_file,
output_dir=tmp_path / "out",
epochs=1,
batch_size=8,
seq_len=64,
resume_from=tmp_path / "nonexistent",
)
```
- [ ] **Step 2: Run test, verify failure**
```bash
uv run pytest tests/test_trainer.py::TestResumeFrom -v
```
Expected: FAIL with `TypeError: train() got an unexpected keyword argument 'resume_from'`.
- [ ] **Step 3: Add resume_from parameter to train()**
In `ai_mouse/trainer.py`, modify the `train` function signature:
```python
def train(
data_path: Path,
output_dir: Path,
epochs: int = 300,
batch_size: int = 64,
lr: float = 3e-4,
seq_len: int = 64,
progress_callback: Callable[[dict], None] | None = None,
config: TrainConfig | None = None,
resume_from: Path | None = None,
) -> None:
```
Update the docstring's Args section to add:
```
resume_from: if given, load model weights from this checkpoint
directory (must contain flow_model.pt). Used for
two-stage training (pretrain → fine-tune).
```
Then, after model construction (find the line `model = TrajectoryFlowModel(...)` around line 341) and BEFORE the optimiser is created, insert:
```python
# ---- Resume from checkpoint if requested ----
if resume_from is not None:
resume_path = Path(resume_from) / "flow_model.pt"
if not resume_path.exists():
raise FileNotFoundError(
f"resume_from checkpoint not found: {resume_path}"
)
logger.info("Resuming from checkpoint: %s", resume_path)
state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_trainer.py::TestResumeFrom -v
```
Expected: 2 tests pass.
- [ ] **Step 5: Run all trainer tests**
```bash
uv run pytest tests/test_trainer.py -v
```
Expected: all tests pass.
- [ ] **Step 6: Commit**
```bash
git add ai_mouse/trainer.py tests/test_trainer.py
git commit -m "feat(trainer): add resume_from for two-stage training"
```
---
## Task 9: Server — Auto-Resume When Pretrained Checkpoint Exists
**Files:**
- Modify: `ai_mouse/server/routes_train.py`
- Modify: `tests/test_server.py` (verify still passes)
- [ ] **Step 1: Update routes_train.py to detect and pass resume_from**
In `ai_mouse/server/routes_train.py`, modify `_paths()` to also return the pretrained dir:
Find:
```python
def _paths() -> tuple[Path, Path]:
data_dir = get_data_dir()
return data_dir / "traces.jsonl", data_dir / "models_v2"
```
Replace with:
```python
def _paths() -> tuple[Path, Path, Path]:
data_dir = get_data_dir()
return (
data_dir / "traces.jsonl",
data_dir / "models_v2",
data_dir / "models_v2_pretrained",
)
```
Update `_trace_count()` and `_model_trained()`:
```python
def _trace_count() -> int:
traces_path, _, _ = _paths()
if not traces_path.exists():
return 0
return sum(
1
for line in traces_path.read_text(encoding="utf-8").splitlines()
if line.strip()
)
def _model_trained() -> bool:
_, models_dir, _ = _paths()
return (models_dir / "flow_model.pt").exists()
```
Then update `_train_sse_generator`'s inner `run_training_async`:
Find:
```python
async def run_training_async() -> None:
from ai_mouse.trainer import train
traces_path, models_dir = _paths()
data_path = Path(req.data_path) if req.data_path else traces_path
output_dir = Path(req.output_dir) if req.output_dir else models_dir
try:
await asyncio.to_thread(
train,
data_path=data_path,
output_dir=output_dir,
epochs=req.epochs,
progress_callback=callback,
)
except Exception as exc: # noqa: BLE001
queue.put_nowait({"error": str(exc)})
```
Replace with:
```python
async def run_training_async() -> None:
from ai_mouse.trainer import train
traces_path, models_dir, pretrained_dir = _paths()
data_path = Path(req.data_path) if req.data_path else traces_path
output_dir = Path(req.output_dir) if req.output_dir else models_dir
# Auto-detect pretrained checkpoint and switch to fine-tune mode
resume_from: Path | None = None
effective_lr = 3e-4
if (pretrained_dir / "flow_model.pt").exists():
resume_from = pretrained_dir
effective_lr = 1e-5 # fine-tune lr
logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
queue.put_nowait({
"info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
f"running fine-tune at lr={effective_lr}",
})
try:
await asyncio.to_thread(
train,
data_path=data_path,
output_dir=output_dir,
epochs=req.epochs,
lr=effective_lr,
progress_callback=callback,
resume_from=resume_from,
)
except Exception as exc: # noqa: BLE001
queue.put_nowait({"error": str(exc)})
```
- [ ] **Step 2: Run server tests, verify still pass**
```bash
uv run pytest tests/test_server.py -v
```
Expected: all tests pass. The auto-resume behaviour is exercised when `models_v2_pretrained` exists, which is not the case in the default test environment, so the existing tests still hit the from-scratch path.
- [ ] **Step 3: Run full test suite**
```bash
uv run pytest -x
```
Expected: all tests pass.
- [ ] **Step 4: Commit**
```bash
git add ai_mouse/server/routes_train.py
git commit -m "feat(server): auto-resume from pretrained checkpoint when available"
```
---
## Task 10: Generator — Remove speed_profile and median±1.1 Hard Clip
**Files:**
- Modify: `ai_mouse/generator.py`
- Modify: `tests/test_generator.py`
- [ ] **Step 1: Update existing tests to be tolerant of the new behaviour**
The current tests use freshly-initialised (untrained) weights and only check structural properties (timestamps monotonic, endpoints close, click events present). These should keep passing — verify.
```bash
uv run pytest tests/test_generator.py -v
```
Expected: all tests pass.
Add a new test that explicitly checks the OLD speed_profile is GONE (regression guard):
Append to `tests/test_generator.py`:
```python
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
# (untrained model produces noisy outputs but post-processing must
# not collapse them to the same template)
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")
```
- [ ] **Step 2: Run test, may pass for untrained model but verifies regression guard**
```bash
uv run pytest tests/test_generator.py::TestPostProcessing -v
```
Expected: PASS (untrained models produce random outputs, so they vary).
- [ ] **Step 3: Remove the speed_profile and median clip blocks**
In `ai_mouse/generator.py`, find this block (lines ~261-272):
```python
# The model tends to produce exaggerated deceleration at the tail
# (last 10 points log_dt ~3-5 vs middle ~1.5).
# Cap the max-to-median ratio to ~3× (i.e., tail Δt ≤ 3× median Δt)
median_ldt = float(np.median(log_dt[1:]))
# Allow max log_dt = median + 1.1 (exp(1.1) ≈ 3× ratio)
max_allowed = median_ldt + 1.1
min_allowed = max(median_ldt - 1.1, 0.0)
for i in range(1, len(log_dt)):
if log_dt[i] > max_allowed:
log_dt[i] = max_allowed
elif log_dt[i] < min_allowed:
log_dt[i] = min_allowed
```
DELETE this entire block.
Then find the speed_profile block immediately after (lines ~273-286):
```python
# Apply asymmetric speed profile: start slow, fast in middle, gentle end
# Mimics natural mouse movement (accelerate → cruise → decelerate)
t_frac = np.linspace(0, 1, len(log_dt))
speed_profile = np.zeros_like(log_dt, dtype=float)
for i in range(1, len(log_dt)):
t = t_frac[i]
if t < 0.15:
# Acceleration phase: start slow (+0.3 at t=0, → 0 at t=0.15)
speed_profile[i] = 0.3 * (1.0 - t / 0.15)
elif t > 0.85:
# Deceleration phase: end slightly slow (+0.2 at t=1)
speed_profile[i] = 0.2 * ((t - 0.85) / 0.15)
# Middle: speed_profile = 0 (fastest, no penalty)
log_dt[1:] = log_dt[1:] + speed_profile[1:]
```
DELETE this entire block too.
- [ ] **Step 4: Update the docstring at the top of the file**
In `ai_mouse/generator.py`, find the module docstring (lines 1-22). Replace it with:
```python
"""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).
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.
"""
```
- [ ] **Step 5: Run tests**
```bash
uv run pytest tests/test_generator.py -v
```
Expected: all tests pass.
- [ ] **Step 6: Commit**
```bash
git add ai_mouse/generator.py tests/test_generator.py
git commit -m "refactor(generator): remove deterministic speed_profile and hard log_dt clip
These post-processing hacks were added to compensate for small-data
training. With Balabit pretraining they suppress the multimodal
timing distribution and cause the template-y Δt curves seen in the
verify UI."
```
---
## Task 11: Generator — Add Lateral Gaussian Smoothing
**Files:**
- Modify: `ai_mouse/generator.py`
- Modify: `tests/test_generator.py`
- [ ] **Step 1: Write failing test for the smoothing helper**
Append to `tests/test_generator.py`:
```python
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)
```
- [ ] **Step 2: Run test, verify failure**
```bash
uv run pytest tests/test_generator.py::TestGaussianSmooth -v
```
Expected: FAIL with `ImportError: cannot import name '_gaussian_smooth'`.
- [ ] **Step 3: Implement _gaussian_smooth**
Add to `ai_mouse/generator.py`, just below the imports (before any function definitions):
```python
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()
smoothed = np.convolve(x, kernel, mode="same")
smoothed[0] = x[0]
smoothed[-1] = x[-1]
return smoothed
```
- [ ] **Step 4: Run smoothing tests, verify pass**
```bash
uv run pytest tests/test_generator.py::TestGaussianSmooth -v
```
Expected: 4 tests pass.
- [ ] **Step 5: Apply smoothing to lateral in generate()**
In `ai_mouse/generator.py`, find the spatial post-processing block. After this code:
```python
# 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
```
Add:
```python
# Lateral 5-point gaussian smoothing (endpoints preserved)
lateral = _gaussian_smooth(lateral, sigma=1.0)
```
- [ ] **Step 6: Update module docstring to document the new smoothing step**
In the module docstring at the top of `ai_mouse/generator.py`, find:
```python
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).
```
Replace with:
```python
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).
```
- [ ] **Step 7: Run all generator tests**
```bash
uv run pytest tests/test_generator.py -v
```
Expected: all tests pass.
- [ ] **Step 8: Commit**
```bash
git add ai_mouse/generator.py tests/test_generator.py
git commit -m "feat(generator): add 5-point gaussian smoothing on lateral"
```
---
## Task 12: Eval — Kinematics Metrics
**Files:**
- Create: `ai_mouse/eval/__init__.py`
- Create: `ai_mouse/eval/metrics.py`
- Create: `tests/test_eval_metrics.py`
- [ ] **Step 1: Write failing tests for kinematics metrics**
Create `tests/test_eval_metrics.py`:
```python
"""Tests for the eval metrics module."""
from __future__ import annotations
import numpy as np
import pytest
class TestKinematics:
def test_compute_speed_constant_velocity(self):
"""Constant-velocity trajectory has constant speed."""
from ai_mouse.eval.metrics import compute_speed
# 10 points, moving 10 px in 100 ms each step → speed = 0.1 px/ms
xs = np.arange(0, 100, 10, dtype=float)
ys = np.zeros(10, dtype=float)
ts = np.arange(0, 1000, 100, dtype=float)
v = compute_speed(xs, ys, ts)
# All speeds should be ≈ 0.1 px/ms
assert v.shape == (9,) # n-1 differences
np.testing.assert_allclose(v, 0.1, rtol=1e-4)
def test_compute_speed_handles_zero_dt(self):
"""Adjacent points with same timestamp must not produce NaN/inf."""
from ai_mouse.eval.metrics import compute_speed
xs = np.array([0.0, 10.0, 20.0])
ys = np.array([0.0, 0.0, 0.0])
ts = np.array([0.0, 0.0, 100.0]) # zero dt between [0] and [1]
v = compute_speed(xs, ys, ts)
assert np.isfinite(v).all()
def test_compute_acceleration(self):
"""Linearly increasing speed → constant acceleration."""
from ai_mouse.eval.metrics import compute_acceleration
# speeds: 0.1, 0.2, 0.3, 0.4 over dt = 100 ms each → a = 0.001 px/ms²
speeds = np.array([0.1, 0.2, 0.3, 0.4])
ts = np.array([100.0, 200.0, 300.0, 400.0])
a = compute_acceleration(speeds, ts)
np.testing.assert_allclose(a, 0.001, rtol=1e-4)
def test_compute_jerk(self):
from ai_mouse.eval.metrics import compute_jerk
# accelerations: 0.001, 0.002, 0.003 over dt = 100 ms → j = 0.00001
accels = np.array([0.001, 0.002, 0.003])
ts = np.array([200.0, 300.0, 400.0])
j = compute_jerk(accels, ts)
np.testing.assert_allclose(j, 1e-5, rtol=1e-4)
class TestStatsSummary:
def test_compute_stats_returns_expected_keys(self):
from ai_mouse.eval.metrics import compute_stats
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
s = compute_stats(x)
assert "mean" in s
assert "std" in s
assert "cv" in s
assert "p25" in s
assert "p50" in s
assert "p75" in s
assert "p95" in s
def test_cv_for_constant_is_zero(self):
from ai_mouse.eval.metrics import compute_stats
x = np.full(10, 3.0)
s = compute_stats(x)
assert s["cv"] == 0.0
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_eval_metrics.py::TestKinematics tests/test_eval_metrics.py::TestStatsSummary -v
```
Expected: FAIL with `ModuleNotFoundError: No module named 'ai_mouse.eval'`.
- [ ] **Step 3: Create the eval package and metrics module**
Create `ai_mouse/eval/__init__.py`:
```python
"""Evaluation module: kinematic metrics and Markdown report generation."""
```
Create `ai_mouse/eval/metrics.py`:
```python
"""Kinematic metrics for mouse trajectory evaluation.
All inputs are 1-D NumPy arrays. Time is in milliseconds, position in pixels.
Velocities are px/ms, accelerations px/ms², jerks px/ms³.
"""
from __future__ import annotations
import numpy as np
def compute_speed(
xs: np.ndarray, ys: np.ndarray, ts: np.ndarray, eps: float = 1e-6
) -> np.ndarray:
"""Compute scalar speed at each step.
Args:
xs: (N,) x coordinates.
ys: (N,) y coordinates.
ts: (N,) timestamps in ms.
eps: minimum dt (ms) to avoid div-by-zero.
Returns:
(N-1,) array of speeds (px/ms).
"""
dx = np.diff(xs)
dy = np.diff(ys)
dt = np.maximum(np.diff(ts), eps)
return np.hypot(dx, dy) / dt
def compute_acceleration(speeds: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Compute scalar acceleration from speeds.
Args:
speeds: (M,) speeds (px/ms). Typically M = N-1 from compute_speed.
ts: (M+1,) timestamps that produced those speeds (ms).
We use the midpoints between adjacent ts.
eps: minimum dt (ms) to avoid div-by-zero.
Returns:
(M-1,) array of accelerations (px/ms²).
"""
if len(speeds) < 2:
return np.array([], dtype=float)
# ts has length M+1; we need M-1 dts between speed points
# speed[i] is between ts[i] and ts[i+1], so it's at ts midpoint (ts[i]+ts[i+1])/2
midpoints = (ts[:-1] + ts[1:]) / 2.0
dt = np.maximum(np.diff(midpoints), eps)
return np.diff(speeds) / dt
def compute_jerk(accels: np.ndarray, ts: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Compute jerk from accelerations.
Args:
accels: (K,) accelerations.
ts: (K+2,) timestamps that produced those accelerations.
Used to derive midpoint-of-midpoint dts.
eps: minimum dt to avoid div-by-zero.
Returns:
(K-1,) array of jerks (px/ms³).
"""
if len(accels) < 2:
return np.array([], dtype=float)
# Approximate dt for jerks as average dt of original ts (good enough for stats)
dt_avg = np.maximum(np.diff(ts).mean(), eps)
return np.diff(accels) / dt_avg
def compute_stats(x: np.ndarray) -> dict[str, float]:
"""Summary statistics for a 1-D distribution.
Returns:
dict with keys: mean, std, cv (coef of variation), p25, p50, p75, p95.
"""
if len(x) == 0:
return {k: 0.0 for k in ("mean", "std", "cv", "p25", "p50", "p75", "p95")}
x = np.asarray(x, dtype=float)
mean = float(x.mean())
std = float(x.std(ddof=1)) if len(x) > 1 else 0.0
cv = std / mean if mean != 0 else 0.0
return {
"mean": mean,
"std": std,
"cv": cv,
"p25": float(np.percentile(x, 25)),
"p50": float(np.percentile(x, 50)),
"p75": float(np.percentile(x, 75)),
"p95": float(np.percentile(x, 95)),
}
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_eval_metrics.py -v
```
Expected: 6 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/eval/ tests/test_eval_metrics.py
git commit -m "feat(eval): kinematics metrics (speed, accel, jerk, stats)"
```
---
## Task 13: Eval — FFT Spectrum + KL Divergence
**Files:**
- Modify: `ai_mouse/eval/metrics.py` (add FFT and KL functions)
- Modify: `tests/test_eval_metrics.py`
- [ ] **Step 1: Write failing tests**
Append to `tests/test_eval_metrics.py`:
```python
class TestFftSpectrum:
def test_finds_dominant_frequency(self):
"""A pure 8 Hz signal should have its peak near 8 Hz."""
from ai_mouse.eval.metrics import fft_spectrum
# Sample at 100 Hz for 1 second
sample_rate_hz = 100.0
ts_ms = np.arange(0, 1000, 1000 / sample_rate_hz)
signal = np.sin(2 * np.pi * 8 * ts_ms / 1000) # 8 Hz sine
freqs, mags = fft_spectrum(signal, sample_rate_hz)
peak_freq = freqs[np.argmax(mags)]
assert abs(peak_freq - 8.0) < 1.0 # within 1 Hz
def test_returns_only_positive_frequencies(self):
from ai_mouse.eval.metrics import fft_spectrum
signal = np.random.randn(64)
freqs, mags = fft_spectrum(signal, 50.0)
assert (freqs >= 0).all()
assert len(freqs) == len(mags)
class TestKlDivergence:
def test_identical_distributions_zero_kl(self):
"""KL(p, p) ≈ 0."""
from ai_mouse.eval.metrics import kl_divergence_histograms
rng = np.random.default_rng(42)
x = rng.normal(0, 1, 5000)
y = rng.normal(0, 1, 5000)
kl = kl_divergence_histograms(x, y, bins=50)
assert kl < 0.05
def test_different_distributions_positive_kl(self):
"""Different means → positive KL."""
from ai_mouse.eval.metrics import kl_divergence_histograms
rng = np.random.default_rng(42)
x = rng.normal(0, 1, 5000)
y = rng.normal(3, 1, 5000)
kl = kl_divergence_histograms(x, y, bins=50)
assert kl > 0.5
def test_handles_disjoint_supports(self):
"""No NaN even when histograms have non-overlapping bins."""
from ai_mouse.eval.metrics import kl_divergence_histograms
x = np.array([1.0, 1.1, 1.2, 1.3, 1.4])
y = np.array([10.0, 10.1, 10.2, 10.3, 10.4])
kl = kl_divergence_histograms(x, y, bins=10)
assert np.isfinite(kl)
```
- [ ] **Step 2: Run tests, verify failure**
```bash
uv run pytest tests/test_eval_metrics.py::TestFftSpectrum tests/test_eval_metrics.py::TestKlDivergence -v
```
Expected: FAIL with `ImportError`.
- [ ] **Step 3: Implement fft_spectrum and kl_divergence_histograms**
Append to `ai_mouse/eval/metrics.py`:
```python
def fft_spectrum(
signal: np.ndarray, sample_rate_hz: float
) -> tuple[np.ndarray, np.ndarray]:
"""Compute one-sided FFT magnitude spectrum.
Args:
signal: 1-D real-valued signal.
sample_rate_hz: Sampling rate in Hz.
Returns:
(freqs, magnitudes) — positive frequencies only.
Magnitudes are absolute values of complex FFT coefficients.
"""
n = len(signal)
if n == 0:
return np.array([]), np.array([])
# Zero-mean to remove DC component which dominates the spectrum
s = signal - signal.mean()
fft = np.fft.rfft(s)
freqs = np.fft.rfftfreq(n, d=1.0 / sample_rate_hz)
return freqs, np.abs(fft)
def kl_divergence_histograms(
x: np.ndarray,
y: np.ndarray,
bins: int = 50,
eps: float = 1e-10,
) -> float:
"""KL divergence KL(P_x || P_y) estimated via shared-bin histograms.
Both arrays are histogrammed over their joint range. Empty bins get
`eps` mass to avoid log(0) — keeps result finite even for disjoint
supports.
Args:
x: samples from distribution P.
y: samples from distribution Q (the "reference").
bins: number of histogram bins.
eps: smoothing constant for empty bins.
Returns:
scalar KL divergence (nats). Always finite, ≥ 0.
"""
if len(x) == 0 or len(y) == 0:
return 0.0
lo = float(min(x.min(), y.min()))
hi = float(max(x.max(), y.max()))
if hi <= lo:
return 0.0
edges = np.linspace(lo, hi, bins + 1)
px, _ = np.histogram(x, bins=edges, density=False)
qy, _ = np.histogram(y, bins=edges, density=False)
px = px.astype(float) + eps
qy = qy.astype(float) + eps
px /= px.sum()
qy /= qy.sum()
return float(np.sum(px * np.log(px / qy)))
```
- [ ] **Step 4: Run tests, verify pass**
```bash
uv run pytest tests/test_eval_metrics.py::TestFftSpectrum tests/test_eval_metrics.py::TestKlDivergence -v
```
Expected: 5 tests pass.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/eval/metrics.py tests/test_eval_metrics.py
git commit -m "feat(eval): add FFT spectrum and KL divergence metrics"
```
---
## Task 14: Eval — Report Generation
**Files:**
- Create: `ai_mouse/eval/report.py`
- Modify: `tests/test_eval_metrics.py` (add light-weight smoke test)
- [ ] **Step 1: Write a smoke test that the report module is importable and produces output**
Append to `tests/test_eval_metrics.py`:
```python
class TestReportGeneration:
def test_generates_report_md(self, tmp_path):
"""Smoke test: build_report writes an MD file with all expected sections."""
from ai_mouse.eval.report import build_report
# Synthetic generated traces (3 traces, 50 points each)
rng = np.random.default_rng(0)
gen_traces = []
for _ in range(3):
xs = np.cumsum(rng.uniform(0, 5, 50))
ys = np.cumsum(rng.uniform(-1, 1, 50))
ts = np.cumsum(rng.uniform(5, 20, 50))
gen_traces.append({"xs": xs, "ys": ys, "ts": ts})
# Synthetic reference
ref_traces = []
for _ in range(5):
xs = np.cumsum(rng.uniform(0, 5, 50))
ys = np.cumsum(rng.uniform(-1, 1, 50))
ts = np.cumsum(rng.uniform(5, 20, 50))
ref_traces.append({"xs": xs, "ys": ys, "ts": ts})
out_md = tmp_path / "report.md"
build_report(
generated_traces=gen_traces,
reference_traces=ref_traces,
output_md=out_md,
tag="smoke-test",
model_dir="/fake/model/dir",
)
assert out_md.exists()
content = out_md.read_text(encoding="utf-8")
assert "# Eval Report" in content
assert "smoke-test" in content
assert "速度" in content or "speed" in content.lower()
assert "FFT" in content.upper()
# PNG plots should exist next to MD
plot_dir = tmp_path / "plots"
assert plot_dir.exists()
assert any(plot_dir.iterdir())
```
- [ ] **Step 2: Run test, verify failure**
```bash
uv run pytest tests/test_eval_metrics.py::TestReportGeneration -v
```
Expected: FAIL with `ImportError: cannot import build_report`.
- [ ] **Step 3: Implement report.py**
Create `ai_mouse/eval/report.py`:
```python
"""Markdown report generation for eval results.
Outputs a self-contained .md file with embedded PNG plots in a sibling
'plots/' directory.
"""
from __future__ import annotations
import logging
from datetime import datetime
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg") # headless
import matplotlib.pyplot as plt # noqa: E402
from ai_mouse.eval.metrics import (
compute_acceleration,
compute_jerk,
compute_speed,
compute_stats,
fft_spectrum,
kl_divergence_histograms,
)
logger = logging.getLogger(__name__)
def _aggregate_kinematics(traces: list[dict]) -> dict[str, np.ndarray]:
"""Concatenate per-trace speed/accel/jerk arrays from a list of traces.
Args:
traces: list of {"xs", "ys", "ts"} dicts (1-D ndarrays).
Returns:
dict with keys "speed", "accel", "jerk", "dt" — each a flat ndarray.
"""
speeds, accels, jerks, dts = [], [], [], []
for tr in traces:
xs, ys, ts = tr["xs"], tr["ys"], tr["ts"]
if len(xs) < 4:
continue
v = compute_speed(xs, ys, ts)
a = compute_acceleration(v, ts)
j = compute_jerk(a, ts)
speeds.append(v)
accels.append(a)
jerks.append(j)
dts.append(np.diff(ts))
return {
"speed": np.concatenate(speeds) if speeds else np.array([]),
"accel": np.concatenate(accels) if accels else np.array([]),
"jerk": np.concatenate(jerks) if jerks else np.array([]),
"dt": np.concatenate(dts) if dts else np.array([]),
}
def _plot_distribution(
gen: np.ndarray,
ref: np.ndarray,
title: str,
output: Path,
xlabel: str,
bins: int = 50,
) -> None:
"""Side-by-side histogram of gen vs ref."""
fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
if len(gen) > 0:
ax.hist(gen, bins=bins, alpha=0.5, label="生成", density=True)
if len(ref) > 0:
ax.hist(ref, bins=bins, alpha=0.5, label="参考 (Balabit)", density=True)
ax.set_title(title)
ax.set_xlabel(xlabel)
ax.set_ylabel("密度")
ax.legend()
fig.tight_layout()
fig.savefig(output)
plt.close(fig)
def _plot_fft_overlay(
gen_traces: list[dict],
ref_traces: list[dict],
output: Path,
sample_rate_hz: float = 100.0,
) -> None:
"""Average FFT magnitude over lateral component for gen vs ref."""
def _avg_spectrum(traces: list[dict]) -> tuple[np.ndarray, np.ndarray]:
all_freqs = None
all_mags = []
for tr in traces:
xs, ys = tr["xs"], tr["ys"]
if len(xs) < 8:
continue
# Use the cross-track ('lateral') signal: project onto perpendicular
# of start→end vector. Approximate by detrended y.
sig = ys - np.linspace(ys[0], ys[-1], len(ys))
f, m = fft_spectrum(sig, sample_rate_hz)
if all_freqs is None:
all_freqs = f
all_mags.append(m)
elif len(m) == len(all_freqs):
all_mags.append(m)
if not all_mags:
return np.array([]), np.array([])
return all_freqs, np.mean(all_mags, axis=0)
fig, ax = plt.subplots(figsize=(8, 4), dpi=100)
f_gen, m_gen = _avg_spectrum(gen_traces)
f_ref, m_ref = _avg_spectrum(ref_traces)
if len(f_gen) > 0:
ax.plot(f_gen, m_gen, label="生成", alpha=0.7)
if len(f_ref) > 0:
ax.plot(f_ref, m_ref, label="参考 (Balabit)", alpha=0.7)
ax.axvspan(4, 12, alpha=0.1, color="green", label="生理震颤区间 4–12 Hz")
ax.set_title("FFT 频谱（横向偏移信号）")
ax.set_xlabel("Hz")
ax.set_ylabel("|FFT|")
ax.set_xlim(0, sample_rate_hz / 2)
ax.legend()
fig.tight_layout()
fig.savefig(output)
plt.close(fig)
def _plot_paths_overlay(traces: list[dict], output: Path, max_traces: int = 5) -> None:
"""Plot up to N generated trajectories on the same axes."""
fig, ax = plt.subplots(figsize=(6, 5), dpi=100)
for i, tr in enumerate(traces[:max_traces]):
ax.plot(tr["xs"], tr["ys"], alpha=0.6, label=f"路径 {i+1}")
ax.invert_yaxis() # screen coords
ax.set_title(f"前 {min(max_traces, len(traces))} 条生成轨迹")
ax.set_xlabel("x (px)")
ax.set_ylabel("y (px)")
ax.set_aspect("equal", adjustable="datalim")
ax.legend()
fig.tight_layout()
fig.savefig(output)
plt.close(fig)
def build_report(
generated_traces: list[dict],
reference_traces: list[dict],
output_md: Path,
tag: str,
model_dir: str,
sample_rate_hz: float = 100.0,
) -> None:
"""Build a Markdown eval report with embedded plots.
Args:
generated_traces: list of {"xs","ys","ts"} from the generator under test.
reference_traces: list of {"xs","ys","ts"} from Balabit (ground truth).
output_md: destination .md path. plots/ created in same dir.
tag: short identifier (e.g. "baseline", "post-finetune").
model_dir: model directory path string (for provenance).
sample_rate_hz: nominal sample rate for FFT (mouse data is irregular —
100 Hz is a sensible nominal).
"""
plot_dir = output_md.parent / "plots"
plot_dir.mkdir(parents=True, exist_ok=True)
gen_kin = _aggregate_kinematics(generated_traces)
ref_kin = _aggregate_kinematics(reference_traces)
# --- KL divergences ---
kl_speed = kl_divergence_histograms(gen_kin["speed"], ref_kin["speed"])
kl_accel = kl_divergence_histograms(gen_kin["accel"], ref_kin["accel"])
kl_jerk = kl_divergence_histograms(gen_kin["jerk"], ref_kin["jerk"])
kl_dt = kl_divergence_histograms(gen_kin["dt"], ref_kin["dt"])
# --- Stats ---
stats_gen = {k: compute_stats(v) for k, v in gen_kin.items()}
stats_ref = {k: compute_stats(v) for k, v in ref_kin.items()}
# --- Plots ---
_plot_distribution(gen_kin["speed"], ref_kin["speed"],
"速度分布", plot_dir / f"{tag}-speed.png", "px/ms")
_plot_distribution(gen_kin["accel"], ref_kin["accel"],
"加速度分布", plot_dir / f"{tag}-accel.png", "px/ms²")
_plot_distribution(gen_kin["jerk"], ref_kin["jerk"],
"Jerk 分布", plot_dir / f"{tag}-jerk.png", "px/ms³")
_plot_distribution(gen_kin["dt"], ref_kin["dt"],
"Δt 分布", plot_dir / f"{tag}-dt.png", "ms")
_plot_fft_overlay(generated_traces, reference_traces,
plot_dir / f"{tag}-fft.png", sample_rate_hz)
_plot_paths_overlay(generated_traces, plot_dir / f"{tag}-paths.png")
# --- Markdown ---
now = datetime.now().strftime("%Y-%m-%d %H:%M")
lines = [
f"# Eval Report: {tag} ({now})",
"",
"## 模型信息",
f"- Checkpoint dir: `{model_dir}`",
f"- 生成样本数: {len(generated_traces)}",
f"- 参考样本数: {len(reference_traces)}",
"",
"## KL 散度（生成 vs 参考，越小越好）",
"| 指标 | KL |",
"|---|---|",
f"| 速度分布 | {kl_speed:.4f} |",
f"| 加速度分布 | {kl_accel:.4f} |",
f"| Jerk 分布 | {kl_jerk:.4f} |",
f"| Δt 分布 | {kl_dt:.4f} |",
"",
"## 摘要统计",
"| 指标 | 生成 mean | 参考 mean | 生成 CV | 参考 CV |",
"|---|---|---|---|---|",
]
for key, label in [("speed", "速度"), ("accel", "加速度"), ("jerk", "jerk"), ("dt", "Δt")]:
lines.append(
f"| {label} | {stats_gen[key]['mean']:.4g} | {stats_ref[key]['mean']:.4g} | "
f"{stats_gen[key]['cv']:.3f} | {stats_ref[key]['cv']:.3f} |"
)
lines += [
"",
"## 直方图",
f"![速度](plots/{tag}-speed.png)",
f"![加速度](plots/{tag}-accel.png)",
f"![Jerk](plots/{tag}-jerk.png)",
f"![Δt](plots/{tag}-dt.png)",
"",
"## FFT 频谱（横向偏移）",
f"![FFT](plots/{tag}-fft.png)",
"",
"## 生成轨迹示例",
f"![轨迹](plots/{tag}-paths.png)",
"",
]
output_md.write_text("\n".join(lines), encoding="utf-8")
logger.info("Report written to %s", output_md)
```
- [ ] **Step 4: Run smoke test, verify pass**
```bash
uv run pytest tests/test_eval_metrics.py::TestReportGeneration -v
```
Expected: 1 test passes. Plots and Markdown file created in `tmp_path`.
- [ ] **Step 5: Commit**
```bash
git add ai_mouse/eval/report.py tests/test_eval_metrics.py
git commit -m "feat(eval): Markdown report builder with matplotlib plots"
```
---
## Task 15: Eval — CLI
**Files:**
- Create: `ai_mouse/eval/__main__.py`
- [ ] **Step 1: Implement the eval CLI**
Create `ai_mouse/eval/__main__.py`:
```python
"""CLI: `python -m ai_mouse.eval --model-dir ... --reference ... --output ...`
Loads N synthetic start/end pairs, calls the generator, loads M reference
traces from a Balabit-format jsonl, and writes a Markdown report.
"""
from __future__ import annotations
import argparse
import json
import logging
import math
import random
import sys
from pathlib import Path
import numpy as np
logger = logging.getLogger(__name__)
def _load_reference_jsonl(path: Path, n_samples: int) -> list[dict]:
"""Load up to n_samples reference traces from a JSONL file.
Returns list of {"xs","ys","ts"} 1-D ndarrays.
"""
out: list[dict] = []
with path.open("r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
rec = json.loads(line)
except json.JSONDecodeError:
continue
moves = [e for e in rec.get("events", []) if e.get("type") == "move"]
if len(moves) < 4:
continue
xs = np.array([e["x"] for e in moves], dtype=float)
ys = np.array([e["y"] for e in moves], dtype=float)
ts = np.array([e["t"] for e in moves], dtype=float)
out.append({"xs": xs, "ys": ys, "ts": ts})
if len(out) >= n_samples:
break
return out
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
rng = random.Random(seed)
out: list[dict] = []
for i in range(n_samples):
# Random start/end on a 800x600 canvas, distance 100..600 px
sx = rng.randint(50, 750)
sy = rng.randint(50, 550)
angle = rng.uniform(0, 2 * math.pi)
dist = rng.randint(100, 600)
ex = int(sx + dist * math.cos(angle))
ey = int(sy + dist * math.sin(angle))
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)
except Exception as exc: # noqa: BLE001
logger.warning("generate() failed at i=%d: %s", i, exc)
continue
# Drop click events (last 2)
moves = pts[:-2]
if len(moves) < 4:
continue
xs = np.array([p[0] for p in moves], dtype=float)
ys = np.array([p[1] for p in moves], dtype=float)
ts = np.array([p[2] for p in moves], dtype=float)
out.append({"xs": xs, "ys": ys, "ts": ts})
return out
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Generate eval report comparing model output to reference traces.")
parser.add_argument("--model-dir", required=True, help="Path to trained model dir (with flow_model.pt)")
parser.add_argument("--reference", type=Path, required=True, help="JSONL reference traces (Balabit holdout)")
parser.add_argument("--n-samples", type=int, default=200, help="Number of generated samples")
parser.add_argument("--n-reference", type=int, default=1000, help="Number of reference samples to load")
parser.add_argument("--output", type=Path, required=True, help="Output Markdown file")
parser.add_argument("--tag", default="eval", help="Tag string used in plot filenames")
parser.add_argument("--seed", type=int, default=0)
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
if not Path(args.model_dir).exists():
logger.error("Model dir not found: %s", args.model_dir)
return 2
if not args.reference.exists():
logger.error("Reference jsonl not found: %s", args.reference)
return 2
logger.info("Loading reference from %s ...", args.reference)
ref_traces = _load_reference_jsonl(args.reference, args.n_reference)
logger.info("Loaded %d reference traces", len(ref_traces))
logger.info("Generating %d samples from %s ...", args.n_samples, args.model_dir)
gen_traces = _generate_n_samples(args.model_dir, args.n_samples, seed=args.seed)
logger.info("Generated %d valid traces", len(gen_traces))
if not gen_traces or not ref_traces:
logger.error("Empty trace sets — aborting")
return 1
from ai_mouse.eval.report import build_report
args.output.parent.mkdir(parents=True, exist_ok=True)
build_report(
generated_traces=gen_traces,
reference_traces=ref_traces,
output_md=args.output,
tag=args.tag,
model_dir=args.model_dir,
)
logger.info("Done. Report at %s", args.output)
return 0
if __name__ == "__main__":
sys.exit(main())
```
- [ ] **Step 2: Smoke-test the CLI help text**
```bash
uv run python -m ai_mouse.eval --help
```
Expected: argparse help text printed.
- [ ] **Step 3: Run all tests**
```bash
uv run pytest -x
```
Expected: all tests pass.
- [ ] **Step 4: Commit**
```bash
git add ai_mouse/eval/__main__.py
git commit -m "feat(eval): CLI for generating evaluation reports"
```
---
## Task 16: Unified Train CLI
**Files:**
- Create: `ai_mouse/__main__.py`
- [ ] **Step 1: Create unified CLI dispatcher**
Create `ai_mouse/__main__.py`:
```python
"""Unified CLI: `python -m ai_mouse {train,eval,balabit-adapter}`
Subcommands dispatch to the underlying modules. This is the recommended
top-level entry; you can also call `python -m ai_mouse.eval` etc. directly.
"""
from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path
def _train_main(args: argparse.Namespace) -> int:
from ai_mouse.trainer import train
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
train(
data_path=Path(args.data),
output_dir=Path(args.output),
epochs=args.epochs,
batch_size=args.batch_size,
lr=args.lr,
seq_len=args.seq_len,
resume_from=Path(args.resume_from) if args.resume_from else None,
)
return 0
def _eval_main(args: argparse.Namespace) -> int:
from ai_mouse.eval.__main__ import main as eval_main
# Reconstruct argv for the sub-CLI
argv = [
"--model-dir", args.model_dir,
"--reference", str(args.reference),
"--n-samples", str(args.n_samples),
"--n-reference", str(args.n_reference),
"--output", str(args.output),
"--tag", args.tag,
"--seed", str(args.seed),
]
return eval_main(argv)
def _balabit_main(args: argparse.Namespace) -> int:
from ai_mouse.data_adapters.balabit import main as bal_main
argv = [
"--input", str(args.input),
"--output", str(args.output),
"--window-ms", str(args.window_ms),
"--min-dist", str(args.min_dist),
"--min-events", str(args.min_events),
"--max-span-ms", str(args.max_span_ms),
"--max-gap-ms", str(args.max_gap_ms),
]
if args.overwrite:
argv.append("--overwrite")
return bal_main(argv)
def main() -> int:
p = argparse.ArgumentParser(prog="ai_mouse", description="AI Mouse trajectory toolkit")
sub = p.add_subparsers(dest="cmd", required=True)
# train
pt = sub.add_parser("train", help="Train (or fine-tune) the Flow Matching model")
pt.add_argument("--data", required=True, help="Path to traces.jsonl")
pt.add_argument("--output", required=True, help="Output checkpoint dir")
pt.add_argument("--epochs", type=int, default=200)
pt.add_argument("--batch-size", type=int, default=64)
pt.add_argument("--lr", type=float, default=3e-4)
pt.add_argument("--seq-len", type=int, default=64)
pt.add_argument("--resume-from", default=None, help="Checkpoint dir to resume from (for fine-tune)")
pt.set_defaults(func=_train_main)
# eval
pe = sub.add_parser("eval", help="Generate evaluation report")
pe.add_argument("--model-dir", required=True)
pe.add_argument("--reference", type=Path, required=True)
pe.add_argument("--n-samples", type=int, default=200)
pe.add_argument("--n-reference", type=int, default=1000)
pe.add_argument("--output", type=Path, required=True)
pe.add_argument("--tag", default="eval")
pe.add_argument("--seed", type=int, default=0)
pe.set_defaults(func=_eval_main)
# balabit-adapter
pb = sub.add_parser("balabit-adapter", help="Convert Balabit dataset to traces.jsonl")
pb.add_argument("--input", type=Path, required=True)
pb.add_argument("--output", type=Path, default=Path("data/pretrain_traces.jsonl"))
pb.add_argument("--window-ms", type=int, default=1200)
pb.add_argument("--min-dist", type=int, default=50)
pb.add_argument("--min-events", type=int, default=5)
pb.add_argument("--max-span-ms", type=int, default=5000)
pb.add_argument("--max-gap-ms", type=int, default=200)
pb.add_argument("--overwrite", action="store_true")
pb.set_defaults(func=_balabit_main)
args = p.parse_args()
return args.func(args)
if __name__ == "__main__":
sys.exit(main())
```
- [ ] **Step 2: Smoke-test all three subcommands**
```bash
uv run python -m ai_mouse train --help
uv run python -m ai_mouse eval --help
uv run python -m ai_mouse balabit-adapter --help
```
Expected: each prints argparse help.
- [ ] **Step 3: Commit**
```bash
git add ai_mouse/__main__.py
git commit -m "feat: unified CLI (python -m ai_mouse {train,eval,balabit-adapter})"
```
---
## Task 17: Run Pre-Refactor Baseline Eval (Optional)
This task captures the "before" snapshot so we can measure improvement quantitatively. **Requires**: Balabit data already converted (Task 18 needs to run first if you don't have Balabit yet, OR you can use the existing 605 `traces.jsonl` as both reference and "real data" — less informative but doable).
- [ ] **Step 1: Capture baseline using existing 605 traces as reference**
If Balabit isn't downloaded yet, use the existing 605 traces as a reference (split-off subset):
```bash
# Take the existing 605 traces as reference (use what you have)
uv run python -m ai_mouse eval \
--model-dir data/models_v2 \
--reference data/traces.jsonl \
--n-samples 100 \
--n-reference 200 \
--output data/eval_reports/2026-05-10-baseline-pre-refactor.md \
--tag baseline-pre-refactor
```
Expected: Markdown report at `data/eval_reports/2026-05-10-baseline-pre-refactor.md` plus PNG plots. **Note the KL values — these are the "before" numbers to beat.**
- [ ] **Step 2: Inspect the report**
Open the file in a Markdown viewer or your editor and look at:
- Speed/accel/jerk KL divergence
- Δt CV (should be very low for the current template-y model)
- The path PNG (should show the high-frequency lateral jitter)
- [ ] **Step 3: Save baseline numbers**
Manually note (or copy into a comment file) the KL values for later comparison:
```bash
echo "Pre-refactor baseline (2026-05-10):" > data/eval_reports/_BASELINE_NOTES.txt
echo " See: 2026-05-10-baseline-pre-refactor.md" >> data/eval_reports/_BASELINE_NOTES.txt
echo " Acceptance: post-refactor KL must be < 50% of these numbers" >> data/eval_reports/_BASELINE_NOTES.txt
```
- [ ] **Step 4: Commit baseline report**
```bash
git add data/eval_reports/2026-05-10-baseline-pre-refactor.md \
data/eval_reports/_BASELINE_NOTES.txt \
data/eval_reports/plots/
git commit -m "docs(eval): pre-refactor baseline report"
```
---
## Task 18: Run Balabit Adapter + Pretraining
**Prerequisites:** Download Balabit dataset to `data/balabit_raw/`:
```bash
# In a separate terminal — user manual step
cd /d/code/python/side/ai_mouse/data
git clone --depth 1 https://github.com/balabit/Mouse-Dynamics-Challenge.git balabit_raw
# OR: download zip from the GitHub page and extract to data/balabit_raw/
```
Verify the structure:
```bash
ls data/balabit_raw/
# Expect to see directories like 'training_files/' or 'test_files/' containing user folders
```
- [ ] **Step 1: Convert Balabit → pretrain_traces.jsonl**
```bash
uv run python -m ai_mouse balabit-adapter \
--input data/balabit_raw/training_files \
--output data/pretrain_traces.jsonl \
--overwrite
```
Expected: log shows "Wrote N segments to data/pretrain_traces.jsonl" with N typically 5,000–50,000 depending on what's in the dataset.
If N < 1000, **stop and investigate**:
- Try `--window-ms 2000` (longer window)
- Try `--min-dist 30` (shorter minimum distance)
- Inspect `data/balabit_raw/` structure — the path may be wrong (try `data/balabit_raw/test_files` instead)
- [ ] **Step 2: Spot-check converted data**
```bash
uv run python -c "
import json
with open('data/pretrain_traces.jsonl') as f:
lines = f.readlines()
print(f'Total: {len(lines)} traces')
sample = json.loads(lines[0])
print('Sample meta:', sample['meta'])
print('Sample event count:', len(sample['events']))
print('First 3 events:', sample['events'][:3])
"
```
Expected: meta has `start`, `end`, `dist`, `angle`, `source: balabit`. Events list looks like real `move` records with sane (x, y, t).
- [ ] **Step 3: Reserve a 5% holdout for eval**
```bash
uv run python -c "
import random
from pathlib import Path
random.seed(42)
src = Path('data/pretrain_traces.jsonl')
lines = src.read_text(encoding='utf-8').splitlines()
random.shuffle(lines)
split = int(len(lines) * 0.95)
Path('data/pretrain_train.jsonl').write_text('\n'.join(lines[:split]) + '\n', encoding='utf-8')
Path('data/balabit_holdout.jsonl').write_text('\n'.join(lines[split:]) + '\n', encoding='utf-8')
print(f'Train: {split} Holdout: {len(lines)-split}')
"
```
- [ ] **Step 4: Run pretraining**
This will take **2 hours – 2 days** depending on hardware. Run in background or overnight.
```bash
uv run python -m ai_mouse train \
--data data/pretrain_train.jsonl \
--output data/models_v2_pretrained \
--epochs 200 \
--batch-size 128 \
--lr 3e-4
```
Expected outputs in `data/models_v2_pretrained/`:
- `flow_model.pt`
- `click_dist.json` (will have default values since Balabit has no click events — this is expected)
- `duration_dist.json`
- `train_config.json`
- [ ] **Step 5: Verify checkpoint loads cleanly**
```bash
uv run python -c "
import torch
from ai_mouse.models import TrajectoryFlowModel
m = TrajectoryFlowModel(seq_len=64)
state = torch.load('data/models_v2_pretrained/flow_model.pt', weights_only=True)
m.load_state_dict(state)
print('Pretrain checkpoint loads OK')
"
```
Expected: `Pretrain checkpoint loads OK`.
- [ ] **Step 6: Commit (do not commit the model weights — they're in .gitignore)**
```bash
# Just commit the holdout split logic via a small note
git add docs/superpowers/plans/2026-05-10-balabit-pretrain-refactor.md
git commit -m "chore: ran Balabit conversion + pretraining (artifacts in data/, not committed)"
```
---
## Task 19: Fine-tune on 605 Traces + Final Eval
- [ ] **Step 1: Run fine-tune**
```bash
uv run python -m ai_mouse train \
--data data/traces.jsonl \
--output data/models_v2 \
--epochs 50 \
--batch-size 64 \
--lr 1e-5 \
--resume-from data/models_v2_pretrained
```
Expected: training completes in 5–30 minutes, `data/models_v2/flow_model.pt` updated.
- [ ] **Step 2: Generate post-refactor eval report**
```bash
uv run python -m ai_mouse eval \
--model-dir data/models_v2 \
--reference data/balabit_holdout.jsonl \
--n-samples 200 \
--n-reference 1000 \
--output data/eval_reports/2026-05-10-final.md \
--tag final
```
- [ ] **Step 3: Compare to baseline**
Open both reports side-by-side:
- `data/eval_reports/2026-05-10-baseline-pre-refactor.md`
- `data/eval_reports/2026-05-10-final.md`
Verify against the spec's acceptance criteria:
1. ✅ **主观**：Δt 曲线在 verify 页面应该不再 5 条重合
2. ✅ **主观**：lateral 轨迹无高频锯齿（看 paths.png）
3. ✅ **量化**：speed KL < 50% of pre-refactor baseline
4. ✅ **量化**：FFT 4–12 Hz 区间出现 peak（看 fft.png）
5. ✅ **回归**：所有非废弃测试通过
- [ ] **Step 4: Run full test suite as final regression check**
```bash
uv run pytest -v
```
Expected: all tests pass.
- [ ] **Step 5: Manual verification in the UI**
```bash
uv run python main.py
```
Open http://127.0.0.1:8765 in browser, navigate to "验证效果" tab, generate 5 paths from (100, 200) to (700, 400), and compare visually to the original screenshot.
Expected:
- 5 Δt curves visibly diverge (no longer overlapping)
- Lateral trajectories smooth, no zigzag
- Average duration similar to before (still in plausible range)
- [ ] **Step 6: Commit final report**
```bash
git add data/eval_reports/2026-05-10-final.md data/eval_reports/plots/
git commit -m "docs(eval): post-refactor final eval report
Acceptance criteria met:
- speed KL <baseline*0.5
- FFT shows 4-12Hz peak
- Δt diversity restored (subjective)
- Lateral jitter eliminated (subjective)
- All tests pass"
```
---
## Self-Review Checklist
After completing all tasks, verify:
- [ ] All 19 tasks committed
- [ ] All tests pass (`uv run pytest -v`)
- [ ] Baseline report exists for comparison
- [ ] Final report shows quantitative improvement
- [ ] UI manually verified
- [ ] No new files outside the spec's "Files Changed" list
- [ ] `data/` artifacts NOT committed (verified via `git status --short`)
Reference in New Issue View Git Blame Copy Permalink