yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Huang Qi
2026-05-12 00:34:05 +08:00
parent ba52c49edf
commit 6c96ab68c8
11 changed files with 11 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

283
tools/scroll/trainer.py Normal file
View File

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