"""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})