feat(actpred): 多帧注意力架构改造（Transformer 替换 concat 融合，支持多帧图像+历史动作序列）

actpred/README.md

@@ -1,19 +1,26 @@
 # actpred — 动作预测模块
 
-输入摄像头图像 + 历史动作序列，预测下一个连续动作（回归）。
+输入多帧摄像头图像 + 历史动作序列，预测下一时间步的连续动作（回归）。
 
-纯神经网络实现（PyTorch），支持轻量 CNN 或 MobileNetV2 图像编码，MLP/GRU 动作历史编码。总参数量 ~140K（默认）至 ~2.3M。
+纯神经网络实现（PyTorch）：轻量 CNN 或 MobileNetV2 逐帧编码，图像/动作 token 进轻量
+Transformer（pre-LN）自注意力融合，取 [PRED] token 回归下一步动作。
+默认配置总参数量 ~450K。
 
 ## 架构
 
 ```
-图像 (3×64×64)  ──►  CNN 编码器  ──►  img_feat (128)
-                                          │
-                                          ├─ concat ──►  融合头  ──►  预测动作 (action_dim)
-                                          │
-动作历史 (H×D)  ──►  MLP/GRU    ──►  act_feat (64)
+F 帧图像 (F×3×64×64) ─► 共享 CNN 编码器 ─► F 个 token ─► proj → d_model ┐
+                                                                          ├─► [img×F, act×H, PRED] + 位置编码
+动作历史 (H×D)      ─► Linear proj → d_model ─► H 个 token ────────┤        │
+可学习 [PRED] token ────────────────────────────────────────┘        ▼
+                                       TransformerEncoder (pre-LN, 2 层, 4 头, d_model=128)
+                                                                                 │
+                                             取 PRED 位置输出 ─► MLP 头 ─► 预测动作 (action_dim)
 ```
 
+token 序列长度固定为 F+H+1（默认 3+5+1=9），除 batch 外无动态 shape，对 NPU（如树莓派
+AI 扩展板）部署友好。
+
 ## 快速开始
 
 ```bash
@@ -23,8 +30,10 @@ cd actpred && bash _set.sh
 # 训练
 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 \
+# 推理（--images 按时间顺序旧→新传入 frame_len 个图像路径，冷启动时可重复传同一帧；
+#      --action-history 格式：冒号分隔各步，每步 action_dim 个值，共 history_len 步）
+bash _run.sh --model checkpoints/best_inference.pt \
+    --images frame_t-2.jpg frame_t-1.jpg frame_t.jpg \
     --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，导出后自动验证一致性）
@@ -39,7 +48,7 @@ python test_actpred.py
 
 ## 预训练（从 .mt 脚本，无需图像）
 
-从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵，用零图像占位预训练动作动力学（history → next action），
+从 bot4cj 的 `.mt` 动作脚本中提取 `@Action` 矩阵，用全零多帧图像占位预训练动作动力学（history → next action），
 保存的 checkpoint 与正常训练格式完全兼容，可直接用于推理或作为微调起点。
 
 ```bash
@@ -74,7 +83,8 @@ data/
         ...
 ```
 
-每个 timestep 的样本：当前图像 + 前 `history_len` 步动作 → 当前步动作。
+每个 timestep 的样本：帧 `[t-frame_len+1, t]`（旧→新）+ 前 `history_len` 步动作 → 当前步动作。
+采样起点为 `max(history_len, frame_len-1)`。
 
 ### 数据采集
 
@@ -96,8 +106,11 @@ recorder.flush()
 | `--image-size` | 64 | 图像缩放尺寸 |
 | `--action-dim` | 18 | 动作向量维度（对齐 bot4cj 18 列：10 舵机 + 2 云台 + 扬声器频率 + 延时 + 4 预留）|
 | `--history-len` | 5 | 历史动作窗口长度 |
+| `--frame-len` | 3 | 输入图像帧数（含当前帧，旧帧在前）|
 | `--backbone` | simple | CNN 类型 (simple/mobilenet) |
-| `--history-encoder` | mlp | 历史编码器 (mlp/gru) |
+| `--attn-dim` | 128 | 注意力 token 维度 d_model（须能被 attn-heads 整除）|
+| `--attn-heads` | 4 | 注意力头数 |
+| `--attn-layers` | 2 | TransformerEncoder 层数 |
 | `--epochs` | 100 | 训练轮数 |
 | `--lr` | 0.001 | 学习率 |
 | `--batch-size` | 32 | 批次大小 |
@@ -124,7 +137,8 @@ recorder.flush()
 from predict import ActionPredictorInference
 
 predictor = ActionPredictorInference("checkpoints/best_inference.pt")
-print(predictor.action_dim, predictor.history_len)  # 公共属性
-action = predictor.predict(image_rgb, action_history)  # np.ndarray (action_dim,)
+print(predictor.action_dim, predictor.history_len, predictor.frame_len)  # 公共属性
+# images: 长度为 frame_len 的 RGB 图像列表（旧→新，末帧为当前帧；冷启动可重复同一帧）
+action = predictor.predict(images, action_history)  # np.ndarray (action_dim,)
 # action_history 形状须为 (history_len, action_dim)
 ```

