feat(actpred): 动作预测模块入库（PyTorch 训练/预测/ONNX 导出）

actpred/.gitignore

@@ -0,0 +1,24 @@
+# Virtual environment
+venv/
+.env/
+
+# Python bytecode
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Checkpoints and trained models
+checkpoints/
+*.pt
+*.pth
+*.onnx
+
+# Dataset (too large for git)
+data/
+
+# IDE
+.vscode/
+.idea/
+
+# OS
+.DS_Store

actpred/README.md

@@ -0,0 +1,83 @@
+# actpred — 动作预测模块
+
+输入摄像头图像 + 历史动作序列，预测下一个连续动作（回归）。
+
+纯神经网络实现（PyTorch），支持轻量 CNN 或 MobileNetV2 图像编码，MLP/GRU 动作历史编码。总参数量 ~140K（默认）至 ~2.3M。
+
+## 架构
+
+```
+图像 (3×64×64)  ──►  CNN 编码器  ──►  img_feat (128)
+                                          │
+                                          ├─ concat ──►  融合头  ──►  预测动作 (action_dim)
+                                          │
+动作历史 (H×D)  ──►  MLP/GRU    ──►  act_feat (64)
+```
+
+## 快速开始
+
+```bash
+# 安装
+cd actpred && bash _set.sh
+
+# 训练
+bash _run.sh train --data-dir ./data --epochs 100 --action-dim 8
+
+# 推理
+bash _run.sh --model checkpoints/best_inference.pt --image frame.jpg \
+    --action-history "45,0,45,0,45,0,45,0:45,10,45,10,45,10,45,10"
+
+# 导出 ONNX
+bash _run.sh export --model checkpoints/best_inference.pt --output model.onnx
+```
+
+## 数据格式
+
+```
+data/
+    episode_001/
+        images/
+            0000.jpg
+            0001.jpg
+            ...
+        actions.npy       # shape: (T, action_dim) float32
+    episode_002/
+        ...
+```
+
+每个 timestep 的样本：当前图像 + 前 `history_len` 步动作 → 当前步动作。
+
+### 数据采集
+
+```python
+from dataset import DataRecorder
+
+recorder = DataRecorder("./data", action_dim=8)
+for image, action in robot_loop():
+    recorder.record(image, action)
+recorder.flush()
+```
+
+## 配置
+
+所有超参数通过 CLI 参数覆盖，默认值见 `config.py`：
+
+| 参数 | 默认值 | 说明 |
+|------|--------|------|
+| `--image-size` | 64 | 图像缩放尺寸 |
+| `--action-dim` | 8 | 动作向量维度 |
+| `--history-len` | 5 | 历史动作窗口长度 |
+| `--backbone` | simple | CNN 类型 (simple/mobilenet) |
+| `--history-encoder` | mlp | 历史编码器 (mlp/gru) |
+| `--epochs` | 100 | 训练轮数 |
+| `--lr` | 0.001 | 学习率 |
+| `--batch-size` | 32 | 批次大小 |
+
+## 编程 API
+
+```python
+from predict import ActionPredictorInference
+
+predictor = ActionPredictorInference("checkpoints/best_inference.pt")
+action = predictor.predict(image_rgb, action_history)  # np.ndarray (action_dim,)
+```

actpred/_run.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+cd "$(dirname "$0")" || exit
+source ./venv/bin/activate
+
+if [ "${1:-}" = "train" ]; then
+    shift
+    exec python3 train.py "$@"
+elif [ "${1:-}" = "export" ]; then
+    shift
+    exec python3 export_onnx.py "$@"
+else
+    exec python3 predict.py "$@"
+fi

actpred/_set.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+cd "$(dirname "$0")" || exit
+python3 -m venv venv
+source venv/bin/activate
+
+# CPU-only PyTorch (no CUDA dependency)
+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
+
+mkdir -p checkpoints data

actpred/config.py

