"""Scroll training pipeline for ScrollCVAE. Pipeline: 1. Load scroll traces from JSONL -> (seq, cond) tensors 2. Apply 4x data augmentation 3. Train ScrollCVAE with: MSE(delta_norm) + 1.5*MSE(log_dt) + beta*KL 4. Save: scroll_model.pt, scroll_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 import torch.nn as nn from torch.distributions import Normal, kl_divergence from torch.utils.data import DataLoader, TensorDataset from tools.config import ScrollTrainConfig from tools.scroll.models import ScrollCVAE logger = logging.getLogger(__name__) # Mode -> one-hot index _MODE_INDEX = {"target": 0, "fast": 1, "precise": 2} # --------------------------------------------------------------------------- # Data loading # --------------------------------------------------------------------------- def load_scroll_data( data_path: Path, seq_len: int = 32, ) -> tuple[np.ndarray, np.ndarray]: """Load scroll JSONL and return (seq, cond) arrays. Args: data_path: path to scroll_traces.jsonl seq_len: number of wheel-event steps to pad/truncate to Returns: seq: (N, seq_len, 2) float32 -- (delta_norm, log_dt) cond: (N, 7) float32 -- [dist/5000, log(dist/500), direction, viewport_norm, mode_onehot*3] """ data_path = Path(data_path) seq_list: list[np.ndarray] = [] cond_list: list[np.ndarray] = [] 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: logger.warning("Skipping line %d: invalid JSON", i) continue if "meta" not in trace or "events" not in trace: logger.warning("Skipping line %d: missing meta or events", i) continue meta = trace["meta"] events = trace["events"] if not events: continue distance = float(meta.get("distance", 0)) if distance <= 0: start_y = float(meta.get("start_scrollY", 0)) target_y = float(meta.get("target_scrollY", 0)) distance = abs(target_y - start_y) if distance <= 0: distance = 1.0 direction_str = meta.get("direction", "down") direction = 1.0 if direction_str == "down" else -1.0 viewport_height = float(meta.get("viewport_height", 900)) viewport_norm = viewport_height / 1000.0 mode_str = meta.get("mode", "target") mode_idx = _MODE_INDEX.get(mode_str, 0) mode_onehot = np.zeros(3, dtype=np.float32) mode_onehot[mode_idx] = 1.0 deltas = np.array([float(e.get("deltaY", 0)) for e in events], dtype=np.float32) times = np.array([float(e.get("t", 0)) for e in events], dtype=np.float32) if len(deltas) == 0: continue delta_norm = deltas / distance dt_raw = np.diff(times).clip(0.0) log_dt = np.log(dt_raw + 1.0) log_dt_padded = np.concatenate([[0.0], log_dt]) seq_raw = np.stack([delta_norm, log_dt_padded], axis=1).astype(np.float32) n = len(seq_raw) if n >= seq_len: seq_out = seq_raw[:seq_len] else: pad = np.zeros((seq_len - n, 2), dtype=np.float32) seq_out = np.concatenate([seq_raw, pad], axis=0) dist_norm = distance / 5000.0 log_dist = math.log(max(distance, 1.0) / 500.0) cond_arr = np.array( [dist_norm, log_dist, direction, viewport_norm, *mode_onehot], dtype=np.float32, ) seq_list.append(seq_out) cond_list.append(cond_arr) if not seq_list: raise ValueError(f"No valid scroll traces found in {data_path}") logger.info("Loaded %d scroll traces from %s", len(seq_list), data_path) return ( np.stack(seq_list, axis=0), np.stack(cond_list, axis=0), ) # --------------------------------------------------------------------------- # Data augmentation (4x) # --------------------------------------------------------------------------- def _augment_scroll( seq: np.ndarray, cond: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: """4x augmentation for scroll sequences. Variants: 0 -- original 1 -- speed x0.8: log_dt[1:] += log(1.25) 2 -- speed x1.2: log_dt[1:] += log(1/1.2) 3 -- temporal noise: log_dt[1:] += N(0, 0.05) """ log_1_25 = math.log(1.25) log_inv_1_2 = math.log(1.0 / 1.2) seqs = [seq] conds = [cond] s1 = seq.copy() s1[:, 1:, 1] += log_1_25 seqs.append(s1) conds.append(cond.copy()) s2 = seq.copy() s2[:, 1:, 1] += log_inv_1_2 seqs.append(s2) conds.append(cond.copy()) s3 = seq.copy() noise = np.random.normal(0.0, 0.05, size=s3[:, 1:, 1].shape).astype(np.float32) s3[:, 1:, 1] += noise seqs.append(s3) conds.append(cond.copy()) return np.concatenate(seqs, axis=0), np.concatenate(conds, axis=0) # --------------------------------------------------------------------------- # Main training function # --------------------------------------------------------------------------- def train_scroll( data_path: Path, output_dir: Path, epochs: int = 100, batch_size: int = 32, lr: float = 5e-4, seq_len: int = 32, progress_callback: Callable[[dict], None] | None = None, config: ScrollTrainConfig | None = None, ) -> None: """Train ScrollCVAE and save artefacts to output_dir.""" cfg = config or ScrollTrainConfig() epochs = epochs or cfg.epochs batch_size = batch_size or cfg.batch_size lr = lr or cfg.lr seq_len = seq_len or cfg.seq_len data_path = Path(data_path) output_dir = Path(output_dir) if not data_path.exists(): raise FileNotFoundError(f"Data file not found: {data_path}") logger.info("Starting scroll training: epochs=%d, data=%s", epochs, data_path) seq_np, cond_np = load_scroll_data(data_path, seq_len=seq_len) seq_np, cond_np = _augment_scroll(seq_np, cond_np) seq_t = torch.from_numpy(seq_np) cond_t = torch.from_numpy(cond_np) output_dir.mkdir(parents=True, exist_ok=True) model = ScrollCVAE(seq_len=seq_len, latent_dim=16, hidden=64, cond_dim=7) optimiser = torch.optim.AdamW(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=max(epochs, 1)) ds = TensorDataset(seq_t, cond_t) loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False) model.train() for epoch in range(epochs): beta = min(cfg.beta_max, cfg.beta_max * epoch / max(cfg.beta_warmup_epochs, 1)) epoch_loss = 0.0 n_batches = 0 for seq_b, cond_b in loader: optimiser.zero_grad() recon, mu, logvar = model(seq_b, cond_b) delta_loss = nn.functional.mse_loss(recon[:, :, 0], seq_b[:, :, 0]) logdt_loss = nn.functional.mse_loss(recon[:, :, 1], seq_b[:, :, 1]) recon_loss = cfg.weight_delta * delta_loss + cfg.weight_log_dt * logdt_loss std = torch.exp(0.5 * logvar) q = Normal(mu, std) p = Normal(torch.zeros_like(mu), torch.ones_like(std)) kl_loss = kl_divergence(q, p).mean() loss = recon_loss + beta * kl_loss loss.backward() 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) if progress_callback is not None: progress_callback({ "epoch": epoch + 1, "total": epochs, "loss": epoch_loss, "stage": "scroll", }) torch.save(model.state_dict(), output_dir / "scroll_model.pt") scroll_cfg = { "seq_len": seq_len, "latent_dim": model.latent_dim, "hidden": model.hidden, "cond_dim": model.cond_dim, "epochs": epochs, "batch_size": batch_size, "lr": lr, "beta_max": cfg.beta_max, "beta_anneal_epochs": cfg.beta_warmup_epochs, "w_delta": cfg.weight_delta, "w_logdt": cfg.weight_log_dt, } (output_dir / "scroll_config.json").write_text(json.dumps(scroll_cfg, indent=2)) logger.info("Scroll training complete. Model saved to %s", output_dir) if progress_callback is not None: progress_callback({"done": True})