actpred/config.py

@@ -8,6 +8,7 @@ class DataConfig:
     image_size: int = 64
     action_dim: int = 18
     history_len: int = 5
+    frame_len: int = 3
     data_dir: str = "./data"
     normalize_actions: bool = False
     action_stats_path: Optional[str] = None
@@ -17,10 +18,11 @@ class DataConfig:
 class ModelConfig:
     cnn_backbone: str = "simple"
     cnn_feature_dim: int = 128
-    action_hidden_dim: int = 64
     fused_hidden_dim: int = 128
     dropout: float = 0.1
-    history_encoder: str = "mlp"
+    attn_dim: int = 128
+    attn_heads: int = 4
+    attn_layers: int = 2
 
 
 @dataclass
@@ -60,15 +62,17 @@ def build_argparser(description: str = "") -> argparse.ArgumentParser:
     parser.add_argument("--image-size", type=int, default=None)
     parser.add_argument("--action-dim", type=int, default=None)
     parser.add_argument("--history-len", type=int, default=None)
+    parser.add_argument("--frame-len", type=int, default=None)
     parser.add_argument("--data-dir", type=str, default=None)
     parser.add_argument("--normalize-actions", action=argparse.BooleanOptionalAction, default=None)
     parser.add_argument("--action-stats-path", type=str, default=None)
     parser.add_argument("--backbone", type=str, default=None, choices=["simple", "mobilenet"])
     parser.add_argument("--cnn-feature-dim", type=int, default=None)
-    parser.add_argument("--action-hidden-dim", type=int, default=None)
     parser.add_argument("--fused-hidden-dim", type=int, default=None)
     parser.add_argument("--dropout", type=float, default=None)
-    parser.add_argument("--history-encoder", type=str, default=None, choices=["mlp", "gru"])
+    parser.add_argument("--attn-dim", type=int, default=None)
+    parser.add_argument("--attn-heads", type=int, default=None)
+    parser.add_argument("--attn-layers", type=int, default=None)
     parser.add_argument("--batch-size", type=int, default=None)
     parser.add_argument("--epochs", type=int, default=None)
     parser.add_argument("--lr", type=float, default=None)
@@ -95,11 +99,14 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
     checks = [
         ("image_size", data_cfg.image_size, 1, 1024, int),
         ("history_len", data_cfg.history_len, 1, 1024, int),
+        ("frame_len", data_cfg.frame_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),
+        ("attn_dim", model_cfg.attn_dim, 1, 4096, int),
+        ("attn_heads", model_cfg.attn_heads, 1, 64, int),
+        ("attn_layers", model_cfg.attn_layers, 1, 64, int),
     ]
     for name, val, lo, hi, typ in checks:
         if typ is int:
@@ -110,6 +117,11 @@ def validate_config_values(data_cfg: DataConfig, model_cfg: ModelConfig,
             raise ValueError(
                 f"checkpoint config {name}={val} 超出合理范围 [{lo}, {hi}]")
 
+    if model_cfg.attn_dim % model_cfg.attn_heads != 0:
+        raise ValueError(
+            f"attn_dim={model_cfg.attn_dim} 必须能被 attn_heads="
+            f"{model_cfg.attn_heads} 整除")
+
     if train_cfg is not None:
         t_checks = [
             ("batch_size", train_cfg.batch_size, 1, 65536, int),
@@ -162,15 +174,17 @@ def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
         "image_size": ("data", "image_size"),
         "action_dim": ("data", "action_dim"),
         "history_len": ("data", "history_len"),
+        "frame_len": ("data", "frame_len"),
         "data_dir": ("data", "data_dir"),
         "normalize_actions": ("data", "normalize_actions"),
         "action_stats_path": ("data", "action_stats_path"),
         "backbone": ("model", "cnn_backbone"),
         "cnn_feature_dim": ("model", "cnn_feature_dim"),
-        "action_hidden_dim": ("model", "action_hidden_dim"),
         "fused_hidden_dim": ("model", "fused_hidden_dim"),
         "dropout": ("model", "dropout"),
-        "history_encoder": ("model", "history_encoder"),
+        "attn_dim": ("model", "attn_dim"),
+        "attn_heads": ("model", "attn_heads"),
+        "attn_layers": ("model", "attn_layers"),
         "batch_size": ("train", "batch_size"),
         "epochs": ("train", "epochs"),
         "lr": ("train", "lr"),

actpred/dataset.py

@@ -17,10 +17,12 @@ def _natural_key(p):
 class ActionDataset(Dataset):
     def __init__(self, data_dir: str, history_len: int, action_dim: int,
                  image_size: int = 64, normalize_actions: bool = False,
-                 action_stats_path: str = None, episode_indices: list = None):
+                 action_stats_path: str = None, episode_indices: list = None,
+                 frame_len: int = 3):
         self.history_len = history_len
         self.action_dim = action_dim
         self.image_size = image_size
+        self.frame_len = frame_len
 
         self.transform = transforms.Compose([
             transforms.Resize((image_size, image_size)),
@@ -51,8 +53,11 @@ class ActionDataset(Dataset):
                 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]))