@@ -0,0 +1,110 @@
+import argparse
+from dataclasses import dataclass, asdict, field
+from typing import Optional
+
+
+@dataclass
+class DataConfig:
+    image_size: int = 64
+    action_dim: int = 8
+    history_len: int = 5
+    data_dir: str = "./data"
+    normalize_actions: bool = False
+    action_stats_path: Optional[str] = None
+
+
+@dataclass
+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"
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 32
+    epochs: int = 100
+    lr: float = 1e-3
+    weight_decay: float = 1e-4
+    val_split: float = 0.15
+    checkpoint_dir: str = "./checkpoints"
+    patience: int = 15
+    num_workers: int = 2
+
+
+@dataclass
+class Config:
+    data: DataConfig = field(default_factory=DataConfig)
+    model: ModelConfig = field(default_factory=ModelConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d):
+        return cls(
+            data=DataConfig(**d.get("data", {})),
+            model=ModelConfig(**d.get("model", {})),
+            train=TrainConfig(**d.get("train", {})),
+        )
+
+
+def build_argparser(description: str = "") -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=description)
+    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("--data-dir", type=str, default=None)
+    parser.add_argument("--normalize-actions", action="store_true", 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("--batch-size", type=int, default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--lr", type=float, default=None)
+    parser.add_argument("--weight-decay", type=float, default=None)
+    parser.add_argument("--val-split", type=float, default=None)
+    parser.add_argument("--checkpoint-dir", type=str, default=None)
+    parser.add_argument("--patience", type=int, default=None)
+    parser.add_argument("--num-workers", type=int, default=None)
+    parser.add_argument("--save-config", type=str, default=None, help="保存配置到 JSON 文件")
+    return parser
+
+
+def apply_overrides(cfg: Config, args: argparse.Namespace) -> Config:
+    mapping = {
+        "image_size": ("data", "image_size"),
+        "action_dim": ("data", "action_dim"),
+        "history_len": ("data", "history_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"),
+        "batch_size": ("train", "batch_size"),
+        "epochs": ("train", "epochs"),
+        "lr": ("train", "lr"),
+        "weight_decay": ("train", "weight_decay"),
+        "val_split": ("train", "val_split"),
+        "checkpoint_dir": ("train", "checkpoint_dir"),
+        "patience": ("train", "patience"),
+        "num_workers": ("train", "num_workers"),
+    }
+    args_dict = vars(args)
+    for arg_name, (section, field) in mapping.items():
+        val = args_dict.get(arg_name)
+        if val is not None:
+            setattr(getattr(cfg, section), field, val)
+    return cfg

actpred/dataset.py

@@ -0,0 +1,104 @@
+import os
+import numpy as np
+from pathlib import Path
+from PIL import Image
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+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):
+        self.history_len = history_len
+        self.action_dim = action_dim
+        self.image_size = image_size
+
+        self.transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+
+        data_path = Path(data_dir)
+        episodes = sorted([d for d in data_path.iterdir() if d.is_dir()])
+        if episode_indices is not None:
+            episodes = [episodes[i] for i in episode_indices]
+
+        self.samples = []
+        for ep_dir in episodes:
+            actions_path = ep_dir / "actions.npy"
+            images_dir = ep_dir / "images"
+            if not actions_path.exists() or not images_dir.is_dir():
+                continue
+            actions = np.load(actions_path).astype(np.float32)
+            n_frames = len(actions)
+            image_files = sorted(images_dir.glob("*.jpg")) + sorted(images_dir.glob("*.png"))
+            n_images = len(image_files)
+            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]))
+
+        self._action_mean = None
+        self._action_std = None
+        if normalize_actions:
+            if action_stats_path and os.path.exists(action_stats_path):
+                stats = np.load(action_stats_path)
+                self._action_mean = stats[0]
+                self._action_std = stats[1]
+            else:
+                all_actions = np.stack([s[2] for s in self.samples])
+                self._action_mean = all_actions.mean(axis=0)
+                self._action_std = all_actions.std(axis=0) + 1e-8
+
+    def __len__(self):
+        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)
+        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}
+
+    def get_action_stats(self):
+        if self._action_mean is not None:
+            return np.stack([self._action_mean, self._action_std])
+        return None
+
+    @staticmethod
+    def get_episode_count(data_dir: str) -> int:
+        data_path = Path(data_dir)
+        return len([d for d in data_path.iterdir() if d.is_dir()])
+
+
+class DataRecorder:
+    def __init__(self, output_dir: str, action_dim: int):
+        self.output_dir = Path(output_dir)
+        self.action_dim = action_dim
+        self._frame_idx = 0
+        self._actions = []
+        ep_idx = 0
+        while (self.output_dir / f"episode_{ep_idx:04d}").exists():
+            ep_idx += 1
+        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)
+
+    def record(self, image: np.ndarray, action: np.ndarray):
+        img = Image.fromarray(image)
+        img.save(self.images_dir / f"{self._frame_idx:04d}.jpg", quality=90)
+        self._actions.append(action.astype(np.float32))
+        self._frame_idx += 1
+
+    def flush(self):
+        if self._actions:
+            actions = np.array(self._actions, dtype=np.float32)
+            np.save(self.episode_dir / "actions.npy", actions)
+        return str(self.episode_dir)

