Huang Qi
4d414fd180
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
100 KiB
100 KiB
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
Prerequisites:
- 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.
- 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.
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:
[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:
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):
# ---------------------------------------------------------------------------
# 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
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
uv run pytest -x
Expected: all tests pass.
- Step 7: Initial commit
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:
"""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
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:
"""Data adapters: convert external datasets to the project's traces.jsonl format."""
Create ai_mouse/data_adapters/balabit.py:
"""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
uv run pytest tests/test_balabit_adapter.py -v
Expected: 3 tests pass.
- Step 5: Commit
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(implementparse_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:
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
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:
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
uv run pytest tests/test_balabit_adapter.py::TestParseSessionCsv -v
Expected: 4 tests pass.
- Step 5: Commit
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(implementsegment_by_clicks) -
Modify:
tests/test_balabit_adapter.py(add tests) -
Step 1: Write failing tests
Append to tests/test_balabit_adapter.py:
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
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:
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
uv run pytest tests/test_balabit_adapter.py::TestSegmentByClicks -v
Expected: 6 tests pass.
- Step 5: Commit
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(implementfilter_segments) -
Modify:
tests/test_balabit_adapter.py -
Step 1: Write failing tests
Append to tests/test_balabit_adapter.py:
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
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:
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
uv run pytest tests/test_balabit_adapter.py::TestFilterSegments -v
Expected: 7 tests pass.
- Step 5: Commit
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(implementprocess_session, addmain) -
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:
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:
import json
- Step 2: Run tests, verify failure
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:
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
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:
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:
"""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:
if __name__ == "__main__":
import sys
sys.exit(main())
- Step 7: Smoke-test the CLI
uv run python -m ai_mouse.data_adapters.balabit --help
Expected: argparse help text printed.
- Step 8: Run all adapter tests
uv run pytest tests/test_balabit_adapter.py -v
Expected: all tests pass.
- Step 9: Commit
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_augmentusage withTrajectoryDataset) -
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:
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
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):
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
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):
# ---- 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:
# ---- 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):
ds = TensorDataset(seq_t, cond_t)
loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
Replace with:
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:
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:
uv run pytest tests/test_trainer.py -v
Expected: all tests pass (including TestAugment, TestLoadAndPrepare, TestTrain, TestTrajectoryDataset).
- Step 7: Commit
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(addresume_fromparameter totrain()) -
Modify:
tests/test_trainer.py -
Step 1: Write failing test for resume_from
Append to tests/test_trainer.py:
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
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:
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:
# ---- 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
uv run pytest tests/test_trainer.py::TestResumeFrom -v
Expected: 2 tests pass.
- Step 5: Run all trainer tests
uv run pytest tests/test_trainer.py -v
Expected: all tests pass.
- Step 6: Commit
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:
def _paths() -> tuple[Path, Path]:
data_dir = get_data_dir()
return data_dir / "traces.jsonl", data_dir / "models_v2"
Replace with:
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():
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:
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:
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
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
uv run pytest -x
Expected: all tests pass.
- Step 4: Commit
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.
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:
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
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):
# 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):
# 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:
"""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
uv run pytest tests/test_generator.py -v
Expected: all tests pass.
- Step 6: Commit
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:
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
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):
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
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:
# 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:
# 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:
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:
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
uv run pytest tests/test_generator.py -v
Expected: all tests pass.
- Step 8: Commit
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:
"""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
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:
"""Evaluation module: kinematic metrics and Markdown report generation."""
Create ai_mouse/eval/metrics.py:
"""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
uv run pytest tests/test_eval_metrics.py -v
Expected: 6 tests pass.
- Step 5: Commit
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:
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
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:
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
uv run pytest tests/test_eval_metrics.py::TestFftSpectrum tests/test_eval_metrics.py::TestKlDivergence -v
Expected: 5 tests pass.
- Step 5: Commit
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:
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
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:
"""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"",
f"",
f"",
f"",
"",
"## FFT 频谱（横向偏移）",
f"",
"",
"## 生成轨迹示例",
f"",
"",
]
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
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
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:
"""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
uv run python -m ai_mouse.eval --help
Expected: argparse help text printed.
- Step 3: Run all tests
uv run pytest -x
Expected: all tests pass.
- Step 4: Commit
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:
"""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
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
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):
# 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:
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
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/:
# 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:
ls data/balabit_raw/
# Expect to see directories like 'training_files/' or 'test_files/' containing user folders
- Step 1: Convert Balabit → pretrain_traces.jsonl
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 (trydata/balabit_raw/test_filesinstead) -
Step 2: Spot-check converted data
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
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.
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
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)
# 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
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
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.mddata/eval_reports/2026-05-10-final.md
Verify against the spec's acceptance criteria:
- ✅ 主观：Δt 曲线在 verify 页面应该不再 5 条重合
- ✅ 主观：lateral 轨迹无高频锯齿（看 paths.png）
- ✅ 量化：speed KL < 50% of pre-refactor baseline
- ✅ 量化：FFT 4–12 Hz 区间出现 peak（看 fft.png）
- ✅ 回归：所有非废弃测试通过
- Step 4: Run full test suite as final regression check
uv run pytest -v
Expected: all tests pass.
- Step 5: Manual verification in the UI
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
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 viagit status --short)