+            # 每个样本：帧 [t-frame_len+1, t]（旧→新）+ 动作历史 [t-history_len, t) → 动作 t
+            start = max(history_len, frame_len - 1)
+            for t in range(start, n):
+                self.samples.append((image_files[t - frame_len + 1: t + 1],
+                                     actions[t - history_len:t], actions[t]))
 
         self._action_mean = None
         self._action_std = None
@@ -77,15 +82,18 @@ class ActionDataset(Dataset):
         return len(self.samples)
 
     def __getitem__(self, idx):
-        image_path, action_history, target = self.samples[idx]
-        image = Image.open(image_path).convert("RGB")
-        image = self.transform(image)
+        image_paths, action_history, target = self.samples[idx]
+        frames = []
+        for p in image_paths:
+            image = Image.open(p).convert("RGB")
+            frames.append(self.transform(image))
+        images = torch.stack(frames)  # (F, 3, S, S)
         action_history = torch.tensor(action_history, dtype=torch.float32)
         target = torch.tensor(target, dtype=torch.float32)
         if self._action_mean is not None:
             action_history = (action_history - self._action_mean) / self._action_std
             target = (target - self._action_mean) / self._action_std
-        return {"image": image, "action_history": action_history, "target": target}
+        return {"images": images, "action_history": action_history, "target": target}
 
     def get_action_stats(self):
         if self._action_mean is not None:

actpred/export_onnx.py

@@ -19,16 +19,17 @@ def main():
     model.load_state_dict(checkpoint["model_state_dict"])
     model.eval()
 
-    dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
+    dummy_images = torch.randn(
+        1, data_cfg.frame_len, 3, data_cfg.image_size, data_cfg.image_size)
     dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
 
     torch.onnx.export(
-        model, (dummy_image, dummy_history),
+        model, (dummy_images, dummy_history),
         args.output,
-        input_names=["image", "action_history"],
+        input_names=["images", "action_history"],
         output_names=["predicted_action"],
         dynamic_axes={
-            "image": {0: "batch"},
+            "images": {0: "batch"},
             "action_history": {0: "batch"},
             "predicted_action": {0: "batch"},
         },
@@ -41,11 +42,11 @@ def main():
         import onnxruntime as ort
         import numpy as np
         sess = ort.InferenceSession(args.output)
-        img_np = dummy_image.numpy()
+        img_np = dummy_images.numpy()
         hist_np = dummy_history.numpy()
-        ort_out = sess.run(None, {"image": img_np, "action_history": hist_np})[0]
+        ort_out = sess.run(None, {"images": img_np, "action_history": hist_np})[0]
         with torch.no_grad():
-            torch_out = model(dummy_image, dummy_history).numpy()
+            torch_out = model(dummy_images, dummy_history).numpy()
         max_diff = np.abs(ort_out - torch_out).max()
         print(f"ONNX 一致性验证通过（最大差异: {max_diff:.6e}）")
     except ImportError:

actpred/model.py

@@ -53,34 +53,7 @@ class MobileNetEncoder(nn.Module):
         return self.proj(x)
 
 
-class MLPHistoryEncoder(nn.Module):
-    def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(inplace=True),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, hidden_dim),
-            nn.ReLU(inplace=True),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
-
-
-class GRUHistoryEncoder(nn.Module):
-    def __init__(self, input_dim: int, hidden_dim: int = 64, dropout: float = 0.1):
-        super().__init__()
-        self.gru = nn.GRU(input_size=input_dim, hidden_size=hidden_dim,
-                          num_layers=1, batch_first=True)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        _, h = self.gru(x)
-        return self.dropout(h.squeeze(0))
-
-
-class FusionHead(nn.Module):
+class PredictionHead(nn.Module):
     def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.1):
         super().__init__()
         self.net = nn.Sequential(
@@ -97,6 +70,13 @@ class FusionHead(nn.Module):
 
 
 class ActionPredictor(nn.Module):
+    """多帧图像 + 历史动作的 Transformer 注意力融合预测器
+
+    token 序列：[img×F, act×H, PRED]，长度固定为 F+H+1，
+    取 PRED 位置输出进预测头回归下一时间步动作。
+    除 batch 维外全部维度静态，保证 ONNX/NPU 部署友好。
+    """
+
     def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig,
                  pretrained: bool = True):
         super().__init__()
