actpred: 两轮代码审计修复与文档合并

- 第 1 轮安全审计修复: F-1~F-10（路径校验、预训练架构修正等）
- 第 2 轮复审修复 11 项: 3 项上轮修复不完整（R01 checkpoint 键校验/
  R02 fused_hidden_dim 值域/R03 pretrained=False）+ 8 项新问题
  （actions.npy 形状校验、帧数警告、CLI 值域校验、stats 形状与路径、
  onnxscript 依赖、归一化语义、解析去重）
- 单元测试由 17 扩展至 51 用例（50 通过 + 1 环境性跳过）
- 审计文档两轮合并至 actpred_CODE_AUDIT_20260808.md，更新 INDEX

actpred/README.md

@@ -34,6 +34,9 @@ bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
 python test_actpred.py
 ```
 
+> **注意**：torch ≥ 2.12 的 ONNX 导出依赖 `onnxscript`，`_set.sh` 已一并安装；
+> 如单独使用 requirements.txt 安装，请额外执行 `pip install onnxscript onnxruntime`。
+
 ## 预训练（从 .mt 脚本，无需图像）
 
 从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵，用零图像占位预训练动作动力学（history → next action），
@@ -43,8 +46,11 @@ python test_actpred.py
 # 预训练（解析 ../bot4cj/scripts/*.mt）
 bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 --noise-std 1.0
 
-# 用预训练权重作为微调起点（需修改 train.py 加载逻辑，或手动加载）
-bash _run.sh train --data-dir ./data --epochs 50
+# 用预训练权重作为微调起点
+# 注意：预训练在非归一化动作空间进行，微调时请勿开启 --normalize-actions，
+# 否则输入分布与预训练权重不匹配，预训练收益会大打折扣
+bash _run.sh train --data-dir ./data --epochs 50 \
+    --init-from checkpoints/pretrain_best_inference.pt
 ```
 
 | 参数 | 说明 |
@@ -105,6 +111,13 @@ recorder.flush()
 | `--num-workers` | 2 | DataLoader 工作进程数 |
 | `--checkpoint-dir` | ./checkpoints | 检查点保存目录 |
 
