refactor(scroll): move trainer/models/collector to tools/scroll/

tools/scroll/collector.py

@@ -0,0 +1,99 @@
+"""Scroll collection state and target generation.
+
+The actual scroll state machine runs in JavaScript (wheel events are
+client-side). This module handles server-side target generation and
+trace persistence.
+"""
+from __future__ import annotations
+
+import logging
+import random
+
+from tools.config import SCROLL_MODES
+
+logger = logging.getLogger(__name__)
+
+
+class ScrollCollector:
+    """Manages scroll collection sessions."""
+
+    def __init__(
+        self,
+        mode: str,
+        count: int,
+        page_height: int = 10000,
+        viewport_height: int = 900,
+    ):
+        if mode not in SCROLL_MODES:
+            raise ValueError(f"Unknown mode: {mode}. Use: {', '.join(SCROLL_MODES)}")
+        self.mode = mode
+        self.count = count
+        self.page_height = page_height
+        self.viewport_height = viewport_height
+        self.collected = 0
+
+        cfg = SCROLL_MODES[mode]
+        self.dist_min = cfg.dist_min
+        self.dist_max = cfg.dist_max
+        self.success_radius = cfg.success_radius
+
+    def next_target(self, current_scrollY: int) -> dict:
+        """Generate next target scroll position.
+
+        The target must be reachable — i.e. the user must be able to scroll
+        such that the target band enters the viewport's success zone (centered).
+
+        Reachability constraint:
+            To hit a target at T, the user needs scrollTop ≈ T - viewportCenter + bandHeight/2.
+            scrollTop must be in [0, pageHeight - viewportHeight].
+
+        So valid target range is:
+            T_min = viewportCenter - successRadius  (scrollTop=0 → target at top of success zone)
+            T_max = maxScrollTop + viewportCenter + successRadius
+
+        Returns:
+            {"target_scrollY": int, "direction": "up"|"down"}
+        """
+        viewport_center = self.viewport_height // 2
+        max_scroll_top = self.page_height - self.viewport_height
+
+        # Valid range for targetScrollY so it's always reachable
+        target_min = viewport_center - self.success_radius
+        target_max = max_scroll_top + viewport_center + self.success_radius
+
+        # Clamp to ensure sane bounds
+        target_min = max(0, target_min)
+        target_max = min(self.page_height, target_max)
+
+        max_down = min(self.page_height - current_scrollY, target_max - current_scrollY)
+        max_up = min(current_scrollY, current_scrollY - target_min)
+
+        for _ in range(100):
+            dist = random.randint(self.dist_min, self.dist_max)
+            direction = random.choice(["up", "down"])
+
+            if direction == "down" and 0 < dist <= max_down:
+                candidate = current_scrollY + dist
+                if target_min <= candidate <= target_max:
+                    return {"target_scrollY": candidate, "direction": "down"}
+            elif direction == "up" and 0 < dist <= max_up:
+                candidate = current_scrollY - dist
+                if target_min <= candidate <= target_max:
+                    return {"target_scrollY": candidate, "direction": "up"}
+
+        # Fallback: generate a valid target within bounds
+        valid_down = min(max_down, self.dist_max)
+        valid_up = min(max_up, self.dist_max)
+
+        if valid_down >= self.dist_min:
+            dist = random.randint(self.dist_min, max(self.dist_min, valid_down))
+            return {"target_scrollY": current_scrollY + dist, "direction": "down"}
+        elif valid_up >= self.dist_min:
+            dist = random.randint(self.dist_min, max(self.dist_min, valid_up))
+            return {"target_scrollY": current_scrollY - dist, "direction": "up"}
+        else:
+            # Edge case: move to center of valid range
+            target = max(target_min, min(target_max, (target_min + target_max) // 2))
+            direction = "down" if target > current_scrollY else "up"
+            return {"target_scrollY": target, "direction": direction}
+

tools/scroll/models.py

@@ -0,0 +1,75 @@
+"""ScrollCVAE — generates realistic scroll wheel event sequences.
+
+Architecture mirrors JointCVAE but smaller (scroll sequences are simpler):
+- Encoder: bidirectional GRU(hidden=64, layers=2)
+- Decoder: unidirectional GRU(hidden=64, layers=2)
+- Input/output: (delta_norm, log_Δt) per time step
+- Condition: [dist_norm, log_dist, direction, viewport_norm, mode_onehot×3] = 7 dims
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+
+
+class ScrollCVAE(nn.Module):
+    def __init__(
+        self,
+        seq_len: int = 32,
+        latent_dim: int = 16,
+        hidden: int = 64,
+        cond_dim: int = 7,
+    ):
+        super().__init__()
+        self.seq_len = seq_len
+        self.latent_dim = latent_dim
+        self.hidden = hidden
+        self.cond_dim = cond_dim
+        self.feat_dim = 2  # (delta_norm, log_Δt)
+
+        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):
+        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_cat = torch.cat([h_n[-2], h_n[-1]], dim=-1)
+        return self.enc_mu(h_cat), self.enc_logvar(h_cat)
+
+    def decode(self, z: torch.Tensor, cond: 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, logvar):
+        std = torch.exp(0.5 * logvar)
+        return Normal(mu, std).rsample()
+
+    def forward(self, seq, cond):
+        mu, logvar = self.encode(seq, cond)
+        z = self.reparameterise(mu, logvar)
+        recon = self.decode(z, cond)
+        return recon, mu, logvar

tools/scroll/trainer.py

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