@@ -106,26 +86,44 @@ class ActionPredictor(nn.Module):
         else:
             self.image_encoder = SimpleCNN(model_cfg.cnn_feature_dim)
 
-        history_input_dim = data_cfg.history_len * data_cfg.action_dim
-        if model_cfg.history_encoder == "gru":
-            self.action_encoder = GRUHistoryEncoder(
-                data_cfg.action_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
-        else:
-            self.action_encoder = MLPHistoryEncoder(
-                history_input_dim, model_cfg.action_hidden_dim, model_cfg.dropout)
+        d = model_cfg.attn_dim
+        self.img_proj = nn.Linear(model_cfg.cnn_feature_dim, d)
+        self.act_proj = nn.Linear(data_cfg.action_dim, d)
 
-        fused_input = model_cfg.cnn_feature_dim + model_cfg.action_hidden_dim
-        self.head = FusionHead(
-            fused_input, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
+        seq_len = data_cfg.frame_len + data_cfg.history_len + 1
+        self.pred_token = nn.Parameter(torch.zeros(1, 1, d))
+        self.pos_embed = nn.Parameter(torch.zeros(1, seq_len, d))
+        nn.init.trunc_normal_(self.pred_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
 
-        self._use_gru = model_cfg.history_encoder == "gru"
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d, nhead=model_cfg.attn_heads,
+            dim_feedforward=d * 2, dropout=model_cfg.dropout,
+            batch_first=True, norm_first=True)
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=model_cfg.attn_layers,
+            enable_nested_tensor=False)
+
+        self.head = PredictionHead(
+            d, model_cfg.fused_hidden_dim, data_cfg.action_dim, model_cfg.dropout)
+
+        self._frame_len = data_cfg.frame_len
         self._history_len = data_cfg.history_len
         self._action_dim = data_cfg.action_dim
+        self._image_size = data_cfg.image_size
+        self._cnn_feature_dim = model_cfg.cnn_feature_dim
 
-    def forward(self, image: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
-        img_feat = self.image_encoder(image)
-        if self._use_gru:
-            act_feat = self.action_encoder(action_history)
-        else:
-            act_feat = self.action_encoder(action_history.flatten(1))
-        return self.head(torch.cat([img_feat, act_feat], dim=-1))
+    def forward(self, images: torch.Tensor, action_history: torch.Tensor) -> torch.Tensor:
+        # images: (B, F, 3, S, S)  action_history: (B, H, D)
+        b = images.size(0)
+        frames = images.reshape(
+            b * self._frame_len, 3, self._image_size, self._image_size)
+        img_feat = self.image_encoder(frames).reshape(
+            b, self._frame_len, self._cnn_feature_dim)
+        img_tokens = self.img_proj(img_feat)
+        act_tokens = self.act_proj(action_history)
+        pred_token = self.pred_token.expand(b, -1, -1)
+        tokens = torch.cat([img_tokens, act_tokens, pred_token], dim=1)
+        tokens = tokens + self.pos_embed
+        out = self.encoder(tokens)
+        return self.head(out[:, -1])

actpred/predict.py

@@ -28,6 +28,7 @@ class ActionPredictorInference:
 
         self.action_dim = self._data_cfg.action_dim
         self.history_len = self._data_cfg.history_len
+        self.frame_len = self._data_cfg.frame_len
 
         self._transform = transforms.Compose([
             transforms.Resize((self._data_cfg.image_size, self._data_cfg.image_size)),
@@ -61,13 +62,23 @@ class ActionPredictorInference:
             self._action_std = stats[1].astype(np.float32)
 
     @torch.no_grad()
-    def predict(self, image: np.ndarray, action_history: np.ndarray) -> np.ndarray:
+    def predict(self, images, action_history: np.ndarray) -> np.ndarray:
+        """images: 长度为 frame_len 的 RGB 图像序列（旧→新，末帧为当前帧）"""
+        if not isinstance(images, (list, tuple)):
+            images = [images]
+        if len(images) != self._data_cfg.frame_len:
+            raise ValueError(
+                f"images 应为 {self._data_cfg.frame_len} 帧，实际为 {len(images)} 帧。"
+                f"冷启动时可重复传入同一帧。")
         expected_shape = (self._data_cfg.history_len, self._data_cfg.action_dim)
         if action_history.shape != expected_shape:
             raise ValueError(
                 f"action_history 形状应为 {expected_shape}，实际为 {action_history.shape}")
-        img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
-        img_tensor = self._transform(img).unsqueeze(0).to(self._device)
+        frames = []
+        for image in images:
+            img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
+            frames.append(self._transform(img))
+        img_tensor = torch.stack(frames).unsqueeze(0).to(self._device)  # (1, F, 3, S, S)
         history = torch.tensor(action_history, dtype=torch.float32)
         if self._action_mean is not None:
             history = (history - self._action_mean) / self._action_std
@@ -93,7 +104,9 @@ def parse_action_history(s: str, action_dim: int) -> np.ndarray:
 def main():
     parser = argparse.ArgumentParser(description="actpred inference")
     parser.add_argument("--model", required=True, help="checkpoint path")
-    parser.add_argument("--image", required=True, help="image file path")
+    parser.add_argument("--images", required=True, nargs="+",
+                        help="按时间顺序（旧→新）传入 frame_len 个图像路径；"
+                             "冷启动时可重复传同一帧路径")
     parser.add_argument("--action-history", required=True,
                         help="冒号分隔各步，每步 action_dim 个值，共 history_len 步。"
                              "e.g. '0,0,0,0:1,0,0,0' (action_dim=4, history_len=2)")
@@ -106,9 +119,13 @@ def main():
     predictor = ActionPredictorInference(args.model, args.device)
     action_dim = args.action_dim or predictor.action_dim
 
-    image = Image.open(args.image).convert("RGB")
+    if len(args.images) != predictor.frame_len:
+        raise SystemExit(
+            f"--images 需传入 {predictor.frame_len} 个路径（旧→新），"
+            f"实际为 {len(args.images)} 个。冷启动时可重复传同一帧路径。")
+    images = [np.array(Image.open(p).convert("RGB")) for p in args.images]
     history = parse_action_history(args.action_history, action_dim)
-    predicted = predictor.predict(np.array(image), history)
+    predicted = predictor.predict(images, history)
 
     if args.output_format == "json":
         print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))

actpred/pretrain_mt.py

@@ -76,24 +76,25 @@ class MTActionDataset(Dataset):
         }
 
 
-def pretrain_one_epoch(model, loader, optimizer, device, image_size,
+def pretrain_one_epoch(model, loader, optimizer, device, image_size, frame_len,
                        max_grad_norm=None, noise_std=0.0):
     model.train()
     total_loss = 0.0
     count = 0
-    zero_image = None
+    zero_images = 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
+        if zero_images is None or zero_images.size(0) != history.size(0):
+            zero_images = torch.zeros(
+                history.size(0), frame_len, 3, image_size, image_size,
+                device=device,
             )
         optimizer.zero_grad()
-        pred = model(zero_image, history)
+        pred = model(zero_images, history)
         loss = torch.nn.functional.mse_loss(pred, target)
         loss.backward()
         if max_grad_norm:
@@ -105,19 +106,20 @@ def pretrain_one_epoch(model, loader, optimizer, device, image_size,
 
 
 @torch.no_grad()
-def validate(model, loader, device, image_size):
+def validate(model, loader, device, image_size, frame_len):
     model.eval()
     total_loss = 0.0
     count = 0
-    zero_image = None
+    zero_images = 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
+        if zero_images is None or zero_images.size(0) != history.size(0):
+            zero_images = torch.zeros(
+                history.size(0), frame_len, 3, image_size, image_size,
+                device=device,
             )
-        pred = model(zero_image, history)
+        pred = model(zero_images, history)
         loss = torch.nn.functional.mse_loss(pred, target)
         total_loss += loss.item() * history.size(0)
         count += history.size(0)
@@ -198,8 +200,10 @@ def main():
     print(f"Model params: {n_params:,}")
     print(f"Config: action_dim={cfg.data.action_dim}, "
           f"history_len={cfg.data.history_len}, "
+          f"frame_len={cfg.data.frame_len}, "
           f"backbone={cfg.model.cnn_backbone}, "
-          f"encoder={cfg.model.history_encoder}")
+          f"attn={cfg.model.attn_layers}x{cfg.model.attn_heads}h"
+          f"@{cfg.model.attn_dim}")
     print(f"Pre-training: noise_std={args.noise_std}")
 
     optimizer = torch.optim.AdamW(
@@ -222,9 +226,11 @@ def main():
               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,
+            cfg.data.image_size, cfg.data.frame_len,
+            cfg.train.max_grad_norm, args.noise_std,
         )
-        val_loss = validate(model, val_loader, device, cfg.data.image_size)
+        val_loss = validate(model, val_loader, device,
+                            cfg.data.image_size, cfg.data.frame_len)
         scheduler.step()
         print(f"  train_loss={train_loss:.6f}  val_loss={val_loss:.6f}")
 

actpred/test_actpred.py

@@ -79,45 +79,57 @@ class TestParseActionHistory(unittest.TestCase):
 
 
 class TestActionPredictor(unittest.TestCase):
-    """测试模型前向传播"""
+    """测试多帧注意力模型前向传播"""
 
-    def test_forward_simple_mlp(self):
-        data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
-        model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
+    def test_forward_multi_frame(self):
+        data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5, frame_len=3)
+        model_cfg = ModelConfig(cnn_backbone="simple")
         model = ActionPredictor(model_cfg, data_cfg)
         model.eval()
 
-        image = torch.randn(2, 3, 64, 64)
+        images = torch.randn(2, 3, 3, 64, 64)
         history = torch.randn(2, 5, 18)
         with torch.no_grad():
-            out = model(image, history)
+            out = model(images, history)
         self.assertEqual(out.shape, (2, 18))
 
-    def test_forward_simple_gru(self):
-        data_cfg = DataConfig(image_size=64, action_dim=18, history_len=5)
-        model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="gru")
+    def test_forward_single_frame_degenerate(self):
+        """frame_len=1 退化为单帧输入也应正常前向"""
+        data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=1)
+        model_cfg = ModelConfig(cnn_backbone="simple")
         model = ActionPredictor(model_cfg, data_cfg)
         model.eval()
 
