initial commit

.gitignore

@@ -0,0 +1,48 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.egg-info/
+dist/
+build/
+
+# Virtual environment
+.venv/
+venv/
+env/
+
+# IDE / Editor
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Data (large binary files)
+bags/
+dataset/
+
+# Model checkpoints / weights
+checkpoints/
+*.pt
+
+# Logs (TensorBoard, etc.)
+logs/
+
+# Benchmark evaluation results
+benchmark/results/
+
+# Evaluation figures
+*.png
+*.jpg
+*.jpeg
+*.pdf
+*.svg
+
+# Shell scripts (optional — uncomment if you want to ignore)
+# *.sh
+
+# ROS bag files
+*.bag
+*.bag.active

DATASET_FORMAT.md

@@ -0,0 +1,95 @@
+# UZH FPV Dataset Format
+
+> 由 `rosbag2wds.py` 从 DAVIS 事件相机 ROS bag 转换生成
+
+## 目录结构
+
+```
+dataset/
+├── <dataset_name>/
+│   ├── shard_0000.tar        # WebDataset shard (图像 + 对齐的 GT)
+│   ├── shard_0001.tar
+│   ├── ...
+│   ├── imu_sequence.npz      # 完整 IMU 序列 (独立存储)
+│   └── metadata.json         # 数据集元信息
+```
+
+## 文件说明
+
+### 1. WebDataset Shard (`shard_*.tar`)
+
+每个 tar 文件包含 `shard_size` 个样本（默认 2000），每个样本的 key 为 `frame_<index>`，包含以下字段：
+
+| Key  | 类型         | 内容                                                              |
+|------|-------------|-------------------------------------------------------------------|
+| `jpg` | JPEG bytes  | 灰度图，尺寸 `320×240`，JPEG quality=85                              |
+| `ts`  | float64     | 图像时间戳（ROS bag 系统时间 `t.to_sec()`）                          |
+| `pose`| float32[7]  | 位姿：`[x, y, z, qx, qy, qz, qw]`（位置 + 单位四元数）              |
+| `vel` | float32[6]  | 速度：`[vx, vy, vz, wx, wy, wz]`（线速度 + 角速度）                 |
+
+**读取示例 (Python):**
+
+```python
+import webdataset as wds
+
+dataset = wds.WebDataset("dataset/<name>/shard_0000.tar")
+for sample in dataset:
+    img = sample["jpg"]       # JPEG bytes
+    ts = sample["ts"]         # bytes -> np.frombuffer(..., dtype=np.float64)
+    pose = sample["pose"]     # bytes -> np.frombuffer(..., dtype=np.float32).reshape(7)
+    vel = sample["vel"]       # bytes -> np.frombuffer(..., dtype=np.float32).reshape(6)
+```
+
+### 2. IMU 序列 (`imu_sequence.npz`)
+
+独立存储的完整 IMU 数据（NPZ 压缩格式），包含三个数组：
+
+| Key                | 类型         | 形状         | 内容                        |
+|--------------------|-------------|-------------|----------------------------|
+| `timestamps`       | float64     | (N,)        | IMU 时间戳                   |
+| `accelerations`    | float32     | (N, 3)      | 线性加速度 `(ax, ay, az)` m/s² |
+| `angular_velocities`| float32    | (N, 3)      | 角速度 `(gx, gy, gz)` rad/s  |
+
+**读取示例:**
+
+```python
+import numpy as np
+data = np.load("dataset/<name>/imu_sequence.npz")
+timestamps = data["timestamps"]
+acc = data["accelerations"]
+gyro = data["angular_velocities"]
+```
+
+### 3. 元信息 (`metadata.json`)
+
+```json
+{
+  "dataset_name": "indoor_forward_7",
+  "source_bag": "/mnt/indoor_forward_7_davis_with_gt.bag",
+  "num_images": 2459,
+  "num_imu_messages": 66632,
+  "num_ground_truth": 33350,
+  "image_size": [320, 240],
+  "imu_frequency_hz": 999.02,
+  "camera_frequency_hz": 36.89,
+  "gt_frequency_hz": 500.01,
+  "coordinate_system": "horizontal (z aligned with gravity, assumed from GT)",
+  "velocity_dimensions": 6
+}
+```
+
+## 数据来源
+
+| Topic                  | 内容              | 频率 (典型值) |
+|------------------------|------------------|-------------|
+| `/dvs/image_raw`       | 灰度图像 (mono8)  | ~30–50 Hz   |
+| `/dvs/imu`             | IMU (加速度+角速度)| ~1000 Hz    |
+| `/groundtruth/odometry`| 位姿真值          | ~500 Hz     |
+
+## 预处理说明
+
+- **时间戳**: 统一使用 ROS bag 系统时间 `t.to_sec()`，而非 `msg.header.stamp`
+- **时间对齐**: 图像与 GT 通过最近邻时间戳匹配，最大允许偏差 0.1s
+- **速度计算**: GT 速度由位姿差分计算（前向/后向有限差分 + 四元数旋转向量），忽略 bag 中原始 twist 数据
+- **时间裁剪**: 所有数据裁剪至 GT 时间范围内，去除首尾无 GT 的片段
+- **图像缩放**: 原始 DAVIS 分辨率 `240×180` → 缩放至 `320×240` (INTER_LINEAR)