+## 部署注意
+
+- 启用 `--normalize-actions` 训练的模型，其 `action_stats.npy` 保存在 checkpoint 同目录；
+  部署时**必须将 stats 文件与 checkpoint 一起拷贝**（推理端优先按 checkpoint 内记录的路径查找，
+  找不到时回退到 checkpoint 同目录的 `action_stats.npy`）。
+- `actions.npy` 的列数必须与 `--action-dim` 一致；图像与动作帧数不一致时训练会打印 WARNING 并按较少者使用。
+
 ## 编程 API
 
 ```python

actpred/_set.sh

@@ -10,4 +10,8 @@ pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 # Remaining deps from Tsinghua mirror
 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
 
+# ONNX 导出依赖（torch>=2.12 的 torch.onnx.export 走 dynamo 路径，onnxscript 为硬依赖）
+pip install onnxscript onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+pip check
 mkdir -p checkpoints data

actpred/config.py

@@ -79,10 +79,84 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
     parser.add_argument("--num-workers", type=int, default=None)
     parser.add_argument("--seed", type=int, default=None)
     parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--init-from", type=str, default=None,
+                        help="预训练 checkpoint 路径，作为微调起点（strict=False 加载兼容权重）")
     parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件")
     return parser
 
 
+def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
+                           train_cfg: TrainConfig = None):
+    """校验从 checkpoint 恢复或 CLI 传入的配置值域合理性
+
+    防止恶意/损坏 checkpoint 注入异常值（如超大维度导致 OOM），
+    也防止训练 CLI 越界参数（如 --history-len 0 / --lr -1）静默进入训练。
+    """
+    checks = [
+        ("image_size", data_cfg.image_size, 1, 1024, int),
+        ("history_len", data_cfg.history_len, 1, 1024, int),
+        ("action_dim", data_cfg.action_dim, 1, 1024, int),
+        ("cnn_feature_dim", model_cfg.cnn_feature_dim, 1, 4096, int),
+        ("action_hidden_dim", model_cfg.action_hidden_dim, 1, 4096, int),
+        ("fused_hidden_dim", model_cfg.fused_hidden_dim, 1, 4096, int),
+        ("dropout", model_cfg.dropout, 0.0, 1.0, float),
+    ]
+    for name, val, lo, hi, typ in checks:
+        if typ is int:
+            ok = isinstance(val, int) and not isinstance(val, bool) and lo <= val <= hi
+        else:
+            ok = isinstance(val, (int, float)) and not isinstance(val, bool) and lo <= val <= hi
+        if not ok:
+            raise ValueError(
+                f"checkpoint config {name}={val} 超出合理范围 [{lo}, {hi}]")
+
+    if train_cfg is not None:
+        t_checks = [
+            ("batch_size", train_cfg.batch_size, 1, 65536, int),
+            ("epochs", train_cfg.epochs, 1, 100000, int),
+            ("lr", train_cfg.lr, 1e-9, 10.0, float),
+            ("weight_decay", train_cfg.weight_decay, 0.0, 1.0, float),
+            ("val_split", train_cfg.val_split, 0.0, 1.0, float),
+            ("patience", train_cfg.patience, 0, 100000, int),
+            ("num_workers", train_cfg.num_workers, 0, 256, int),
+            ("max_grad_norm", train_cfg.max_grad_norm, 0.0, 1e6, float),
+        ]
+        for name, val, lo, hi, typ in t_checks:
+            if typ is int:
+                ok = isinstance(val, int) and not isinstance(val, bool) and lo <= val <= hi
+            else:
+                ok = isinstance(val, (int, float)) and not isinstance(val, bool) and lo <= val <= hi
+            if not ok:
+                raise ValueError(
+                    f"配置 {name}={val} 超出合理范围 [{lo}, {hi}]")
+
+
+def parse_checkpoint_config(cfg_dict: dict) -> Config:
+    """从 checkpoint 恢复 Config，校验键集合与值域
+
+    - 顶层缺 data/model 键 → ValueError（带明确信息）
+    - data/model 段含未知键 → ValueError（列出未知键，避免裸 TypeError）
+    - 值域非法 → ValueError（复用 validate_config_values）
+    """
+    if not isinstance(cfg_dict, dict):
+        raise ValueError(f"checkpoint config 应为 dict，实际为 {type(cfg_dict).__name__}")
+    for section in ("data", "model"):
+        if section not in cfg_dict:
+            raise ValueError(f"checkpoint config 缺少 '{section}' 段")
+        section_dict = cfg_dict[section]
+        if not isinstance(section_dict, dict):
+            raise ValueError(f"checkpoint config['{section}'] 应为 dict，实际为 {type(section_dict).__name__}")
+        known = set(DataConfig.__dataclass_fields__) if section == "data" \
+            else set(ModelConfig.__dataclass_fields__)
+        unknown = set(section_dict) - known
+        if unknown:
+            raise ValueError(
+                f"checkpoint config['{section}'] 含未知键: {sorted(unknown)}")
+    cfg = Config.from_dict(cfg_dict)
+    validate_config_values(cfg.data, cfg.model, cfg.train)
+    return cfg
+
+
 def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
     mapping = {
         "image_size": ("data", "image_size"),

actpred/dataset.py

@@ -39,10 +39,17 @@ class ActionDataset(Dataset):
             if not actions_path.exists() or not images_dir.is_dir():
                 continue
             actions = np.load(actions_path).astype(np.float32)
+            if actions.ndim != 2 or actions.shape[1] != action_dim:
+                raise ValueError(
+                    f"{actions_path} 形状应为 (T, {action_dim})，实际为 {actions.shape}。"
+                    f"请检查该 episode 的动作维度与 --action-dim 是否一致。")
             n_frames = len(actions)
             image_files = sorted(list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.png")),
                                  key=_natural_key)
             n_images = len(image_files)
+            if n_frames != n_images:
+                print(f"WARNING: {ep_dir} 动作帧数 {n_frames} 与图像数 {n_images} 不一致，"
+                      f"按较少者（{min(n_frames, n_images)}）使用，请检查录制完整性")
             n = min(n_frames, n_images)
             for t in range(history_len, n):
                 self.samples.append((image_files[t], actions[t - history_len:t], actions[t]))
@@ -109,11 +116,20 @@ class DataRecorder:
                 ep_idx = max(ep_idx, int(d.name.split("_")[1]) + 1)
             except (ValueError, IndexError):
                 pass
-        self.episode_dir = self.output_dir / f"episode_{ep_idx:04d}"
-        self.images_dir = self.episode_dir / "images"
-        self.images_dir.mkdir(parents=True, exist_ok=True)
+        # 原子创建目录，避免并发 Recorder 分配到同一 episode
+        while True:
+            self.episode_dir = self.output_dir / f"episode_{ep_idx:04d}"
+            self.images_dir = self.episode_dir / "images"
+            try:
+                self.images_dir.mkdir(parents=True, exist_ok=False)
+                break
+            except FileExistsError:
+                ep_idx += 1
 
     def record(self, image: np.ndarray, action: np.ndarray):
+        if action.shape != (self.action_dim,):
+            raise ValueError(
+                f"action 维度应为 ({self.action_dim},)，实际为 {action.shape}")
         img = Image.fromarray(image)
         img.save(self.images_dir / f"{self._frame_idx:04d}.jpg", quality=90)
         self._actions.append(action.astype(np.float32))

actpred/export_onnx.py

@@ -1,6 +1,6 @@
 import argparse
 import torch
-from config import DataConfig, ModelConfig
+from config import parse_checkpoint_config
 from model import ActionPredictor
 
 
@@ -11,11 +11,11 @@ def main():
     args = parser.parse_args()
 
     checkpoint = torch.load(args.model, map_location="cpu", weights_only=True)
-    cfg_dict = checkpoint["config"]
-    data_cfg = DataConfig(**cfg_dict["data"])
-    model_cfg = ModelConfig(**cfg_dict["model"])
+    cfg = parse_checkpoint_config(checkpoint["config"])
+    data_cfg = cfg.data
+    model_cfg = cfg.model
 
-    model = ActionPredictor(model_cfg, data_cfg)
+    model = ActionPredictor(model_cfg, data_cfg, pretrained=False)
     model.load_state_dict(checkpoint["model_state_dict"])
     model.eval()
 

actpred/model.py

@@ -27,12 +27,15 @@ class SimpleCNN(nn.Module):
 
 
 class MobileNetEncoder(nn.Module):
-    def __init__(self, feature_dim: int = 128):
+    def __init__(self, feature_dim: int = 128, pretrained: bool = True):
         super().__init__()
         from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
+        weights = MobileNet_V2_Weights.DEFAULT if pretrained else None
         try:
-            backbone = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+            backbone = mobilenet_v2(weights=weights)
         except Exception as e:
+            if not pretrained:
+                raise RuntimeError("加载 MobileNetV2 架构失败") from e
             raise RuntimeError(
                 "MobileNetV2 预训练权重下载失败，可能是网络受限。"
                 "请在联网环境中预下载权重"
@@ -94,10 +97,12 @@ class FusionHead(nn.Module):
 
 
 class ActionPredictor(nn.Module):
-    def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig):
+    def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig,
+                 pretrained: bool = True):
         super().__init__()
         if model_cfg.cnn_backbone == "mobilenet":
-            self.image_encoder = MobileNetEncoder(model_cfg.cnn_feature_dim)
+            self.image_encoder = MobileNetEncoder(
+                model_cfg.cnn_feature_dim, pretrained=pretrained)
         else:
             self.image_encoder = SimpleCNN(model_cfg.cnn_feature_dim)
 

actpred/mt_parser.py

@@ -1,7 +1,7 @@
 """轻量级 .mt 脚本 @Action 矩阵解析器
 
 从 .mt 文件中提取 @Action 块，解析为 (n_rows, n_cols) 的 float32 矩阵。
-支持 `-` 占位符（沿用上一行同列的值）。
+支持 `-` 占位符（沿用上一行同列的值）；首行 `-` 从 @Servo 块的初始角度填充。
 """
 
 import re
@@ -15,6 +15,27 @@ def _natural_key(p: Path):
     return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", p.name)]
 
 
+def _parse_servo_init(text: str) -> dict:
+    """从 @Servo 块提取 {servo_name: initial_angle}（max 列）"""
+    match = re.search(r"@Servo\s+(.*?)\s*@", text, re.DOTALL)
+    if not match:
+        return {}
+    servo_init = {}
+    for line in match.group(1).split("\n"):
+        line = re.sub(r"#.*$", "", line).strip()
+        if line.endswith(";"):
+            line = line[:-1].strip()
+        if not line:
+            continue
+        cols = re.split(r"\s+", line)
+        if len(cols) >= 2 and cols[0] != "name":
+            try:
+                servo_init[cols[0]] = float(cols[1])
+            except ValueError:
+                pass
+    return servo_init
+
+
 def parse_mt_action(path) -> np.ndarray:
     """解析 .mt 文件中的 @Action 矩阵
 
@@ -22,15 +43,19 @@ def parse_mt_action(path) -> np.ndarray:
         path: .mt 文件路径
 
     Returns:
-        (n_data_rows, n_cols) float32 矩阵，表头行已跳过，`-` 已前向填充
+        (n_data_rows, n_cols) float32 矩阵，表头行已跳过，`-` 已前向填充。
+        首行 `-` 从 @Servo 块的初始角度填充。
 
     Raises:
-        ValueError: 文件中无 @Action 块或数据行不足
+        ValueError: 文件中无 @Action 块、数据行不足、行列不一致或数值非法
     """
     text = Path(path).read_text(encoding="utf-8")
     # 先去除整行注释（# 开头），避免注释中的 @Action 干扰正则匹配
     text = re.sub(r"(?m)^[ \t]*#.*$", "", text)
 
+    # 解析 @Servo 块的初始角度，用于首行 `-` 前向填充
+    servo_init = _parse_servo_init(text)
+
     match = re.search(r"@Action\s+(.*?)\s*@", text, re.DOTALL)
     if not match:
         raise ValueError(f"未找到 @Action 块: {path}")
@@ -55,16 +80,32 @@ def parse_mt_action(path) -> np.ndarray:
         )
 
     # 第一行是表头（列名如 s1.A, s2.A, ...），跳过
+    header = rows[0]
     data_rows = rows[1:]
-    n_cols = len(rows[0])
+    n_cols = len(header)
+
+    # 为每列计算初始值：s1.A → @Servo 中 s1 的 max 角度，其余为 0.0
+    col_init = []
+    for col_name in header:
+        servo_name = col_name.split(".")[0] if "." in col_name else None
+        col_init.append(servo_init.get(servo_name, 0.0))
 
     matrix = np.zeros((len(data_rows), n_cols), dtype=np.float32)
     for i, row in enumerate(data_rows):
+        if len(row) != n_cols:
+            raise ValueError(
+                f"@Action 第 {i + 2} 行列数 {len(row)} 与表头 {n_cols} "
+                f"不一致: {path}")
         for j, val in enumerate(row):
             if val == "-":
-                matrix[i, j] = matrix[i - 1, j] if i > 0 else 0.0
+                matrix[i, j] = matrix[i - 1, j] if i > 0 else col_init[j]
             else:
-                matrix[i, j] = float(val)
+                try:
+                    matrix[i, j] = float(val)
+                except ValueError:
+                    raise ValueError(
+                        f"@Action 第 {i + 2} 行第 {j + 1} 列值 "
+                        f"'{val}' 无法转为数值: {path}")
 
     return matrix
 
@@ -86,6 +127,6 @@ def parse_mt_directory(mt_dir, pattern="*.mt"):
         try:
             matrix = parse_mt_action(f)
             results.append((str(f), matrix))
-        except ValueError:
-            pass
+        except Exception as e:
+            print(f"WARNING: 跳过 .mt 文件 {f}: {e}")
     return results

actpred/predict.py

@@ -6,22 +6,23 @@ import torch
 from PIL import Image
 from torchvision import transforms
 
-from config import Config, ModelConfig, DataConfig
+from config import parse_checkpoint_config
 from model import ActionPredictor
 
 
 class ActionPredictorInference:
     def __init__(self, checkpoint_path: str, device: str = None):
         checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
-        cfg_dict = checkpoint["config"]
-        self._data_cfg = DataConfig(**cfg_dict["data"])
-        self._model_cfg = ModelConfig(**cfg_dict["model"])
+        cfg = parse_checkpoint_config(checkpoint["config"])
+        self._data_cfg = cfg.data
+        self._model_cfg = cfg.model
 
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
         self._device = torch.device(device)
 
-        self._model = ActionPredictor(self._model_cfg, self._data_cfg).to(self._device)
+        self._model = ActionPredictor(
+            self._model_cfg, self._data_cfg, pretrained=False).to(self._device)
         self._model.load_state_dict(checkpoint["model_state_dict"])
         self._model.eval()
 
@@ -36,15 +37,28 @@ class ActionPredictorInference:
 
         self._action_mean = None
         self._action_std = None
-        stats_path = os.path.join(os.path.dirname(os.path.abspath(checkpoint_path)),
-                                  "action_stats.npy")
         if self._data_cfg.normalize_actions:
+            # 优先使用 checkpoint config 记录的 stats 路径（相对路径按 checkpoint 目录解析），
+            # 旧 checkpoint 无该字段时回退到 checkpoint 同目录的 action_stats.npy
+            stats_path = self._data_cfg.action_stats_path
+            if stats_path:
+                if not os.path.isabs(stats_path):
+                    stats_path = os.path.join(
+                        os.path.dirname(os.path.abspath(checkpoint_path)), stats_path)
+            else:
+                stats_path = os.path.join(
+                    os.path.dirname(os.path.abspath(checkpoint_path)),
+                    "action_stats.npy")
             if not os.path.exists(stats_path):
                 raise FileNotFoundError(
                     f"模型训练时启用了动作归一化，但未找到统计文件: {stats_path}")
             stats = np.load(stats_path)
-            self._action_mean = stats[0]
-            self._action_std = stats[1]
+            if stats.shape != (2, self._data_cfg.action_dim):
+                raise ValueError(
+                    f"action_stats.npy 形状应为 (2, {self._data_cfg.action_dim})，"
+                    f"实际为 {stats.shape}，请检查文件: {stats_path}")
+            self._action_mean = stats[0].astype(np.float32)
+            self._action_std = stats[1].astype(np.float32)
 
     @torch.no_grad()
     def predict(self, image: np.ndarray, action_history: np.ndarray) -> np.ndarray:
@@ -69,8 +83,9 @@ def parse_action_history(s: str, action_dim: int) -> np.ndarray:
     history = []
     for step in steps:
         values = [float(v) for v in step.split(",")]
-        assert len(values) == action_dim, \
-            f"Expected {action_dim} values per step, got {len(values)}"
+        if len(values) != action_dim:
+            raise ValueError(
+                f"Expected {action_dim} values per step, got {len(values)}")
         history.append(values)
     return np.array(history, dtype=np.float32)
 

actpred/pretrain_mt.py

@@ -20,10 +20,10 @@ import random
 
 import numpy as np
 import torch
-from torch.utils.data import Dataset, DataLoader, Subset
+from torch.utils.data import Dataset, DataLoader
 from tqdm import tqdm
 
-from config import Config, build_argparser, apply_overrides
+from config import Config, build_argparser, apply_overrides, validate_config_values
 from model import ActionPredictor
 from mt_parser import parse_mt_directory
 
@@ -34,20 +34,26 @@ class MTActionDataset(Dataset):
     每个样本: (action_history[history_len, action_dim], target[action_dim])
     """
 
-    def __init__(self, mt_dir: str, history_len: int, action_dim: int):
+    def __init__(self, mt_dir: str, history_len: int, action_dim: int,
+                 file_indices: list = None, files: list = None):
         self.history_len = history_len
         self.action_dim = action_dim
         self.samples = []
 
-        for path, matrix in parse_mt_directory(mt_dir):
-            n_rows = matrix.shape[1]
-            if n_rows < action_dim:
+        # 复用调用方已解析的结果（避免 train/val 各自重复解析同一目录）
+        all_files = files if files is not None else parse_mt_directory(mt_dir)
+        if file_indices is not None:
+            all_files = [all_files[i] for i in file_indices]
+
+        for path, matrix in all_files:
+            n_cols = matrix.shape[1]
+            if n_cols < action_dim:
                 matrix = np.pad(
                     matrix,
-                    ((0, 0), (0, action_dim - n_rows)),
+                    ((0, 0), (0, action_dim - n_cols)),
                     mode="constant",
                 )
-            elif n_rows > action_dim:
+            elif n_cols > action_dim:
                 matrix = matrix[:, :action_dim]
 
             for i in range(len(matrix) - history_len):
@@ -56,7 +62,8 @@ class MTActionDataset(Dataset):
                 self.samples.append((history.copy(), target.copy()))
 
         print(f"MTActionDataset: {len(self.samples)} samples "
-              f"from .mt files (history_len={history_len}, action_dim={action_dim})")
+              f"from {len(all_files)} files "
+              f"(history_len={history_len}, action_dim={action_dim})")
 
     def __len__(self):
         return len(self.samples)
@@ -128,6 +135,7 @@ def main():
     args = parser.parse_args()
 
     cfg = apply_overrides(Config(), args)
+    validate_config_values(cfg.data, cfg.model, cfg.train)
     cfg.data.data_dir = args.mt_dir
 
     os.makedirs(cfg.train.checkpoint_dir, exist_ok=True)
@@ -141,26 +149,32 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Device: {device}")
 
-    full_ds = MTActionDataset(
-        args.mt_dir, cfg.data.history_len, cfg.data.action_dim
-    )
-    if len(full_ds) == 0:
+    # 按文件级划分 train/val，避免滑动窗口样本重叠导致的数据泄漏
+    all_files = parse_mt_directory(args.mt_dir)
+    if len(all_files) == 0:
         raise ValueError(
             f"数据集为空。请检查 mt-dir={args.mt_dir} "
             f"及 .mt 文件中 @Action 行数 > history_len={cfg.data.history_len}"
         )
 
-    n = len(full_ds)
-    indices = list(range(n))
+    n_files = len(all_files)
+    file_indices = list(range(n_files))
     rng = np.random.RandomState(cfg.train.seed)
-    rng.shuffle(indices)
-    val_count = max(1, int(n * cfg.train.val_split))
-    val_indices = indices[:val_count]
-    train_indices = indices[val_count:]
+    rng.shuffle(file_indices)
+    val_file_count = max(1, int(n_files * cfg.train.val_split))
+    val_file_indices = file_indices[:val_file_count]
+    train_file_indices = file_indices[val_file_count:]
 
-    train_ds = Subset(full_ds, train_indices)
-    val_ds = Subset(full_ds, val_indices)
+    train_ds = MTActionDataset(
+        args.mt_dir, cfg.data.history_len, cfg.data.action_dim,
+        file_indices=train_file_indices, files=all_files,
+    )
+    val_ds = MTActionDataset(
+        args.mt_dir, cfg.data.history_len, cfg.data.action_dim,
+        file_indices=val_file_indices, files=all_files,
+    )
 
+    print(f"Files: {len(train_file_indices)} train, {len(val_file_indices)} val")
     print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
     if len(train_ds) == 0:
         raise ValueError("训练集为空（样本数不足）")
@@ -179,7 +193,7 @@ def main():
         shuffle=False, num_workers=cfg.train.num_workers,
     )
 
-    model = ActionPredictor(cfg.model, cfg.data).to(device)
+    model = ActionPredictor(cfg.model, cfg.data, pretrained=False).to(device)
     n_params = sum(p.numel() for p in model.parameters())
     print(f"Model params: {n_params:,}")
     print(f"Config: action_dim={cfg.data.action_dim}, "
@@ -256,7 +270,7 @@ def main():
     else:
         print(f"Best pre-trained model: {inference_path}")
         print("Fine-tune: bash _run.sh train --data-dir ./data "
-              f"--model {inference_path}")
+              f"--init-from {inference_path}")
 
     if args.save_config:
         with open(args.save_config, "w") as f:

actpred/requirements.txt

@@ -2,7 +2,7 @@
 torch>=2.0,<3.0
 torchvision>=0.15,<1.0
 numpy>=1.24,<3.0
-Pillow>=9.0,<12.0
+Pillow>=9.0,<14.0  # torchvision 依赖，上限放宽避免与 torch 源安装冲突
 
 # Training utilities
 tqdm>=4.60,<5.0

actpred/test_actpred.py

@@ -6,19 +6,41 @@
     cd actpred && python test_actpred.py
 """
 
+import contextlib
+import importlib.util
+import io
 import os
 import tempfile
 import unittest
 from pathlib import Path
+from unittest import mock
 
 import numpy as np
 from PIL import Image
 import torch
 
-from config import Config, DataConfig, ModelConfig, build_argparser, apply_overrides
+from config import (Config, DataConfig, ModelConfig, build_argparser,
+                    apply_overrides, validate_config_values, parse_checkpoint_config)
 from dataset import _natural_key, ActionDataset, DataRecorder
 from model import ActionPredictor
-from predict import parse_action_history
+from mt_parser import parse_mt_action, parse_mt_directory
+from predict import ActionPredictorInference, parse_action_history
+from pretrain_mt import MTActionDataset
+
+
+def _write_mt(dir_path, name, content):
+    path = Path(dir_path) / name
+    path.write_text(content, encoding="utf-8")
+    return path
+
+
+def _make_seq_mt(n_rows):
+    """生成 3 列 @Action 块，第 i 行为 [i, i*10, i*100]"""
+    lines = ["@Action", "a.A  b.A  c.A;"]
+    for i in range(n_rows):
+        lines.append(f"{i}  {i * 10}  {i * 100};")
+    lines.append("@")
+    return "\n".join(lines) + "\n"
 
 
 class TestNaturalKey(unittest.TestCase):
@@ -52,7 +74,7 @@ class TestParseActionHistory(unittest.TestCase):
         self.assertEqual(history.shape, (3, 4))
 
     def test_wrong_dim_raises(self):
-        with self.assertRaises(AssertionError):
+        with self.assertRaises(ValueError):
             parse_action_history("1,2,3", action_dim=4)
 
 
@@ -179,6 +201,37 @@ class TestActionDataset(unittest.TestCase):
                               normalize_actions=True,
                               action_stats_path="/nonexistent/path.npy")
 
+    def test_action_dim_mismatch_raises(self):
+        """actions.npy 列数与 action_dim 不一致应在加载时抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            ep_dir = base / "episode_0001"
+            images_dir = ep_dir / "images"
+            images_dir.mkdir(parents=True, exist_ok=True)
+            for i in range(8):
+                Image.fromarray(np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
+                                ).save(images_dir / f"{i:04d}.jpg")
+            np.save(ep_dir / "actions.npy", np.random.randn(8, 6).astype(np.float32))
+            with self.assertRaisesRegex(ValueError, "actions.npy"):
+                ActionDataset(tmpdir, history_len=3, action_dim=4)
+
+    def test_frame_mismatch_warns(self):
+        """图像与动作帧数不一致时应打印 WARNING 并截断"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base = Path(tmpdir)
+            ep_dir = base / "episode_0001"
+            images_dir = ep_dir / "images"
+            images_dir.mkdir(parents=True, exist_ok=True)
+            for i in range(8):
+                Image.fromarray(np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8)
+                                ).save(images_dir / f"{i:04d}.jpg")
+            np.save(ep_dir / "actions.npy", np.random.randn(12, 4).astype(np.float32))
+            buf = io.StringIO()
+            with contextlib.redirect_stdout(buf):
+                ds = ActionDataset(tmpdir, history_len=3, action_dim=4)
+            self.assertIn("WARNING", buf.getvalue())
+            self.assertEqual(len(ds), 8 - 3)  # 按较少者（8 帧图像）截断
+
 
 class TestConfig(unittest.TestCase):
     """测试配置和 CLI 参数"""
@@ -215,5 +268,387 @@ class TestConfig(unittest.TestCase):
         self.assertIsNone(args.normalize_actions)
 
 
+class TestMtParser(unittest.TestCase):
+    """测试 .mt 脚本 @Action 解析器"""
+
+    def test_parse_basic_and_forward_fill(self):
+        """正常解析：跳过表头/注释，`-` 沿用上一行同列值"""
+        content = (
+            "@Action\n"
+            "s1.A  s2.A  spk.F;\t# column name\n"
+            "  10    20    100;\n"
+            "  11    21    101;\n"
+            "   -    22      -;\n"
+            "@\n"
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "basic.mt", content)
+            matrix = parse_mt_action(path)
+        self.assertEqual(matrix.shape, (3, 3))
+        self.assertEqual(matrix.dtype, np.float32)
+        np.testing.assert_array_equal(matrix[0], [10, 20, 100])
+        np.testing.assert_array_equal(matrix[2], [11, 22, 101])
+
+    def test_first_row_dash_uses_servo_init(self):
+        """首行 `-` 从 @Servo 初始角度（max 列）填充，无映射列填 0"""
+        content = (
+            "@Servo\n"
+            " name   max  min;\n"
+            "  s1    90     -;\n"
+            "  s2   180     -;\n"
+            "@\n"
+            "\n"
+            "@Action\n"
+            "s1.A  s2.A  spk.F;\n"
+            "   -     -      -;\n"
+            "  45     -     10;\n"
+            "@\n"
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "servo.mt", content)
+            matrix = parse_mt_action(path)
+        np.testing.assert_array_equal(matrix[0], [90, 180, 0])
+        np.testing.assert_array_equal(matrix[1], [45, 180, 10])
+
+    def test_no_action_block_raises(self):
+        """无 @Action 块应抛 ValueError"""
+        content = "@Servo\nname  max;\n  s1   90;\n@\n"
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "noaction.mt", content)
+            with self.assertRaisesRegex(ValueError, "@Action"):
+                parse_mt_action(path)
+
+    def test_header_only_raises(self):
+        """只有表头没有数据行应抛 ValueError"""
+        content = "@Action\na.A  b.A;\n@\n"
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "empty.mt", content)
+            with self.assertRaises(ValueError):
+                parse_mt_action(path)
+
+    def test_row_col_mismatch_raises(self):
+        """数据行列数与表头不一致应抛带行号的 ValueError"""
+        content = "@Action\na.A  b.A;\n  1    2;\n  1    2    3;\n@\n"
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "mismatch.mt", content)
+            with self.assertRaisesRegex(ValueError, "第 3 行"):
+                parse_mt_action(path)
+
+    def test_invalid_value_raises(self):
+        """非数值且非 `-` 的单元格应抛带行列号的 ValueError"""
+        content = "@Action\na.A  b.A;\n  1  xyz;\n@\n"
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _write_mt(tmpdir, "badval.mt", content)
+            with self.assertRaisesRegex(ValueError, "xyz"):
+                parse_mt_action(path)
+
+    def test_directory_skips_bad_files_with_warning(self):
+        """parse_mt_directory 跳过坏文件并打印 WARNING，保留好文件"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _write_mt(tmpdir, "good.mt", _make_seq_mt(4))
+            _write_mt(tmpdir, "bad.mt", "@Action\na.A  b.A;\n  1  xyz;\n@\n")
+            buf = io.StringIO()
+            with contextlib.redirect_stdout(buf):
+                results = parse_mt_directory(tmpdir)
+        self.assertEqual(len(results), 1)
+        self.assertTrue(results[0][0].endswith("good.mt"))
+        self.assertIn("WARNING", buf.getvalue())
+        self.assertIn("bad.mt", buf.getvalue())
+
+
+class TestMTActionDataset(unittest.TestCase):
+    """测试预训练滑动窗口数据集"""
+
+    @staticmethod
+    def _build(mt_dir, history_len, action_dim, file_indices=None):
+        with contextlib.redirect_stdout(io.StringIO()):
+            return MTActionDataset(mt_dir, history_len, action_dim,
+                                   file_indices=file_indices)
+
+    def test_sliding_window_count_and_shapes(self):
+        """N 行数据、窗口 H → N-H 个样本，形状与 target 值正确"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _write_mt(tmpdir, "a.mt", _make_seq_mt(7))
+            ds = self._build(tmpdir, history_len=3, action_dim=3)
+            self.assertEqual(len(ds), 4)  # 7 - 3 = 4 samples
+            sample = ds[0]
+            self.assertEqual(sample["action_history"].shape, (3, 3))
+            self.assertEqual(sample["target"].shape, (3,))
+            # 样本 0 的 target 是第 4 行数据 [3, 30, 300]
+            np.testing.assert_array_equal(sample["target"].numpy(), [3, 30, 300])
+
+    def test_pad_columns_to_action_dim(self):
+        """列数 < action_dim 时右侧补零"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _write_mt(tmpdir, "a.mt", _make_seq_mt(5))
+            ds = self._build(tmpdir, history_len=3, action_dim=5)
+            sample = ds[0]
+            self.assertEqual(sample["action_history"].shape, (3, 5))
+            np.testing.assert_array_equal(
+                sample["action_history"].numpy()[:, 3:], np.zeros((3, 2)))
+            np.testing.assert_array_equal(sample["target"].numpy()[3:], [0, 0])
+
+    def test_truncate_columns_to_action_dim(self):
+        """列数 > action_dim 时截断多余列"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _write_mt(tmpdir, "a.mt", _make_seq_mt(5))
+            ds = self._build(tmpdir, history_len=3, action_dim=2)
+            sample = ds[0]
+            self.assertEqual(sample["action_history"].shape, (3, 2))
+            np.testing.assert_array_equal(sample["target"].numpy(), [3, 30])
+
+    def test_file_indices_filter(self):
+        """file_indices 只加载指定文件，实现文件级 train/val 划分"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            _write_mt(tmpdir, "a.mt", _make_seq_mt(7))  # 4 samples
+            _write_mt(tmpdir, "b.mt", _make_seq_mt(4))  # 1 sample
+            ds_all = self._build(tmpdir, history_len=3, action_dim=3)
+            ds_a = self._build(tmpdir, history_len=3, action_dim=3,
+                               file_indices=[0])
+            ds_b = self._build(tmpdir, history_len=3, action_dim=3,
+                               file_indices=[1])
+            self.assertEqual(len(ds_all), 5)
+            self.assertEqual(len(ds_a), 4)
+            self.assertEqual(len(ds_b), 1)
+
+
+class TestValidateConfigValues(unittest.TestCase):
+    """测试 checkpoint config 值域校验"""
+
+    def test_default_config_passes(self):
+        cfg = Config()
+        validate_config_values(cfg.data, cfg.model)  # 不应抛异常
+
+    def test_out_of_range_or_non_int_raises(self):
+        for field, value in [("image_size", 99999), ("image_size", 0),
+                             ("history_len", -1), ("action_dim", "18")]:
+            cfg = Config()
+            setattr(cfg.data, field, value)
+            with self.assertRaises(ValueError, msg=f"{field}={value}"):
+                validate_config_values(cfg.data, cfg.model)
+
+    def test_model_field_out_of_range_raises(self):
+        """model 字段越界（fused_hidden_dim 超限、dropout 越界）应抛 ValueError"""
+        for field, value in [("fused_hidden_dim", 10 ** 9),
+                             ("fused_hidden_dim", 0),
+                             ("cnn_feature_dim", 99999),
+                             ("action_hidden_dim", -1),
+                             ("dropout", 1.5),
+                             ("dropout", -0.1)]:
+            cfg = Config()
+            setattr(cfg.model, field, value)
+            with self.assertRaises(ValueError, msg=f"{field}={value}"):
+                validate_config_values(cfg.data, cfg.model)
+
+    def test_train_field_out_of_range_raises(self):
+        """train 字段越界（lr 为负、val_split>1 等）应抛 ValueError"""
+        for field, value in [("lr", -1.0), ("lr", 0.0),
+                             ("val_split", 2.0), ("val_split", -0.1),
+                             ("batch_size", 0), ("epochs", 0),
+                             ("patience", -1), ("num_workers", -1),
+                             ("max_grad_norm", -1.0), ("weight_decay", -1.0)]:
+            cfg = Config()
+            setattr(cfg.train, field, value)
+            with self.assertRaises(ValueError, msg=f"{field}={value}"):
+                validate_config_values(cfg.data, cfg.model, cfg.train)
+
+    def test_default_train_passes(self):
+        cfg = Config()
+        validate_config_values(cfg.data, cfg.model, cfg.train)  # 不应抛异常
+
+
+class TestParseCheckpointConfig(unittest.TestCase):
+    """测试 checkpoint config 结构解析（未知键/缺段应给明确错误而非裸 TypeError）"""
+
+    def test_valid_config_roundtrip(self):
+        cfg = Config()
+        parsed = parse_checkpoint_config(cfg.to_dict())
+        self.assertEqual(parsed.data.action_dim, cfg.data.action_dim)
+        self.assertEqual(parsed.model.cnn_backbone, cfg.model.cnn_backbone)
+        self.assertEqual(parsed.train.epochs, cfg.train.epochs)
+
+    def test_unknown_key_raises(self):
+        cfg = Config()
+        d = cfg.to_dict()
+        d["data"]["evil_extra_key"] = 123
+        with self.assertRaisesRegex(ValueError, "evil_extra_key"):
+            parse_checkpoint_config(d)
+
+    def test_missing_section_raises(self):
+        cfg = Config()
+        d = cfg.to_dict()
+        del d["data"]
+        with self.assertRaisesRegex(ValueError, "data"):
+            parse_checkpoint_config(d)
+
+    def test_out_of_range_raises(self):
+        cfg = Config()
+        d = cfg.to_dict()
+        d["model"]["fused_hidden_dim"] = 10 ** 9
+        with self.assertRaisesRegex(ValueError, "fused_hidden_dim"):
+            parse_checkpoint_config(d)
+
+    def test_non_dict_raises(self):
+        with self.assertRaises(ValueError):
+            parse_checkpoint_config("not a dict")
+
+
+class TestDataRecorderValidation(unittest.TestCase):
+    """测试 DataRecorder 的输入校验与目录原子分配"""
+
+    def test_wrong_action_dim_raises(self):
+        """record 传入错误维度的 action 应抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            recorder = DataRecorder(tmpdir, action_dim=4)
+            img = np.zeros((32, 32, 3), dtype=np.uint8)
+            with self.assertRaisesRegex(ValueError, "维度"):
+                recorder.record(img, np.zeros(3, dtype=np.float32))
+
+    def test_dir_collision_retries_next_index(self):
+        """目标 episode 目录已被占用时应重试递增索引（模拟并发竞态）"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            r1 = DataRecorder(tmpdir, action_dim=4)
+            self.assertIn("episode_0000", str(r1.episode_dir))
+            # 让 glob 返回空 → r2 从 0 开始分配，与 r1 冲突后应重试到 0001
+            with mock.patch.object(Path, "glob", return_value=iter([])):
+                r2 = DataRecorder(tmpdir, action_dim=4)
+            self.assertIn("episode_0001", str(r2.episode_dir))
+
+
+class TestActionPredictorInference(unittest.TestCase):
+    """测试推理封装端到端：checkpoint 保存 → 加载 → 推理"""
+
+    @staticmethod
+    def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32):
+        cfg = Config()
+        cfg.data.image_size = image_size
+        cfg.data.action_dim = 4
+        cfg.data.history_len = 3
+        cfg.data.normalize_actions = normalize_actions
+        model = ActionPredictor(cfg.model, cfg.data)
+        model.eval()
+        path = os.path.join(tmpdir, "best_inference.pt")
+        torch.save({"model_state_dict": model.state_dict(),
+                    "config": cfg.to_dict()}, path)
+        return path, model, cfg
+
+    def test_end_to_end_predict(self):
+        """加载 checkpoint 后推理输出形状正确、数值有限且可复现"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = self._make_checkpoint(tmpdir)
+            predictor = ActionPredictorInference(path, device="cpu")
+            self.assertEqual(predictor.action_dim, 4)
+            self.assertEqual(predictor.history_len, 3)
+
+            image = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
+            history = np.random.randn(3, 4).astype(np.float32)
+            pred1 = predictor.predict(image, history)
+            pred2 = predictor.predict(image, history)
+            self.assertEqual(pred1.shape, (4,))
+            self.assertTrue(np.all(np.isfinite(pred1)))
+            np.testing.assert_array_equal(pred1, pred2)  # eval 模式确定性
+
+    def test_loaded_weights_match_saved(self):
+        """加载后的模型权重应与保存时逐张量一致"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, model, _ = self._make_checkpoint(tmpdir)
+            predictor = ActionPredictorInference(path, device="cpu")
+            loaded_sd = predictor._model.state_dict()
+            for key, tensor in model.state_dict().items():
+                self.assertTrue(torch.equal(loaded_sd[key].cpu(), tensor),
+                                msg=f"权重不一致: {key}")
+
+    def test_wrong_history_shape_raises(self):
+        """action_history 形状错误应抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = self._make_checkpoint(tmpdir)
+            predictor = ActionPredictorInference(path, device="cpu")
+            image = np.zeros((32, 32, 3), dtype=np.uint8)
+            with self.assertRaisesRegex(ValueError, "action_history"):
+                predictor.predict(image, np.zeros((2, 4), dtype=np.float32))
+
+    def test_missing_action_stats_raises(self):
+        """启用归一化但缺 action_stats.npy 应抛 FileNotFoundError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = self._make_checkpoint(tmpdir, normalize_actions=True)
+            with self.assertRaises(FileNotFoundError):
+                ActionPredictorInference(path, device="cpu")
+
+    def test_normalize_identity_stats_matches_raw(self):
+        """mean=0/std=1 的归一化路径输出应与非归一化一致"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            raw_dir = os.path.join(tmpdir, "raw")
+            norm_dir = os.path.join(tmpdir, "norm")
+            os.makedirs(raw_dir)
+            os.makedirs(norm_dir)
+            torch.manual_seed(0)
+            raw_path, model, cfg = self._make_checkpoint(raw_dir)
+            # 同一份权重另存为启用归一化的 checkpoint
+            cfg.data.normalize_actions = True
+            norm_path = os.path.join(norm_dir, "best_inference.pt")
+            torch.save({"model_state_dict": model.state_dict(),
+                        "config": cfg.to_dict()}, norm_path)
+            stats = np.stack([np.zeros(4), np.ones(4)]).astype(np.float32)
+            np.save(os.path.join(norm_dir, "action_stats.npy"), stats)
+
+            image = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
+            history = np.random.randn(3, 4).astype(np.float32)
+            pred_raw = ActionPredictorInference(
+                raw_path, device="cpu").predict(image, history)
+            pred_norm = ActionPredictorInference(
+                norm_path, device="cpu").predict(image, history)
+            np.testing.assert_allclose(pred_norm, pred_raw, rtol=1e-5, atol=1e-6)
+
+    def test_stats_shape_mismatch_raises(self):
+        """action_stats.npy 形状与 action_dim 不一致应抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = self._make_checkpoint(tmpdir, normalize_actions=True)
+            np.save(os.path.join(tmpdir, "action_stats.npy"), np.zeros((2, 7)))
+            with self.assertRaisesRegex(ValueError, "action_stats.npy"):
+                ActionPredictorInference(path, device="cpu")
+
+    def test_invalid_config_in_checkpoint_raises(self):
+        """checkpoint 中 config 值域非法应在加载时抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, model, cfg = self._make_checkpoint(tmpdir)
+            cfg_dict = cfg.to_dict()
+            cfg_dict["data"]["image_size"] = 99999
+            bad_path = os.path.join(tmpdir, "bad.pt")
+            torch.save({"model_state_dict": model.state_dict(),
+                        "config": cfg_dict}, bad_path)
+            with self.assertRaisesRegex(ValueError, "image_size"):
+                ActionPredictorInference(bad_path, device="cpu")
+
+    def test_unknown_config_key_in_checkpoint_raises(self):
+        """checkpoint config 含未知键应在加载时抛带键名的 ValueError 而非裸 TypeError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, model, cfg = self._make_checkpoint(tmpdir)
+            cfg_dict = cfg.to_dict()
+            cfg_dict["data"]["evil_extra_key"] = 123
+            bad_path = os.path.join(tmpdir, "bad.pt")
+            torch.save({"model_state_dict": model.state_dict(),
+                        "config": cfg_dict}, bad_path)
+            with self.assertRaisesRegex(ValueError, "evil_extra_key"):
+                ActionPredictorInference(bad_path, device="cpu")
+
+
+@unittest.skipUnless(importlib.util.find_spec("onnxscript") is not None,
+                     "torch.onnx.export 需要 onnxscript（当前环境未安装）")
+class TestExportOnnx(unittest.TestCase):
+    """测试 ONNX 导出（环境缺 onnxscript 时自动跳过）"""
+
+    def test_export_creates_onnx_file(self):
+        """从 checkpoint 导出 ONNX 文件应成功生成"""
+        from export_onnx import main as export_onnx_main
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = TestActionPredictorInference._make_checkpoint(tmpdir)
+            out_path = os.path.join(tmpdir, "model.onnx")
+            argv = ["export_onnx.py", "--model", path, "--output", out_path]
+            with mock.patch("sys.argv", argv):
+                export_onnx_main()
+            self.assertTrue(os.path.exists(out_path))
+            self.assertGreater(os.path.getsize(out_path), 0)
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)

actpred/train.py

@@ -6,7 +6,7 @@ import torch
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
-from config import Config, build_argparser, apply_overrides
+from config import Config, build_argparser, apply_overrides, validate_config_values
 from model import ActionPredictor
 from dataset import ActionDataset
 
@@ -62,6 +62,7 @@ def main():
     parser = build_argparser("actpred training")
     args = parser.parse_args()
     cfg = apply_overrides(Config(), args)
+    validate_config_values(cfg.data, cfg.model, cfg.train)
 
     os.makedirs(cfg.train.checkpoint_dir, exist_ok=True)
 
@@ -76,7 +77,7 @@ def main():
     print(f"Device: {device}")
 
     train_indices, val_indices = split_episodes(
-        cfg.data.data_dir, cfg.train.val_split)
+        cfg.data.data_dir, cfg.train.val_split, seed=cfg.train.seed)
     print(f"Episodes: {len(train_indices)} train, {len(val_indices)} val")
 
     train_ds = ActionDataset(
@@ -115,6 +116,16 @@ def main():
     n_params = sum(p.numel() for p in model.parameters())
     print(f"Model params: {n_params:,}")
 
+    if args.init_from:
+        init_ckpt = torch.load(args.init_from, map_location="cpu", weights_only=True)
+        missing, unexpected = model.load_state_dict(
+            init_ckpt["model_state_dict"], strict=False)
+        if missing:
+            print(f"  init-from missing keys: {missing}")
+        if unexpected:
+            print(f"  init-from unexpected keys: {unexpected}")
+        print(f"Loaded pretrained weights from: {args.init_from}")
+
     optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.train.lr,
                                   weight_decay=cfg.train.weight_decay)
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.train.epochs)