-        image = torch.randn(2, 3, 64, 64)
-        history = torch.randn(2, 5, 18)
+        images = torch.randn(2, 1, 3, 32, 32)
+        history = torch.randn(2, 3, 4)
         with torch.no_grad():
-            out = model(image, history)
-        self.assertEqual(out.shape, (2, 18))
+            out = model(images, history)
+        self.assertEqual(out.shape, (2, 4))
 
     def test_forward_dynamic_batch(self):
-        data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3)
-        model_cfg = ModelConfig(cnn_backbone="simple", history_encoder="mlp")
+        data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=2)
+        model_cfg = ModelConfig(cnn_backbone="simple")
         model = ActionPredictor(model_cfg, data_cfg)
         model.eval()
 
         for batch_size in [1, 4, 8]:
-            image = torch.randn(batch_size, 3, 32, 32)
+            images = torch.randn(batch_size, 2, 3, 32, 32)
             history = torch.randn(batch_size, 3, 4)
             with torch.no_grad():
-                out = model(image, history)
+                out = model(images, history)
             self.assertEqual(out.shape, (batch_size, 4))
 
+    def test_custom_attn_config(self):
+        """自定义 attn_dim/heads/layers 应可构建并前向"""
+        data_cfg = DataConfig(image_size=32, action_dim=4, history_len=3, frame_len=2)
+        model_cfg = ModelConfig(attn_dim=64, attn_heads=2, attn_layers=1)
+        model = ActionPredictor(model_cfg, data_cfg)
+        model.eval()
+
+        with torch.no_grad():
+            out = model(torch.randn(1, 2, 3, 32, 32), torch.randn(1, 3, 4))
+        self.assertEqual(out.shape, (1, 4))
+
 
 class TestDataRecorder(unittest.TestCase):
     """测试数据录制器"""