batch_convert.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+# batch_convert.sh
+# 在已运行的 ROS 容器内执行，批量转换尚未转换的数据集
+#
+# 用法（容器内）:
+#   cd /mnt && bash batch_convert.sh
+#
+# 它会自动检测:
+#   - 哪些 .bag 文件尚未转换（通过检查 dataset/<name>/metadata.json）
+#   - 跳过已转换的数据集
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+OUTPUT_DIR="${SCRIPT_DIR}/dataset"
+BAG_DIR="${SCRIPT_DIR}/bags"
+
+# 已转换的数据集列表（通过检查 metadata.json）
+echo "=========================================="
+echo "批量转换 UZH FPV 数据集"
+echo "=========================================="
+
+# 收集所有 bag 文件
+BAG_FILES=()
+while IFS= read -r -d '' f; do
+    BAG_FILES+=("$f")
+done < <(find "$BAG_DIR" -maxdepth 2 -name "*_davis_with_gt.bag" -print0)
+
+if [ ${#BAG_FILES[@]} -eq 0 ]; then
+    echo "❌ 未找到任何 .bag 文件"
+    exit 1
+fi
+
+echo "找到 ${#BAG_FILES[@]} 个 bag 文件"
+
+# 逐个检查并转换
+CONVERTED=0
+SKIPPED=0
+FAILED=0
+
+for bag_path in "${BAG_FILES[@]}"; do
+    bag_name="$(basename "$bag_path")"
+    dataset_name="${bag_name%_davis_with_gt.bag}"
+
+    # 检查是否已转换
+    metadata_file="${OUTPUT_DIR}/${dataset_name}/metadata.json"
+    if [ -f "$metadata_file" ]; then
+        echo "  ⏭️  跳过 ${dataset_name} (已转换)"
+        SKIPPED=$((SKIPPED + 1))
+        continue
+    fi
+
+    echo ""
+    echo "  🔄 转换: ${dataset_name}"
+
+    # 检查依赖
+    if ! python3 -c "import webdataset" 2>/dev/null; then
+        echo "  ⚠️   正在安装 webdataset..."
+        pip3 install webdataset tqdm scipy 2>/dev/null || {
+            echo "  ❌ pip install 失败，跳过 ${dataset_name}"
+            FAILED=$((FAILED + 1))
+            continue
+        }
+    fi
+
+    # 执行转换
+    if python3 "${SCRIPT_DIR}/rosbag2wds.py" \
+        --bag "$bag_path" \
+        --output "$OUTPUT_DIR" \
+        --name "$dataset_name" \
+        --shard_size 2000 \
+        --width 320 --height 240; then
+        echo "  ✅ ${dataset_name} 转换完成"
+        CONVERTED=$((CONVERTED + 1))
+    else
+        echo "  ❌ ${dataset_name} 转换失败"
+        FAILED=$((FAILED + 1))
+    fi
+done
+
+echo ""
+echo "=========================================="
+echo "批量转换结束"
+echo "  已转换: ${CONVERTED}"
+echo "  已跳过: ${SKIPPED}"
+echo "  失败:   ${FAILED}"
+echo "=========================================="

benchmark/benchmark.py

@@ -0,0 +1,309 @@
+"""
+benchmark.py — Unified evaluation entry point.
+
+Two modes:
+  1. Single-model eval:  python -m benchmark.benchmark --checkpoint <path>
+  2. Compare mode:       python -m benchmark.benchmark --compare <checkpoint_dir>
+
+Results are saved to benchmark/results/<exp_name>/.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+import torch
+
+from benchmark.config import eval_cfg, TEST_SCENE_GROUPS
+from benchmark.evaluate import run_full_evaluation, save_results
+
+# Project root (two levels up from benchmark/benchmark.py)
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+RESULTS_DIR = PROJECT_ROOT / "benchmark" / "results"
+
+
+def load_checkpoint(
+    checkpoint_path: Path,
+    device: torch.device,
+) -> torch.nn.Module:
+    """Load a VelocityPredictionModel from a checkpoint file."""
+    from src.velocity_prediction.model import VelocityPredictionModel
+
+    model = VelocityPredictionModel()
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    state_dict = ckpt.get("model_state_dict", ckpt)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    return model
+
+
+def run_single_eval(
+    checkpoint_path: Path,
+    output_dir: Optional[Path] = None,
+    device: torch.device = None,
+    seq_len: Optional[int] = None,
+    batch_size: Optional[int] = None,
+    num_workers: Optional[int] = None,
+    save_plots: bool = True,
+) -> Path:
+    """Evaluate a single checkpoint and save results.
+
+    Returns the output directory path.
+    """
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    seq_len = seq_len or eval_cfg.seq_len
+    batch_size = batch_size or eval_cfg.batch_size
+    num_workers = num_workers or eval_cfg.num_workers
+
+    # Derive experiment name from checkpoint filename (strip extension)
+    exp_name = checkpoint_path.stem  # e.g. "best" or "epoch_050_val_1.827390"
+
+    if output_dir is None:
+        output_dir = RESULTS_DIR / exp_name
+
+    print(f"{'=' * 60}")
+    print(f"Benchmark — Single Model Evaluation")
+    print(f"{'=' * 60}")
+    print(f"  Checkpoint: {checkpoint_path}")
+    print(f"  Device:     {device}")
+    print(f"  Seq len:    {seq_len}")
+    print(f"  Batch size: {batch_size}")
+    print(f"  Output:     {output_dir}")
+    print()
+
+    model = load_checkpoint(checkpoint_path, device)
+
+    results = run_full_evaluation(
+        model=model,
+        device=device,
+        seq_len=seq_len,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        event_threshold=eval_cfg.event_threshold,
+        event_use_log=eval_cfg.event_use_log,
+        scene_groups=TEST_SCENE_GROUPS,
+    )
+
+    save_results(results, save_dir=output_dir, checkpoint_name=exp_name)
+
+    return output_dir
+
+
+def run_compare(
+    checkpoint_dir: Path,
+    output_dir: Optional[Path] = None,
+    device: torch.device = None,
+    seq_len: Optional[int] = None,
+    batch_size: Optional[int] = None,
+    num_workers: Optional[int] = None,
+    pattern: str = "*.pt",
+) -> Path:
+    """Evaluate all checkpoints in a directory and produce a comparison table.
+
+    Returns the output directory path.
+    """
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    seq_len = seq_len or eval_cfg.seq_len
+    batch_size = batch_size or eval_cfg.batch_size
+    num_workers = num_workers or eval_cfg.num_workers
+
+    checkpoint_paths = sorted(Path(checkpoint_dir).glob(pattern))
+    if not checkpoint_paths:
+        print(f"No checkpoints found matching '{pattern}' in {checkpoint_dir}")
+        sys.exit(1)
+
+    if output_dir is None:
+        output_dir = RESULTS_DIR / "compare"
+
+    print(f"{'=' * 60}")
+    print(f"Benchmark — Compare Mode ({len(checkpoint_paths)} checkpoints)")
+    print(f"{'=' * 60}")
+    print(f"  Checkpoint dir: {checkpoint_dir}")
+    print(f"  Device:         {device}")
+    print(f"  Seq len:        {seq_len}")
+    print(f"  Batch size:     {batch_size}")
+    print(f"  Output:         {output_dir}")
+    print()
+
+    all_global_metrics = []
+
+    for ckpt_path in checkpoint_paths:
+        exp_name = ckpt_path.stem
+        print(f"\n── Evaluating {exp_name} ──")
+
+        model = load_checkpoint(ckpt_path, device)
+
+        results = run_full_evaluation(
+            model=model,
+            device=device,
+            seq_len=seq_len,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            event_threshold=eval_cfg.event_threshold,
+            event_use_log=eval_cfg.event_use_log,
+            scene_groups=TEST_SCENE_GROUPS,
+        )
+
+        # Save individual results
+        ckpt_output_dir = output_dir / exp_name
+        save_results(results, save_dir=ckpt_output_dir, checkpoint_name=exp_name)
+
+        all_global_metrics.append((exp_name, results["global"]))
+
+    # ── Comparison table ──
+    print(f"\n\n{'=' * 60}")
+    print("Comparison Summary")
+    print(f"{'=' * 60}")
+
+    header = f"{'Checkpoint':<30} {'RMSE vx':>10} {'RMSE vy':>10} {'RMSE xy':>10} {'MAE vx':>10} {'MAE vy':>10} {'R² vx':>8} {'R² vy':>8}"
+    sep = "-" * len(header)
+    print(header)
+    print(sep)
+
+    rows = []
+    for name, metrics in all_global_metrics:
+        row = (
+            f"{name:<30} "
+            f"{metrics.get('rmse_vx', 0):>10.4f} "
+            f"{metrics.get('rmse_vy', 0):>10.4f} "
+            f"{metrics.get('rmse_xy', 0):>10.4f} "
+            f"{metrics.get('mae_vx', 0):>10.4f} "
+            f"{metrics.get('mae_vy', 0):>10.4f} "
+            f"{metrics.get('r2_vx', 0):>8.4f} "
+            f"{metrics.get('r2_vy', 0):>8.4f}"
+        )
+        print(row)
+        rows.append(row)
+
+    # Save comparison CSV
+    import csv
+    csv_path = output_dir / "comparison.csv"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    fieldnames = ["checkpoint", "rmse_vx", "rmse_vy", "rmse_xy", "mae_vx", "mae_vy",
+                  "mae_xy", "r2_vx", "r2_vy", "count"]
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for name, metrics in all_global_metrics:
+            row = {"checkpoint": name, **metrics}
+            writer.writerow(row)
+    print(f"\nComparison CSV: {csv_path}")
+
+    # Save comparison text
+    txt_path = output_dir / "comparison.txt"
+    with open(txt_path, "w") as f:
+        f.write("Benchmark Comparison\n")
+        f.write(f"{'=' * 60}\n\n")
+        f.write(header + "\n")
+        f.write(sep + "\n")
+        for row in rows:
+            f.write(row + "\n")
+    print(f"Comparison TXT: {txt_path}")
+
+    return output_dir
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Unified benchmark for velocity prediction models.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "Examples:\n"
+            "  # Single model evaluation\n"
+            "  python -m benchmark.benchmark --checkpoint checkpoints/best.pt\n\n"
+            "  # Compare all checkpoints in a directory\n"
+            "  python -m benchmark.benchmark --compare checkpoints/\n\n"
+            "  # Custom output directory\n"
+            "  python -m benchmark.benchmark --checkpoint checkpoints/best.pt --output my_results/\n"
+        ),
+    )
+
+    # Mutually exclusive mode selection
+    mode = parser.add_mutually_exclusive_group(required=True)
+    mode.add_argument(
+        "--checkpoint", type=str, default=None,
+        help="Path to a single checkpoint .pt file for single-model evaluation.",
+    )
+    mode.add_argument(
+        "--compare", type=str, default=None,
+        help="Directory containing multiple .pt checkpoints for comparison.",
+    )
+
+    # Optional overrides
+    parser.add_argument(
+        "--output", type=str, default=None,
+        help="Output directory for results (default: benchmark/results/<exp_name>/).",
+    )
+    parser.add_argument(
+        "--device", type=str, default=None,
+        help="Device override, e.g. 'cuda:0' or 'cpu' (default: auto-detect).",
+    )
+    parser.add_argument(
+        "--seq-len", type=int, default=None,
+        help=f"Sequence length override (default: {eval_cfg.seq_len}).",
+    )
+    parser.add_argument(
+        "--batch-size", type=int, default=None,
+        help=f"Batch size override (default: {eval_cfg.batch_size}).",
+    )
+    parser.add_argument(
+        "--num-workers", type=int, default=None,
+        help=f"DataLoader workers override (default: {eval_cfg.num_workers}).",
+    )
+    parser.add_argument(
+        "--pattern", type=str, default="*.pt",
+        help="Glob pattern for --compare mode (default: '*.pt').",
+    )
+    parser.add_argument(
+        "--no-plots", action="store_true",
+        help="Skip generating per-scene plots.",
+    )
+
+    args = parser.parse_args()
+
+    # Resolve device
+    device = None
+    if args.device is not None:
+        device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
+
+    # Resolve output directory
+    output_dir = Path(args.output) if args.output else None
+
+    if args.checkpoint:
+        ckpt_path = Path(args.checkpoint)
+        if not ckpt_path.exists():
+            print(f"Error: checkpoint not found: {ckpt_path}")
+            sys.exit(1)
+        run_single_eval(
+            checkpoint_path=ckpt_path,
+            output_dir=output_dir,
+            device=device,
+            seq_len=args.seq_len,
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+            save_plots=not args.no_plots,
+        )
+    elif args.compare:
+        ckpt_dir = Path(args.compare)
+        if not ckpt_dir.is_dir():
+            print(f"Error: checkpoint directory not found: {ckpt_dir}")
+            sys.exit(1)
+        run_compare(
+            checkpoint_dir=ckpt_dir,
+            output_dir=output_dir,
+            device=device,
+            seq_len=args.seq_len,
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+            pattern=args.pattern,
+        )
+
+
+if __name__ == "__main__":
+    main()

benchmark/config.py

@@ -0,0 +1,98 @@
+"""
+Benchmark configuration — evaluation-only scene splits and metric definitions.
+
+This config is independent from src.velocity_prediction.config so that
+evaluation scenarios can be changed without touching training config.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Dict
+
+
+# ──────────────────────────── Dataset root ────────────────────────────
+
+DATASET_ROOT = Path(__file__).resolve().parents[1] / "dataset"
+
+# ──────────────────────────── Scene splits ────────────────────────────
+
+# Each scene group has a name, a list of scene dirs, and a difficulty label.
+# The test scenes are the primary evaluation set; val scenes are for
+# checkpoint selection reference.
+
+
+@dataclass
+class SceneGroup:
+    name: str
+    scenes: List[str]
+    difficulty: str = "medium"  # easy / medium / hard
+
+
+# ── Validation scenes (for checkpoint selection reference) ──
+VAL_SCENE_GROUPS: List[SceneGroup] = [
+    SceneGroup("indoor_forward_7", ["indoor_forward_7"], "hard"),
+    SceneGroup("outdoor_forward_1", ["outdoor_forward_1"], "easy"),
+    # SceneGroup("indoor_forward_6", ["indoor_forward_6"], "medium"),
+    # SceneGroup("indoor_forward_9", ["indoor_forward_9"], "easy"),
+    # SceneGroup("indoor_forward_10", ["indoor_forward_10"], "easy"),
+    # SceneGroup("indoor_forward_5", ["indoor_forward_5"], "medium"),
+]
+
+# ── Test scenes (primary evaluation) ──
+TEST_SCENE_GROUPS: List[SceneGroup] = [
+    SceneGroup("indoor_forward_7", ["indoor_forward_7"], "hard"),
+    SceneGroup("outdoor_forward_1", ["outdoor_forward_1"], "easy"),
+    SceneGroup("outdoor_forward_5", ["outdoor_forward_5"], "hard"),
+    SceneGroup("indoor_forward_6", ["indoor_forward_6"], "medium"),
+    SceneGroup("indoor_forward_9", ["indoor_forward_9"], "easy"),
+    SceneGroup("indoor_forward_10", ["indoor_forward_10"], "easy"),
+    SceneGroup("indoor_forward_5", ["indoor_forward_5"], "medium"),
+]
+
+# Flat lists for convenience
+VAL_SCENES: List[str] = [s for g in VAL_SCENE_GROUPS for s in g.scenes]
+TEST_SCENES: List[str] = [s for g in TEST_SCENE_GROUPS for s in g.scenes]
+
+# Difficulty grouping
+DIFFICULTY_GROUPS: Dict[str, List[str]] = {}
+for g in TEST_SCENE_GROUPS:
+    DIFFICULTY_GROUPS.setdefault(g.difficulty, []).extend(g.scenes)
+
+
+# ──────────────────────────── Evaluation parameters ────────────────────────────
+
+
+@dataclass
+class EvalConfig:
+    """Parameters used when running evaluation."""
+
+    # Sequence length (must match what the model was trained with)
+    seq_len: int = 8
+
+    # Batch size for evaluation (can be larger than training)
+    batch_size: int = 64
+
+    # Data loading
+    num_workers: int = 2
+
+    # Event simulation (must match training config)
+    event_threshold: float = 0.1
+    event_use_log: bool = True
+
+    # Output directory (relative to benchmark/results/)
+    output_dir: str = "results"
+
+    # Whether to generate per-scene plots
+    save_plots: bool = True
+
+    # Device override (None = auto-detect)
+    device: str = "cuda"
+
+
+# ──────────────────────────── Metrics definition ────────────────────────────
+
+# Metrics computed per-axis and overall
+METRICS = ["rmse", "mae", "r2"]
+
+# Singleton
+eval_cfg = EvalConfig()

benchmark/evaluate.py

@@ -0,0 +1,458 @@
+"""
+Core evaluation logic: run model on one or more scenes, compute metrics,
+generate visualizations, and save structured results.
+
+This module is called by benchmark.py (the user-facing entry point).
+"""
+
+import numpy as np
+from pathlib import Path
+from typing import List, Optional, Dict, Tuple
+
+import torch
+import torch.nn as nn
+
+from src.velocity_prediction.model import VelocityPredictionModel
+from src.velocity_prediction.dataset import create_val_loader
+from src.velocity_prediction.config import VELOCITY_MEAN, VELOCITY_STD
+
+from benchmark.config import (
+    eval_cfg,
+    TEST_SCENE_GROUPS,
+    VAL_SCENE_GROUPS,
+    DIFFICULTY_GROUPS,
+    DATASET_ROOT,
+)
+
+
+# ──────────────────────────── Metrics ────────────────────────────
+
+
+def compute_metrics(
+    pred: np.ndarray,
+    target: np.ndarray,
+) -> Dict[str, float]:
+    """
+    Compute RMSE, MAE, R² for each axis and overall.
+
+    Args:
+        pred:    (N, 2)  denormalized predictions
+        target:  (N, 2)  denormalized ground truth
+
+    Returns:
+        dict with keys like rmse_vx, rmse_vy, rmse_xy, mae_vx, ...
+    """
+    # Per-axis
+    rmse_x = float(np.sqrt(np.mean((pred[:, 0] - target[:, 0]) ** 2)))
+    rmse_y = float(np.sqrt(np.mean((pred[:, 1] - target[:, 1]) ** 2)))
+    rmse_xy = float(np.sqrt(np.mean(np.sum((pred - target) ** 2, axis=1))))
+
+    mae_x = float(np.mean(np.abs(pred[:, 0] - target[:, 0])))
+    mae_y = float(np.mean(np.abs(pred[:, 1] - target[:, 1])))
+    mae_xy = float(np.mean(np.sqrt(np.sum((pred - target) ** 2, axis=1))))
+
+    # R² per axis
+    def r2(p, t):
+        ss_res = np.sum((t - p) ** 2)
+        ss_tot = np.sum((t - np.mean(t)) ** 2)
+        return float(1 - ss_res / ss_tot) if ss_tot > 1e-12 else 0.0
+
+    r2_x = r2(pred[:, 0], target[:, 0])
+    r2_y = r2(pred[:, 1], target[:, 1])
+
+    return {
+        "rmse_vx": rmse_x,
+        "rmse_vy": rmse_y,
+        "rmse_xy": rmse_xy,
+        "mae_vx": mae_x,
+        "mae_vy": mae_y,
+        "mae_xy": mae_xy,
+        "r2_vx": r2_x,
+        "r2_vy": r2_y,
+        "count": len(pred),
+    }
+
+
+# ──────────────────────────── Per-scene evaluation ────────────────────────────
+
+
+@torch.no_grad()
+def evaluate_scene(
+    model: nn.Module,
+    scene_names: List[str],
+    device: torch.device,
+    seq_len: int = 8,
+    batch_size: int = 64,
+    num_workers: int = 2,
+    event_threshold: float = 0.1,
+    event_use_log: bool = True,
+) -> Dict:
+    """
+    Evaluate model on one or more scenes.
+
+    Returns:
+        dict with keys:
+            preds:      (N, 2) denormalized predictions
+            targets:    (N, 2) denormalized ground truth
+            metrics:    dict of scalar metrics
+    """
+    loader = create_val_loader(
+        scene_names=scene_names,
+        seq_len=seq_len,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        event_threshold=event_threshold,
+        event_use_log=event_use_log,
+    )
+
+    model.eval()
+    all_preds = []
+    all_targets = []
+
+    for batch in loader:
+        events = batch["events"].to(device)
+        tilt = batch["tilt"].to(device)
+        target = batch["v_body_target"].to(device)  # (B, S, 2) normalized
+
+        pred = model(events, tilt)                   # (B, 2) normalized
+        target_last = target[:, -1, :]               # (B, 2) normalized
+
+        all_preds.append(pred.cpu().numpy())
+        all_targets.append(target_last.cpu().numpy())
+
+    if not all_preds:
+        return {"preds": np.zeros((0, 2)), "targets": np.zeros((0, 2)), "metrics": {}}
+
+    preds = np.concatenate(all_preds, axis=0)
+    targets = np.concatenate(all_targets, axis=0)
+
+    # Denormalize
+    mean = np.array(VELOCITY_MEAN, dtype=np.float32)
+    std = np.array(VELOCITY_STD, dtype=np.float32)
+    preds_denorm = preds * std + mean
+    targets_denorm = targets * std + mean
+
+    metrics = compute_metrics(preds_denorm, targets_denorm)
+
+    return {
+        "preds": preds_denorm,
+        "targets": targets_denorm,
+        "metrics": metrics,
+    }
+
+
+# ──────────────────────────── Full evaluation suite ────────────────────────────
+
+
+def run_full_evaluation(
+    model: nn.Module,
+    device: torch.device,
+    seq_len: int = 8,
+    batch_size: int = 64,
+    num_workers: int = 2,
+    event_threshold: float = 0.1,
+    event_use_log: bool = True,
+    scene_groups=None,
+) -> Dict:
+    """
+    Run evaluation on all scene groups.
+
+    Returns nested dict:
+        {
+            "global": { metrics... },
+            "per_scene": {
+                "indoor_forward_7": { metrics..., "preds": ..., "targets": ... },
+                ...
+            },
+            "by_difficulty": {
+                "easy": { metrics... },
+                "hard": { metrics... },
+            }
+        }
+    """
+    if scene_groups is None:
+        from benchmark.config import TEST_SCENE_GROUPS
+        scene_groups = TEST_SCENE_GROUPS
+
+    per_scene = {}
+    all_preds = []
+    all_targets = []
+
+    for group in scene_groups:
+        for scene_name in group.scenes:
+            result = evaluate_scene(
+                model, [scene_name], device,
+                seq_len=seq_len, batch_size=batch_size,
+                num_workers=num_workers,
+                event_threshold=event_threshold,
+                event_use_log=event_use_log,
+            )
+            per_scene[scene_name] = result
+            if result["preds"].shape[0] > 0:
+                all_preds.append(result["preds"])
+                all_targets.append(result["targets"])
+
+    # Global metrics (all scenes combined)
+    if all_preds:
+        global_preds = np.concatenate(all_preds, axis=0)
+        global_targets = np.concatenate(all_targets, axis=0)
+        global_metrics = compute_metrics(global_preds, global_targets)
+    else:
+        global_preds = np.zeros((0, 2))
+        global_targets = np.zeros((0, 2))
+        global_metrics = {}
+
+    # By difficulty
+    by_difficulty = {}
+    for diff, scenes in DIFFICULTY_GROUPS.items():
+        diff_preds = []
+        diff_targets = []
+        for s in scenes:
+            if s in per_scene and per_scene[s]["preds"].shape[0] > 0:
+                diff_preds.append(per_scene[s]["preds"])
+                diff_targets.append(per_scene[s]["targets"])
+        if diff_preds:
+            by_difficulty[diff] = compute_metrics(
+                np.concatenate(diff_preds, axis=0),
+                np.concatenate(diff_targets, axis=0),
+            )
+        else:
+            by_difficulty[diff] = {}
+
+    return {
+        "global": global_metrics,
+        "per_scene": per_scene,
+        "by_difficulty": by_difficulty,
+    }
+
+
+# ──────────────────────────── Visualization ────────────────────────────
+
+
+def plot_scene_comparison(
+    preds: np.ndarray,
+    targets: np.ndarray,
+    scene_name: str,
+    save_dir: Path,
+    metrics: Optional[Dict] = None,
+):
+    """Generate time-series and scatter plots for a single scene."""
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    time = np.arange(len(preds))
+
+    # ── Time-series plot ──
+    fig, axes = plt.subplots(2, 1, figsize=(14, 5), sharex=True)
+
+    axes[0].plot(time, targets[:, 0], label="GT vx", color="C0", alpha=0.7, linewidth=0.8)
+    axes[0].plot(time, preds[:, 0], label="Pred vx", color="C1", alpha=0.7, linewidth=0.8)
+    axes[0].set_ylabel("vx (m/s)")
+    axes[0].legend(fontsize=9)
+    axes[0].grid(True, alpha=0.3)
+
+    axes[1].plot(time, targets[:, 1], label="GT vy", color="C0", alpha=0.7, linewidth=0.8)
+    axes[1].plot(time, preds[:, 1], label="Pred vy", color="C1", alpha=0.7, linewidth=0.8)
+    axes[1].set_ylabel("vy (m/s)")
+    axes[1].set_xlabel("Frame index")
+    axes[1].legend(fontsize=9)
+    axes[1].grid(True, alpha=0.3)
+
+    title = f"Body-frame Velocity — {scene_name}"
+    if metrics:
+        title += f"  |  RMSE vx={metrics['rmse_vx']:.3f}  vy={metrics['rmse_vy']:.3f}"
+    fig.suptitle(title)
+    try:
+        plt.tight_layout()
+        plt.savefig(save_dir / f"{scene_name}_timeseries.png", dpi=150, bbox_inches="tight")
+    except Exception as e:
+        print(f"  [WARN] Failed to save {scene_name}_timeseries.png: {e}")
+    plt.close()
+
+    # ── Scatter plot ──
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4.5))
+
+    for ax, pred, target, label in zip(
+        axes, [preds[:, 0], preds[:, 1]], [targets[:, 0], targets[:, 1]], ["vx", "vy"]
+    ):
+        ax.scatter(target, pred, s=3, alpha=0.4, c="C1", edgecolors="none")
+        lim_min = min(target.min(), pred.min())
+        lim_max = max(target.max(), pred.max())
+        margin = (lim_max - lim_min) * 0.05
+        ax.plot([lim_min - margin, lim_max + margin],
+                [lim_min - margin, lim_max + margin], "r--", alpha=0.5, linewidth=1)
+        ax.set_xlabel(f"GT {label} (m/s)")
+        ax.set_ylabel(f"Pred {label} (m/s)")
+        ax.set_aspect("equal")
+        ax.grid(True, alpha=0.3)
+        if metrics:
+            ax.set_title(f"{label} — RMSE: {metrics[f'rmse_{label}']:.4f}")
+
+    fig.suptitle(f"Scatter — {scene_name}")
+    try:
+        plt.tight_layout()
+        plt.savefig(save_dir / f"{scene_name}_scatter.png", dpi=150, bbox_inches="tight")
+    except Exception as e:
+        print(f"  [WARN] Failed to save {scene_name}_scatter.png: {e}")
+    plt.close()
+
+
+def plot_global_comparison(
+    results: Dict,
+    save_dir: Path,
+):
+    """Generate a summary figure comparing all scenes."""
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    scenes = list(results["per_scene"].keys())
+    if not scenes:
+        return
+
+    # Bar chart: RMSE vx and vy per scene
+    rmse_vx = [results["per_scene"][s]["metrics"].get("rmse_vx", 0) for s in scenes]
+    rmse_vy = [results["per_scene"][s]["metrics"].get("rmse_vy", 0) for s in scenes]
+    rmse_xy = [results["per_scene"][s]["metrics"].get("rmse_xy", 0) for s in scenes]
+
+    x = np.arange(len(scenes))
+    width = 0.25
+
+    fig, ax = plt.subplots(figsize=(10, 4.5))
+    bars1 = ax.bar(x - width, rmse_vx, width, label="RMSE vx", alpha=0.8)
+    bars2 = ax.bar(x, rmse_vy, width, label="RMSE vy", alpha=0.8)
+    bars3 = ax.bar(x + width, rmse_xy, width, label="RMSE xy", alpha=0.8)
+
+    ax.set_xticks(x)
+    ax.set_xticklabels(scenes, rotation=15, ha="right")
+    ax.set_ylabel("RMSE (m/s)")
+    ax.set_title("Per-Scene RMSE Comparison")
+    ax.legend(fontsize=9)
+    ax.grid(True, alpha=0.3, axis="y")
+
+    # Annotate values
+    for bars in [bars1, bars2, bars3]:
+        for bar in bars:
+            height = bar.get_height()
+            ax.annotate(f"{height:.3f}",
+                        xy=(bar.get_x() + bar.get_width() / 2, height),
+                        xytext=(0, 2), textcoords="offset points",
+                        ha="center", va="bottom", fontsize=7)
+
+    try:
+        plt.tight_layout()
+        plt.savefig(save_dir / "per_scene_rmse.png", dpi=150, bbox_inches="tight")
+    except Exception as e:
+        print(f"  [WARN] Failed to save per_scene_rmse.png: {e}")
+    plt.close()
+
+
+# ──────────────────────────── Results serialization ────────────────────────────
+
+
+def save_results(
+    results: Dict,
+    save_dir: Path,
+    checkpoint_name: str = "model",
+):
+    """Save all evaluation results to disk."""
+    save_dir = Path(save_dir)
+    plots_dir = save_dir / "plots"
+    save_dir.mkdir(parents=True, exist_ok=True)
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    # ── 1. Global metrics ──
+    global_m = results["global"]
+    lines = [
+        f"Benchmark Results: {checkpoint_name}",
+        f"{'=' * 50}",
+        f"Total samples: {global_m.get('count', '?')}",
+        "",
+        "── Global Metrics ──",
+        f"  RMSE vx:  {global_m.get('rmse_vx', 'N/A'):.4f}  m/s",
+        f"  RMSE vy:  {global_m.get('rmse_vy', 'N/A'):.4f}  m/s",
+        f"  RMSE xy:  {global_m.get('rmse_xy', 'N/A'):.4f}  m/s",
+        f"  MAE  vx:  {global_m.get('mae_vx', 'N/A'):.4f}  m/s",
+        f"  MAE  vy:  {global_m.get('mae_vy', 'N/A'):.4f}  m/s",
+        f"  MAE  xy:  {global_m.get('mae_xy', 'N/A'):.4f}  m/s",
+        f"  R²   vx:  {global_m.get('r2_vx', 'N/A'):.4f}",
+        f"  R²   vy:  {global_m.get('r2_vy', 'N/A'):.4f}",
+        "",
+    ]
+
+    # ── 2. Per-scene metrics ──
+    lines.append("── Per-Scene Metrics ──")
+    lines.append(f"  {'Scene':<22} {'RMSE vx':>10} {'RMSE vy':>10} {'RMSE xy':>10} "
+                 f"{'MAE vx':>10} {'MAE vy':>10} {'R² vx':>8} {'R² vy':>8} {'Samples':>8}")
+    lines.append("  " + "-" * 96)
+
+    for scene_name, scene_result in results["per_scene"].items():
+        m = scene_result["metrics"]
+        lines.append(
+            f"  {scene_name:<22} {m.get('rmse_vx', 0):>10.4f} {m.get('rmse_vy', 0):>10.4f} "
+            f"{m.get('rmse_xy', 0):>10.4f} {m.get('mae_vx', 0):>10.4f} "
+            f"{m.get('mae_vy', 0):>10.4f} {m.get('r2_vx', 0):>8.4f} "
+            f"{m.get('r2_vy', 0):>8.4f} {m.get('count', 0):>8}"
+        )
+
+    lines.append("")
+
+    # ── 3. By difficulty ──
+    lines.append("── By Difficulty ──")
+    for diff, metrics in results["by_difficulty"].items():
+        lines.append(f"  {diff:<10}  RMSE vx={metrics.get('rmse_vx', 0):.4f}  "
+                     f"RMSE vy={metrics.get('rmse_vy', 0):.4f}  "
+                     f"RMSE xy={metrics.get('rmse_xy', 0):.4f}  "
+                     f"(samples={metrics.get('count', 0)})")
+
+    summary_text = "\n".join(lines)
+    with open(save_dir / "summary.txt", "w") as f:
+        f.write(summary_text)
+    print(summary_text)
+
+    # ── 4. CSV: per-scene metrics ──
+    import csv
+    csv_path = save_dir / "per_scene_metrics.csv"
+    fieldnames = ["scene", "rmse_vx", "rmse_vy", "rmse_xy", "mae_vx", "mae_vy",
+                  "mae_xy", "r2_vx", "r2_vy", "count"]
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        for scene_name, scene_result in results["per_scene"].items():
+            row = {"scene": scene_name, **scene_result["metrics"]}
+            writer.writerow(row)
+    print(f"Per-scene CSV: {csv_path}")
+
+    # ── 5. Global metrics CSV (single row) ──
+    csv_path_global = save_dir / "metrics.csv"
+    with open(csv_path_global, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=["checkpoint"] + list(global_m.keys()))
+        writer.writeheader()
+        row = {"checkpoint": checkpoint_name, **global_m}
+        writer.writerow(row)
+    print(f"Global metrics CSV: {csv_path_global}")
+
+    # ── 6. Plots ──
+    for scene_name, scene_result in results["per_scene"].items():
+        if scene_result["preds"].shape[0] > 0:
+            plot_scene_comparison(
+                scene_result["preds"],
+                scene_result["targets"],
+                scene_name,
+                plots_dir,
+                metrics=scene_result["metrics"],
+            )
+
+    plot_global_comparison(results, plots_dir)
+
+    import matplotlib.pyplot as plt
+    plt.close("all")
+
+    print(f"\nAll results saved to: {save_dir.resolve()}")

