refactor: move trainer/models/utils/config to tools/

tools/config.py

@@ -0,0 +1,160 @@
+"""Centralised configuration for ai_mouse.
+
+All magic numbers and hyperparameters live here so they can be tuned
+from one place, overridden per-instance, or serialised to JSON.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+# ---------------------------------------------------------------------------
+# Training configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TrainConfig:
+    """Hyperparameters for Flow Matching training."""
+
+    epochs: int = 300
+    batch_size: int = 64
+    lr: float = 3e-4
+    seq_len: int = 64
+    # Transformer backbone
+    d_model: int = 128
+    nhead: int = 4
+    num_layers: int = 4
+    dim_feedforward: int = 256
+    dropout: float = 0.1
+    # Flow matching
+    sigma_min: float = 1e-4
+    # Data augmentation
+    augment: bool = True
+    # Duration distribution bins
+    dist_bins: list[float] = field(
+        default_factory=lambda: [0, 50, 100, 200, 400, 600, 800, 1200, float("inf")]
+    )
+
+
+# ---------------------------------------------------------------------------
+# Generation configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class GenerateConfig:
+    """Tuneable knobs for Flow Matching inference."""
+
+    n_steps: int = 10           # Euler ODE steps
+    seq_len: int = 64
+    dt_clip_min_ms: float = 2.0
+    dt_clip_max_ms: float = 150.0
+
+
+# ---------------------------------------------------------------------------
+# Scroll subsystem configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ScrollModeConfig:
+    """Parameters for a single scroll collection mode."""
+
+    dist_min: int
+    dist_max: int
+    success_radius: int
+
+
+SCROLL_MODES: dict[str, ScrollModeConfig] = {
+    "target": ScrollModeConfig(dist_min=500, dist_max=3000, success_radius=80),
+    "fast": ScrollModeConfig(dist_min=3000, dist_max=8000, success_radius=120),
+    "precise": ScrollModeConfig(dist_min=200, dist_max=800, success_radius=40),
+}
+
+
+@dataclass
+class ScrollTrainConfig:
+    """Hyperparameters for ScrollCVAE training."""
+
+    epochs: int = 100
+    batch_size: int = 32
+    lr: float = 5e-4
+    seq_len: int = 32
+    beta_max: float = 0.3
+    beta_warmup_epochs: int = 15
+    weight_delta: float = 1.0
+    weight_log_dt: float = 1.5
+
+
+# ---------------------------------------------------------------------------
+# Server configuration
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ServerConfig:
+    """Web server and data directory settings."""
+
+    host: str = "127.0.0.1"
+    port: int = 8765
+    data_dir: Path = field(default_factory=lambda: Path("data"))
+    canvas_width: int = 800
+    canvas_height: int = 600
+    open_browser: bool = True
+
+
+# ---------------------------------------------------------------------------
+# 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

tools/models.py

@@ -0,0 +1,236 @@
+"""Conditional Flow Matching model with Transformer backbone.
+
+Predicts velocity field v_θ(x_t, t, cond) that transports noise to data.
+- x_t: (B, T, 3) noisy trajectory at time t
+- t: (B,) interpolation time ∈ [0,1]
+- cond: (B, 3) condition [dist_norm, log_dist, log_dur]
+- Output: (B, T, 3) velocity field
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+
+
+# ---------------------------------------------------------------------------
+# Sinusoidal time embedding
+# ---------------------------------------------------------------------------
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    """Map scalar timestep t ∈ [0,1] to a d_model-dimensional vector."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            t: (B,) scalar timesteps
+
+        Returns:
+            (B, d_model) embedding vectors
+        """
+        half = self.d_model // 2
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half
+        )
+        args = t.unsqueeze(-1) * freqs.unsqueeze(0)  # (B, half)
+        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # (B, d_model)
+
+
+# ---------------------------------------------------------------------------
+# TrajectoryFlowModel
+# ---------------------------------------------------------------------------
+
+
+class TrajectoryFlowModel(nn.Module):
+    """Conditional Flow Matching model for mouse trajectory generation.
+
+    Architecture:
+    - SinusoidalTimeEmbedding: t scalar → d_model vector
+    - Input projection: 3 → d_model
+    - Learned positional embedding: (1, seq_len, d_model)
+    - Time embed + Condition embed added as bias to all tokens
+    - 4-layer TransformerEncoder (pre-norm, GELU, batch_first=True)
+    - Output projection: d_model → 3
+
+    Args:
+        seq_len: Number of time steps (default 64).
+        d_model: Transformer hidden dimension (default 128).
+        nhead: Number of attention heads (default 4).
+        num_layers: Number of transformer layers (default 4).
+        dim_feedforward: Feedforward hidden size (default 256).
+        dropout: Dropout rate (default 0.1).
+        cond_dim: Condition vector size (default 3).
+    """
+
+    def __init__(
+        self,
+        seq_len: int = 64,
+        d_model: int = 128,
+        nhead: int = 4,
+        num_layers: int = 4,
+        dim_feedforward: int = 256,
+        dropout: float = 0.1,
+        cond_dim: int = 3,
+    ):
+        super().__init__()
+        self.seq_len = seq_len
+        self.d_model = d_model
+        self.cond_dim = cond_dim
+
+        # Input projection: (forward, lateral, log_dt) → d_model
+        self.input_proj = nn.Linear(3, d_model)
+
+        # Learned positional embedding
+        self.pos_embed = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
+
+        # Time embedding
+        self.time_embed = SinusoidalTimeEmbedding(d_model)
+        self.time_mlp = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+
+        # Condition embedding
+        self.cond_mlp = nn.Sequential(
+            nn.Linear(cond_dim, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+        )
+
+        # Transformer encoder (pre-norm via norm_first=True)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer, num_layers=num_layers, enable_nested_tensor=False
+        )
+
+        # Output projection: d_model → 3
+        self.output_proj = nn.Linear(d_model, 3)
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        cond: torch.Tensor,
+    ) -> torch.Tensor:
+        """Predict velocity field.
+
+        Args:
+            x_t: (B, T, 3) noisy trajectory at interpolation time t
+            t: (B,) interpolation timestep ∈ [0,1]
+            cond: (B, 3) condition vector [dist_norm, log_dist, log_dur]
+
+        Returns:
+            (B, T, 3) predicted velocity field
+        """
+        # Project input tokens
+        h = self.input_proj(x_t)  # (B, T, d_model)
+
+        # Add positional embedding
+        h = h + self.pos_embed  # (B, T, d_model)
+
+        # Time embedding → bias added to all tokens
+        t_emb = self.time_mlp(self.time_embed(t))  # (B, d_model)
+        h = h + t_emb.unsqueeze(1)  # broadcast over T
+
+        # Condition embedding → bias added to all tokens
+        c_emb = self.cond_mlp(cond)  # (B, d_model)
+        h = h + c_emb.unsqueeze(1)  # broadcast over T
+
+        # Transformer
+        h = self.transformer(h)  # (B, T, d_model)
+
+        # Output projection
+        return self.output_proj(h)  # (B, T, 3)
+
+
+# ---------------------------------------------------------------------------
+# Legacy JointCVAE — kept for backward compatibility with generator.py
+# ---------------------------------------------------------------------------
+
+
+class JointCVAE(nn.Module):
+    """Joint Conditional VAE for mouse trajectory generation (legacy).
+
+    Kept for backward compatibility with the existing generator.
+    See TrajectoryFlowModel for the new approach.
+    """
+
+    def __init__(
+        self,
+        seq_len: int = 64,
+        latent_dim: int = 32,
+        hidden: int = 128,
+        cond_dim: int = 3,
+    ):
+        super().__init__()
+        self.seq_len = seq_len
+        self.latent_dim = latent_dim
+        self.hidden = hidden
+        self.cond_dim = cond_dim
+        self.feat_dim = 3
+
+        self.enc_gru = nn.GRU(
+            input_size=self.feat_dim + cond_dim,
+            hidden_size=hidden,
+            num_layers=2,
+            batch_first=True,
+            bidirectional=True,
+        )
+        self.enc_mu = nn.Linear(hidden * 2, latent_dim)
+        self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
+
+        self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
+        self.dec_gru = nn.GRU(
+            input_size=latent_dim + cond_dim,
+            hidden_size=hidden,
+            num_layers=2,
+            batch_first=True,
+        )
+        self.dec_out = nn.Linear(hidden, self.feat_dim)
+
+    def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        B, T, _ = seq.shape
+        c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
+        x_in = torch.cat([seq, c_exp], dim=-1)
+        _, h_n = self.enc_gru(x_in)
+        h_fwd = h_n[-2]
+        h_bwd = h_n[-1]
+        h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
+        return self.enc_mu(h_cat), self.enc_logvar(h_cat)
+
+    def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        B = z.shape[0]
+        zc = torch.cat([z, cond], dim=-1)
+        h0_flat = self.dec_h0(zc)
+        h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
+        inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
+        out, _ = self.dec_gru(inp, h0)
+        return self.dec_out(out)
+
+    def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        std = torch.exp(0.5 * logvar)
+        return Normal(mu, std).rsample()
+
+    def forward(
+        self, seq: torch.Tensor, cond: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        mu, logvar = self.encode(seq, cond)
+        z = self.reparameterise(mu, logvar)
+        recon = self.decode(z, cond)
+        return recon, mu, logvar

tools/trainer.py

@@ -0,0 +1,509 @@
+"""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})

tools/utils.py

@@ -0,0 +1,25 @@
+# sites/ai_mouse/ai_mouse/_utils.py
+"""Shared utility functions used by both _generator.py and _trainer.py."""
+from __future__ import annotations
+
+import numpy as np
+
+
+def resample_arc(xy: np.ndarray, n_points: int) -> np.ndarray:
+    """Resample a 2-D polyline to exactly n_points via arc-length interpolation.
+
+    Args:
+        xy: (M, 2) array of (x, y) coordinates.
+        n_points: desired number of output points.
+
+    Returns:
+        (n_points, 2) array uniformly spaced along cumulative arc length.
+    """
+    arc = np.concatenate(
+        [[0], np.cumsum(np.linalg.norm(np.diff(xy, axis=0), axis=1))]
+    )
+    s_new = np.linspace(0, arc[-1], n_points)
+    return np.stack(
+        [np.interp(s_new, arc, xy[:, 0]), np.interp(s_new, arc, xy[:, 1])],
+        axis=1,
+    )