@@ -169,17 +181,30 @@ class TestActionDataset(unittest.TestCase):
     def test_dataset_loading(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4)
-            ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32)
-            self.assertEqual(len(ds), 7)  # 10 - 3 = 7 samples
+            ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32,
+                               frame_len=3)
+            # 采样起点 max(history_len, frame_len-1) = 3 → 10 - 3 = 7 samples
+            self.assertEqual(len(ds), 7)
 
             sample = ds[0]
-            self.assertIn("image", sample)
+            self.assertIn("images", sample)
             self.assertIn("action_history", sample)
             self.assertIn("target", sample)
-            self.assertEqual(sample["image"].shape, (3, 32, 32))
+            self.assertEqual(sample["images"].shape, (3, 3, 32, 32))
             self.assertEqual(sample["action_history"].shape, (3, 4))
             self.assertEqual(sample["target"].shape, (4,))
 
+    def test_frame_len_gt_history_shifts_start(self):
+        """frame_len-1 > history_len 时采样起点由帧窗口决定"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self._create_dummy_episode(Path(tmpdir), n_frames=10, action_dim=4)
+            ds = ActionDataset(tmpdir, history_len=3, action_dim=4, image_size=32,
+                               frame_len=6)
+            # 起点 max(3, 5) = 5 → 10 - 5 = 5 samples
+            self.assertEqual(len(ds), 5)
+            sample = ds[0]
+            self.assertEqual(sample["images"].shape, (6, 3, 32, 32))
+
     def test_episode_count(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             self._create_dummy_episode(Path(tmpdir))
@@ -241,17 +266,24 @@ class TestConfig(unittest.TestCase):
         self.assertEqual(cfg.data.image_size, 64)
         self.assertEqual(cfg.data.action_dim, 18)
         self.assertEqual(cfg.data.history_len, 5)
+        self.assertEqual(cfg.data.frame_len, 3)
         self.assertEqual(cfg.model.cnn_backbone, "simple")
-        self.assertEqual(cfg.model.history_encoder, "mlp")
+        self.assertEqual(cfg.model.attn_dim, 128)
+        self.assertEqual(cfg.model.attn_heads, 4)
+        self.assertEqual(cfg.model.attn_layers, 2)
         self.assertEqual(cfg.train.epochs, 100)
 
     def test_apply_overrides(self):
         parser = build_argparser("test")
-        args = parser.parse_args(["--epochs", "50", "--lr", "0.01", "--backbone", "mobilenet"])
+        args = parser.parse_args(["--epochs", "50", "--lr", "0.01",
+                                  "--backbone", "mobilenet",
+                                  "--frame-len", "4", "--attn-dim", "64"])
         cfg = apply_overrides(Config(), args)
         self.assertEqual(cfg.train.epochs, 50)
         self.assertEqual(cfg.train.lr, 0.01)
         self.assertEqual(cfg.model.cnn_backbone, "mobilenet")
+        self.assertEqual(cfg.data.frame_len, 4)
+        self.assertEqual(cfg.model.attn_dim, 64)
 
     def test_normalize_actions_flag(self):
         """--normalize-actions 和 --no-normalize-actions 都应可用"""
@@ -421,18 +453,23 @@ class TestValidateConfigValues(unittest.TestCase):
 
     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")]:
+                             ("history_len", -1), ("frame_len", 0),
+                             ("frame_len", 99999), ("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"""
+        """model 字段越界（fused_hidden_dim 超限、dropout 越界、attn_* 越界）应抛 ValueError"""
         for field, value in [("fused_hidden_dim", 10 ** 9),
                              ("fused_hidden_dim", 0),
                              ("cnn_feature_dim", 99999),
-                             ("action_hidden_dim", -1),
+                             ("attn_dim", 0),
+                             ("attn_dim", 99999),
+                             ("attn_heads", 0),
+                             ("attn_heads", 128),
+                             ("attn_layers", 0),
                              ("dropout", 1.5),
                              ("dropout", -0.1)]:
             cfg = Config()