download_davis_gt_rosbags.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Download script for UZH FPV DAVIS rosbags with ground truth
+# Skips files that already exist in the current directory.
+# Uses the final download.ifi.uzh.ch URLs directly (avoids 301 redirect chain).
+set -euo pipefail
+
+BASE="https://download.ifi.uzh.ch/rpg/web/datasets/uzh-fpv-newer-versions/v3"
+
+URLS=(
+    # Indoor forward facing
+    "${BASE}/indoor_forward_3_davis_with_gt.bag"
+    "${BASE}/indoor_forward_5_davis_with_gt.bag"
+    "${BASE}/indoor_forward_6_davis_with_gt.bag"
+    "${BASE}/indoor_forward_9_davis_with_gt.bag"
+    "${BASE}/indoor_forward_10_davis_with_gt.bag"
+
+    # Indoor 45 degree downward
+    "${BASE}/indoor_45_2_davis_with_gt.bag"
+    "${BASE}/indoor_45_4_davis_with_gt.bag"
+    "${BASE}/indoor_45_9_davis_with_gt.bag"
+    "${BASE}/indoor_45_12_davis_with_gt.bag"
+    "${BASE}/indoor_45_13_davis_with_gt.bag"
+    "${BASE}/indoor_45_14_davis_with_gt.bag"
+
+    # Outdoor forward facing
+    "${BASE}/outdoor_forward_1_davis_with_gt.bag"
+    "${BASE}/outdoor_forward_5_davis_with_gt.bag"
+
+    # Outdoor 45 degree downward
+    "${BASE}/outdoor_45_1_davis_with_gt.bag"
+)
+
+for url in "${URLS[@]}"; do
+    filename=$(basename "$url")
+    if [ -f "$filename" ]; then
+        echo "SKIP: $filename already exists"
+    else
+        echo "DOWNLOAD: $filename"
+        wget --continue --show-progress "$url" -O "$filename"
+    fi
+done
+
+echo ""
+echo "Done. All DAVIS with-GT rosbags downloaded."

requirements.txt

@@ -0,0 +1,17 @@
+# ===== Core =====
+numpy
+scipy
+
+# ===== PyTorch (CUDA 12.4) =====
+# Install via: pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
+torch>=2.4
+torchvision
+
+# ===== Data loading =====
+webdataset
+opencv-python
+
+# ===== Training / logging =====
+tensorboard
+tqdm
+matplotlib

rosbag2wds.py

