ai_mouse/docs/superpowers/specs/2026-05-10-balabit-pretrain-refactor-design.md

Balabit 预训练 + Fine-tune 重构设计

日期：2026-05-10 作者：Claude + 用户 状态：待用户评审

1. 背景与目标

1.1 问题

当前 TrajectoryFlowModel 仅用 605 条用户采集数据训练，生成质量差：

• 空间锯齿：lateral 方向高频抖动（见 verify 页面截图）
• 时间模板化：5 条生成轨迹的 Δt 曲线几乎完全重合（多样性丧失）
• 反检测能力弱

1.2 根本原因

1. 数据量不足（605 条），模型欠拟合 + 记噪声
2. 后处理过度：generator.py 中的 speed_profile（确定性钟形曲线）和 median±1.1 硬 clip 把多样性压扁

1.3 目标

• 用 Balabit Mouse Dynamics Challenge 数据集（公开真实用户会话）做预训练
• 用现有 605 条做 fine-tune 适配本任务分布
• 重写后处理，让真人数据本身的速度模式自然显现
• 输出量化评估报表，便于迭代判断进步

1.4 不在范围（YAGNI）

• ❌ 滚轮模型 ai_mouse/scroll/（独立子系统，本次冻结不动）
• ❌ 模型架构变更（仍是 TrajectoryFlowModel，不换 Diffusion——CPU 推理约束）
• ❌ 对抗训练 / Discriminator（验证完基础方案再说）
• ❌ 浏览器插件、其他公开数据集
• ❌ 前端 UI 变更

2. 整体架构

Balabit raw sessions (CSV)
        │
        ▼   [新] ai_mouse/data_adapters/balabit.py
   click 锚定切分（Pressed 事件前 W ms 的 Move）
        │
        ▼
  data/pretrain_traces.jsonl   ← 与现有 traces.jsonl 格式 100% 兼容
        │
        ▼   [改] ai_mouse/trainer.py（streaming dataloader、resume_from）
  data/models_v2_pretrained/   ← 预训练 checkpoint（Balabit only）
        │
        ▼   [改] trainer.py 加 --resume-from 支持加载已有权重
  data/models_v2/   ← 605 条 fine-tune 后的最终权重（部署用）
        │
        ▼   [改] ai_mouse/generator.py（砍硬模板后处理）
  推理：CPU < 200ms

  评估：[新] ai_mouse/eval/
        ▼
  data/eval_reports/YYYY-MM-DD-<tag>.md   ← 量化指标 + 直方图 + FFT

关键设计原则：

• 数据格式不变：Balabit 转换后产出和现有 traces.jsonl 格式 100% 一致。load_and_prepare_data / 旋转坐标系 / 6× 增强逻辑全部不动
• 训练流程加一个阶段：原来"605 → train → models_v2"，现在变成"Balabit → pretrain → models_v2_pretrained → 605 fine-tune → models_v2"
• 滚轮子系统完全不动

3. 文件变更清单

3.1 新增

路径	用途
ai_mouse/data_adapters/__init__.py	包初始化
ai_mouse/data_adapters/balabit.py	Balabit CSV → traces.jsonl 适配器 + CLI
ai_mouse/eval/__init__.py	包初始化
ai_mouse/eval/metrics.py	运动学指标计算（速度/加速度/jerk/FFT）
ai_mouse/eval/report.py	Markdown 报表生成
ai_mouse/eval/__main__.py	CLI: python -m ai_mouse.eval
ai_mouse/__main__.py	CLI: python -m ai_mouse train ...（统一入口）
tests/test_balabit_adapter.py	适配器单元测试
tests/test_eval_metrics.py	指标计算正确性测试

3.2 修改