@@ -440,6 +477,14 @@ class TestValidateConfigValues(unittest.TestCase):
             with self.assertRaises(ValueError, msg=f"{field}={value}"):
                 validate_config_values(cfg.data, cfg.model)
 
+    def test_attn_dim_not_divisible_by_heads_raises(self):
+        """attn_dim 不能被 attn_heads 整除应抛 ValueError"""
+        cfg = Config()
+        cfg.model.attn_dim = 130
+        cfg.model.attn_heads = 4
+        with self.assertRaisesRegex(ValueError, "整除"):
+            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),
@@ -474,6 +519,15 @@ class TestParseCheckpointConfig(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "evil_extra_key"):
             parse_checkpoint_config(d)
 
+    def test_legacy_checkpoint_config_rejected(self):
+        """旧架构 checkpoint（含 history_encoder/action_hidden_dim）应拒载"""
+        cfg = Config()
+        d = cfg.to_dict()
+        d["model"]["history_encoder"] = "mlp"
+        d["model"]["action_hidden_dim"] = 64
+        with self.assertRaisesRegex(ValueError, "history_encoder"):
+            parse_checkpoint_config(d)
+
     def test_missing_section_raises(self):
         cfg = Config()
         d = cfg.to_dict()
@@ -519,11 +573,13 @@ class TestActionPredictorInference(unittest.TestCase):
     """测试推理封装端到端：checkpoint 保存 → 加载 → 推理"""
 
     @staticmethod
-    def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32):
+    def _make_checkpoint(tmpdir, normalize_actions=False, image_size=32,
+                         frame_len=3):
         cfg = Config()
         cfg.data.image_size = image_size
         cfg.data.action_dim = 4
         cfg.data.history_len = 3
+        cfg.data.frame_len = frame_len
         cfg.data.normalize_actions = normalize_actions
         model = ActionPredictor(cfg.model, cfg.data)
         model.eval()
@@ -533,17 +589,19 @@ class TestActionPredictorInference(unittest.TestCase):
         return path, model, cfg
 
     def test_end_to_end_predict(self):