@@ -0,0 +1,396 @@
+#!/usr/bin/env python3
+"""
+ROS bag to WebDataset converter for DAVIS dataset
+Extracts: grayscale images, IMU sequence, ground truth poses and velocities
+
+Usage:
+    python convert_bag_to_webdataset.py --bag <path_to.bag> --output <output_dir> --name <dataset_name>
+"""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+import numpy as np
+import rosbag
+from cv_bridge import CvBridge
+import cv2
+import webdataset as wds
+from tqdm import tqdm
+from scipy.spatial.transform import Rotation as R
+
+
+class BagToWebDataset:
+    def __init__(self, bag_path: str, output_dir: str, dataset_name: str, 
+                 shard_size: int = 2000, image_width: int = 320, image_height: int = 240):
+        self.bag_path = Path(bag_path)
+        self.output_dir = Path(output_dir) / dataset_name
+        self.dataset_name = dataset_name
+        self.shard_size = shard_size
+        self.image_width = image_width
+        self.image_height = image_height
+        
+        self.bridge = CvBridge()
+        
+        # Data containers
+        self.images: List[Tuple[float, np.ndarray]] = []      # (timestamp, image)
+        self.imu_timestamps: List[float] = []
+        self.imu_acc: List[np.ndarray] = []                   # (ax, ay, az)
+        self.imu_gyro: List[np.ndarray] = []                  # (gx, gy, gz)
+        self.gt_timestamps: List[float] = []
+        self.gt_poses: List[np.ndarray] = []                  # (x, y, z, qx, qy, qz, qw)
+        self.gt_velocities: List[np.ndarray] = []             # (vx, vy, vz, wx, wy, wz)
+        
+    def extract_all_data(self):
+        """Extract all data from ROS bag"""
+        print(f"Opening bag: {self.bag_path}")
+        bag = rosbag.Bag(str(self.bag_path), 'r')
+        
+        # Count messages for progress bar
+        topic_counts = {topic: bag.get_message_count(topic) for topic in 
+                       ['/dvs/image_raw', '/dvs/imu', '/groundtruth/odometry']}
+        total_msgs = sum(topic_counts.values())
+        
+        print(f"Topics: {topic_counts}")
+        
+        with tqdm(total=total_msgs, desc="Extracting messages") as pbar:
+            for topic, msg, t in bag.read_messages(topics=['/dvs/image_raw', '/dvs/imu', '/groundtruth/odometry']):
+                
+                if topic == '/dvs/image_raw':
+                    self._process_image(msg, t)  # 传入 t
+                    
+                elif topic == '/dvs/imu':
+                    self._process_imu(msg, t)    # 传入 t
+                    
+                elif topic == '/groundtruth/odometry':
+                    self._process_odometry(msg, t)  # 传入 t
+                
+                pbar.update(1)
+        
+        bag.close()
+        
+        # Post-processing: compute velocities from poses if not directly available
+        self._ensure_velocities()
+        
+        # Print statistics
+        print(f"\nExtraction completed:")
+        print(f"  Images: {len(self.images)}")
+        print(f"  IMU messages: {len(self.imu_timestamps)}")
+        print(f"  Ground truth poses: {len(self.gt_timestamps)}")
+        print(f"  Ground truth velocities: {len(self.gt_velocities)}")
+        
+    def crop_to_gt_time_range(self):
+        """裁剪所有数据，只保留 GT 时间范围内的部分"""
+        
+        if len(self.gt_timestamps) == 0:
+            print("Warning: No GT data found, skipping crop")
+            return
+        
+        gt_start = min(self.gt_timestamps)
+        gt_end = max(self.gt_timestamps)
+        
+        print(f"\nCropping to GT time range: {gt_start:.3f} - {gt_end:.3f} ({gt_end - gt_start:.1f}s)")
+        
+        # 裁剪图像
+        original_img_count = len(self.images)
+        self.images = [(ts, img) for ts, img in self.images if gt_start <= ts <= gt_end]
+        print(f"  Images: {original_img_count} -> {len(self.images)}")
+        
+        # 裁剪 IMU
+        original_imu_count = len(self.imu_timestamps)
+        imu_filtered = [(ts, acc, gyro) for ts, acc, gyro 
+                        in zip(self.imu_timestamps, self.imu_acc, self.imu_gyro)
+                        if gt_start <= ts <= gt_end]
+        
+        if imu_filtered:
+            self.imu_timestamps = [item[0] for item in imu_filtered]
+            self.imu_acc = [item[1] for item in imu_filtered]
+            self.imu_gyro = [item[2] for item in imu_filtered]
+        print(f"  IMU: {original_imu_count} -> {len(self.imu_timestamps)}")
+        
+        # GT 数据本身已经在范围内，不需要裁剪
+        print(f"  GT: {len(self.gt_timestamps)} (unchanged)")
+
+    def _process_image(self, msg, t):
+        """Process grayscale image message using system time"""
+        try:
+            # Convert ROS image to OpenCV format
+            cv_img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='mono8')
+            
+            # Resize if needed
+            if self.image_width and self.image_height:
+                cv_img = cv2.resize(cv_img, (self.image_width, self.image_height), 
+                                   interpolation=cv2.INTER_LINEAR)
+            
+            # 使用系统物理时间，而不是 msg.header.stamp
+            timestamp = t.to_sec()
+            self.images.append((timestamp, cv_img))
+            
+        except Exception as e:
+            print(f"Error processing image: {e}")
+    
+    def _process_imu(self, msg, t):
+        """Process IMU message using system time"""
+        timestamp = t.to_sec()  # 使用系统物理时间
+        
+        # Linear acceleration (m/s^2)
+        acc = np.array([msg.linear_acceleration.x,
+                        msg.linear_acceleration.y,
+                        msg.linear_acceleration.z], dtype=np.float32)
+        
+        # Angular velocity (rad/s)
+        gyro = np.array([msg.angular_velocity.x,
+                         msg.angular_velocity.y,
+                         msg.angular_velocity.z], dtype=np.float32)
+        
+        self.imu_timestamps.append(timestamp)
+        self.imu_acc.append(acc)
+        self.imu_gyro.append(gyro)
+    
+    def _process_odometry(self, msg, t):
+        """Process ground truth odometry using system time"""
+        timestamp = t.to_sec()  # 使用系统物理时间
+        
+        # Position (x, y, z)
+        pos = np.array([msg.pose.pose.position.x,
+                        msg.pose.pose.position.y,
+                        msg.pose.pose.position.z], dtype=np.float32)
+        
+        # Orientation (qx, qy, qz, qw) - already normalized
+        quat = np.array([msg.pose.pose.orientation.x,
+                         msg.pose.pose.orientation.y,
+                         msg.pose.pose.orientation.z,
+                         msg.pose.pose.orientation.w], dtype=np.float32)
+        
+        pose = np.concatenate([pos, quat])
+        self.gt_timestamps.append(timestamp)
+        self.gt_poses.append(pose)
+        
+        # Velocity: always compute from pose differences in post-processing
+        vel = None
+        
+        self.gt_velocities.append(vel)
+    
+    def _ensure_velocities(self):
+        # 数据集中 twist 数据为 0 直接利用时间戳差值
+        # """Compute velocities from pose differences if not directly available"""
+        # # Check if velocities are missing
+        # missing_velocities = any(v is None for v in self.gt_velocities)
+        
+        # if not missing_velocities:
+        #     return
+        
+        print("Computing velocities from pose differences...")
+        
+        computed_velocities = []
+        for i in range(len(self.gt_timestamps)):
+            if i == 0:
+                # Use forward difference for first frame
+                if len(self.gt_timestamps) > 1:
+                    dt = self.gt_timestamps[1] - self.gt_timestamps[0]
+                    if dt > 0:
+                        # Linear velocity
+                        v_lin = (self.gt_poses[1][:3] - self.gt_poses[0][:3]) / dt
+                        # Angular velocity (from quaternion difference)
+                        q0 = self.gt_poses[0][3:7]
+                        q1 = self.gt_poses[1][3:7]
+                        dq = R.from_quat(q1) * R.from_quat(q0).inv()
+                        v_ang = dq.as_rotvec() / dt
+                        computed_velocities.append(np.concatenate([v_lin, v_ang]))
+                    else:
+                        computed_velocities.append(np.zeros(6, dtype=np.float32))
+                else:
+                    computed_velocities.append(np.zeros(6, dtype=np.float32))
+            else:
+                # Use backward difference
+                dt = self.gt_timestamps[i] - self.gt_timestamps[i-1]
+                if dt > 0:
+                    v_lin = (self.gt_poses[i][:3] - self.gt_poses[i-1][:3]) / dt
+                    q0 = self.gt_poses[i-1][3:7]
+                    q1 = self.gt_poses[i][3:7]
+                    dq = R.from_quat(q1) * R.from_quat(q0).inv()
+                    v_ang = dq.as_rotvec() / dt
+                    computed_velocities.append(np.concatenate([v_lin, v_ang]))
+                else:
+                    computed_velocities.append(np.zeros(6, dtype=np.float32))
+        
+        # Replace missing velocities
+        for i in range(len(self.gt_velocities)):
+            if self.gt_velocities[i] is None:
+                self.gt_velocities[i] = computed_velocities[i]
+    
+    def save_imu_sequence(self):
+        """Save IMU sequence as NPZ file"""
+        imu_data = {
+            'timestamps': np.array(self.imu_timestamps, dtype=np.float64),
+            'accelerations': np.array(self.imu_acc, dtype=np.float32),
+            'angular_velocities': np.array(self.imu_gyro, dtype=np.float32)
+        }
+        
+        imu_path = self.output_dir / 'imu_sequence.npz'
+        imu_path.parent.mkdir(parents=True, exist_ok=True)
+        np.savez_compressed(imu_path, **imu_data)
+        
+        print(f"Saved IMU sequence: {imu_path}")
+        return imu_path
+    
+    def align_ground_truth_to_images(self) -> List[Tuple[float, np.ndarray, np.ndarray, np.ndarray]]:
+        """Align ground truth (pose + velocity) to each image using nearest timestamp"""
+        aligned_gt = []
+        
+        gt_timestamps = np.array(self.gt_timestamps)
+        gt_poses = np.array(self.gt_poses)
+        gt_vels = np.array(self.gt_velocities)
+        
+        for img_ts, img in tqdm(self.images, desc="Aligning ground truth to images"):
+            idx = np.argmin(np.abs(gt_timestamps - img_ts))
+            time_diff = abs(gt_timestamps[idx] - img_ts)
+            
+            if time_diff < 0.1:
+                aligned_gt.append((img_ts, img, gt_poses[idx], gt_vels[idx]))  # 保存图像
+        
+        return aligned_gt
+    
+    def save_as_webdataset(self, aligned_gt: List[Tuple[float, np.ndarray, np.ndarray, np.ndarray]]):
+        """Save images and aligned ground truth as WebDataset tar files"""
+        
+        num_shards = (len(aligned_gt) + self.shard_size - 1) // self.shard_size
+        
+        print(f"Saving {len(aligned_gt)} samples into {num_shards} shards...")
+        
+        for shard_idx in range(num_shards):
+            start_idx = shard_idx * self.shard_size
+            end_idx = min((shard_idx + 1) * self.shard_size, len(aligned_gt))
+            
+            tar_path = self.output_dir / f'shard_{shard_idx:04d}.tar'
+            
+            with wds.TarWriter(str(tar_path)) as sink:
+                for local_idx, (img_ts, img, pose, vel) in enumerate(aligned_gt):
+
+                    # Encode image as JPEG
+                    _, img_encoded = cv2.imencode('.jpg', img, 
+                                                   [cv2.IMWRITE_JPEG_QUALITY, 85])
+                    img_bytes = img_encoded.tobytes()
+                    
+                    # Prepare metadata
+                    sample_key = f'frame_{local_idx:08d}'
+                    
+                    # Write to tar
+                    sink.write({
+                        '__key__': sample_key,
+                        'jpg': img_bytes,
+                        'ts': np.array([img_ts], dtype=np.float64).tobytes(),
+                        'pose': pose.astype(np.float32).tobytes(),
+                        'vel': vel.astype(np.float32).tobytes()
+                    })
+            
+            print(f"  Saved {tar_path} ({end_idx - start_idx} samples)")
+    
+    def save_metadata(self):
+        """Save dataset metadata"""
+        metadata = {
+            'dataset_name': self.dataset_name,
+            'source_bag': str(self.bag_path),
+            'num_images': len(self.images),
+            'num_imu_messages': len(self.imu_timestamps),
+            'num_ground_truth': len(self.gt_timestamps),
+            'image_size': [self.image_width, self.image_height],
+            'imu_frequency_hz': len(self.imu_timestamps) / (self.imu_timestamps[-1] - self.imu_timestamps[0]) if len(self.imu_timestamps) > 1 else 0,
+            'camera_frequency_hz': len(self.images) / (self.images[-1][0] - self.images[0][0]) if len(self.images) > 1 else 0,
+            'gt_frequency_hz': len(self.gt_timestamps) / (self.gt_timestamps[-1] - self.gt_timestamps[0]) if len(self.gt_timestamps) > 1 else 0,
+            'coordinate_system': 'horizontal (z aligned with gravity, assumed from GT)',
+            'velocity_dimensions': 6,  # (vx, vy, vz, wx, wy, wz)
+        }
+        
+        metadata_path = self.output_dir / 'metadata.json'
+        with open(metadata_path, 'w') as f:
+            json.dump(metadata, f, indent=2)
+        
+        print(f"Saved metadata: {metadata_path}")
+    
+    def convert(self):
+        """Main conversion pipeline"""
+        print(f"\n{'='*60}")
+        print(f"Converting: {self.bag_path.name}")
+        print(f"Output: {self.output_dir}")
+        print(f"{'='*60}\n")
+        
+        # Step 1: Extract all data from bag
+        self.extract_all_data()
+        
+        # 裁剪掉无 GT 的时间段
+        self.crop_to_gt_time_range()
+
+        # Step 2: Save IMU sequence
+        self.save_imu_sequence()
+        
+        # # Step 3: Align ground truth to images
+        aligned_gt = self.align_ground_truth_to_images()
+        
+        if len(aligned_gt) == 0:
+            print("Error: No aligned ground truth found!")
+            sys.exit(1)
+        
+        # # Step 4: Save as WebDataset
+        self.save_as_webdataset(aligned_gt)
+        
+        # # Step 5: Save metadata
+        self.save_metadata()
+
+        self.diagnose_timestamps()
+        
+        print(f"\n✅ Conversion completed for {self.bag_path.name}")
+
+
+
+    def diagnose_timestamps(self):
+        """Print timestamp ranges for debugging"""
+        img_timestamps = [t for t, _ in self.images]
+        gt_timestamps = self.gt_timestamps
+        
+        print(f"Image timestamps: {min(img_timestamps):.3f} - {max(img_timestamps):.3f}")
+        print(f"GT timestamps:    {min(gt_timestamps):.3f} - {max(gt_timestamps):.3f}")
+        print(f"Image duration: {max(img_timestamps) - min(img_timestamps):.3f}s")
+        print(f"GT duration:    {max(gt_timestamps) - min(gt_timestamps):.3f}s")
+        
+        # Check if there's a constant offset
+        if len(img_timestamps) > 0 and len(gt_timestamps) > 0:
+            offset = gt_timestamps[0] - img_timestamps[0]
+            print(f"Initial offset (first GT - first image): {offset:.3f}s")
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert ROS bag to WebDataset format')
+    parser.add_argument('--bag', type=str, required=True, help='Path to ROS bag file')
+    parser.add_argument('--output', type=str, default='./dataset', help='Output directory')
+    parser.add_argument('--name', type=str, default=None, help='Dataset name (default: bag filename without extension)')
+    parser.add_argument('--shard_size', type=int, default=2000, help='Number of samples per shard')
+    parser.add_argument('--width', type=int, default=320, help='Image width (resize)')
+    parser.add_argument('--height', type=int, default=240, help='Image height (resize)')
+    
+    args = parser.parse_args()
+    
+    # Validate inputs
+    if not os.path.exists(args.bag):
+        print(f"Error: Bag file not found: {args.bag}")
+        sys.exit(1)
+    
+    # Set dataset name
+    if args.name is None:
+        args.name = Path(args.bag).stem
+    
+    # Run conversion
+    converter = BagToWebDataset(
+        bag_path=args.bag,
+        output_dir=args.output,
+        dataset_name=args.name,
+        shard_size=args.shard_size,
+        image_width=args.width,
+        image_height=args.height
+    )
+    
+    converter.convert()
+
+if __name__ == '__main__':
+    main()

src/event_utils.py

