Huang Qi
31fd884dfd
Drop the pre-migration PyTorch inference pipeline now that the ONNX-backed
MouseModel/ScrollModel in mouse.py and scroll.py are wired up through the
public ai_mouse API.
Deleted:
* src/ai_mouse/generator.py (legacy torch flow ODE + post-processing)
* src/ai_mouse/coord.py (legacy public coord transforms,
superseded by ai_mouse._coord)
* src/ai_mouse/_scroll_legacy.py (legacy torch scroll VAE inference)
* scripts/build_golden_*.py (one-shot capture scripts, no longer
needed once goldens are committed)
* tests/unit/test_generator.py (legacy module gone)
* tests/unit/test_scroll_generator.py (legacy module gone)
* tests/unit/test_coord.py (legacy module gone; ai_mouse._coord is
tested by test__coord.py)
* scripts/ (empty, removed)
Tools migrations:
* tools/trainer.py: import encode_trajectory from ai_mouse._coord
instead of the deleted ai_mouse.coord
* tools/server/routes_verify.py, tools/server/routes_scroll.py: route to
the public ai_mouse.generate / generate_scroll. They no longer accept
a model_dir override — the bundled ONNX is the source of truth, and a
fresh export goes through `python -m tools.export_onnx`.
* tools/eval/__main__.py: same migration; model_dir CLI arg retained as
a deprecation shim but ignored.
Final src/ai_mouse/ layout (matches plan):
__init__.py, _assets.py, _coord.py, _model_cache.py, _postprocess.py,
errors.py, mouse.py, py.typed, scroll.py, assets/
Test suite: 188 passed (was 188 before deletion; obsolete suites cleaned
out alongside the modules they covered).
510 lines
17 KiB
Python
510 lines
17 KiB
Python
"""Training pipeline for Conditional Flow Matching mouse trajectory model.
|
||
|
||
Pipeline:
|
||
1. Load traces from JSONL, convert to rotated coordinate frame
|
||
2. Apply 6× data augmentation
|
||
3. Train TrajectoryFlowModel with OT-Conditional Flow Matching:
|
||
- x1 = real data, x0 = randn_like(x1), t = rand(B)
|
||
- x_t = (1-t)*x0 + t*x1
|
||
- v_target = x1 - x0
|
||
- v_pred = model(x_t, t, cond)
|
||
- loss = MSE(v_pred, v_target)
|
||
4. Save: flow_model.pt, click_dist.json, duration_dist.json, train_config.json
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import math
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import DataLoader
|
||
|
||
from ai_mouse._coord import encode_trajectory
|
||
from tools.config import TrainConfig
|
||
from tools.models import TrajectoryFlowModel
|
||
from tools.utils import resample_arc
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Distance bins for duration distribution (in pixels)
|
||
_DIST_BINS: list[float] = [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Data loading
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def load_and_prepare_data(
|
||
data_path: Path,
|
||
seq_len: int = 64,
|
||
) -> tuple[np.ndarray, np.ndarray, list[float]]:
|
||
"""Load JSONL traces and convert to rotated-frame tensors.
|
||
|
||
Args:
|
||
data_path: path to traces.jsonl
|
||
seq_len: number of time steps to resample each trajectory to
|
||
|
||
Returns:
|
||
seq: (N, seq_len, 3) float32 — (forward, lateral, log_Δt)
|
||
cond: (N, 3) float32 — [dist_norm, log_dist, log_dur]
|
||
click_durs: list of float click durations in ms
|
||
"""
|
||
data_path = Path(data_path)
|
||
seq_list: list[np.ndarray] = []
|
||
cond_list: list[np.ndarray] = []
|
||
click_durs: list[float] = []
|
||
|
||
for i, raw_line in enumerate(data_path.read_text(encoding="utf-8").splitlines(), 1):
|
||
line = raw_line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
trace = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
meta = trace["meta"]
|
||
events = trace["events"]
|
||
|
||
if "start" not in meta or "end" not in meta:
|
||
logger.warning("Skipping line %d: missing start/end in meta", i)
|
||
continue
|
||
|
||
sx, sy = meta["start"]
|
||
ex, ey = meta["end"]
|
||
|
||
# Extract move events
|
||
moves = [(e["x"], e["y"], e["t"]) for e in events if e["type"] == "move"]
|
||
if len(moves) < 2:
|
||
continue
|
||
|
||
xs = np.array([m[0] for m in moves], dtype=float)
|
||
ys = np.array([m[1] for m in moves], dtype=float)
|
||
ts = np.array([m[2] for m in moves], dtype=float)
|
||
|
||
xy_raw = np.stack([xs, ys], axis=1)
|
||
|
||
# Reject degenerate (zero-length) trajectories
|
||
total_arc = float(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1).sum())
|
||
if total_arc < 1.0:
|
||
continue
|
||
|
||
# Resample spatial positions to seq_len via arc-length
|
||
xy_resampled = resample_arc(xy_raw, seq_len) # (seq_len, 2)
|
||
|
||
# Resample timestamps along the same arc-length grid
|
||
arc_dist = np.concatenate(
|
||
[[0.0], np.cumsum(np.linalg.norm(np.diff(xy_raw, axis=0), axis=1))]
|
||
)
|
||
s_uniform = np.linspace(0.0, arc_dist[-1], seq_len)
|
||
ts_resampled = np.interp(s_uniform, arc_dist, ts)
|
||
|
||
# Convert spatial coords to rotated frame (forward, lateral)
|
||
fl = encode_trajectory(xy_resampled, (sx, sy), (ex, ey)) # (seq_len, 2)
|
||
|
||
# Compute Δt intervals (length seq_len-1) → log(Δt+1), pad 0 at front
|
||
dt_raw = np.diff(ts_resampled).clip(0.0)
|
||
log_dt = np.log(dt_raw + 1.0) # (seq_len-1,)
|
||
log_dt_padded = np.concatenate([[0.0], log_dt]) # (seq_len,) — first step has no interval
|
||
|
||
# Stack into (seq_len, 3)
|
||
seq_arr = np.stack([fl[:, 0], fl[:, 1], log_dt_padded], axis=1).astype(np.float32)
|
||
|
||
# Condition vector
|
||
dist = float(meta["dist"]) if meta["dist"] > 0 else float(
|
||
math.hypot(ex - sx, ey - sy)
|
||
)
|
||
dist = max(dist, 1.0)
|
||
total_dur = float(ts_resampled[-1] - ts_resampled[0])
|
||
total_dur = max(total_dur, 1.0)
|
||
|
||
cond_arr = np.array(
|
||
[
|
||
dist / 2000.0, # dist_norm
|
||
math.log(dist / 100.0), # log_dist
|
||
math.log(total_dur / 500.0), # log_dur
|
||
],
|
||
dtype=np.float32,
|
||
)
|
||
|
||
seq_list.append(seq_arr)
|
||
cond_list.append(cond_arr)
|
||
|
||
# Click duration (down→up)
|
||
downs = [e for e in events if e["type"] == "down"]
|
||
ups = [e for e in events if e["type"] == "up"]
|
||
if downs and ups:
|
||
click_durs.append(float(ups[-1]["t"] - downs[-1]["t"]))
|
||
|
||
if not seq_list:
|
||
raise ValueError(f"No valid traces found in {data_path}")
|
||
|
||
return (
|
||
np.stack(seq_list, axis=0), # (N, seq_len, 3)
|
||
np.stack(cond_list, axis=0), # (N, 3)
|
||
click_durs,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Data augmentation (6×)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _augment(
|
||
seq: np.ndarray,
|
||
cond: np.ndarray,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""6× augmentation operating in the rotated (forward, lateral, log_dt) frame.
|
||
|
||
Variants:
|
||
0 — original
|
||
1 — lateral flip: lateral → −lateral
|
||
2 — speed ×0.8: log_Δt[1:] += log(1.25)
|
||
3 — speed ×1.2: log_Δt[1:] += log(1/1.2)
|
||
4 — temporal noise: log_Δt[1:] += N(0, 0.05)
|
||
5 — combined: lateral flip + speed ×0.9
|
||
|
||
Args:
|
||
seq: (N, T, 3) — (forward, lateral, log_dt)
|
||
cond: (N, 3) — [dist_norm, log_dist, log_dur]
|
||
|
||
Returns:
|
||
seq_aug: (6N, T, 3)
|
||
cond_aug: (6N, 3)
|
||
"""
|
||
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)
|
||
|
||
seqs = [seq]
|
||
conds = [cond]
|
||
|
||
# 1. Lateral flip
|
||
s1 = seq.copy()
|
||
s1[:, :, 1] = -s1[:, :, 1]
|
||
seqs.append(s1)
|
||
conds.append(cond.copy())
|
||
|
||
# 2. Speed ×0.8 (longer duration: log_dt += log(1.25))
|
||
s2 = seq.copy()
|
||
s2[:, 1:, 2] += log_1_25
|
||
c2 = cond.copy()
|
||
c2[:, 2] += log_1_25 # log_dur updated
|
||
seqs.append(s2)
|
||
conds.append(c2)
|
||
|
||
# 3. Speed ×1.2 (shorter duration: log_dt += log(1/1.2))
|
||
s3 = seq.copy()
|
||
s3[:, 1:, 2] += log_inv_1_2
|
||
c3 = cond.copy()
|
||
c3[:, 2] += log_inv_1_2
|
||
seqs.append(s3)
|
||
conds.append(c3)
|
||
|
||
# 4. Temporal noise
|
||
s4 = seq.copy()
|
||
noise = np.random.normal(0.0, 0.05, size=s4[:, 1:, 2].shape).astype(np.float32)
|
||
s4[:, 1:, 2] += noise
|
||
seqs.append(s4)
|
||
conds.append(cond.copy())
|
||
|
||
# 5. Lateral flip + speed ×0.9 (i.e. log_dt += log(1/0.9))
|
||
s5 = seq.copy()
|
||
s5[:, :, 1] = -s5[:, :, 1]
|
||
s5[:, 1:, 2] += log_1_1
|
||
c5 = cond.copy()
|
||
c5[:, 2] += log_1_1
|
||
seqs.append(s5)
|
||
conds.append(c5)
|
||
|
||
return np.concatenate(seqs, axis=0), np.concatenate(conds, axis=0)
|
||
|
||
|
||
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_INV_0_9 = 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_INV_0_9
|
||
c[2] += self._LOG_INV_0_9
|
||
|
||
return torch.from_numpy(s), torch.from_numpy(c)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Duration distribution (per distance bin)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _compute_duration_dist(data_path: Path) -> dict:
|
||
"""Compute per-distance-bin log-normal parameters for trace duration.
|
||
|
||
Args:
|
||
data_path: path to traces.jsonl
|
||
|
||
Returns:
|
||
dict with keys "bins" (list of floats) and "params" (list of dicts
|
||
with "mu_log" and "sigma_log" per bin).
|
||
"""
|
||
bin_durations: list[list[float]] = [[] for _ in range(len(_DIST_BINS) - 1)]
|
||
|
||
for raw_line in Path(data_path).read_text(encoding="utf-8").splitlines():
|
||
line = raw_line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
trace = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
meta = trace["meta"]
|
||
events = trace["events"]
|
||
dist = float(meta.get("dist", 0))
|
||
|
||
moves = [e for e in events if e["type"] == "move"]
|
||
if len(moves) < 2:
|
||
continue
|
||
dur = float(moves[-1]["t"] - moves[0]["t"])
|
||
if dur <= 0:
|
||
continue
|
||
|
||
# Find bin
|
||
for i in range(len(_DIST_BINS) - 1):
|
||
if _DIST_BINS[i] <= dist < _DIST_BINS[i + 1]:
|
||
bin_durations[i].append(dur)
|
||
break
|
||
|
||
params = []
|
||
for durs in bin_durations:
|
||
if len(durs) >= 2:
|
||
log_durs = np.log(np.array(durs, dtype=float))
|
||
mu_log = float(np.mean(log_durs))
|
||
sigma_log = float(np.std(log_durs, ddof=1))
|
||
else:
|
||
mu_log = float(np.log(500.0))
|
||
sigma_log = 0.5
|
||
params.append({"mu_log": mu_log, "sigma_log": max(sigma_log, 0.05)})
|
||
|
||
return {"bins": _DIST_BINS, "params": params}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main training function
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
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:
|
||
"""Train TrajectoryFlowModel with OT-Conditional Flow Matching.
|
||
|
||
Args:
|
||
data_path: path to traces.jsonl
|
||
output_dir: directory where artefacts are written
|
||
epochs: training epochs
|
||
batch_size: mini-batch size
|
||
lr: AdamW learning rate
|
||
seq_len: number of time steps per trajectory
|
||
progress_callback: optional callable invoked each epoch with
|
||
{"epoch": n, "total": N, "loss": f}.
|
||
Called once at the end with {"done": True, "mu", "sigma"}.
|
||
config: optional TrainConfig for model hyperparameters
|
||
resume_from: if given, load model weights from this checkpoint
|
||
directory (must contain flow_model.pt). Used for
|
||
two-stage training (pretrain → fine-tune).
|
||
"""
|
||
data_path = Path(data_path)
|
||
output_dir = Path(output_dir)
|
||
|
||
if config is None:
|
||
config = TrainConfig(
|
||
epochs=epochs, batch_size=batch_size, lr=lr, seq_len=seq_len
|
||
)
|
||
|
||
if not data_path.exists():
|
||
raise FileNotFoundError(f"Data file not found: {data_path}")
|
||
|
||
# ---- Load & prepare ----
|
||
logger.info("Loading data from %s", data_path)
|
||
seq_np, cond_np, click_durs = load_and_prepare_data(data_path, seq_len=seq_len)
|
||
logger.info("Loaded %d traces", len(seq_np))
|
||
|
||
# ---- 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))
|
||
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# ---- Model & optimiser ----
|
||
model = TrajectoryFlowModel(
|
||
seq_len=seq_len,
|
||
d_model=config.d_model,
|
||
nhead=config.nhead,
|
||
num_layers=config.num_layers,
|
||
dim_feedforward=config.dim_feedforward,
|
||
dropout=config.dropout,
|
||
cond_dim=3,
|
||
)
|
||
|
||
# ---- 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)
|
||
|
||
optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
|
||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)
|
||
|
||
loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
|
||
|
||
# ---- Training loop: OT-Conditional Flow Matching ----
|
||
model.train()
|
||
for epoch in range(epochs):
|
||
epoch_loss = 0.0
|
||
n_batches = 0
|
||
|
||
for x1_batch, cond_batch in loader:
|
||
B = x1_batch.shape[0]
|
||
|
||
# Sample noise x0 ~ N(0, I)
|
||
x0 = torch.randn_like(x1_batch)
|
||
|
||
# Sample random timestep t ~ U[0, 1]
|
||
t = torch.rand(B)
|
||
|
||
# Interpolate: x_t = (1-t)*x0 + t*x1
|
||
t_expand = t[:, None, None] # (B, 1, 1) for broadcasting
|
||
x_t = (1.0 - t_expand) * x0 + t_expand * x1_batch
|
||
|
||
# Target velocity: v = x1 - x0
|
||
v_target = x1_batch - x0
|
||
|
||
# Predict velocity
|
||
v_pred = model(x_t, t, cond_batch)
|
||
|
||
# MSE loss
|
||
loss = torch.nn.functional.mse_loss(v_pred, v_target)
|
||
|
||
optimiser.zero_grad()
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
||
optimiser.step()
|
||
|
||
epoch_loss += loss.item()
|
||
n_batches += 1
|
||
|
||
scheduler.step()
|
||
|
||
epoch_loss /= max(n_batches, 1)
|
||
logger.debug("Epoch %d/%d loss=%.6f", epoch + 1, epochs, epoch_loss)
|
||
|
||
if progress_callback is not None:
|
||
progress_callback({"epoch": epoch + 1, "total": epochs, "loss": epoch_loss})
|
||
|
||
# ---- Save model ----
|
||
torch.save(model.state_dict(), output_dir / "flow_model.pt")
|
||
logger.info("Saved flow_model.pt to %s", output_dir)
|
||
|
||
# ---- Click duration distribution ----
|
||
if click_durs:
|
||
arr = np.array(click_durs)
|
||
arr = arr[(arr >= 20) & (arr <= 500)]
|
||
if len(arr) >= 2:
|
||
mu_c = float(arr.mean())
|
||
sigma_c = max(float(arr.std()), 1.0)
|
||
else:
|
||
mu_c, sigma_c = 80.0, 30.0
|
||
else:
|
||
mu_c, sigma_c = 80.0, 30.0
|
||
|
||
click_dist = {"mu": mu_c, "sigma": sigma_c, "low": 20.0, "high": 500.0}
|
||
(output_dir / "click_dist.json").write_text(json.dumps(click_dist, indent=2))
|
||
|
||
# ---- Duration distribution (per distance bin) ----
|
||
dur_dist = _compute_duration_dist(data_path)
|
||
(output_dir / "duration_dist.json").write_text(json.dumps(dur_dist, indent=2))
|
||
|
||
# ---- Train config ----
|
||
train_cfg = {
|
||
"seq_len": seq_len,
|
||
"epochs": epochs,
|
||
"batch_size": batch_size,
|
||
"lr": lr,
|
||
"d_model": config.d_model,
|
||
"nhead": config.nhead,
|
||
"num_layers": config.num_layers,
|
||
"dim_feedforward": config.dim_feedforward,
|
||
"dropout": config.dropout,
|
||
"cond_dim": 3,
|
||
}
|
||
(output_dir / "train_config.json").write_text(json.dumps(train_cfg, indent=2))
|
||
|
||
if progress_callback is not None:
|
||
progress_callback({"done": True, "mu": mu_c, "sigma": sigma_c})
|