284 lines
8.6 KiB
Python
284 lines
8.6 KiB
Python
"""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})
|