@@ -0,0 +1,103 @@
+"""
+Event camera simulation utilities for ML preprocessing.
+
+Core logic extracted from EventCameraSimulator (test.py).
+Designed for frame-by-frame preprocessing in training pipelines.
+
+Output:
+    events_binary:  (-1, 0, +1)  hard threshold decision
+    events_strength:  [-1, 1]    continuous change intensity (clipped & normalized)
+"""
+
+import numpy as np
+from collections import deque
+
+
+class EventProcessor:
+    """Lightweight event computation module. No visualization, no OpenCV dependency."""
+
+    def __init__(self, threshold=0.1, use_log=True, auto_threshold=False):
+        self.threshold = threshold
+        self.use_log = use_log
+        self.auto_threshold = auto_threshold
+
+        self.prev_brightness = None
+        self.change_history = deque(maxlen=100)
+        self.threshold_scale = 1.5
+
+    def reset(self):
+        """Clear temporal state (call on video/reset)."""
+        self.prev_brightness = None
+        self.change_history.clear()
+
+    def _to_grayscale(self, frame):
+        """Convert frame to grayscale float32."""
+        if frame.ndim == 3:
+            # RGB/HWC -> gray via luminance weights
+            gray = 0.299 * frame[..., 0] + 0.587 * frame[..., 1] + 0.114 * frame[..., 2]
+        else:
+            gray = frame
+        return gray.astype(np.float32)
+
+    def _compute_change(self, brightness):
+        """Compute log or linear brightness change."""
+        if self.use_log:
+            eps = 1e-3
+            return np.log(brightness + eps) - np.log(self.prev_brightness + eps)
+        else:
+            return brightness - self.prev_brightness
+
+    def _update_auto_threshold(self, change):
+        """Adapt threshold based on global change statistics."""
+        abs_change = np.abs(change)
+        mean_change = np.mean(abs_change)
+        self.change_history.append(mean_change)
+
+        if len(self.change_history) > 10:
+            avg_change = np.mean(self.change_history)
+            new_threshold = max(avg_change * self.threshold_scale, 0.01)
+            self.threshold = self.threshold * 0.9 + new_threshold * 0.1
+
+            if self.use_log:
+                self.threshold = np.clip(self.threshold, 0.01, 0.5)
+            else:
+                self.threshold = np.clip(self.threshold, 1, 50)
+
+    def __call__(self, frame):
+        """
+        Process a single frame.
+
+        Args:
+            frame: np.ndarray, shape (H, W) or (H, W, C), uint8 or float.
+
+        Returns:
+            events_binary:   np.ndarray (H, W), values in {-1, 0, +1}
+            events_strength: np.ndarray (H, W), values in [-1, 1]
+            event_count:     int, number of non-zero events
+        """
+        brightness = self._to_grayscale(frame)
+
+        # First frame — initialise, no events
+        if self.prev_brightness is None:
+            self.prev_brightness = brightness
+            h, w = brightness.shape
+            return np.zeros((h, w), dtype=np.int8), np.zeros((h, w), dtype=np.float32), 0
+
+        change = self._compute_change(brightness)
+
+        if self.auto_threshold:
+            self._update_auto_threshold(change)
+
+        # Binary events
+        events_binary = np.zeros_like(brightness, dtype=np.int8)
+        events_binary[change > self.threshold] = 1
+        events_binary[change < -self.threshold] = -1
+
+        # Continuous strength: clip to [-threshold, threshold] then normalise to [-1, 1]
+        events_strength = np.clip(change, -self.threshold, self.threshold) / self.threshold
+
+        event_count = int(np.count_nonzero(events_binary))
+
+        self.prev_brightness = brightness
+
+        return events_binary, events_strength, event_count

src/velocity_prediction/README.md

@@ -0,0 +1,198 @@
+# Velocity Prediction from Event Frames + Attitude
+
+基于 UZH-FPV 数据集，通过模拟事件帧 + 姿态输入预测机体速度（机体系 vx, vy）。
+
+---
+
+## 项目结构
+
+```
+uzh_fpv/
+├── dataset/                          # UZH-FPV 数据集（WebDataset shards）
+│   ├── indoor_forward_3/
+│   ├── indoor_forward_5/
+│   ├── ...
+│   └── outdoor_45_1/
+├── src/
+│   ├── event_utils.py                # 模拟事件帧生成（已有模块）
+│   └── velocity_prediction/          # 本工程
+│       ├── __init__.py
+│       ├── config.py                 # 全局配置
+│       ├── utils.py                  # 四元数运算、坐标变换
+│       ├── transforms.py             # 数据预处理管线
+│       ├── dataset.py                # WebDataset 加载 + 序列采样
+│       ├── model.py                  # 网络模型定义
+│       ├── train.py                  # 训练入口
+│       └── evaluate.py               # 评估与可视化
+├── DATASET_FORMAT.md                 # 数据集格式说明
+└── requirements.txt
+```
+
+---
+
+## 网络架构
+
+```
+事件帧 (1, 240, 320) ──► CNN (4层 Conv+Pool+GAP, 256-d)
+                                                      │
+姿态 tilt_angles (3,) ──► PoseMLP (3→32→64, 64-d) ───┤
+                                                      │
+                                              concat (320-d)  ← 每帧融合
+                                                      │
+                                              GRU (hidden=128)
+                                                      │
+                                              Head MLP (128→64→2)
+                                                      │
+                                              [vx_body, vy_body]
+```
+
+### 各模块参数
+
+| 模块 | 参数量 | 说明 |
+|------|--------|------|
+| CNN Encoder | 387,840 | 4 层 Conv2D(3×3) + BN + ReLU + MaxPool(2×2) + GAP，通道 1→32→64→128→256 |
+| PoseMLP | 2,240 | 3→32→64，两层全连接 |
+| GRU | 172,800 | 单层，input=320, hidden=128 |
+| Head MLP | 8,386 | 128→64→2 |
+| **总计** | **~571K** | FP32 约 2.3 MB |
+
+---
+
+## 数据预处理
+
+### 输入变换管线（transforms.py）
+
+每个 WebDataset 样本依次经过：
+
+1. **DecodeSample** — JPEG 解码为灰度图 (H, W)，pose/vel 字节转 numpy
+2. **SimulateEvents** — `EventProcessor` 计算帧间亮度变化，输出二值事件帧 (1, H, W)，值域 {-1, 0, 1}
+3. **ComputeTilt** — 从四元数 `[qx, qy, qz, qw]` 中提取偏航角 yaw，移除后得到 tilt 旋转向量 (3,)
+4. **ComputeBodyVelocity** — 世界系速度 `[vx, vy, vz]` → 补偿偏航 → 转到机体系，取 `[vx_body, vy_body]` (2,)
+
+### 坐标系变换逻辑（utils.py）
+
+```
+输入: q_world_to_body (四元数), v_world (3,)
+
+Step 1: 从四元数分解偏航角 yaw
+Step 2: 构造纯偏航四元数 q_yaw
+Step 3: q_tilt = q_yaw^{-1} * q_world_to_body  →  tilt_angles (旋转向量)
+Step 4: v_yaw_comp = q_yaw^{-1} * v_world       →  偏航补偿
+Step 5: v_body = q_tilt^{-1} * v_yaw_comp       →  转到机体系
+Step 6: 取 v_body[:2] 作为回归目标
+```
+
+### 序列采样（dataset.py）
+
+- 从 shard 中取连续 `seq_len` 帧（默认 8 帧）构成一个训练样本
+- 输出 batch 维度：`(B, S, 1, H, W)` 事件帧, `(B, S, 3)` tilt, `(B, S, 2)` 速度 GT
+- 模型预测最后一帧的速度
+
+### 数据集划分
+
+| 集 | 场景 |
+|----|------|
+| **训练** | indoor_forward_3/5/6/7/9/10, indoor_45_2/4/9/12 |
+| **验证** | indoor_45_13/14 |
+| **测试** | outdoor_forward_1/3/5, outdoor_45_1 |
+
+---
+
+## 运行方式
+
+### 1. 安装依赖
+
+```bash
+# 使用 uv（推荐）
+uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+uv pip install webdataset opencv-python matplotlib tensorboard numpy scipy
+
+# 或使用 pip
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install webdataset opencv-python matplotlib tensorboard numpy scipy
+```
+
+### 2. 训练
+
+```bash
+# 从项目根目录执行
+cd /home/hexone/Workplace/ws_asmo/uzh_fpv
+
+# 激活虚拟环境后
+python -m src.velocity_prediction.train
+```
+
+训练参数在 `config.py` 的 `TrainConfig` 中配置，关键参数：
+
+| 参数 | 默认值 | 说明 |
+|------|--------|------|
+| seq_len | 8 | 每序列帧数 |
+| batch_size | 32 | 批次大小 |
+| epochs | 100 | 训练轮数 |
+| lr | 1e-3 | 学习率 |
+| event_threshold | 0.1 | 事件模拟阈值 |
+
+训练输出：
+- `logs/` — TensorBoard 日志
+- `checkpoints/` — 模型 checkpoint（每 10 轮保存 + 最优模型 `best.pt`）
+
+### 3. 评估
+
+```bash
+python -m src.velocity_prediction.evaluate --checkpoint checkpoints/best.pt
+```
+
+输出：
+- 控制台打印 RMSE（vx, vy, xy）
+- `eval_velocity.png` — 预测 vs GT 时序对比图
+- `eval_scatter.png` — 散点图
+
+### 4. 模型参数量检查
+
+```bash
+python -m src.velocity_prediction.model
+```
+
+输出类似：
+```
+Total trainable parameters: 571,266 (0.571 M)
+```
+
+---
+
+## 模型导出与部署
+
+### 导出 TorchScript
+
+```python
+import torch
+from src.velocity_prediction.model import VelocityPredictionModel
+
+model = VelocityPredictionModel()
+model.load_state_dict(torch.load("checkpoints/best.pt")["model_state_dict"])
+model.eval()
+
+# 导出
+traced = torch.jit.trace(model, (torch.randn(1, 8, 1, 240, 320), torch.randn(1, 8, 3)))
+traced.save("velocity_model.pt")
+```
+
+### RV1106 部署注意事项
+
+- 模型 ~0.57M 参数，FP32 ~2.3 MB，INT8 量化后 ~0.6 MB
+- CNN 部分可跑 NPU（0.5 TOPS），GRU 需 ARM CPU 执行
+- 若需裁剪：`config.py` 中 `CNNConfig.channels` 减半（32→16, 64→32, 128→64, 256→128），或 `GRUConfig.hidden_size` 从 128 降至 64
+
+---
+
+## 文件职责速查
+
+| 文件 | 职责 | 关键类/函数 |
+|------|------|------------|
+| `config.py` | 所有可配置参数 | `ModelConfig`, `TrainConfig`, `TRAIN_SCENES` |
+| `utils.py` | 四元数运算、坐标变换 | `decompose_tilt`, `world_vel_to_body` |
+| `transforms.py` | 数据预处理管线 | `DecodeSample`, `SimulateEvents`, `ComputeTilt`, `ComputeBodyVelocity` |
+| `dataset.py` | WebDataset 加载 + 序列采样 | `create_train_loader`, `create_val_loader` |
+| `model.py` | 网络模型 | `VelocityPredictionModel`, `CNNEncoder`, `PoseMLP` |
+| `train.py` | 训练循环 | `train_one_epoch`, `validate`, `main` |
+| `evaluate.py` | 评估与可视化 | `evaluate`, `plot_results`, `plot_scatter` |

src/velocity_prediction/__init__.py

@@ -0,0 +1,7 @@
+"""
+Velocity prediction from simulated event frames + attitude.
+
+Pipeline:
+    Event frame (1, H, W) ──► CNN ──┐
+    Tilt angles    (3,)    ──► MLP ──┤──► concat ──► GRU ──► Head ──► [vx_body, vy_body]
+"""

src/velocity_prediction/config.py