actpred/export_onnx.py

@@ -0,0 +1,44 @@
+import argparse
+import torch
+from config import DataConfig, ModelConfig
+from model import ActionPredictor
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Export actpred model to ONNX")
+    parser.add_argument("--model", required=True, help="checkpoint path (best_inference.pt)")
+    parser.add_argument("--output", default="model.onnx", help="output ONNX file path")
+    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"])
+
+    model = ActionPredictor(model_cfg, data_cfg)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    dummy_image = torch.randn(1, 3, data_cfg.image_size, data_cfg.image_size)
+    if model_cfg.history_encoder == "gru":
+        dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
+    else:
+        dummy_history = torch.randn(1, data_cfg.history_len, data_cfg.action_dim)
+
+    torch.onnx.export(
+        model, (dummy_image, dummy_history),
+        args.output,
+        input_names=["image", "action_history"],
+        output_names=["predicted_action"],
+        dynamic_axes={
+            "image": {0: "batch"},
+            "action_history": {0: "batch"},
+            "predicted_action": {0: "batch"},
+        },
+        opset_version=17,
+    )
+    print(f"ONNX model exported to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

actpred/model.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from config import DataConfig, ModelConfig
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self, feature_dim: int = 128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(3, 16, 3, stride=2, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(16, 32, 3, stride=2, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 64, 3, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, feature_dim, 3, stride=2, padding=1),
+            nn.BatchNorm2d(feature_dim),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x).flatten(1)
+
+
+class MobileNetEncoder(nn.Module):
+    def __init__(self, feature_dim: int = 128):
+        super().__init__()
+        from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
+        backbone = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+        self.features = backbone.features
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.proj = nn.Linear(1280, feature_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        x = self.pool(x).flatten(1)
+        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):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, 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 // 2),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_dim // 2, output_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class ActionPredictor(nn.Module):
+    def __init__(self, model_cfg: ModelConfig, data_cfg: DataConfig):
+        super().__init__()
+        if model_cfg.cnn_backbone == "mobilenet":
+            self.image_encoder = MobileNetEncoder(model_cfg.cnn_feature_dim)
+        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)
+
+        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)
+
+        self._use_gru = model_cfg.history_encoder == "gru"
+        self._history_len = data_cfg.history_len
+        self._action_dim = data_cfg.action_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))

actpred/predict.py

@@ -0,0 +1,94 @@
+import argparse
+import json
+import numpy as np
+import torch
+from PIL import Image
+from torchvision import transforms
+
+from config import Config, ModelConfig, DataConfig
+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"])
+
+        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.load_state_dict(checkpoint["model_state_dict"])
+        self._model.eval()
+
+        self._transform = transforms.Compose([
+            transforms.Resize((self._data_cfg.image_size, self._data_cfg.image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ])
+
+        self._action_mean = None
+        self._action_std = None
+        stats_path = checkpoint_path.rsplit("/", 1)[0] + "/action_stats.npy"
+        try:
+            stats = np.load(stats_path)
+            self._action_mean = stats[0]
+            self._action_std = stats[1]
+        except (FileNotFoundError, ValueError):
+            pass
+
+    @torch.no_grad()
+    def predict(self, image: np.ndarray, action_history: np.ndarray) -> np.ndarray:
+        img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
+        img_tensor = self._transform(img).unsqueeze(0).to(self._device)
+        history = torch.tensor(action_history, dtype=torch.float32)
+        if self._action_mean is not None:
+            history = (history - self._action_mean) / self._action_std
+        history = history.unsqueeze(0).to(self._device)
+        pred = self._model(img_tensor, history).squeeze(0).cpu().numpy()
+        if self._action_mean is not None:
+            pred = pred * self._action_std + self._action_mean
+        return pred
+
+
+def parse_action_history(s: str, action_dim: int) -> np.ndarray:
+    steps = s.split(":")
+    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)}"
+        history.append(values)
+    return np.array(history, dtype=np.float32)
+
+
+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("--action-history", required=True,
+                        help="colon-separated action vectors, e.g. '45,0,45,0:45,10,45,10'")
+    parser.add_argument("--action-dim", type=int, default=None,
+                        help="override action dim (auto-detected from checkpoint if omitted)")
+    parser.add_argument("--output-format", choices=["json", "raw"], default="json")
+    parser.add_argument("--device", type=str, default=None)
+    args = parser.parse_args()
+
+    predictor = ActionPredictorInference(args.model, args.device)
+    action_dim = args.action_dim or predictor._data_cfg.action_dim
+
+    image = Image.open(args.image).convert("RGB")
+    history = parse_action_history(args.action_history, action_dim)
+    predicted = predictor.predict(np.array(image), history)
+
+    if args.output_format == "json":
+        print(json.dumps({"predicted_action": predicted.tolist()}, indent=2))
+    else:
+        print(" ".join(f"{v:.4f}" for v in predicted))
+
+
+if __name__ == "__main__":
+    main()

actpred/requirements.txt

@@ -0,0 +1,8 @@
+# Core
+torch>=2.0
+torchvision>=0.15
+numpy>=1.24
+Pillow>=9.0
+
+# Training utilities
+tqdm>=4.60

actpred/train.py

@@ -0,0 +1,143 @@
+import json
+import os
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Subset
+from tqdm import tqdm
+
+from config import Config, build_argparser, apply_overrides
+from model import ActionPredictor
+from dataset import ActionDataset
+
+
+def split_episodes(data_dir: str, val_ratio: float, seed: int = 42):
+    n = ActionDataset.get_episode_count(data_dir)
+    indices = list(range(n))
+    rng = np.random.RandomState(seed)
+    rng.shuffle(indices)
+    val_count = max(1, int(n * val_ratio))
+    return indices[val_count:], indices[:val_count]
+
+
+def train_one_epoch(model, loader, optimizer, device):
+    model.train()
+    total_loss = 0.0
+    count = 0
+    for batch in tqdm(loader, desc="  train", leave=False):
+        image = batch["image"].to(device)
+        history = batch["action_history"].to(device)
+        target = batch["target"].to(device)
+        optimizer.zero_grad()
+        pred = model(image, history)
+        loss = torch.nn.functional.mse_loss(pred, target)
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item() * image.size(0)
+        count += image.size(0)
+    return total_loss / max(count, 1)
+
+
+@torch.no_grad()
+def validate(model, loader, device):
+    model.eval()
+    total_loss = 0.0
+    count = 0
+    for batch in tqdm(loader, desc="  val  ", leave=False):
+        image = batch["image"].to(device)
+        history = batch["action_history"].to(device)
+        target = batch["target"].to(device)
+        pred = model(image, history)
+        loss = torch.nn.functional.mse_loss(pred, target)
+        total_loss += loss.item() * image.size(0)
+        count += image.size(0)
+    return total_loss / max(count, 1)
+
+
+def main():
+    parser = build_argparser("actpred training")
+    args = parser.parse_args()
+    cfg = apply_overrides(Config(), args)
+
+    os.makedirs(cfg.train.checkpoint_dir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    train_indices, val_indices = split_episodes(
+        cfg.data.data_dir, cfg.train.val_split)
+    print(f"Episodes: {len(train_indices)} train, {len(val_indices)} val")
+
+    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)
+    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)
+
+    print(f"Samples: {len(train_ds)} train, {len(val_ds)} val")
+
+    train_loader = DataLoader(train_ds, batch_size=cfg.train.batch_size,
+                              shuffle=True, num_workers=cfg.train.num_workers)
+    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:,}")
+
+    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, "best.pt")
+
+    for epoch in range(1, cfg.train.epochs + 1):
+        print(f"\nEpoch {epoch}/{cfg.train.epochs}  lr={optimizer.param_groups[0]['lr']:.6f}")
+        train_loss = train_one_epoch(model, train_loader, optimizer, device)
+        val_loss = validate(model, val_loader, device)
+        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)
+            inference_path = os.path.join(cfg.train.checkpoint_dir, "best_inference.pt")
+            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} (patience={cfg.train.patience})")
+                break
+
+    print(f"\nTraining done. Best val_loss={best_val_loss:.6f}")
+    print(f"Best model saved to: {best_path}")
+
+    action_stats = train_ds.get_action_stats()
+    if action_stats is not None:
+        stats_path = os.path.join(cfg.train.checkpoint_dir, "action_stats.npy")
+        np.save(stats_path, action_stats)
+        print(f"Action stats saved to: {stats_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()