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
This commit is contained in:
qing committed
1 parent
23a855963b
commit
63cfdd2066
15 files changed
+960
-65
No files matched your search
+15
-2
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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"),
|
||||
|
||||
+19
-3
@@ -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))
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
+9
-4
@@ -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)
|
||||
|
||||
|
||||
+49
-8
@@ -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
|
||||
+26
-11
@@ -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)
|
||||
|
||||
|
||||
+37
-23
@@ -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:
|
||||
|
||||
@@ -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
|
||||
+438
-3
@@ -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)
|
||||
+13
-2
@@ -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)
|
||||
|
||||
+3
-2
@@ -7,7 +7,7 @@
|
||||
```
|
||||
docs/
|
||||
├── INDEX.md # 本文件（文档总索引）
|
||||
├── security/ # 安全审计与加固（25 篇）
|
||||
├── security/ # 安全审计与加固（26 篇）
|
||||
├── reviews/ # 代码审查日志（4 篇）
|
||||
└── reports/ # 分析报告与开发规划（9 篇）
|
||||
```
|
||||
@@ -57,6 +57,7 @@ docs/
|
||||
| libth | [libth_SECURITY_AUDIT_20260804.md](security/libth_SECURITY_AUDIT_20260804.md) | C++ 线程库（Thread、Cv_wait、Fsm） |
|
||||
| aibox | [aibox_SECURITY_AUDIT_20260801.md](security/aibox_SECURITY_AUDIT_20260801.md) | Python AI 智能体系统（WebSocket） |
|
||||
| BotOps | [BotOps_SECURITY_AUDIT_20260801.md](security/BotOps_SECURITY_AUDIT_20260801.md) | BotOps 运维控制面板 |
|
||||
| actpred | [actpred_SECURITY_AUDIT_20260808.md](security/actpred_SECURITY_AUDIT_20260808.md) | Python 动作预测模块（PyTorch 训练/推理、.mt 解析、ONNX 导出） |
|
||||
|
||||
---
|
||||
|
||||
@@ -66,7 +67,7 @@ docs/
|
||||
|------|------|
|
||||
| [core_CODE_REVIEW_LOG_20260730.md](reviews/core_CODE_REVIEW_LOG_20260730.md) | Core 模块代码审查日志（2026-07-30，自动化审查） |
|
||||
| [DEPLOY_DEFECT_AUDIT_20260806.md](reviews/DEPLOY_DEFECT_AUDIT_20260806.md) | 部署缺陷审查（2026-08-06，首次原生部署实录，含修复记录） |
|
||||
| [actpred_CODE_AUDIT_20260808.md](reviews/actpred_CODE_AUDIT_20260808.md) | actpred 模块代码审计（2026-08-08，PyTorch 动作预测，16 项发现，全部修复） |
|
||||
| [actpred_CODE_AUDIT_20260808.md](reviews/actpred_CODE_AUDIT_20260808.md) | actpred 模块代码审计（2026-08-08，两轮合并：第 1 轮 16 项 + 第 2 轮复审 11 项，共 27 项全部修复） |
|
||||
| [bot4cj_CODE_AUDIT_20260808.md](reviews/bot4cj_CODE_AUDIT_20260808.md) | bot4cj 模块代码审计（2026-08-08，仓颉机器人硬件脚本解析器，15 项发现，待修复） |
|
||||
|
||||
---
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
# actpred 模块代码审计报告
|
||||
# actpred 模块代码审计报告（第 1 轮 + 第 2 轮复审）
|
||||
|
||||
> 审查日期：2026-08-08 (Asia/Shanghai)
|
||||
> 审查人：Automated Review
|
||||
> 审查范围：`actpred/` 全模块（Python，PyTorch 动作预测神经网络）
|
||||
>
|
||||
> **文档合并说明**：本文档为两轮审计的合并版——第 1 轮（16 项发现，全部修复）见 §二~§七；
|
||||
> 第 2 轮复审（11 项发现，含 3 项上轮修复不完整，全部修复）见 §八。
|
||||
|
||||
---
|
||||
|
||||
@@ -345,3 +348,135 @@ OK
|
||||
---
|
||||
|
||||
*本报告由自动化代码审计生成，供开发团队参考。修复建议需结合实际场景评估。*
|
||||
|
||||
---
|
||||
|
||||
## 八、第 2 轮复审（2026-08-08）
|
||||
|
||||
> 前置文档：[第 1 轮安全审计](../security/actpred_SECURITY_AUDIT_20260808.md)（10 项发现，全部标记已修复 ✅）
|
||||
|
||||
### 8.1 复审方法
|
||||
|
||||
本轮采用**静态通读 + 动态验证**：除逐文件复查全部源码外，对疑点编写了针对性复现脚本（checkpoint 畸形 config、值域校验边界、数据形状不匹配、CLI 越界参数等），**所有结论均以实测输出为准**，不采信审计文档的"已修复"标记。
|
||||
|
||||
实测环境：Intel Xeon E3-1200 v2, Linux 24.04, PyTorch CPU-only。第 1 轮审计后测试已扩展至 **39 用例（38 通过 + 1 环境性跳过）**。
|
||||
|
||||
### 8.2 复审结论摘要
|
||||
|
||||
| 类别 | 数量 | 说明 |
|
||||
|------|------|------|
|
||||
| A 类：上轮"已修复"标记与实现不符（修复不完整/未覆盖） | 3 | ACT-R01 ~ ACT-R03，均涉及安全审计 F-3 / F-8 的修复遗漏 |
|
||||
| B 类：新发现问题 | 8 | ACT-R04 ~ ACT-R11，数据校验、CLI 校验、部署约定、文档缺口 |
|
||||
| **合计** | **11** | 无 P0/P1，全部为 P3 级；安全评级维持 **B+**；**全部已修复 ✅** |
|
||||
|
||||
模块仍为离线 CLI 工具（不监听端口、不联网、输入均来自本机可信路径），无远程可利用漏洞。本轮问题集中在**输入健壮性、修复覆盖完整性、部署约定文档化**。
|
||||
|
||||
### 8.3 A 类：上轮修复项复审（3 项）
|
||||
|
||||
#### [ACT-R01] checkpoint config 多余键仍抛裸 TypeError（F-8 声称已修复，实际未修复）🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `predict.py`、`export_onnx.py`（`DataConfig(**cfg_dict["data"])` / `ModelConfig(**cfg_dict["model"])`）
|
||||
- **上轮声明**: 安全审计 F-8 描述为"多余键触发裸 TypeError"，修复项为新增 `validate_config_values()`，结论"已修复 ✅"
|
||||
- **实测**: 向 checkpoint config 的 `data` 段注入未知键 `evil_extra_key`，加载时抛 `TypeError: DataConfig.__init__() got an unexpected keyword argument 'evil_extra_key'`。`validate_config_values` 在 `DataConfig(**...)` **之后**才被调用，未知键在构造阶段即崩溃，校验函数根本无法执行。
|
||||
- **影响**: 畸形/损坏 checkpoint（或未来 config 结构演进产生的多余键）报错信息不含文件名与上下文；上轮"已修复"结论与实现不符。
|
||||
- **修复**: 新增 `parse_checkpoint_config()`，构造前校验键集合，未知键/缺段/非 dict 抛带键名的 ValueError。
|
||||
|
||||
#### [ACT-R02] validate_config_values 遗漏 fused_hidden_dim，恶意 checkpoint 可致 OOM（F-8 修复不完整）🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `config.py`
|
||||
- **实测**: 设 `model.fused_hidden_dim = 10**9`，`validate_config_values` **通过**；构造 `ActionPredictor` 时 `FusionHead` 分配内存抛 `RuntimeError ... DefaultCPUAllocator: can't allocate`。另有 `dropout`（可 <0 或 >1）也未被校验。
|
||||
- **影响**: 与 F-8 防护目标（恶意/损坏 checkpoint 防内存耗尽）直接冲突——攻击面同类、防护漏了 `fused_hidden_dim`。
|
||||
- **修复**: 校验表补充 `fused_hidden_dim ∈ [1, 4096]`、`dropout ∈ [0, 1]`，并扩展 train 字段校验（lr/val_split/batch_size/epochs/patience/num_workers/max_grad_norm/weight_decay）。
|
||||
|
||||
#### [ACT-R03] pretrain_mt.py 未传 pretrained=False，mobilenet 离线预训练仍失败（F-3 修复不完整）🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `pretrain_mt.py`
|
||||
- **实测**: `model = ActionPredictor(cfg.model, cfg.data)` **未传 `pretrained=False`**。`--backbone mobilenet` 时构造即尝试下载 ImageNet 权重（~14MB）：① 离线/受限网络环境下预训练直接 `RuntimeError`；② 零图像预训练会用 `load_state_dict` 全程覆盖这些权重，下载纯属浪费。
|
||||
- **影响**: F-3 的修复只覆盖了"推理/导出"两个入口，遗漏了第三个模型构造入口（预训练）。当前默认 backbone=simple 未触发，属潜伏问题。
|
||||
- **修复**: 构造模型传 `pretrained=False`。
|
||||
|
||||
### 8.4 B 类：新发现问题（8 项）
|
||||
|
||||
#### [ACT-R04] ActionDataset 不校验 actions.npy 列数与 action_dim 一致性 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `dataset.py`
|
||||
- **实测**: 构造 6 列 `actions.npy` 配 `action_dim=4`，数据集**加载成功**，样本 `action_history` 形状 (3, 6)；直到训练 `DataLoader` collate 阶段才因形状不一致报错，错误远离数据源头且信息不友好。
|
||||
- **修复**: 加载后校验 `actions.ndim == 2 and actions.shape[1] == action_dim`，否则抛带 episode 路径的 ValueError。
|
||||
|
||||
#### [ACT-R05] 图像/动作帧数不匹配时静默截断，无任何警告 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `dataset.py`（`n = min(n_frames, n_images)`）
|
||||
- **实测**: 12 帧动作 + 8 张图像，数据集静默按 8 帧使用，**尾部 4 帧动作被丢弃且无提示**。
|
||||
- **修复**: `n_frames != n_images` 时打印 WARNING（含 episode 路径与数量）。
|
||||
|
||||
#### [ACT-R06] 训练 CLI 参数无值域校验，越界值静默接受 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `config.py`（`apply_overrides`）、`train.py`
|
||||
- **实测**: `--history-len 0`、`--val-split 2.0`、`--lr -1` **全部静默接受**并进入训练：`history_len=0` 产生空 history 样本；`val_split=2.0` 验证集吞掉全部 episode；`lr=-1` loss 发散。
|
||||
- **对比**: `validate_config_values()` 只被推理/导出路径调用，**训练 CLI 完全不校验**——同一份参数两套待遇。
|
||||
- **修复**: train.py / pretrain_mt.py 在 `apply_overrides` 后调用 `validate_config_values(cfg.data, cfg.model, cfg.train)`。
|
||||
|
||||
#### [ACT-R07] predict.py 加载 action_stats.npy 不校验形状 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `predict.py`
|
||||
- **实测**: `action_stats.npy` 形状错误（如 (2, 7) 配 action_dim=4）时加载成功，推理阶段抛难以定位的广播错误 `ValueError: operands could not be broadcast together with shapes (3,4) (7,)`。
|
||||
- **修复**: 加载后校验 `stats.shape == (2, action_dim)`，并强制转 float32（兼容 float64 stats）。
|
||||
|
||||
#### [ACT-R08] 推理端硬编码 stats 路径约定，config 中 action_stats_path 为死字段 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `predict.py` vs `train.py`、`config.py`
|
||||
- **描述**: 训练时 `action_stats_path` 被写入 checkpoint config；但推理端**从不读取该字段**，硬编码从 checkpoint 同目录找 `action_stats.npy`。部署约定（stats 必须与 checkpoint 同目录、同名）未在任何文档说明。
|
||||
- **修复**: 推理端优先解析 config 记录的 `action_stats_path`（相对路径按 checkpoint 目录解析），旧 checkpoint 回退同目录约定；README 新增「部署注意」章节。
|
||||
|
||||
#### [ACT-R09] README 未说明 ONNX 导出硬依赖 onnxscript，按文档执行直接失败 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `README.md`、`export_onnx.py`、`_set.sh`
|
||||
- **实测**: 当前 venv（torch 2.13）执行 `bash _run.sh export` 抛 `ModuleNotFoundError: No module named 'onnxscript'`——torch ≥2.12 的 `torch.onnx.export` 走 dynamo 路径，onnxscript 为硬依赖。
|
||||
- **修复**: `_set.sh` 增加 onnxscript/onnxruntime 安装；README 补充依赖说明。
|
||||
|
||||
#### [ACT-R10] 预训练-微调归一化语义不匹配（方法论）🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `pretrain_mt.py`（全程无归一化）vs `train.py`（`--normalize-actions` 支持）
|
||||
- **描述**: 预训练始终在**非归一化**动作空间训练；若微调时开启 `--normalize-actions`，`--init-from` 加载的预训练权重将在**归一化空间**继续训练，输入分布突变。
|
||||
- **修复**: README 微调示例注明「预训练权重搭配非归一化微调使用」。
|
||||
|
||||
#### [ACT-R11] MTActionDataset 命名误导 + parse_mt_directory 重复解析 3 次 🟡 P3 ✅ 已修复
|
||||
|
||||
- **位置**: `pretrain_mt.py`（`n_rows = matrix.shape[1]` 实为**列数**）
|
||||
- **描述**: ① 变量名 `n_rows` 取自 `shape[1]`（列数），与直觉相反，易在后续维护中引入行列颠倒 bug；② `main()` 先 `parse_mt_directory` 一次做文件级划分，train_ds / val_ds 构造时又各解析一次全部文件，**同一目录共解析 3 次**。
|
||||
- **修复**: 改名 `n_cols`；`MTActionDataset` 新增 `files` 参数复用预解析结果，解析次数 3 次 → 1 次。
|
||||
|
||||
### 8.5 测试覆盖缺口（复审确认）
|
||||
|
||||
复审时 `test_actpred.py` 39 用例，以下路径无覆盖：checkpoint config 异常结构（多余键/缺段）、`validate_config_values` 的 model/train 字段、`ActionDataset` 归一化正常路径、`actions.npy` 列数不匹配/帧数不匹配、`split_episodes` 边界、`train.py`/`pretrain_mt.py` 的 `main()` 端到端、stats 形状错误。
|
||||
|
||||
### 8.6 修复记录（全部 ✅）
|
||||
|
||||
> 修复时间：2026-08-08（复审当日）
|
||||
> 全部 11 项问题已修复并验证，单元测试由 39 用例扩展至 **51 用例（50 通过 + 1 环境性跳过）**。
|
||||
|
||||
| 问题 | 修复内容 | 文件 |
|
||||
|------|----------|------|
|
||||
| ACT-R01 | 新增 `parse_checkpoint_config()`：校验 config 键集合，未知键/缺段/非 dict 抛带键名的 ValueError | config.py、predict.py、export_onnx.py |
|
||||
| ACT-R02 | 校验表补充 `fused_hidden_dim ∈ [1,4096]`、`dropout ∈ [0,1]`，并扩展 train 字段校验 | config.py |
|
||||
| ACT-R03 | `pretrain_mt.py` 构造模型传 `pretrained=False`，消除零图像预训练的无谓下载与离线失败 | pretrain_mt.py |
|
||||
| ACT-R04 | `actions.npy` 加载后校验 `ndim==2 且 shape[1]==action_dim`，否则抛带 episode 路径的 ValueError | dataset.py |
|
||||
| ACT-R05 | 帧数不匹配打印 WARNING（含 episode 路径与数量） | dataset.py |
|
||||
| ACT-R06 | train.py / pretrain_mt.py 在 `apply_overrides` 后调用 `validate_config_values(cfg.data, cfg.model, cfg.train)` | train.py、pretrain_mt.py |
|
||||
| ACT-R07 | stats 加载后校验形状 `(2, action_dim)`，并强制转 float32（兼容 float64 stats） | predict.py |
|
||||
| ACT-R08 | 推理端优先解析 checkpoint config 记录的 `action_stats_path`（相对路径按 checkpoint 目录解析），旧 checkpoint 回退同目录约定；README 新增「部署注意」 | predict.py、README.md |
|
||||
| ACT-R09 | `_set.sh` 增加 onnxscript/onnxruntime 安装；README 补充依赖说明 | _set.sh、README.md |
|
||||
| ACT-R10 | README 微调示例注明「预训练权重搭配非归一化微调使用」 | README.md |
|
||||
| ACT-R11 | `n_rows` 改名 `n_cols`；`MTActionDataset` 新增 `files` 参数复用预解析结果，解析次数 3 次 → 1 次 | pretrain_mt.py |
|
||||
|
||||
### 8.7 验证记录
|
||||
|
||||
- 全部 51 项单元测试通过（`Ran 51 tests ... OK (skipped=1)`）
|
||||
- 复现脚本确认：多余键 → ValueError（含键名）；fused_hidden_dim=10⁹ → 加载期拦截（不再 OOM）；stats 形状错误 → ValueError（含路径与期望形状）；`--lr -1` → CLI 拦截
|
||||
- 预训练端到端（3 个 .mt、3 epochs）正常，空验证集保底保存路径正确
|
||||
- 推理端 stats 双路径（config 相对路径 / 同目录回退）均验证通过
|
||||
- 上轮声称的亮点机制（weights_only、文件级划分、stats 隔离、保底保存等）经复查**均属实**
|
||||
|
||||
---
|
||||
|
||||
*本文档为两轮审计合并版：第 1 轮（§二~§七）+ 第 2 轮复审（§八），全部 27 项发现均已修复，修复项均经回归验证。*
|
||||
@@ -0,0 +1,131 @@
|
||||
# actpred 安全审计报告
|
||||
|
||||
**审计日期**: 2026-08-08
|
||||
**审计范围**: `actpred/`（config.py / model.py / dataset.py / train.py / pretrain_mt.py / mt_parser.py / predict.py / export_onnx.py / test_actpred.py / _run.sh / _set.sh / requirements.txt，共 ~1100 行）
|
||||
**审计轮次**: 第 1 轮
|
||||
**安全评级**: B+
|
||||
|
||||
---
|
||||
|
||||
## 1. 审计概述
|
||||
|
||||
actpred 是动作预测模块（图像 + 动作历史 → 下一动作回归，PyTorch），为**离线 CLI 工具**：不监听端口、不接入 pymsg/HTTP，无对外网络暴露面。攻击面收敛为三类输入：
|
||||
|
||||
1. **模型 checkpoint**（`torch.load`，pickle 反序列化风险点）
|
||||
2. **训练数据**（`actions.npy` / 图像文件 / `.mt` 脚本）
|
||||
3. **依赖供应链**（pip 安装、MobileNetV2 预训练权重下载）
|
||||
|
||||
本轮发现 **10 项问题（0 P0 / 0 P1 / 3 P2 / 7 P3）**，全部已修复；另确认 8 项已有良好机制。由于模块不联网、输入均来自本机可信路径，未发现可远程利用的漏洞，问题集中在**功能正确性与方法论**层面。
|
||||
|
||||
---
|
||||
|
||||
## 2. 发现问题
|
||||
|
||||
### F-1: 微调闭环缺失 —— pretrain 完成后打印的命令无法执行 ✅ 已修复
|
||||
|
||||
- **等级**: P2（功能性）
|
||||
- **位置**: `pretrain_mt.py` L258-259、`train.py`（全文）、`README.md` L46-47
|
||||
- **问题**: 预训练结束时提示 `Fine-tune: bash _run.sh train --data-dir ./data --model <path>`，但 `train.py` 的 argparser（`config.build_argparser`）**没有 `--model` 参数**，也没有任何加载预训练权重的逻辑——照提示执行会直接 argparse 报错退出。README 也承认"需修改 train.py 加载逻辑，或手动加载"。预训练产物（`pretrain_best_inference.pt`）目前只能用于直接推理，无法进入微调流程，预训练的核心价值（初始化 action_encoder / fusion_head）落不了地。
|
||||
- **修复**: `config.build_argparser` 新增 `--init-from` 参数；`train.py` 模型构造后以 `torch.load(weights_only=True)` + `load_state_dict(strict=False)` 加载预训练权重（打印 missing/unexpected keys）；`pretrain_mt.py` 完成提示 `--model` → `--init-from`；README 微调示例同步更新。验证：现有 pretrain checkpoint 加载后 missing=0、unexpected=0，完全兼容。
|
||||
|
||||
### F-2: 预训练 train/val 按样本级随机划分，滑动窗口重叠导致数据泄漏 ✅ 已修复
|
||||
|
||||
- **等级**: P2（方法论）
|
||||
- **位置**: `pretrain_mt.py` L153-162
|
||||
- **问题**: `MTActionDataset` 用滑动窗口采样（步长 1），相邻样本共享 `history_len-1 = 4` 行数据；随后按**样本索引**随机划分 train/val。验证集中几乎每个样本都与某个训练样本高度重叠 → **val_loss 系统性偏乐观**，早停与模型选择基于失真信号（上次预训练的 best val_loss=190.58 可信度打折）。对照 `train.py` 是按 **episode 级**划分（正确做法），预训练未沿用。
|
||||
- **修复**: `MTActionDataset` 新增 `file_indices` 参数，只加载指定索引的 .mt 文件；`main()` 改为先 `parse_mt_directory` 获取文件列表，按文件级随机划分（train/val 文件不重叠），再分别构造两个 dataset——每个文件的全部滑动窗口只进一个集合。移除不再使用的 `Subset` 导入。验证：13 个有效 .mt 文件 → 12 train / 1 val（20+2 样本），无跨集窗口重叠。
|
||||
|
||||
### F-3: 推理加载 mobilenet checkpoint 会触发 ImageNet 权重网络下载（离线失败 + 无谓供应链面） ✅ 已修复
|
||||
|
||||
- **等级**: P2
|
||||
- **位置**: `model.py` L29-45（`MobileNetEncoder.__init__`）、`predict.py` L24 / `export_onnx.py` L18
|
||||
- **问题**: `MobileNetEncoder` 构造时无条件下载 `MobileNet_V2_Weights.DEFAULT`（ImageNet 预训练权重）。推理/导出路径（`ActionPredictorInference`、`export_onnx`）构造模型后紧接着 `load_state_dict` **完全覆盖**这些权重——下载纯属浪费，且引入两个问题：① 离线/受限网络环境下，backbone=mobilenet 的 checkpoint **无法推理**（构造即抛 RuntimeError）；② 推理路径引入不必要的外部权重下载（供应链面扩大）。当前默认 backbone=simple 未触发，属潜伏问题。
|
||||
- **修复**: `MobileNetEncoder` 与 `ActionPredictor` 新增 `pretrained: bool = True` 参数；`pretrained=False` 时传 `weights=None`（不下载 ImageNet 权重，仅初始化架构）。`predict.py` 和 `export_onnx.py` 构造模型时传 `pretrained=False`，训练路径保持默认 `True` 不变。
|
||||
|
||||
### F-4: mt_parser 行列不一致静默容错，解析失败静默跳过整个文件 ✅ 已修复
|
||||
|
||||
- **等级**: P3（数据完整性）
|
||||
- **位置**: `mt_parser.py` L61-67、L85-90
|
||||
- **问题**: ① 数据行比表头**短**时，缺失列静默保持 0.0（错误数据混入训练集无告警）；数据行比表头**长**时 `matrix[i, j]` 直接 IndexError；② `parse_mt_directory` 对 `ValueError` 一律 `pass`（无 @Action 块、数值非法等），文件被**静默丢弃**——训练数据缺失无任何日志，且 IndexError 不在捕获范围会直接崩溃。
|
||||
- **修复**: 行列不一致时抛带行号的 ValueError；数值解析失败抛带行列号的 ValueError；`parse_mt_directory` 捕获 `Exception`（含 IndexError）并打印 `WARNING: 跳过 .mt 文件 <path>: <reason>`。
|
||||
|
||||
### F-5: mt_parser 首行 `-` 占位符填 0.0，与 bot4cj 语义不符 ✅ 已修复
|
||||
|
||||
- **等级**: P3（数据保真度）
|
||||
- **位置**: `mt_parser.py` L64-65
|
||||
- **问题**: bot4cj 中 `-` 表示"沿用上一状态"，首行的上一状态是 `@Servo` 块的初始角度；actpred 解析器对首行 `-` 一律填 0.0。实测 `jump.mt` 首行 `-` 恰好对应 0 值舵机未失真，但这是巧合——初始角非 0 的脚本首行样本会引入错误的(history, target)对。
|
||||
- **修复**: 新增 `_parse_servo_init()` 解析 `@Servo` 块的 `max` 列（舵机初始角度）；`parse_mt_action` 根据列名（如 `s1.A` → 查 `s1`）为每列计算初始值，首行 `-` 从对应舵机初始角填充，无映射的列（如 `spk.F`）填 0.0。验证：`jump.mt` 首行 s2.A=180、s4.A=-180 与 @Servo 定义一致。
|
||||
|
||||
### F-6: parse_action_history 用 assert 做输入校验，`python -O` 下失效 ✅ 已修复
|
||||
|
||||
- **等级**: P3
|
||||
- **位置**: `predict.py` L72-73
|
||||
- **问题**: 维度校验依赖 `assert`，优化模式（`python -O`）下被剥离，畸形输入将以错误形状进入 `np.array`（可能产生 object 数组或错误广播）。
|
||||
- **修复**: 改为显式 `raise ValueError`；同步更新 `test_actpred.py` 中 `test_wrong_dim_raises` 的断言类型（AssertionError → ValueError）。
|
||||
|
||||
### F-7: DataRecorder 不校验 action 维度，episode 目录分配非原子 ✅ 已修复
|
||||
|
||||
- **等级**: P3
|
||||
- **位置**: `dataset.py` L100-120
|
||||
- **问题**: ① 构造时接收 `action_dim` 但 `record()` 从不校验 `action.shape`——录错维度的动作会在 `flush()` 时静默写出畸形 `actions.npy`（或 `np.array` 构造 object 数组），下游训练才暴露；② episode 索引通过 glob 扫描 + max 计算，两个并发 Recorder 会分配到同一目录互相覆盖。
|
||||
- **修复**: `record()` 增加 `action.shape != (action_dim,)` 校验抛 ValueError；目录创建改为 `mkdir(exist_ok=False)` + `FileExistsError` 重试递增 ep_idx，保证原子性。
|
||||
|
||||
### F-8: checkpoint 内嵌 config 未做值域校验 ✅ 已修复
|
||||
|
||||
- **等级**: P3（低风险）
|
||||
- **位置**: `predict.py` L16-18、`export_onnx.py` L14-16
|
||||
- **问题**: `torch.load(weights_only=True)` 已阻断 pickle 代码执行（✅），但 checkpoint 中的 config 字典仍被直接用于建模：恶意/损坏 checkpoint 可设置超大 `image_size`/`history_len`/`cnn_feature_dim` 造成内存耗尽，或多余键触发裸 TypeError。本模块 checkpoint 来源为本机训练产物，实际风险低。
|
||||
- **修复**: `config.py` 新增 `validate_config_values()` 函数，对 image_size/history_len/action_dim/cnn_feature_dim/action_hidden_dim 做合理性上限校验（[1,1024] 或 [1,4096]）；`predict.py` 和 `export_onnx.py` 加载 config 后调用校验。
|
||||
|
||||
### F-9: venv 实际依赖违反 requirements.txt 版本约束 ✅ 已修复
|
||||
|
||||
- **等级**: P3（供应链/一致性）
|
||||
- **位置**: `_set.sh` L8-11、`requirements.txt` L5
|
||||
- **问题**: 实测 venv 中 `Pillow==12.2.0`，违反 `Pillow>=9.0,<12.0` 上限。原因是 `_set.sh` 先从 PyTorch 源装 torch/torchvision（torchvision 无约束地拉入最新 Pillow），再装 requirements 时未强制降级（或安装顺序导致约束未生效）。声明与实际环境脱节，版本上限形同虚设；且无 hash 锁定。
|
||||
- **修复**: `requirements.txt` 中 Pillow 上限放宽至 `<14.0`（torchvision 依赖管理，加注释说明）；`_set.sh` 末尾新增 `pip check` 验证依赖一致性。
|
||||
|
||||
### F-10: `--seed` 对 episode 划分不生效 ✅ 已修复
|
||||
|
||||
- **等级**: P3（可复现性）
|
||||
- **位置**: `train.py` L78-79
|
||||
- **问题**: `split_episodes(cfg.data.data_dir, cfg.train.val_split)` 未传 `cfg.train.seed`，函数内硬编码默认 `seed=42`。用户指定 `--seed 123` 时权重初始化、DataLoader 洗牌均变化，唯独 train/val episode 划分不变，可复现性语义不完整（也可视为"划分稳定"的特性，但与参数预期不符）。
|
||||
- **修复**: `split_episodes()` 调用处显式传入 `seed=cfg.train.seed`。
|
||||
|
||||
---
|
||||
|
||||
## 3. 已确认的良好机制
|
||||
|
||||
| # | 机制 | 位置 |
|
||||
|---|------|------|
|
||||
| 1 | `torch.load(..., weights_only=True)`，阻断 checkpoint pickle 反序列化 RCE | predict.py L15、export_onnx.py L13 |
|
||||
| 2 | `np.load` 默认 `allow_pickle=False`（numpy ≥1.16.3），npy 文件无 pickle 风险 | dataset.py L41、predict.py L45 |
|
||||
| 3 | 归一化统计量只用训练集计算再供验证集复用，避免统计泄漏 | train.py L86-91 |
|
||||
| 4 | 空验证集返回 `inf` 不参与 best 比较，避免假 best；并有保底保存 | train.py L56-57、L158-164 |
|
||||
| 5 | 图像文件自然排序（`1.jpg < 10.jpg`），避免帧序错乱 | dataset.py L12-14 |
|
||||
| 6 | 全局种子固定（random/numpy/torch/cuda + DataLoader generator）、梯度裁剪、早停 | train.py L69-73、L106-110 |
|
||||
| 7 | ONNX 导出后自动做数值一致性验证（如安装 onnxruntime） | export_onnx.py L40-52 |
|
||||
| 8 | 推理/训练双 checkpoint 分离（best.pt 含 optimizer / best_inference.pt 纯权重），格式互相兼容 | train.py、pretrain_mt.py |
|
||||
|
||||
---
|
||||
|
||||
## 4. 测试覆盖缺口 ✅ 已补齐
|
||||
|
||||
首轮审计时 `test_actpred.py` 仅 17 用例，覆盖自然排序、动作历史解析、模型前向（simple + mlp/gru）、DataRecorder、ActionDataset、config 覆盖机制。以下缺口已全部补齐（现 **39 用例，38 通过 + 1 环境性跳过**）：
|
||||
|
||||
- ✅ `mt_parser.py`（新增 7 用例）：正常解析 + `-` 前向填充、首行 `-` 从 @Servo 初始角填充、无 @Action 块 / 仅表头 / 行列不一致 / 数值非法各自抛 ValueError、`parse_mt_directory` 跳过坏文件并打印 WARNING
|
||||
- ✅ `pretrain_mt.py`（新增 4 用例）：MTActionDataset 滑动窗口样本数与 target 值、列数 < action_dim 补零、> 截断、`file_indices` 文件级过滤
|
||||
- ✅ `ActionPredictorInference` 端到端（新增 6 用例）：checkpoint 保存 → 加载 → 推理（形状/有限性/确定性）、加载权重与保存权重逐张量一致、history 形状错误抛 ValueError、归一化缺 action_stats.npy 抛 FileNotFoundError、恒等统计量（mean=0/std=1）归一化路径与非归一化输出一致、checkpoint config 值域非法抛 ValueError
|
||||
- ✅ `export_onnx.py`（新增 1 用例）：端到端导出并校验产物非空；因当前 venv 缺 `onnxscript`（torch 2.13 的 `torch.onnx.export` 走 dynamo 路径的硬依赖），用 `skipUnless` 环境检测跳过——安装 onnxscript 后自动生效
|
||||
- 另补 `validate_config_values`（2 用例）与 DataRecorder 校验/目录原子分配（2 用例）
|
||||
|
||||
> 注：当前环境 `torch.onnx.export` 实测报 `ModuleNotFoundError: No module named 'onnxscript'`，即 F-8 之外的独立环境问题——ONNX 导出功能在该 venv 不可用，如需导出请 `pip install onnxscript onnxruntime`。
|
||||
|
||||
---
|
||||
|
||||
## 5. 结论与修复优先级
|
||||
|
||||
actpred 无网络暴露面，无 P0/P1 级安全漏洞；`weights_only=True` 与 `np.load` 默认行为已覆盖两个主要反序列化风险点。核心问题是**功能正确性**：
|
||||
|
||||
1. **已修复**（P2）: ✅ F-1 微调加载、✅ F-2 预训练 val 泄漏、✅ F-3 mobilenet 推理离线不可用
|
||||
2. **已修复**（P3）: ✅ F-4 解析器静默容错、✅ F-5 首行 @Servo 填充、✅ F-6 assert 校验、✅ F-7 DataRecorder 校验、✅ F-8 config 值域校验、✅ F-9 依赖一致性、✅ F-10 seed 传递
|
||||
3. **全部 10 项问题已修复，无遗留项**
|
||||
Reference in new issue
Block a user