@@ -0,0 +1,118 @@
+"""
+Global configuration for velocity prediction.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+# ──────────────────────────── Dataset paths ────────────────────────────
+
+DATASET_ROOT = Path(__file__).resolve().parents[2] / "dataset"
+
+# Velocity normalization stats (computed from training set)
+VELOCITY_MEAN = [0.859184, -0.783945]   # [vx, vy]
+VELOCITY_STD  = [2.244513,  1.088335]   # [vx, vy]
+
+# TRAIN_SCENES = [
+#     "indoor_forward_3", "indoor_forward_5", "indoor_forward_6",
+#     "indoor_forward_7", "indoor_forward_9", "indoor_forward_10",
+#     "indoor_45_2", "indoor_45_4", "indoor_45_9", "indoor_45_12",
+# ]
+# VAL_SCENES = [
+#     "indoor_45_13", "indoor_45_14",
+# ]
+# TEST_SCENES = [
+#     "outdoor_forward_1", "outdoor_forward_3", "outdoor_forward_5",
+#     "outdoor_45_1",
+# ]
+
+TRAIN_SCENES = [
+    "indoor_forward_3", "indoor_forward_9", "indoor_forward_10",  # Easy
+    "indoor_forward_5", "indoor_forward_6",                       # Medium
+    "outdoor_forward_3"                                           # Medium 室外
+]
+VAL_SCENES = [
+    "indoor_forward_7",      # Hard 室内
+    "outdoor_forward_1"      # Easy 室外
+    # "indoor_forward_3", "indoor_forward_9", "indoor_forward_10",  # Easy
+]
+TEST_SCENES = [
+    "indoor_forward_7",      # Hard 室内
+    "outdoor_forward_1",     # Easy 室外
+    "outdoor_forward_5"      # Hard 室外
+    # "indoor_forward_3", "indoor_forward_9", "indoor_forward_10",  # Easy
+]
+
+
+# ──────────────────────────── Model architecture ────────────────────────────
+
+@dataclass
+class CNNConfig:
+    in_channels: int = 1
+    channels: tuple = (32, 64, 128, 256)       # per-layer output channels
+    kernel_size: int = 3
+    pool_size: int = 2
+    use_bn: bool = True
+
+
+@dataclass
+class PoseMLPConfig:
+    input_dim: int = 3
+    hidden_dim: int = 32
+    output_dim: int = 64
+
+
+@dataclass
+class GRUConfig:
+    input_size: int = 320       # CNN(256) + PoseMLP(64)
+    hidden_size: int = 128
+    num_layers: int = 1
+    dropout: float = 0.0
+
+
+@dataclass
+class HeadConfig:
+    input_dim: int = 128        # GRU hidden_size
+    hidden_dim: int = 64
+    output_dim: int = 2         # [vx_body, vy_body]
+
+
+@dataclass
+class ModelConfig:
+    cnn: CNNConfig = field(default_factory=CNNConfig)
+    pose_mlp: PoseMLPConfig = field(default_factory=PoseMLPConfig)
+    gru: GRUConfig = field(default_factory=GRUConfig)
+    head: HeadConfig = field(default_factory=HeadConfig)
+
+
+# ──────────────────────────── Training ────────────────────────────
+
+@dataclass
+class TrainConfig:
+    seq_len: int = 8                    # frames per training sequence
+    batch_size: int = 32
+    epochs: int = 100
+    lr: float = 1e-3
+    weight_decay: float = 1e-5
+    lr_scheduler_step: int = 30
+    lr_scheduler_gamma: float = 0.5
+    num_workers: int = 4
+    seed: int = 42
+
+    # Event simulation
+    event_threshold: float = 0.1
+    event_use_log: bool = True
+    event_auto_threshold: bool = False
+
+    # Logging / checkpoint
+    log_dir: str = "logs"
+    checkpoint_dir: str = "checkpoints"
+    log_interval: int = 10
+    save_interval: int = 10
+
+
+# ──────────────────────────── Singleton instances ────────────────────────────
+
+model_cfg = ModelConfig()
+train_cfg = TrainConfig()

src/velocity_prediction/dataset.py

@@ -0,0 +1,120 @@
+"""
+WebDataset-based dataset with sequence sampling.
+
+Each sample from the dataset is a dict:
+    {
+        "events": np.ndarray (1, H, W),   # simulated event frame
+        "tilt":   np.ndarray (3,),         # tilt rotation vector
+        "v_body_target": np.ndarray (2,),  # body-frame [vx, vy]
+    }
+
+The dataset groups consecutive frames into sequences of length seq_len.
+"""
+
+import webdataset as wds
+from pathlib import Path
+from typing import List, Callable, Optional
+
+from src.velocity_prediction.config import DATASET_ROOT, train_cfg
+from src.velocity_prediction.transforms import build_train_transform, build_val_transform
+
+
+def _scene_urls(scene_names: List[str], root: Path = DATASET_ROOT) -> List[str]:
+    """Build list of actual shard file paths for the given scene names."""
+    urls = []
+    for name in scene_names:
+        scene_dir = root / name
+        if not scene_dir.exists():
+            print(f"Warning: scene directory not found: {scene_dir}")
+            continue
+        # List all shard_*.tar files in the scene directory
+        shard_files = sorted(scene_dir.glob("shard_*.tar"))
+        if not shard_files:
+            print(f"Warning: no shard files found in {scene_dir}")
+            continue
+        urls.extend(str(f) for f in shard_files)
+    return urls
+
+
+def _build_pipeline(
+    urls: List[str],
+    transform: Callable,
+    seq_len: int,
+    shuffle: int = 1000,
+    deterministic: bool = False,
+):
+    """
+    Build a WebDataset pipeline that:
+      1. Reads tar shards
+      2. Decodes and transforms individual samples
+      3. Groups consecutive samples into sequences
+    """
+    dataset = wds.WebDataset(urls, shardshuffle=shuffle if not deterministic else 0, empty_check=False)
+
+    if not deterministic:
+        dataset = dataset.shuffle(shuffle)
+
+    dataset = dataset.decode().map(transform)
+
+    # Group into sequences of seq_len consecutive frames
+    dataset = dataset.batched(seq_len, partial=False)
+
+    return dataset
+
+
+def create_train_loader(
+    scene_names: Optional[List[str]] = None,
+    seq_len: int = 8,
+    batch_size: int = 32,
+    num_workers: int = 4,
+    event_threshold: float = 0.1,
+    event_use_log: bool = True,
+):
+    """Create a DataLoader for training."""
+    if scene_names is None:
+        from src.velocity_prediction.config import TRAIN_SCENES
+        scene_names = TRAIN_SCENES
+
+    urls = _scene_urls(scene_names)
+    transform = build_train_transform(
+        event_threshold=event_threshold,
+        event_use_log=event_use_log,
+    )
+    pipeline = _build_pipeline(urls, transform, seq_len=seq_len, shuffle=1000)
+
+    loader = wds.WebLoader(
+        pipeline,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        shuffle=False,  # already shuffled in pipeline
+    )
+    return loader
+
+
+def create_val_loader(
+    scene_names: Optional[List[str]] = None,
+    seq_len: int = 8,
+    batch_size: int = 32,
+    num_workers: int = 4,
+    event_threshold: float = 0.1,
+    event_use_log: bool = True,
+):
+    """Create a DataLoader for validation (deterministic order)."""
+    if scene_names is None:
+        from src.velocity_prediction.config import VAL_SCENES
+        scene_names = VAL_SCENES
+
+    urls = _scene_urls(scene_names)
+    transform = build_val_transform(
+        event_threshold=event_threshold,
+        event_use_log=event_use_log,
+    )
+    pipeline = _build_pipeline(urls, transform, seq_len=seq_len, shuffle=0, deterministic=True)
+
+    loader = wds.WebLoader(
+        pipeline,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        shuffle=False,
+    )
+    return loader

src/velocity_prediction/evaluate.py

@@ -0,0 +1,212 @@
+"""
+Evaluation and visualization for VelocityPredictionModel.
+
+Usage:
+    python -m src.velocity_prediction.evaluate --checkpoint checkpoints/best.pt
+"""
+
+import argparse
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+from src.velocity_prediction.model import VelocityPredictionModel
+from src.velocity_prediction.dataset import create_val_loader
+from src.velocity_prediction.config import train_cfg, VELOCITY_MEAN, VELOCITY_STD
+
+
+@torch.no_grad()
+def evaluate(
+    model: nn.Module,
+    loader,
+    device: torch.device,
+) -> dict:
+    """
+    Run evaluation on a dataloader.
+
+    Returns:
+        dict with keys:
+            preds:  np.ndarray (N, 2)  predicted [vx, vy]
+            targets: np.ndarray (N, 2) ground truth [vx, vy]
+    """
+    model.eval()
+    all_preds = []
+    all_targets = []
+
+    for batch in loader:
+        events = batch["events"].to(device)
+        tilt = batch["tilt"].to(device)
+        target = batch["v_body_target"].to(device)  # (B, S, 2)
+
+        pred = model(events, tilt)                   # (B, 2)
+        target_last = target[:, -1, :]               # (B, 2)
+
+        all_preds.append(pred.cpu().numpy())
+        all_targets.append(target_last.cpu().numpy())
+
+    preds = np.concatenate(all_preds, axis=0)
+    targets = np.concatenate(all_targets, axis=0)
+
+    # Denormalize predictions back to original velocity space
+    mean = np.array(VELOCITY_MEAN, dtype=np.float32)
+    std = np.array(VELOCITY_STD, dtype=np.float32)
+    preds_denorm = preds * std + mean
+    targets_denorm = targets * std + mean
+
+    # ── Diagnostics (in normalized space) ────────────────────────
+    print("\n========== Evaluation Diagnostics (normalized space) ==========")
+    print(f"Total samples: {len(preds)}")
+    print(f"\n--- Targets (normalized) ---")
+    print(f"  vx: mean={targets[:, 0].mean():.6f}, std={targets[:, 0].std():.6f}")
+    print(f"  vy: mean={targets[:, 1].mean():.6f}, std={targets[:, 1].std():.6f}")
+    print(f"\n--- Predictions (normalized) ---")
+    print(f"  vx: mean={preds[:, 0].mean():.6f}, std={preds[:, 0].std():.6f}, "
+          f"min={preds[:, 0].min():.6f}, max={preds[:, 0].max():.6f}")
+    print(f"  vy: mean={preds[:, 1].mean():.6f}, std={preds[:, 1].std():.6f}, "
+          f"min={preds[:, 1].min():.6f}, max={preds[:, 1].max():.6f}")
+    print(f"\n--- Unique prediction values ---")
+    print(f"  vx unique: {len(np.unique(preds[:, 0]))} / {len(preds)}")
+    print(f"  vy unique: {len(np.unique(preds[:, 1]))} / {len(preds)}")
+    vx_range = preds[:, 0].max() - preds[:, 0].min()
+    vy_range = preds[:, 1].max() - preds[:, 1].min()
+    print(f"\n  vx range: {vx_range:.8f}  (constant if near 0)")
+    print(f"  vy range: {vy_range:.8f}  (constant if near 0)")
+    print(f"\n--- Constant prediction check ---")
+    print(f"  pred vx mean ≈ 0?  {abs(preds[:, 0].mean()):.6f} diff from zero")
+    print(f"  pred vy mean ≈ 0?  {abs(preds[:, 1].mean()):.6f} diff from zero")
+    print("=============================================\n")
+
+    # Per-axis and overall RMSE (in original velocity space)
+    rmse_x = np.sqrt(np.mean((preds_denorm[:, 0] - targets_denorm[:, 0]) ** 2))
+    rmse_y = np.sqrt(np.mean((preds_denorm[:, 1] - targets_denorm[:, 1]) ** 2))
+    rmse_xy = np.sqrt(np.mean(np.sum((preds_denorm - targets_denorm) ** 2, axis=1)))
+
+    return {
+        "preds": preds_denorm,      # denormalized for plotting
+        "targets": targets_denorm,  # denormalized for plotting
+        "rmse_x": rmse_x,
+        "rmse_y": rmse_y,
+        "rmse_xy": rmse_xy,
+    }
+
+
+def plot_results(
+    preds: np.ndarray,
+    targets: np.ndarray,
+    save_path: str = "eval_plot.png",
+):
+    """Plot predicted vs ground truth velocity traces."""
+    fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
+
+    time = np.arange(len(preds))
+
+    axes[0].plot(time, targets[:, 0], label="GT vx", color="C0", alpha=0.8)
+    axes[0].plot(time, preds[:, 0], label="Pred vx", color="C1", alpha=0.8)
+    axes[0].set_ylabel("vx (m/s)")
+    axes[0].legend()
+    axes[0].grid(True, alpha=0.3)
+
+    axes[1].plot(time, targets[:, 1], label="GT vy", color="C0", alpha=0.8)
+    axes[1].plot(time, preds[:, 1], label="Pred vy", color="C1", alpha=0.8)
+    axes[1].set_ylabel("vy (m/s)")
+    axes[1].set_xlabel("Frame index")
+    axes[1].legend()
+    axes[1].grid(True, alpha=0.3)
+
+    fig.suptitle("Body-frame Velocity Prediction")
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150)
+    print(f"Plot saved: {save_path}")
+    plt.close()
+
+
+def plot_scatter(
+    preds: np.ndarray,
+    targets: np.ndarray,
+    save_path: str = "eval_scatter.png",
+):
+    """Scatter plot: predicted vs ground truth."""
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
+
+    for ax, pred, target, label in zip(
+        axes, [preds[:, 0], preds[:, 1]], [targets[:, 0], targets[:, 1]], ["vx", "vy"]
+    ):
+        ax.scatter(target, pred, s=2, alpha=0.5)
+        lim_min = min(target.min(), pred.min())
+        lim_max = max(target.max(), pred.max())
+        ax.plot([lim_min, lim_max], [lim_min, lim_max], "r--", alpha=0.5)
+        ax.set_xlabel(f"GT {label} (m/s)")
+        ax.set_ylabel(f"Pred {label} (m/s)")
+        ax.set_aspect("equal")
+        ax.grid(True, alpha=0.3)
+        ax.set_title(f"{label} — RMSE: {np.sqrt(np.mean((pred - target)**2)):.4f}")
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150)
+    print(f"Scatter saved: {save_path}")
+    plt.close()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, default="checkpoints/best.pt",
+                        help="Path to model checkpoint")
+    parser.add_argument("--device", type=str, default="cuda",
+                        help="Device to use (e.g. 'cuda:2', 'cpu')")
+    parser.add_argument("--plot", action="store_true", default=True,
+                        help="Generate evaluation plots")
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
+    print(f"Device: {device}")
+
+    # Load model
+    model = VelocityPredictionModel()
+    ckpt = torch.load(args.checkpoint, map_location="cpu")
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.to(device)
+    print(f"Loaded checkpoint from {args.checkpoint} (epoch={ckpt.get('epoch', '?')})")
+
+    # Validation loader (use test scenes for final eval)
+    from src.velocity_prediction.config import TEST_SCENES
+    loader = create_val_loader(
+        scene_names=TEST_SCENES,
+        seq_len=train_cfg.seq_len,
+        batch_size=train_cfg.batch_size,
+        num_workers=2,
+        event_threshold=train_cfg.event_threshold,
+        event_use_log=train_cfg.event_use_log,
+    )
+
+    # # ── Quick event diagnostics: inspect one batch ───────────────
+    # print("\n========== Event Frame Diagnostics ==========")
+    # sample_batch = next(iter(loader))
+    # ev = sample_batch["events"]  # (B, S, 1, H, W)
+    # print(f"Events shape: {ev.shape}")
+    # print(f"Events dtype: {ev.dtype}")
+    # print(f"Events value counts:  -1: {(ev == -1).sum().item()}, "
+    #       f"0: {(ev == 0).sum().item()}, +1: {(ev == 1).sum().item()}")
+    # total_el = ev.numel()
+    # nonzero = (ev != 0).sum().item()
+    # print(f"Non-zero ratio: {nonzero / total_el:.6f} ({nonzero}/{total_el})")
+    # print(f"Per-sample non-zero: {[(ev[b] != 0).sum().item() for b in range(min(4, ev.shape[0]))]}")
+    # print("=============================================\n")
+
+    # Evaluate
+    results = evaluate(model, loader, device)
+    print(f"\nEvaluation results on test scenes: {TEST_SCENES}")
+    print(f"  RMSE vx: {results['rmse_x']:.4f} m/s")
+    print(f"  RMSE vy: {results['rmse_y']:.4f} m/s")
+    print(f"  RMSE xy: {results['rmse_xy']:.4f} m/s")
+
+    # Plots
+    if args.plot:
+        plot_results(results["preds"], results["targets"], "eval_velocity.png")
+        plot_scatter(results["preds"], results["targets"], "eval_scatter.png")
+
+
+if __name__ == "__main__":
+    main()

src/velocity_prediction/model.py

@@ -0,0 +1,170 @@
+"""
+VelocityPredictionModel: CNN + PoseMLP → concat → GRU → Head → [vx_body, vy_body].
+
+Architecture:
+    Event frame (1, H, W) ──► CNN ──┐
+    Tilt angles    (3,)    ──► MLP ──┤──► concat ──► GRU ──► Head ──► [vx, vy]
+"""
+
+import torch
+import torch.nn as nn
+
+from src.velocity_prediction.config import model_cfg
+
+
+class CNNEncoder(nn.Module):
+    """
+    4-layer ConvNet with BatchNorm, ReLU, MaxPool, ending with Global Avg Pool.
+
+    Input:  (B, S, 1, H, W)  — processed per-frame (flattened to (B*S, 1, H, W))
+    Output: (B, S, C_out)     — per-frame feature vectors
+    """
+
+    def __init__(self, cfg=model_cfg.cnn):
+        super().__init__()
+        channels = cfg.channels
+        in_ch = cfg.in_channels
+
+        layers = []
+        for out_ch in channels:
+            layers.extend([
+                nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2),
+                # nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
+                nn.Identity(),
+                nn.LeakyReLU(inplace=True),
+                nn.MaxPool2d(cfg.pool_size),
+            ])
+            in_ch = out_ch
+
+        self.conv = nn.Sequential(*layers)
+        self.gap = nn.AdaptiveAvgPool2d(1)
+        self.out_dim = channels[-1]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, S, 1, H, W)  event frame sequence
+        Returns:
+            features: (B, S, C_out)
+        """
+        B, S, C, H, W = x.shape
+        x = x.view(B * S, C, H, W)          # (B*S, 1, H, W)
+        x = self.conv(x)                     # (B*S, C_out, H', W')
+        x = self.gap(x)                      # (B*S, C_out, 1, 1)
+        x = x.view(B, S, self.out_dim)       # (B, S, C_out)
+        return x
+
+
+class PoseMLP(nn.Module):
+    """
+    Encode tilt rotation vector (3,) into a compact feature vector.
+
+    Input:  (B, S, 3)
+    Output: (B, S, output_dim)
+    """
+
+    def __init__(self, cfg=model_cfg.pose_mlp):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(cfg.input_dim, cfg.hidden_dim),
+            nn.LeakyReLU(inplace=True),
+            nn.Linear(cfg.hidden_dim, cfg.output_dim),
+            nn.LeakyReLU(inplace=True),
+        )
+        self.out_dim = cfg.output_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (B, S, 3)  →  (B, S, output_dim)
+        """
+        B, S, D = x.shape
+        x = x.view(B * S, D)
+        x = self.net(x)
+        x = x.view(B, S, self.out_dim)
+        return x
+
+
+class VelocityPredictionModel(nn.Module):
+    """
+    Full model: CNN + PoseMLP → concat → GRU → Head → [vx, vy].
+
+    Input:
+        events: (B, S, 1, H, W)
+        tilt:   (B, S, 3)
+    Output:
+        v_body: (B, 2)  — body-frame [vx, vy] for the last frame in the sequence
+    """
+
+    def __init__(self, cnn_cfg=model_cfg.cnn, pose_cfg=model_cfg.pose_mlp,
+                 gru_cfg=model_cfg.gru, head_cfg=model_cfg.head):
+        super().__init__()
+
+        self.cnn = CNNEncoder(cnn_cfg)
+        self.pose_mlp = PoseMLP(pose_cfg)
+
+        fused_dim = self.cnn.out_dim + self.pose_mlp.out_dim  # 256 + 64 = 320
+
+        self.gru = nn.GRU(
+            input_size=fused_dim,
+            hidden_size=gru_cfg.hidden_size,
+            num_layers=gru_cfg.num_layers,
+            dropout=gru_cfg.dropout if gru_cfg.num_layers > 1 else 0.0,
+            batch_first=True,
+        )
+
+        self.head = nn.Sequential(
+            nn.Linear(gru_cfg.hidden_size, head_cfg.hidden_dim),
+            nn.LeakyReLU(inplace=True),
+            nn.Linear(head_cfg.hidden_dim, head_cfg.output_dim),
+        )
+
+        # # Small init for the final layer: start from near-zero output
+        # self.head[-1].weight.data.mul_(0.01)
+        # self.head[-1].bias.data.zero_()
+
+    def forward(self, events: torch.Tensor, tilt: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            events: (B, S, 1, H, W)
+            tilt:   (B, S, 3)
+
+        Returns:
+            v_body: (B, 2)  predicted body-frame [vx, vy] at the last timestep
+        """
+        # Per-frame encoding
+        cnn_feat = self.cnn(events)           # (B, S, 256)
+        pose_feat = self.pose_mlp(tilt)       # (B, S, 64)
+
+        # Fuse per frame
+        fused = torch.cat([cnn_feat, pose_feat], dim=-1)  # (B, S, 320)
+
+        # GRU temporal modelling
+        gru_out, h_n = self.gru(fused)        # gru_out: (B, S, 128), h_n: (1, B, 128)
+
+        # Use last hidden state
+        last_hidden = h_n[-1]                 # (B, 128)
+
+        # Head regression
+        v_body = self.head(last_hidden)       # (B, 2)
+        return v_body
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Count trainable parameters."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+if __name__ == "__main__":
+    # Quick sanity check
+    model = VelocityPredictionModel()
+    total = count_parameters(model)
+    print(f"Total trainable parameters: {total:,} ({total/1e6:.3f} M)")
+
+    # Forward pass test
+    B, S, H, W = 4, 8, 240, 320
+    events = torch.randn(B, S, 1, H, W)
+    tilt = torch.randn(B, S, 3)
+    out = model(events, tilt)
+    print(f"Input events: {events.shape}")
+    print(f"Input tilt:   {tilt.shape}")
+    print(f"Output:       {out.shape}  (should be [4, 2])")

src/velocity_prediction/train.py

@@ -0,0 +1,213 @@
+"""
+Training loop for VelocityPredictionModel.
+
+Usage:
+    python -m src.velocity_prediction.train [--device cuda:0]
+"""
+
+import argparse
+import os
+import time
+import numpy as np
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.tensorboard import SummaryWriter
+
+from src.velocity_prediction.config import train_cfg, model_cfg
+from src.velocity_prediction.model import VelocityPredictionModel, count_parameters
+from src.velocity_prediction.dataset import create_train_loader, create_val_loader
+
+
+def set_seed(seed: int):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def train_one_epoch(
+    model: nn.Module,
+    loader,
+    optimizer: torch.optim.Optimizer,
+    criterion: nn.Module,
+    device: torch.device,
+    epoch: int,
+    writer: SummaryWriter,
+    log_interval: int = 50,
+) -> float:
+    """Train for one epoch. Returns average loss."""
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+    start_time = time.time()
+
+    for batch_idx, batch in enumerate(loader):
+        events = batch["events"].to(device)          # (B, S, 1, H, W)
+        tilt = batch["tilt"].to(device)              # (B, S, 3)
+        target = batch["v_body_target"].to(device)   # (B, S, 2)
+
+        # Predict velocity for the last frame in the sequence
+        pred = model(events, tilt)                   # (B, 2)
+        target_last = target[:, -1, :]               # (B, 2)
+
+        loss = criterion(pred, target_last)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item()
+        num_batches += 1
+
+        if batch_idx % log_interval == 0:
+            elapsed = time.time() - start_time
+            print(f"  Epoch {epoch} | Batch {batch_idx} | Loss: {loss.item():.6f} | {elapsed:.1f}s")
+            writer.add_scalar("train/loss_batch", loss.item(), batch_idx)
+
+    avg_loss = total_loss / max(num_batches, 1)
+    return avg_loss
+
+
+@torch.no_grad()
+def validate(
+    model: nn.Module,
+    loader,
+    criterion: nn.Module,
+    device: torch.device,
+) -> float:
+    """Validate. Returns average loss."""
+    model.eval()
+    total_loss = 0.0
+    num_batches = 0
+
+    for batch in loader:
+        events = batch["events"].to(device)
+        tilt = batch["tilt"].to(device)
+        target = batch["v_body_target"].to(device)
+
+        pred = model(events, tilt)
+        target_last = target[:, -1, :]
+
+        loss = criterion(pred, target_last)
+        total_loss += loss.item()
+        num_batches += 1
+
+    avg_loss = total_loss / max(num_batches, 1)
+    return avg_loss
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", type=str, default="cuda",
+                        help="CUDA device, e.g. 'cuda:0', 'cuda:1' (default: 'cuda')")
+    args = parser.parse_args()
+
+    set_seed(train_cfg.seed)
+    device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
+    print(f"Device: {device}")
+
+    # Create model
+    model = VelocityPredictionModel()
+    model.to(device)
+    total_params = count_parameters(model)
+    print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
+
+    # Data loaders
+    train_loader = create_train_loader(
+        seq_len=train_cfg.seq_len,
+        batch_size=train_cfg.batch_size,
+        num_workers=train_cfg.num_workers,
+        event_threshold=train_cfg.event_threshold,
+        event_use_log=train_cfg.event_use_log,
+    )
+    val_loader = create_val_loader(
+        seq_len=train_cfg.seq_len,
+        batch_size=train_cfg.batch_size,
+        num_workers=train_cfg.num_workers,
+        event_threshold=train_cfg.event_threshold,
+        event_use_log=train_cfg.event_use_log,
+    )
+
+    # Optimizer & scheduler
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=train_cfg.lr,
+        weight_decay=train_cfg.weight_decay,
+    )
+    scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=train_cfg.lr_scheduler_step,
+        gamma=train_cfg.lr_scheduler_gamma,
+    )
+    # criterion = nn.SmoothL1Loss()
+    criterion = nn.MSELoss()
+
+    # Logging
+    log_dir = Path(train_cfg.log_dir)
+    log_dir.mkdir(parents=True, exist_ok=True)
+    writer = SummaryWriter(log_dir=str(log_dir))
+
+    ckpt_dir = Path(train_cfg.checkpoint_dir)
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+
+    best_val_loss = float("inf")
+
+    print(f"\nStarting training for {train_cfg.epochs} epochs...")
+    print(f"  seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}")
+    print(f"  lr={train_cfg.lr}, weight_decay={train_cfg.weight_decay}")
+    print(f"  log_dir={log_dir}, checkpoint_dir={ckpt_dir}\n")
+
+    for epoch in range(1, train_cfg.epochs + 1):
+        epoch_start = time.time()
+
+        train_loss = train_one_epoch(
+            model, train_loader, optimizer, criterion, device, epoch, writer,
+            log_interval=train_cfg.log_interval,
+        )
+        val_loss = validate(model, val_loader, criterion, device)
+        scheduler.step()
+
+        epoch_time = time.time() - epoch_start
+        current_lr = scheduler.get_last_lr()[0]
+
+        print(f"Epoch {epoch:3d}/{train_cfg.epochs} | "
+              f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | "
+              f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")
+
+        writer.add_scalar("train/loss_epoch", train_loss, epoch)
+        writer.add_scalar("val/loss", val_loss, epoch)
+        writer.add_scalar("lr", current_lr, epoch)
+
+        # Save checkpoint
+        if epoch % train_cfg.save_interval == 0:
+            ckpt_path = ckpt_dir / f"epoch_{epoch:03d}_val_{val_loss:.6f}.pt"
+            torch.save({
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "scheduler_state_dict": scheduler.state_dict(),
+                "train_loss": train_loss,
+                "val_loss": val_loss,
+            }, ckpt_path)
+            print(f"  Checkpoint saved: {ckpt_path}")
+
+        # Save best model
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_path = ckpt_dir / "best.pt"
+            torch.save({
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "val_loss": val_loss,
+            }, best_path)
+            print(f"  Best model updated: {best_path} (val_loss={val_loss:.6f})")
+
+    writer.close()
+    print(f"\nTraining complete. Best val loss: {best_val_loss:.6f}")
+    print(f"Best checkpoint: {ckpt_dir / 'best.pt'}")
+
+
+if __name__ == "__main__":
+    main()

src/velocity_prediction/transforms.py

@@ -0,0 +1,142 @@
+"""
+Transforms: event frame generation + coordinate transforms for pose/velocity.
+
+Each transform operates on a single decoded sample dict:
+    {
+        "jpg":  bytes,          # JPEG-encoded grayscale image
+        "ts":   bytes,          # float64 timestamp
+        "pose": bytes,          # float32[7]  [x, y, z, qx, qy, qz, qw]
+        "vel":  bytes,          # float32[6]  [vx, vy, vz, wx, wy, wz]
+    }
+"""
+
+import io
+import numpy as np
+import cv2
+
+from src.event_utils import EventProcessor
+from src.velocity_prediction.utils import decompose_tilt_np, world_vel_to_body_np
+from src.velocity_prediction.config import VELOCITY_MEAN, VELOCITY_STD
+
+
+class DecodeSample:
+    """Decode raw bytes from WebDataset tar entry into numpy arrays."""
+
+    def __call__(self, sample: dict) -> dict:
+        # Image: JPEG bytes → grayscale uint8 (H, W)
+        img = cv2.imdecode(np.frombuffer(sample["jpg"], np.uint8), cv2.IMREAD_GRAYSCALE)
+
+        # Timestamp
+        ts = np.frombuffer(sample["ts"], dtype=np.float64).item()
+
+        # Pose: [x, y, z, qx, qy, qz, qw]
+        pose = np.frombuffer(sample["pose"], dtype=np.float32).copy()
+
+        # Velocity: [vx, vy, vz, wx, wy, wz]
+        vel = np.frombuffer(sample["vel"], dtype=np.float32).copy()
+
+        return {"img": img, "ts": ts, "pose": pose, "vel": vel}
+
+
+class SimulateEvents:
+    """Convert grayscale frame to binary event frame using EventProcessor."""
+
+    def __init__(self, threshold=0.1, use_log=True, auto_threshold=False, verbose=False):
+        self.processor = EventProcessor(
+            threshold=threshold,
+            use_log=use_log,
+            auto_threshold=auto_threshold,
+        )
+        self.verbose = verbose
+        self._frame_count = 0
+
+    def __call__(self, sample: dict) -> dict:
+        img = sample["img"]
+        events_binary, events_strength, event_count = self.processor(img)
+        # Use binary events as network input: shape (1, H, W), values in {-1, 0, 1}
+        sample["events"] = events_binary.astype(np.float32)[np.newaxis, ...]
+
+        if self.verbose:
+            self._frame_count += 1
+            total_pixels = events_binary.shape[0] * events_binary.shape[1]
+            nonzero_ratio = event_count / total_pixels
+            pos_ratio = (events_binary > 0).sum() / total_pixels
+            neg_ratio = (events_binary < 0).sum() / total_pixels
+            if self._frame_count <= 5 or self._frame_count % 100 == 0:
+                print(f"  [EventDiagnostics] frame={self._frame_count} | "
+                      f"nonzero={nonzero_ratio:.4f} (+{pos_ratio:.4f}/-{neg_ratio:.4f}) | "
+                      f"count={event_count}")
+
+        return sample
+
+    def reset(self):
+        self.processor.reset()
+        self._frame_count = 0
+
+
+class ComputeTilt:
+    """Extract tilt rotation vector from pose quaternion (discard position, discard yaw)."""
+
+    def __call__(self, sample: dict) -> dict:
+        q = sample["pose"][3:7]  # [qx, qy, qz, qw]
+        tilt = decompose_tilt_np(q)  # (3,)  rotation vector
+        sample["tilt"] = tilt.astype(np.float32)
+        return sample
+
+
+class ComputeBodyVelocity:
+    """Transform world-frame velocity to body-frame (yaw-compensated)."""
+
+    def __call__(self, sample: dict) -> dict:
+        v_world = sample["vel"][:3]  # [vx, vy, vz]  world frame
+        q = sample["pose"][3:7]      # [qx, qy, qz, qw]
+        v_body = world_vel_to_body_np(v_world, q)  # (3,)
+        # Only predict forward (x) and lateral (y) body velocity
+        sample["v_body_target"] = v_body[:2].astype(np.float32)  # (2,)
+        return sample
+
+
+class NormalizeVelocity:
+    """Normalize body-frame velocity to zero mean, unit variance."""
+
+    def __init__(self):
+        self.mean = np.array(VELOCITY_MEAN, dtype=np.float32)
+        self.std = np.array(VELOCITY_STD, dtype=np.float32)
+
+    def __call__(self, sample: dict) -> dict:
+        sample["v_body_target"] = (sample["v_body_target"] - self.mean) / self.std
+        return sample
+
+
+class Compose:
+    """Chain multiple transforms."""
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, sample: dict) -> dict:
+        for t in self.transforms:
+            sample = t(sample)
+        return sample
+
+
+def build_train_transform(event_threshold=0.1, event_use_log=True):
+    """Build the full transform pipeline for training samples."""
+    return Compose([
+        DecodeSample(),
+        SimulateEvents(threshold=event_threshold, use_log=event_use_log),
+        ComputeTilt(),
+        ComputeBodyVelocity(),
+        NormalizeVelocity(),
+    ])
+
+
+def build_val_transform(event_threshold=0.1, event_use_log=True):
+    """Same as train but with a fresh EventProcessor per sample (no cross-contamination)."""
+    return Compose([
+        DecodeSample(),
+        SimulateEvents(threshold=event_threshold, use_log=event_use_log),
+        ComputeTilt(),
+        ComputeBodyVelocity(),
+        NormalizeVelocity(),
+    ])

