From 45e019952e534815c82c2a07fc452a99d21e438a Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Sat, 8 Aug 2026 17:09:14 +0000 Subject: [PATCH] =?UTF-8?q?actpred:=20action=5Fdim=E6=94=B918=20+=20.mt?= =?UTF-8?q?=E8=84=9A=E6=9C=AC=E9=A2=84=E8=AE=AD=E7=BB=83=E5=85=A5=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - config.py: action_dim默认值 8→18（对齐bot4cj 18列） - mt_parser.py: 轻量级.mt @Action解析器（支持-占位符前向填充） - pretrain_mt.py: 预训练入口（零图像占位，checkpoint兼容现有推理） - _run.sh: 新增pretrain子命令 - README.md: 新增预训练文档 + action_dim说明 - test_actpred.py: 测试断言更新为18 --- actpred/README.md | 28 ++++- actpred/_run.sh | 3 + actpred/config.py | 2 +- actpred/mt_parser.py | 91 ++++++++++++++ actpred/pretrain_mt.py | 268 ++++++++++++++++++++++++++++++++++++++++ actpred/test_actpred.py | 14 +-- 6 files changed, 394 insertions(+), 12 deletions(-) create mode 100644 actpred/mt_parser.py create mode 100644 actpred/pretrain_mt.py diff --git a/actpred/README.md b/actpred/README.md index 4d94aac..7e2d946 100644 --- a/actpred/README.md +++ b/actpred/README.md @@ -21,11 +21,11 @@ cd actpred && bash _set.sh # 训练 -bash _run.sh train --data-dir ./data --epochs 100 --action-dim 8 +bash _run.sh train --data-dir ./data --epochs 100 --action-dim 18 # 推理（--action-history 格式：冒号分隔各步，每步 action_dim 个值，共 history_len 步） bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \ - --action-history "0,0,0,0,0,0,0,0:1,0,0,0,0,0,0,0:0,1,0,0,0,0,0,0:0,0,1,0,0,0,0,0:0,0,0,1,0,0,0,0" + --action-history "45,0,45,0,45,0,45,0,90,90,10,200,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,20,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0:45,0,45,0,45,0,45,0,90,90,0,0,0,0,0,0,0,0" # 导出 ONNX（如安装了 onnxruntime，导出后自动验证一致性） bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx @@ -34,6 +34,26 @@ bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx python test_actpred.py ``` +## 预训练（从 .mt 脚本，无需图像） + +从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵，用零图像占位预训练动作动力学（history → next action）， +保存的 checkpoint 与正常训练格式完全兼容，可直接用于推理或作为微调起点。 + +```bash +# 预训练（解析 ../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 +``` + +| 参数 | 说明 | +|------|------| +| `--mt-dir` | .mt 脚本目录（如 `../bot4cj/scripts`） | +| `--noise-std` | 高斯噪声标准差（数据增强，建议 0.5~2.0） | + +> **注意**：.mt 脚本列数（12）与 `action_dim`（18）不一致时自动补零。预训练保存到 `checkpoints/pretrain_best_inference.pt`。 + ## 数据格式 ``` @@ -55,7 +75,7 @@ data/ ```python from dataset import DataRecorder -recorder = DataRecorder("./data", action_dim=8) +recorder = DataRecorder("./data", action_dim=18) for image, action in robot_loop(): recorder.record(image, action) recorder.flush() @@ -68,7 +88,7 @@ recorder.flush() | 参数 | 默认值 | 说明 | |------|--------|------| | `--image-size` | 64 | 图像缩放尺寸 | -| `--action-dim` | 8 | 动作向量维度 | +| `--action-dim` | 18 | 动作向量维度（对齐 bot4cj 18 列：10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留）| | `--history-len` | 5 | 历史动作窗口长度 | | `--backbone` | simple | CNN 类型 (simple/mobilenet) | | `--history-encoder` | mlp | 历史编码器 (mlp/gru) | diff --git a/actpred/_run.sh b/actpred/_run.sh index 161729a..cd99923 100755 --- a/actpred/_run.sh +++ b/actpred/_run.sh @@ -5,6 +5,9 @@ source ./venv/bin/activate if [ "${1:-}" = "train" ]; then shift exec python3 train.py "$@" +elif [ "${1:-}" = "pretrain" ]; then + shift + exec python3 pretrain_mt.py "$@" elif [ "${1:-}" = "export" ]; then shift exec python3 export_onnx.py "$@" diff --git a/actpred/config.py b/actpred/config.py index fa22dde..0f4bda7 100644 --- a/actpred/config.py +++ b/actpred/config.py @@ -6,7 +6,7 @@ from typing import Optional @dataclass class DataConfig: image_size: int = 64 - action_dim: int = 8 + action_dim: int = 18 history_len: int = 5 data_dir: str = "./data" normalize_actions: bool = False diff --git a/actpred/mt_parser.py b/actpred/mt_parser.py new file mode 100644 index 0000000..3b83f71 --- /dev/null +++ b/actpred/mt_parser.py @@ -0,0 +1,91 @@ +"""轻量级 .mt 脚本 @Action 矩阵解析器 + +从 .mt 文件中提取 @Action 块，解析为 (n_rows, n_cols) 的 float32 矩阵。 +支持 `-` 占位符（沿用上一行同列的值）。 +""" + +import re +from pathlib import Path + +import numpy as np + + +def _natural_key(p: Path): + """自然排序键，确保 s1.mt < s2.mt < s10.mt""" + return [int(t) if t.isdigit() else t for t in re.split(r"(\d+)", p.name)] + + +def parse_mt_action(path) -> np.ndarray: + """解析 .mt 文件中的 @Action 矩阵 + + Args: + path: .mt 文件路径 + + Returns: + (n_data_rows, n_cols) float32 矩阵，表头行已跳过，`-` 已前向填充 + + Raises: + ValueError: 文件中无 @Action 块或数据行不足 + """ + text = Path(path).read_text(encoding="utf-8") + # 先去除整行注释（# 开头），避免注释中的 @Action 干扰正则匹配 + text = re.sub(r"(?m)^[ \t]*#.*$", "", text) + + match = re.search(r"@Action\s+(.*?)\s*@", text, re.DOTALL) + if not match: + raise ValueError(f"未找到 @Action 块: {path}") + + block = match.group(1) + + rows = [] + for line in block.split("\n"): + line = re.sub(r"#.*$", "", line).strip() + if not line: + continue + if line.endswith(";"): + line = line[:-1].strip() + if not line: + continue + cols = re.split(r"\s+", line) + rows.append(cols) + + if len(rows) < 2: + raise ValueError( + f"@Action 至少需要表头 + 1 行数据，实际 {len(rows)} 行: {path}" + ) + + # 第一行是表头（列名如 s1.A, s2.A, ...），跳过 + data_rows = rows[1:] + n_cols = len(rows[0]) + + matrix = np.zeros((len(data_rows), n_cols), dtype=np.float32) + for i, row in enumerate(data_rows): + for j, val in enumerate(row): + if val == "-": + matrix[i, j] = matrix[i - 1, j] if i > 0 else 0.0 + else: + matrix[i, j] = float(val) + + return matrix + + +def parse_mt_directory(mt_dir, pattern="*.mt"): + """解析目录下所有 .mt 文件的 @Action 矩阵 + + Args: + mt_dir: .mt 文件目录 + pattern: 文件匹配模式 + + Returns: + list of (filepath, np.ndarray) 元组，按文件名自然排序 + """ + dir_path = Path(mt_dir) + files = sorted(dir_path.glob(pattern), key=_natural_key) + results = [] + for f in files: + try: + matrix = parse_mt_action(f) + results.append((str(f), matrix)) + except ValueError: + pass + return results diff --git a/actpred/pretrain_mt.py b/actpred/pretrain_mt.py new file mode 100644 index 0000000..7cd4d9f --- /dev/null +++ b/actpred/pretrain_mt.py @@ -0,0 +1,268 @@ +"""actpred 预训练入口：从 .mt 脚本学习动作动力学（无需图像） + +用法: + bash _run.sh pretrain --mt-dir ../bot4cj/scripts --epochs 200 + +流程: + 1. 解析 .mt 文件中的 @Action 矩阵 → 动作序列 + 2. 滑动窗口采样: action_history[H] → target action（下一行） + 3. 用零图像占位训练完整 ActionPredictor（保持 checkpoint 兼容性） + 4. 保存 checkpoint，格式与 train.py 一致，可直接用于推理或微调 + +数据增强: + --noise-std 控制高斯噪声标准差。因 .mt 脚本数据量有限， + 适度噪声（如 0.5~2.0）可有效防止过拟合。 +""" + +import json +import os +import random + +import numpy as np +import torch +from torch.utils.data import Dataset, DataLoader, Subset +from tqdm import tqdm + +from config import Config, build_argparser, apply_overrides +from model import ActionPredictor +from mt_parser import parse_mt_directory + + +class MTActionDataset(Dataset): + """从 .mt @Action 矩阵构建的滑动窗口数据集（无图像） + + 每个样本: (action_history[history_len, action_dim], target[action_dim]) + """ + + def __init__(self, mt_dir: str, history_len: int, action_dim: int): + 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: + matrix = np.pad( + matrix, + ((0, 0), (0, action_dim - n_rows)), + mode="constant", + ) + elif n_rows > action_dim: + matrix = matrix[:, :action_dim] + + for i in range(len(matrix) - history_len): + history = matrix[i : i + history_len] + target = matrix[i + history_len] + 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})") + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + history, target = self.samples[idx] + return { + "action_history": torch.from_numpy(history.copy()), + "target": torch.from_numpy(target.copy()), + } + + +def pretrain_one_epoch(model, loader, optimizer, device, image_size, + max_grad_norm=None, noise_std=0.0): + model.train() + total_loss = 0.0 + count = 0 + zero_image = None + for batch in tqdm(loader, desc=" pretrain", leave=False): + history = batch["action_history"].to(device) + target = batch["target"].to(device) + if noise_std > 0: + history = history + torch.randn_like(history) * noise_std + target = target + torch.randn_like(target) * noise_std + if zero_image is None or zero_image.size(0) != history.size(0): + zero_image = torch.zeros( + history.size(0), 3, image_size, image_size, device=device + ) + optimizer.zero_grad() + pred = model(zero_image, history) + loss = torch.nn.functional.mse_loss(pred, target) + loss.backward() + if max_grad_norm: + torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) + optimizer.step() + total_loss += loss.item() * history.size(0) + count += history.size(0) + return total_loss / max(count, 1) + + +@torch.no_grad() +def validate(model, loader, device, image_size): + model.eval() + total_loss = 0.0 + count = 0 + zero_image = None + for batch in tqdm(loader, desc=" val ", leave=False): + history = batch["action_history"].to(device) + target = batch["target"].to(device) + if zero_image is None or zero_image.size(0) != history.size(0): + zero_image = torch.zeros( + history.size(0), 3, image_size, image_size, device=device + ) + pred = model(zero_image, history) + loss = torch.nn.functional.mse_loss(pred, target) + total_loss += loss.item() * history.size(0) + count += history.size(0) + if count == 0: + return float("inf") + return total_loss / count + + +def main(): + parser = build_argparser("actpred pre-training from .mt scripts") + parser.add_argument("--mt-dir", type=str, required=True, + help=".mt 脚本目录（如 ../bot4cj/scripts）") + parser.add_argument("--noise-std", type=float, default=0.0, + help="高斯噪声标准差（数据增强，建议 0.5~2.0）") + args = parser.parse_args() + + cfg = apply_overrides(Config(), args) + cfg.data.data_dir = args.mt_dir + + os.makedirs(cfg.train.checkpoint_dir, exist_ok=True) + + random.seed(cfg.train.seed) + np.random.seed(cfg.train.seed) + torch.manual_seed(cfg.train.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(cfg.train.seed) + + 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: + 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)) + 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:] + + train_ds = Subset(full_ds, train_indices) + val_ds = Subset(full_ds, val_indices) + + print(f"Samples: {len(train_ds)} train, {len(val_ds)} val") + if len(train_ds) == 0: + raise ValueError("训练集为空（样本数不足）") + if len(val_ds) == 0: + print("WARNING: 验证集为空，模型选择禁用") + + loader_gen = torch.Generator() + loader_gen.manual_seed(cfg.train.seed) + train_loader = DataLoader( + train_ds, batch_size=cfg.train.batch_size, + shuffle=True, num_workers=cfg.train.num_workers, + generator=loader_gen, + ) + val_loader = DataLoader( + val_ds, batch_size=cfg.train.batch_size, + shuffle=False, num_workers=cfg.train.num_workers, + ) + + model = ActionPredictor(cfg.model, cfg.data).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}, " + f"history_len={cfg.data.history_len}, " + f"backbone={cfg.model.cnn_backbone}, " + f"encoder={cfg.model.history_encoder}") + print(f"Pre-training: noise_std={args.noise_std}") + + 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 + ) + + best_val_loss = float("inf") + patience_counter = 0 + best_path = os.path.join(cfg.train.checkpoint_dir, "pretrain_best.pt") + inference_path = os.path.join( + cfg.train.checkpoint_dir, "pretrain_best_inference.pt" + ) + + for epoch in range(1, cfg.train.epochs + 1): + print(f"\nEpoch {epoch}/{cfg.train.epochs} " + f"lr={optimizer.param_groups[0]['lr']:.6f}") + train_loss = pretrain_one_epoch( + model, train_loader, optimizer, device, + cfg.data.image_size, cfg.train.max_grad_norm, args.noise_std, + ) + val_loss = validate(model, val_loader, device, cfg.data.image_size) + scheduler.step() + print(f" train_loss={train_loss:.6f} val_loss={val_loss:.6f}") + + if val_loss < best_val_loss: + best_val_loss = val_loss + patience_counter = 0 + torch.save( + { + "epoch": epoch, + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "val_loss": val_loss, + "config": cfg.to_dict(), + }, + best_path, + ) + torch.save( + { + "model_state_dict": model.state_dict(), + "config": cfg.to_dict(), + }, + inference_path, + ) + print(f" -> saved best (val_loss={val_loss:.6f})") + else: + patience_counter += 1 + if patience_counter >= cfg.train.patience: + print(f"\nEarly stopping at epoch {epoch} " + f"(patience={cfg.train.patience})") + break + + print(f"\nPre-training done. Best val_loss={best_val_loss:.6f}") + + if best_val_loss == float("inf"): + torch.save( + { + "model_state_dict": model.state_dict(), + "config": cfg.to_dict(), + }, + inference_path, + ) + print("WARNING: saved final weights (no valid validation loss)") + else: + print(f"Best pre-trained model: {inference_path}") + print("Fine-tune: bash _run.sh train --data-dir ./data " + f"--model {inference_path}") + + if args.save_config: + with open(args.save_config, "w") as f: + json.dump(cfg.to_dict(), f, indent=2, ensure_ascii=False) + print(f"Config saved to: {args.save_config}") + + +if __name__ == "__main__": + main() diff --git a/actpred/test_actpred.py b/actpred/test_actpred.py index 07a5b92..c5903a8 100644 --- a/actpred/test_actpred.py +++ b/actpred/test_actpred.py @@ -60,28 +60,28 @@ class TestActionPredictor(unittest.TestCase): """测试模型前向传播""" def test_forward_simple_mlp(self): - data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5) + data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5) model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp") model = ActionPredictor(model_cfg, data_cfg) model.eval() image = torch.randn(2, 3, 64, 64) - history = torch.randn(2, 5, 8) + history = torch.randn(2, 5, 18) with torch.no_grad(): out = model(image, history) - self.assertEqual(out.shape, (2, 8)) + self.assertEqual(out.shape, (2, 18)) def test_forward_simple_gru(self): - data_cfg = DataConfig(image_size=64, action_dim=8, history_len=5) + data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5) model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru") model = ActionPredictor(model_cfg, data_cfg) model.eval() image = torch.randn(2, 3, 64, 64) - history = torch.randn(2, 5, 8) + history = torch.randn(2, 5, 18) with torch.no_grad(): out = model(image, history) - self.assertEqual(out.shape, (2, 8)) + self.assertEqual(out.shape, (2, 18)) def test_forward_dynamic_batch(self): data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3) @@ -186,7 +186,7 @@ class TestConfig(unittest.TestCase): def test_default_config(self): cfg = Config() self.assertEqual(cfg.data.image_size, 64) - self.assertEqual(cfg.data.action_dim, 8) + self.assertEqual(cfg.data.action_dim, 18) self.assertEqual(cfg.data.history_len, 5) self.assertEqual(cfg.model.cnn_backbone, "simple") self.assertEqual(cfg.model.history_encoder, "mlp")