From 575c2a27a4b2be0bca07bcd64a5268592de0b1b8 Mon Sep 17 00:00:00 2001 From: qing <1776177662@qq.com> Date: Sat, 8 Aug 2026 06:25:06 +0000 Subject: [PATCH] =?UTF-8?q?feat(actpred):=20=E5=8A=A8=E4=BD=9C=E9=A2=84?= =?UTF-8?q?=E6=B5=8B=E6=A8=A1=E5=9D=97=E5=85=A5=E5=BA=93=EF=BC=88PyTorch?= =?UTF-8?q?=20=E8=AE=AD=E7=BB=83/=E9=A2=84=E6=B5=8B/ONNX=20=E5=AF=BC?= =?UTF-8?q?=E5=87=BA=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- actpred/.gitignore | 24 +++++++ actpred/README.md | 83 +++++++++++++++++++++++ actpred/_run.sh | 13 ++++ actpred/_set.sh | 12 ++++ actpred/config.py | 110 ++++++++++++++++++++++++++++++ actpred/dataset.py | 104 ++++++++++++++++++++++++++++ actpred/export_onnx.py | 44 ++++++++++++ actpred/model.py | 117 ++++++++++++++++++++++++++++++++ actpred/predict.py | 94 +++++++++++++++++++++++++ actpred/requirements.txt | 8 +++ actpred/train.py | 143 +++++++++++++++++++++++++++++++++++++++ 11 files changed, 752 insertions(+) create mode 100644 actpred/.gitignore create mode 100644 actpred/README.md create mode 100755 actpred/_run.sh create mode 100755 actpred/_set.sh create mode 100644 actpred/config.py create mode 100644 actpred/dataset.py create mode 100644 actpred/export_onnx.py create mode 100644 actpred/model.py create mode 100644 actpred/predict.py create mode 100644 actpred/requirements.txt create mode 100644 actpred/train.py diff --git a/actpred/.gitignore b/actpred/.gitignore new file mode 100644 index 0000000..b74f2a3 --- /dev/null +++ b/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 diff --git a/actpred/README.md b/actpred/README.md new file mode 100644 index 0000000..b88ac54 --- /dev/null +++ b/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,) +``` diff --git a/actpred/_run.sh b/actpred/_run.sh new file mode 100755 index 0000000..161729a --- /dev/null +++ b/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 diff --git a/actpred/_set.sh b/actpred/_set.sh new file mode 100755 index 0000000..5753b7d --- /dev/null +++ b/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 diff --git a/actpred/config.py b/actpred/config.py new file mode 100644 index 0000000..ee5f890 --- /dev/null +++ b/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 diff --git a/actpred/dataset.py b/actpred/dataset.py new file mode 100644 index 0000000..3a73d91 --- /dev/null +++ b/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) diff --git a/actpred/export_onnx.py b/actpred/export_onnx.py new file mode 100644 index 0000000..1a19fb4 --- /dev/null +++ b/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() diff --git a/actpred/model.py b/actpred/model.py new file mode 100644 index 0000000..49f03d4 --- /dev/null +++ b/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)) diff --git a/actpred/predict.py b/actpred/predict.py new file mode 100644 index 0000000..05d24ac --- /dev/null +++ b/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() diff --git a/actpred/requirements.txt b/actpred/requirements.txt new file mode 100644 index 0000000..1de5255 --- /dev/null +++ b/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 diff --git a/actpred/train.py b/actpred/train.py new file mode 100644 index 0000000..fc8cae4 --- /dev/null +++ b/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()