src/velocity_prediction/utils.py

@@ -0,0 +1,165 @@
+"""
+Quaternion and coordinate-frame utility functions.
+
+All quaternions follow [x, y, z, w] convention (matching dataset pose field).
+"""
+
+import torch
+import numpy as np
+
+
+# ──────────────────────────── Quaternion operations ────────────────────────────
+
+def quat_normalize(q: torch.Tensor) -> torch.Tensor:
+    """Normalize quaternion.  q: (..., 4)"""
+    return q / torch.norm(q, dim=-1, keepdim=True).clamp(min=1e-12)
+
+
+def quat_conjugate(q: torch.Tensor) -> torch.Tensor:
+    """Conjugate (inverse for unit quaternion).  q: (..., 4)"""
+    return q * torch.tensor([-1, -1, -1, 1], device=q.device)
+
+
+def quat_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Hamilton product.  a, b: (..., 4)"""
+    x1, y1, z1, w1 = a.unbind(-1)
+    x2, y2, z2, w2 = b.unbind(-1)
+    return torch.stack([
+        w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+        w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+        w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+    ], dim=-1)
+
+
+def quat_rotate(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Rotate vector v by quaternion q.  q: (..., 4), v: (..., 3)"""
+    q_conj = quat_conjugate(q)
+    v_pad = torch.zeros_like(v[..., :1])  # (..., 1)
+    v_q = torch.cat([v, v_pad], dim=-1)   # (..., 4)  pure quaternion
+    return quat_mul(quat_mul(q, v_q), q_conj)[..., :3]
+
+
+# ──────────────────────────── Yaw decomposition ────────────────────────────
+
+def quat_to_yaw(q: torch.Tensor) -> torch.Tensor:
+    """
+    Extract yaw (heading) angle from a quaternion in gravity-aligned frame.
+
+    Gravity axis is +z.  Yaw is the rotation around z-axis.
+    Returns angle in radians, shape (...,).
+    """
+    x, y, z, w = q.unbind(-1)
+    # From quaternion to Euler: yaw = atan2(2(w*z + x*y), 1 - 2(y² + z²))
+    siny = 2.0 * (w * z + x * y)
+    cosy = 1.0 - 2.0 * (y * y + z * z)
+    return torch.atan2(siny, cosy)
+
+
+def quat_from_yaw(yaw: torch.Tensor) -> torch.Tensor:
+    """
+    Build a pure-yaw quaternion (rotation around +z).
+    yaw: (...,)  →  quat: (..., 4)
+    """
+    half = yaw * 0.5
+    cos = torch.cos(half)
+    sin = torch.sin(half)
+    z = torch.zeros_like(cos)
+    return torch.stack([z, z, sin, cos], dim=-1)  # [0, 0, sin(yaw/2), cos(yaw/2)]
+
+
+def decompose_tilt(q: torch.Tensor) -> torch.Tensor:
+    """
+    Remove yaw from a quaternion, returning the residual tilt rotation vector.
+
+    Given q = q_yaw * q_tilt (z-yaw first, then body tilt),
+    we compute q_tilt = q_yaw^{-1} * q, then convert to rotation vector.
+
+    Args:
+        q: (..., 4)  unit quaternion in world→body convention.
+
+    Returns:
+        tilt_angles: (..., 3)  rotation vector [rx, ry, rz] representing
+                     the body's deviation from the heading direction.
+    """
+    yaw = quat_to_yaw(q)
+    q_yaw = quat_from_yaw(yaw)
+    q_yaw_inv = quat_conjugate(q_yaw)
+    q_tilt = quat_mul(q_yaw_inv, q)          # remove yaw
+    q_tilt = quat_normalize(q_tilt)
+    return quat_to_rotvec(q_tilt)
+
+
+def quat_to_rotvec(q: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
+    """
+    Convert unit quaternion to rotation vector (axis * angle).
+    q: (..., 4)  →  rotvec: (..., 3)
+    """
+    q = quat_normalize(q)
+    x, y, z, w = q.unbind(-1)
+    angle = 2.0 * torch.acos(w.clamp(-1.0, 1.0))
+    sin_half = torch.sqrt((1.0 - w * w).clamp(min=eps))
+    scale = angle / sin_half
+    # Avoid division by zero for small angles
+    mask = sin_half > eps
+    rx = torch.where(mask, x * scale, torch.zeros_like(x))
+    ry = torch.where(mask, y * scale, torch.zeros_like(y))
+    rz = torch.where(mask, z * scale, torch.zeros_like(z))
+    return torch.stack([rx, ry, rz], dim=-1)
+
+
+# ──────────────────────────── Velocity transformation ────────────────────────────
+
+def world_vel_to_body(
+    v_world: torch.Tensor,
+    q_world_to_body: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Transform world-frame velocity to body-frame velocity.
+
+    Steps:
+        1. Extract yaw from q_world_to_body.
+        2. Build pure-yaw quaternion q_yaw.
+        3. Remove yaw from velocity: v_yaw_compensated = q_yaw^{-1} * v_world
+        4. Rotate to body frame: v_body = q_tilt^{-1} * v_yaw_compensated
+           where q_tilt = q_yaw^{-1} * q_world_to_body
+
+    Args:
+        v_world: (..., 3)  world-frame linear velocity [vx, vy, vz]
+        q_world_to_body: (..., 4)  world→body unit quaternion
+
+    Returns:
+        v_body: (..., 3)  body-frame linear velocity
+    """
+    yaw = quat_to_yaw(q_world_to_body)
+    q_yaw = quat_from_yaw(yaw)
+    q_yaw_inv = quat_conjugate(q_yaw)
+
+    # Step 1: remove yaw from velocity (rotate to yaw-aligned intermediate frame)
+    v_yaw_comp = quat_rotate(q_yaw_inv, v_world)
+
+    # Step 2: compute tilt quaternion
+    q_tilt = quat_mul(q_yaw_inv, q_world_to_body)
+    q_tilt = quat_normalize(q_tilt)
+    q_tilt_inv = quat_conjugate(q_tilt)
+
+    # Step 3: rotate to body frame
+    v_body = quat_rotate(q_tilt_inv, v_yaw_comp)
+    return v_body
+
+
+# ──────────────────────────── NumPy wrappers (for transforms.py) ────────────────────────────
+
+def decompose_tilt_np(q: np.ndarray) -> np.ndarray:
+    """NumPy version of decompose_tilt."""
+    q_t = torch.from_numpy(q)
+    tilt = decompose_tilt(q_t)
+    return tilt.numpy()
+
+
+def world_vel_to_body_np(v_world: np.ndarray, q: np.ndarray) -> np.ndarray:
+    """NumPy version of world_vel_to_body."""
+    v_t = torch.from_numpy(v_world)
+    q_t = torch.from_numpy(q)
+    vb = world_vel_to_body(v_t, q_t)
+    return vb.numpy()

start_ros_container.sh

@@ -0,0 +1,143 @@
+#!/bin/bash
+
+# ROS Podman容器启动脚本
+# 作者自动生成
+# 功能：启动ROS容器，挂载当前目录，设置环境变量，安装pip3，错误不退出
+
+set +e  # 遇到错误时不退出
+
+# 颜色定义（用于美化输出）
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+# 配置参数
+CONTAINER_IMAGE="69a38b2c0905"  # ROS Noetic Desktop Full镜像ID
+WORKSPACE_DIR="$(pwd)"          # 当前工作目录
+MOUNT_POINT="/mnt"               # 容器内挂载点
+
+# ROS环境变量设置
+ROS_ENV_VARS=(
+    "ROS_DISTRO=noetic"
+    "ROS_PYTHON_VERSION=3"
+    "ROS_VERSION=1"
+    "ROS_ETC_DIR=/opt/ros/noetic/etc/ros"
+    "ROS_ROOT=/opt/ros/noetic/share/ros"
+    "ROS_PACKAGE_PATH=/opt/ros/noetic/share"
+    "PYTHONIOENCODING=utf-8"
+    "TZ=Asia/Shanghai"
+)
+
+echo -e "${GREEN}========================================${NC}"
+echo -e "${GREEN}ROS Podman容器启动脚本${NC}"
+echo -e "${GREEN}========================================${NC}"
+
+# 检查Podman是否安装
+if ! command -v podman &> /dev/null; then
+    echo -e "${RED}错误: Podman未安装，请先安装Podman${NC}"
+    echo -e "${YELLOW}Ubuntu/Debian: sudo apt-get install podman${NC}"
+    echo -e "${YELLOW}CentOS/RHEL: sudo yum install podman${NC}"
+    exit 1
+fi
+
+# 检查镜像是否存在
+echo -e "${YELLOW}检查容器镜像...${NC}"
+if ! podman image exists $CONTAINER_IMAGE; then
+    echo -e "${RED}警告: 镜像 $CONTAINER_IMAGE 不存在${NC}"
+    echo -e "${YELLOW}可用的ROS镜像:${NC}"
+    podman image list | grep ros
+    echo -e "${YELLOW}继续使用指定镜像，如果启动失败请更换镜像ID${NC}"
+fi
+
+# 构建环境变量参数
+ENV_ARGS=""
+for env_var in "${ROS_ENV_VARS[@]}"; do
+    ENV_ARGS="$ENV_ARGS -e $env_var"
+done
+
+# 添加额外的环境变量（如果用户需要）
+ENV_ARGS="$ENV_ARGS -e DISPLAY=$DISPLAY"  # 支持GUI应用
+ENV_ARGS="$ENV_ARGS -e QT_X11_NO_MITSHM=1"
+
+# X11支持（用于GUI应用）
+if [ -n "$DISPLAY" ]; then
+    echo -e "${YELLOW}启用X11支持，用于GUI应用...${NC}"
+    xhost +local:root 2>/dev/null
+    ENV_ARGS="$ENV_ARGS --volume /tmp/.X11-unix:/tmp/.X11-unix:rw"
+fi
+
+echo -e "${GREEN}配置信息:${NC}"
+echo -e "  工作目录: $WORKSPACE_DIR"
+echo -e "  挂载点: $MOUNT_POINT"
+echo -e "  镜像ID: $CONTAINER_IMAGE"
+echo -e "  ROS发行版: noetic"
+
+echo -e "${YELLOW}正在启动容器...${NC}"
+
+# 启动容器的命令
+# 使用bash -c来执行多个命令，确保pip3安装即使失败也不会退出容器
+PODMAN_CMD="podman run -it \
+    --rm \
+    --name ros_noetic_container_$(date +%s) \
+    -v $WORKSPACE_DIR:$MOUNT_POINT:rw \
+    $ENV_ARGS \
+    $CONTAINER_IMAGE \
+    /bin/bash -c \"\
+        echo '========================================' && \
+        echo 'ROS容器已启动' && \
+        echo '========================================' && \
+        echo '工作目录已挂载到: $MOUNT_POINT' && \
+        echo '当前ROS版本: ' && \
+        echo \\\$ROS_DISTRO && \
+        echo '' && \
+        echo '正在检查并安装pip3...' && \
+        if command -v pip3 &> /dev/null; then \
+            echo 'pip3已安装，版本: ' && \
+            pip3 --version; \
+        else \
+            echo 'pip3未安装，正在安装...' && \
+            apt-get update 2>/dev/null && \
+            apt-get install -y python3-pip 2>/dev/null; \
+            if [ \\\$? -eq 0 ]; then \
+                echo 'pip3安装成功！'; \
+                pip3 --version; \
+            else \
+                echo '警告: pip3安装失败，请手动安装'; \
+            fi; \
+        fi && \
+        echo '' && \
+        echo '========================================' && \
+        echo '环境变量已设置:' && \
+        env | grep ROS_ && \
+        echo '========================================' && \
+        echo '容器已准备就绪，进入交互式shell...' && \
+        echo '提示: 输入exit退出容器' && \
+        echo '========================================' && \
+        pip config set global.index-url https://mirrors.ustc.edu.cn/pypi/simple && \
+        source /ros_entrypoint.sh && \
+        cd $MOUNT_POINT && \
+        exec /bin/bash -l\""
+
+# 执行命令
+echo -e "${GREEN}执行启动命令...${NC}"
+# echo -e "${YELLOW}命令详情:${NC}"
+# echo "$PODMAN_CMD"
+echo -e "${YELLOW}========================================${NC}"
+
+# 运行容器
+eval $PODMAN_CMD
+
+# 检查退出状态
+if [ $? -ne 0 ]; then
+    echo -e "${RED}容器已退出（退出码非0）${NC}"
+else
+    echo -e "${GREEN}容器正常退出${NC}"
+fi
+
+# 清理X11权限
+if [ -n "$DISPLAY" ]; then
+    xhost -local:root 2>/dev/null
+fi
+
+echo -e "${GREEN}脚本执行完成${NC}"