-        """加载 checkpoint 后推理输出形状正确、数值有限且可复现"""
+        """加载 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)
+            self.assertEqual(predictor.frame_len, 3)
 
-            image = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
+            images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
+                      for _ in range(3)]
             history = np.random.randn(3, 4).astype(np.float32)
-            pred1 = predictor.predict(image, history)
-            pred2 = predictor.predict(image, history)
+            pred1 = predictor.predict(images, history)
+            pred2 = predictor.predict(images, history)
             self.assertEqual(pred1.shape, (4,))
             self.assertTrue(np.all(np.isfinite(pred1)))
             np.testing.assert_array_equal(pred1, pred2)  # eval 模式确定性
@@ -563,9 +621,18 @@ class TestActionPredictorInference(unittest.TestCase):
         with tempfile.TemporaryDirectory() as tmpdir:
             path, _, _ = self._make_checkpoint(tmpdir)
             predictor = ActionPredictorInference(path, device="cpu")
-            image = np.zeros((32, 32, 3), dtype=np.uint8)
+            images = [np.zeros((32, 32, 3), dtype=np.uint8)] * 3
             with self.assertRaisesRegex(ValueError, "action_history"):
-                predictor.predict(image, np.zeros((2, 4), dtype=np.float32))
+                predictor.predict(images, np.zeros((2, 4), dtype=np.float32))
+
+    def test_wrong_frame_count_raises(self):
+        """传入帧数与 frame_len 不一致应抛 ValueError"""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path, _, _ = self._make_checkpoint(tmpdir)
+            predictor = ActionPredictorInference(path, device="cpu")
+            images = [np.zeros((32, 32, 3), dtype=np.uint8)] * 2
+            with self.assertRaisesRegex(ValueError, "帧"):
+                predictor.predict(images, np.zeros((3, 4), dtype=np.float32))
 
     def test_missing_action_stats_raises(self):
         """启用归一化但缺 action_stats.npy 应抛 FileNotFoundError"""
@@ -594,9 +661,9 @@ class TestActionPredictorInference(unittest.TestCase):
             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)
+                raw_path, device="cpu").predict([image] * 3, history)
             pred_norm = ActionPredictorInference(
-                norm_path, device="cpu").predict(image, history)
+                norm_path, device="cpu").predict([image] * 3, history)
             np.testing.assert_allclose(pred_norm, pred_raw, rtol=1e-5, atol=1e-6)
 
     def test_stats_shape_mismatch_raises(self):

actpred/train.py

@@ -25,18 +25,18 @@ def train_one_epoch(model, loader, optimizer, device, max_grad_norm: float = Non
     total_loss = 0.0
     count = 0
     for batch in tqdm(loader, desc="  train", leave=False):
-        image = batch["image"].to(device)
+        images = batch["images"].to(device)
         history = batch["action_history"].to(device)
         target = batch["target"].to(device)
         optimizer.zero_grad()
-        pred = model(image, history)
+        pred = model(images, 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() * image.size(0)
-        count += image.size(0)
+        total_loss += loss.item() * images.size(0)
+        count += images.size(0)
     return total_loss / max(count, 1)
 
 
@@ -46,13 +46,13 @@ def validate(model, loader, device):
     total_loss = 0.0
     count = 0
     for batch in tqdm(loader, desc="  val  ", leave=False):
-        image = batch["image"].to(device)
+        images = batch["images"].to(device)
         history = batch["action_history"].to(device)
         target = batch["target"].to(device)
-        pred = model(image, history)
+        pred = model(images, history)
         loss = torch.nn.functional.mse_loss(pred, target)
-        total_loss += loss.item() * image.size(0)
-        count += image.size(0)
+        total_loss += loss.item() * images.size(0)
+        count += images.size(0)
     if count == 0:
         return float("inf")  # 空验证集不参与比较，避免假 best
     return total_loss / count
@@ -83,7 +83,7 @@ def main():
     train_ds = ActionDataset(
         cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
         cfg.data.image_size, cfg.data.normalize_actions,
-        cfg.data.action_stats_path, train_indices)
+        cfg.data.action_stats_path, train_indices, cfg.data.frame_len)
     if cfg.data.normalize_actions and not cfg.data.action_stats_path:
         # 归一化统计量必须只用训练集计算，再供验证集复用，否则 val loss 不可信
         stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy")
@@ -93,7 +93,7 @@ def main():
     val_ds = ActionDataset(
         cfg.data.data_dir, cfg.data.history_len, cfg.data.action_dim,
         cfg.data.image_size, cfg.data.normalize_actions,
-        cfg.data.action_stats_path, val_indices)
+        cfg.data.action_stats_path, val_indices, cfg.data.frame_len)
 
     print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
     if len(train_ds) == 0: