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

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

## 1. 背景与目标

### 1.1 问题
当前 `TrajectoryFlowModel` 仅用 605 条用户采集数据训练，生成质量差：
- **空间锯齿**：lateral 方向高频抖动（见 verify 页面截图）
- **时间模板化**：5 条生成轨迹的 Δt 曲线几乎完全重合（多样性丧失）
- 反检测能力弱

### 1.2 根本原因
1. **数据量不足**（605 条），模型欠拟合 + 记噪声
2. **后处理过度**：[generator.py](../../../ai_mouse/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`，每行：
```json
{"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
```bash
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 新增参数
```python
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 后内存占用大。改为：

```python
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
```python
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 完整两阶段训练流程
```bash
# 阶段 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 点高斯平滑
```python
# 在 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`，结构：

```markdown
# 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
```bash
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`）

新增：

```python
@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 数据集合规问题 | 仅本地使用，不分发，不商用 |