路径	改动
ai_mouse/trainer.py	加 resume_from 参数；增强从一次性 numpy 改成 Dataset 内 on-the-fly；可选 AMP（有 GPU 时）
ai_mouse/generator.py	砍 speed_profile 与 median±1.1 clip；新增 5 点高斯 lateral 平滑
ai_mouse/config.py	新增 PretrainConfig、FinetuneConfig、BalabitAdapterConfig、EvalConfig
tests/test_generator.py	适配新后处理
tests/test_trainer.py	加 resume_from 测试

3.3 冻结不动

ai_mouse/scroll/*、ai_mouse/coord.py、ai_mouse/models.py、ai_mouse/utils.py、ai_mouse/collector.py、ai_mouse/server/*、static/*、tests/test_scroll_*.py、tests/test_coord.py、tests/test_models.py、tests/test_server.py

4. Balabit 适配器（ai_mouse/data_adapters/balabit.py）

4.1 输入格式

Balabit Mouse Dynamics Challenge 每个 session 是 CSV 文件，列：

record timestamp, client timestamp, button, state, x, y

• state ∈ {Move, Pressed, Released, Drag, Scroll}
• button ∈ {NoButton, Left, Right}

4.2 切分逻辑（click 锚定）

1. 扫描 session 所有事件
2. 找每个 Pressed 事件 P
3. 回溯前 W ms（默认 W=1200）内的所有 Move 事件，构成一段 trace
4. 段的 start = 第一个 Move 的 (x,y)，end = P 的 (x,y)
5. 时间戳归零（第一个 Move 的 t=0）

4.3 过滤规则

丢弃整段（不修复，直接丢弃这一条 trace），满足任一条件即丢：

• Move 事件数 < 5
• dist(start, end) < 50 px
• 时间跨度 > 5000ms（避免长停顿）
• start 或 end 任一坐标 < 0 或 > 5000（避免跨屏瞬移异常值）
• 总弧长 < 50 px（避免抖动残留）
• 段内任意相邻 Move 之间时间差 > 200ms（采样断档，宁可整段丢也不要含断点的样本）

4.4 输出

追加到 data/pretrain_traces.jsonl，每行：

{"meta":{"start":[x,y],"end":[x,y],"dist":int,"angle":float,"source":"balabit","session_id":"user7_session_42"},"events":[{"type":"move","x":int,"y":int,"t":int}, ...]}

关于 click events 的兼容性：Balabit 转换后的 trace 不包含 down/up 事件（预训练只学移动）。这与现有 trainer.py:139-142 的逻辑兼容——load_and_prepare_data 用 if downs and ups 检查，没有就跳过。后果：

• 预训练阶段产出的 click_dist.json 会基于零样本，写入默认值（mu=80, sigma=30）
• Fine-tune 阶段重新基于 605 条产出真正的 click_dist.json，覆盖默认值
• flow_model.pt 和 train_config.json 是预训练真正要保留的产物

4.5 CLI

uv run python -m ai_mouse.data_adapters.balabit \
  --input /path/to/balabit/sessions/ \
  --output data/pretrain_traces.jsonl \
  --window-ms 1200

5. 训练 pipeline 改造（ai_mouse/trainer.py）

5.1 新增参数

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,   # 新增：加载已有权重
    use_amp: bool = False,              # 新增：mixed precision
) -> None:

5.2 Dataset 改造

当前 _augment 是一次性把全部数据 6× 复制到内存。Balabit 几万条样本 × 6 后内存占用大。改为：

class TrajectoryDataset(torch.utils.data.Dataset):
    def __init__(self, seq, cond, augment: bool = True):
        self.seq = seq
        self.cond = cond
        self.augment = augment
        self._n_aug = 6 if augment else 1
    
    def __len__(self):
        return len(self.seq) * self._n_aug
    
    def __getitem__(self, idx):
        base_idx = idx // self._n_aug
        aug_id = idx % self._n_aug
        seq, cond = self.seq[base_idx], self.cond[base_idx]
        return self._apply_augment(seq, cond, aug_id)

每次 __getitem__ 即时增强，内存占用 = 原始数据 1×。

5.3 Resume from checkpoint

if resume_from is not None:
    state = torch.load(resume_from / "flow_model.pt", weights_only=True)
    model.load_state_dict(state)

如果 fine-tune 阶段的 cond_dim 与预训练不一致，需要明确报错。

5.4 完整两阶段训练流程

# 阶段 1：预训练
uv run python -m ai_mouse train \
  --data data/pretrain_traces.jsonl \
  --output data/models_v2_pretrained \
  --epochs 200 --lr 3e-4 --batch-size 128

# 阶段 2：微调
uv run python -m ai_mouse train \
  --data data/traces.jsonl \
  --output data/models_v2 \
  --epochs 50 --lr 1e-5 --batch-size 64 \
  --resume-from data/models_v2_pretrained

5.5 服务端 API

POST /api/train 当前不支持 resume。本次不改服务端 API——CLI 是主要预训练入口，UI 上的"训练"按钮仍用于 605 条的 fine-tune（默认会自动从 models_v2_pretrained resume，如果存在）。

逻辑：

• models_v2_pretrained/flow_model.pt 存在 → fine-tune 模式（lr=1e-5, epochs=50）
• 不存在 → 走原逻辑（from scratch）

6. Generator 改造（ai_mouse/generator.py）

6.1 砍掉

• lines 263–271 max_allowed = median + 1.1、min_allowed = median - 1.1 硬 clip 整段
• lines 273–286 speed_profile 整段（acceleration/deceleration phase）

6.2 保留

• 端点 snap (lines 220–229)
• 起点 lateral 衰减 (lines 232–236)
• forward 单调性强制 (lines 239–246)
• log_dt 安全 clip [0, 5] (line 255)
• click duration 采样

6.3 新增：lateral 5 点高斯平滑

# 在 lateral monotonic fix 之后、decode_trajectory 之前
def _gaussian_smooth(x: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """5-point gaussian smoothing, preserving endpoints."""
    kernel = np.exp(-0.5 * (np.arange(-2, 3) / sigma) ** 2)
    kernel /= kernel.sum()
    smoothed = np.convolve(x, kernel, mode="same")
    smoothed[0] = x[0]   # preserve endpoint
    smoothed[-1] = x[-1]
    return smoothed

lateral = _gaussian_smooth(lateral, sigma=1.0)

只对 lateral 平滑，不对 forward 平滑（forward 已经被单调性约束保护，再平滑会损害准确性）。

6.4 预期效果

• 解决高频锯齿（lateral 平滑）
• Δt 多样性恢复（不再被压扁到 median ± 1.1）
• 速度模式由模型自己决定（learned from Balabit real data）

7. 评估模块（ai_mouse/eval/）

7.1 指标（metrics.py）

每条轨迹计算：

• 速度 v[i] = sqrt(dx² + dy²) / dt（单位 px/ms）
• 加速度 a[i] = (v[i+1] - v[i]) / dt
• Jerk j[i] = (a[i+1] - a[i]) / dt
• Δt 序列

聚合到样本集：

• 速度/加速度/jerk 的均值、std、变异系数 CV、p25/p50/p75/p95
• Δt 分布的 CV
• FFT 频谱：每条轨迹 lateral 信号做 FFT，看 4–12Hz 频段是否有 peak（生理震颤）
• 多样性：N 条样本之间的 PCA 方差贡献（衡量是否模板化）

对比"参考分布"（从 pretrain_traces.jsonl 随机抽 1000 条作为 holdout）：

• 速度/加速度分布的 KL 散度（直方图离散化估计）
• jerk 分布的 KL 散度

7.2 报表（report.py）

输出 Markdown 到 data/eval_reports/YYYY-MM-DD-<tag>.md，结构：

# Eval Report: <tag> (2026-05-10 18:30)

## 模型信息
- Checkpoint: data/models_v2/flow_model.pt
- 训练参数: ...

## 摘要
| 指标 | 生成 | 参考 | 评价 |
|---|---|---|---|
| 速度 KL | 0.12 | 0 | OK |
| FFT 4-12Hz peak | 8.3Hz @ 0.04 | 7.1Hz @ 0.05 | OK |
| Δt CV | 0.45 | 0.52 | 接近 |
| ...

## 直方图（PNG 嵌入）
![速度分布](plots/2026-05-10-baseline-speed.png)
...

## 5 条生成轨迹示例
![](plots/2026-05-10-baseline-paths.png)

## vs 上次报表
- 速度 KL 从 0.34 → 0.12（提升）
- ...

PNG 用 matplotlib 输出到 data/eval_reports/plots/。

7.3 CLI

uv run python -m ai_mouse.eval \
  --model-dir data/models_v2 \
  --reference data/pretrain_traces.jsonl \
  --n-samples 1000 \
  --output data/eval_reports/2026-05-10-baseline.md \
  --tag baseline

8. 配置变更（ai_mouse/config.py）

新增：

@dataclass
class PretrainConfig:
    """Hyperparameters for Balabit pretraining."""
    epochs: int = 200
    batch_size: int = 128
    lr: float = 3e-4
    seq_len: int = 64

@dataclass
class FinetuneConfig:
    """Hyperparameters for fine-tuning on user data."""
    epochs: int = 50
    batch_size: int = 64
    lr: float = 1e-5  # 比预训练小一个数量级
    seq_len: int = 64

@dataclass
class BalabitAdapterConfig:
    """Settings for Balabit data conversion."""
    window_ms: int = 1200
    min_dist: int = 50
    min_events: int = 5
    max_span_ms: int = 5000
    max_gap_ms: int = 200

@dataclass
class EvalConfig:
    """Settings for evaluation report generation."""
    n_samples: int = 1000
    fft_freq_band: tuple[float, float] = (4.0, 12.0)
    kl_bins: int = 50

TrainConfig 保持不变（向后兼容现有训练脚本）。

9. 测试策略

9.1 新增测试

• test_balabit_adapter.py：用合成 CSV 测试切分逻辑、过滤规则、边界条件（空 session、无 click、坐标溢出）
• test_eval_metrics.py：固定输入下指标计算的正确性

9.2 更新测试

• test_generator.py：移除对 speed_profile 的断言；新增 lateral 平滑断言
• test_trainer.py：新增 resume_from 测试

9.3 保持通过

所有非 test_generator.py、非 test_trainer.py 的测试保持通过。test_server.py 不变（服务端 API 未改）。

10. 验收标准

最终评估报表（data/eval_reports/<final>.md）应显示：

1. 主观：5 条生成轨迹的 Δt 曲线明显多样化（不再重合）
2. 主观：lateral 无高频锯齿
3. 量化：速度分布 KL（vs Balabit holdout）< 当前实现（pre-refactor baseline，第一次跑评估时记下）的 50%
4. 量化：FFT 频谱在 4–12Hz 区间出现 peak
5. 回归：所有非废弃测试保持通过

11. 工作量估计

阶段	时间
Balabit 适配器（含测试）	1–2 天
Trainer 改造（streaming + resume）	1 天
Generator 后处理改造	0.5 天
评估模块	1–2 天
跑预训练 + fine-tune（CPU/GPU 视情况）	0.5–2 天
调参迭代 + 报表对比	1–3 天

最小可行版本（搭起来跑通第一版）：3–5 天 完整调到验收：1–2 周

12. 风险与备选

风险	缓解
Balabit 切分后样本数不足（< 5000）	放宽 min_dist 到 30，或扩大窗口到 2000ms
预训练后 fine-tune 出现灾难性遗忘	lr 调更小（1e-6），epochs 减到 20
模型架构 cond_dim 在预训练/fine-tune 不一致	强制相同；不一致时直接 raise
评估报表实现工作量过大	第一版只做基础指标（速度/Δt CV），FFT 和 PCA 后置
Balabit 数据集合规问题	仅本地使用，不分发，不商用