fix: remove denormalization in eval scripts

NormalizeVelocity transform is disabled, so model outputs are already in
original m/s space. Denormalizing inflates RMSE by ~3.5x.

- evaluate.py: compute RMSE in model output space directly
- benchmark/evaluate.py: same, plus fix misleading "normalized" comment

Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>

TODO.md

@@ -0,0 +1,109 @@
+# TODO — 训练/评估改进计划
+
+---
+
+## P1 · 截断反传(TBPTT):修 GRU 隐状态长程失配(新)
+
+### 问题诊断(2025-07-09)
+
+`diag_state_mismatch.py` 确认核心问题是 **GRU 隐状态训练/推理分布失配**:
+
+| checkpoint | A 无状态(h=0 每 8 帧重置) | B 有状态(h 持续滚动) | B/A |
+|---|---|---|---|
+| epoch 4 (val 7.91) | RMSE 0.14~4.0 | RMSE 1.6~4.3 | 1.1~1.4 |
+| epoch 180 (val 9.17) | RMSE 0.10~0.98 | RMSE 2.5~5.4 | 2.8~38.9 |
+
+训练越久,分叉越大。epoch 180 模型在窗口内几乎背下数据(A=0.14),但一滚动就崩(B=5.4)。**训练 loss 越低,长程越差**。
+
+**根因**:训练时每个序列从 `h=0` 开始(seq_len=8),GRU 只见「零后 ≤8 步」的状态分布。推理时 h 滚动数百帧,进入 OOD 区域,预测崩溃。val_loss 从 7.91 升到 9.17 进一步印证过拟合。
+
+### 修法:TBPTT
+
+跨 batch 传递 detach 后的隐状态,让 GRU 暴露于长程状态:
+
+```
+batch_k:     frames [k,  k+7]  → GRU(fused, h_prev) → h_curr  (detach → batch_{k+1})
+batch_{k+1}: frames [k+8, k+15] → GRU(fused, h_detached) → h_next  (detach → ...)
+```
+
+### 具体改动
+
+1. **`model.py::VelocityPredictionModel.forward`**
+   - 签名: `forward(self, events, tilt, h=None) -> (v_body, h_new)`
+   - `gru_out, h_new = self.gru(fused, h)`
+   - `model.step` 保持不变
+
+2. **`train.py::train_one_epoch`**
+   - 接受 `h_state` 参数,初值为 `None`
+   - `pred_seq, h_new = model(events, tilt, h_state)`
+   - `h_state = h_new.detach()` 传出供下一 batch
+   - 每个 epoch 重置 `h_state = None`
+
+3. **`train.py::validate`**
+   - 保持原样传 `h=None`(每 batch 独立窗口,与当前评估一致)
+
+4. **`dataset.py`**
+   - 当前 `create_train_loader` 使用 `shuffle=1000` → 序列打乱 → 无法跨 batch 传 h
+   - 方案:新增 `create_tbptt_loader`:
+     - `shuffle=0, deterministic=True`
+     - `stride=seq_len`(非重叠)
+     - `num_workers=0`(确保严格时序)
+     - 场景级 shuffle:在加载前打乱场景列表,但每个场景内帧序严格
+   - 直接将原始 `_build_pipeline` 参数由 `stride=1` 改为 `stride=seq_len`,关闭 shuffle
+
+### 注意事项
+
+- 梯度只回传到当前 batch 的 seq_len 步,不回溯更早 batch(detach 保证)
+- 首 batch `h=None` 仍从零初始化
+- BN 不受影响(CNN 在 per-frame 维度独立运行)
+- TBPTT loader 帧数 = floor(scene_frames / seq_len) * seq_len,末尾不足部分丢弃
+- 需起新 run_id(旧 checkpoint 全部在 h=0 重置模式下训练,不兼容)
+
+### 验证
+
+- `diag_state_mismatch.py` 重新评估:预期 B/A < 1.5
+- 同一场景下 stateful RMSE 接近 stateless RMSE
+
+---
+
+## P2 · 考虑中(未纳入本次)
+
+### 2.1 归一化被注释
+
+确认训练/评估两端 self-consistent(输出+target 都是原始 m/s),暂保留。
+
+### 2.2 验证集/测试集划分
+
+当前划分可调,涉及 config.py 场景列表调整和重训。
+
+### 2.3 速度平滑正则
+
+`((pred[:, 1:] - pred[:, :-1]) ** 2).mean()` 加权 `lambda=0.01`。建议 TBPTT 收敛后再加。
+
+### 2.4 Loss mask 前几步
+
+TBPTT 后 GRU 状态连贯,此问题可能自然缓解,暂不动。
+
+---
+
+## 落地顺序
+
+1. 改 `model.py::forward` 签名,返回 `(pred, h_new)`
+2. 改 `train.py` train + validate 适配新签名
+3. 新增/修改 loader 实现 TBPTT 模式(不 shuffle,非重叠)
+4. 起新 run_id 训练,观察 val_loss 走势
+5. 用 `diag_state_mismatch.py` 验证 B/A 比值
+
+---
+
+## 已完成
+
+以下为旧版 TODO 中已完成的项目,保留归档。
+
+### ~~P1 · 序列级损失监督~~(已完成,2025-07)
+
+已在 `train.py:54-59` + `model.py:149-153` 实现:
+- `model.forward` 返回 `(B, S, 2)`,对整段 `gru_out` 过 head
+- 全序列逐 step MSE loss `(B, S, 2) → (B, S) → scalar`
+- TensorBoard `loss_step_{i:02d}` 分位数记录
+- 验证通过:输出 shape `(B, 8, 2)`,TensorBoard 观察 loss_step 分布正常

analyse.md

@@ -0,0 +1,210 @@
+# UZH-FPV 速度预测 — 问题分析
+
+## 一、训练方案问题
+
+### 1.1 归一化被注释掉，训练与评估不一致
+
+`transforms.py:142,153` 中 `NormalizeVelocity()` 在 train 和 val 管线中均被注释。这意味着：
+
+- **训练时**模型直接回归原始速度值（`v_right` 约 -3~+3 m/s，`v_forward` 约 0~+8 m/s），数值范围大，梯度尺度不稳定，收敛慢。
+- **评估时** `evaluate.py:58-61` 和 `benchmark/evaluate.py:130-133` 却假设输出是归一化的，做了 `preds * std + mean` 反归一化。如果模型输出的是原始速度，反归一化后结果完全错误。
+- **结论**：要么启用 `NormalizeVelocity()` 并保持评估一致，要么去掉评估中的反归一化。当前状态是两边对不上。
+
+### 1.2 验证集与测试集重叠
+
+`config.py:40-53`：
+
+```python
+VAL_SCENES = ["indoor_forward_3"]
+TEST_SCENES = ["indoor_forward_3"]
+```
+
+验证集和测试集都是同一个场景 `indoor_forward_3`。这意味着：
+
+- 早停选择的 checkpoint 已经在这个场景上过拟合了，测试指标无意义。
+- 无法衡量泛化能力。
+
+### 1.3 训练集包含测试场景
+
+`config.py:31-38` 的 `TRAIN_SCENES` 包含 `outdoor_forward_1` 和 `outdoor_forward_5`，而 `TEST_SCENES` 中也有它们（注释中）。虽然当前 `TEST_SCENES` 只写了 `indoor_forward_3`，但注释里残留的测试场景与训练集重叠，容易误用。
+
+### 1.4 滑窗 stride = seq_len，无重叠
+
+`config.py:112`：
+
+```python
+sliding_window_stride: int = 32
+```
+
+`seq_len` 也是 32，所以滑窗不重叠。对于 1000 帧的场景，只产生约 31 个序列（1000/32），数据利用率低。通常 stride=1（全重叠）或 stride=seq_len//2（50% 重叠）能大幅增加训练样本量。
+
+### 1.5 训练时使用滑窗，评估时使用 stateful 逐帧推理，两者不一致
+
+- **训练**：`model(events, tilt)` 一次输入整个序列 (B, S, 1, H, W)，GRU 内部状态在序列内传播，但序列间重置。
+- **评估**：`model.step(events, tilt, h)` 逐帧推理，手动维护 GRU hidden state 跨帧传播。
+- 两种模式下的 GRU 行为不同：训练时每个序列从零状态开始，评估时状态持续累积。如果训练时序列长度不够覆盖相关时间尺度，评估时累积的长程状态可能产生分布偏移。
+
+### 1.6 损失函数只监督最后一帧
+
+`train.py:57-58`：
+
+```python
+pred = model(events, tilt)               # (B, 2)
+target_last = target[:, -1, :]           # (B, 2)
+loss = criterion(pred, target_last)
+```
+
+序列中前 S-1 帧完全没有损失信号。GRU 只在最后一帧收到梯度，前序时间步的隐藏状态更新缺乏直接监督。这可能导致 GRU 的中间状态退化。
+
+### 1.7 学习率调度器步长与 epoch 数不匹配
+
+`config.py:103,106`：
+
+```python
+epochs: int = 1000
+lr_scheduler_step: int = 30
+```
+
+StepLR 每 30 个 epoch 衰减一次 gamma=0.5。1000 epoch 内衰减约 33 次，最终学习率约为 `1e-3 * 0.5^33 ≈ 1.16e-13`，几乎为零。模型在后几百个 epoch 基本停止学习。
+
+---
+
+## 二、输入输出问题
+
+### 2.1 输出维度命名不一致
+
+多个地方对输出维度的命名互相矛盾：
+
+| 位置 | 命名 |
+|------|------|
+| `model.py:95` docstring | `[v_forward, v_lateral]` |
+| `model.py:134` docstring | `[v_forward, v_lateral]` |
+| `model.py:172` docstring | `[v_right, v_forward]` |
+| `config.py:86` | `[v_right, v_forward]` |
+| `transforms.py:107` | `[v_right, v_forward]` |
+| `evaluate.py:36` | `[v_right, v_forward]` |
+| `AGENTS.md` | `[v_right, v_forward]` |
+
+`model.py` 的 docstring 写的是 `[v_forward, v_lateral]`（前向、侧向），但实际代码和配置文件都使用 `[v_right, v_forward]`（右向、前向）。docstring 与实现不一致，容易误导。
+
+### 2.2 输入 tilt 的定义与使用存在歧义
+
+`transforms.py:77-90` 中 `ComputeTilt` 计算的是 **body up 向量**（世界坐标系下的机体上方向量），维度 (3,)，单位向量。
+
+但 `model.py` 中 PoseMLP 的输入标注为 `tilt`，docstring 写的是 "tilt rotation vector"。实际上输入是 body up 向量（编码 pitch/roll），不是 rotation vector（轴角表示）。命名和文档有误导性。
+
+### 2.3 速度归一化统计量可能不准确
+
+`config.py:15-16`：
+
+```python
+VELOCITY_MEAN = [-2.902497, 3.837231]
+VELOCITY_STD  = [3.453774,  3.722085]
+```
+
+注释说 "computed from forward scenes only, 28363 frames"。但：
+
+- 归一化被注释掉了，这些统计量从未被使用。
+- 如果未来启用，需要确认这些统计量是否覆盖了所有训练场景的分布，特别是 45° 飞行场景（侧向速度更大）。
+
+---
+
+## 三、模型架构问题
+
+### 3.1 CNN 编码器被禁用（输出全零）
+
+`model.py:138-139` 注释掉的代码：
+
+```python
+# cnn_feat = events.new_zeros(B, S, self.cnn.out_dim) # 全零替代
+```
+
+当前 CNN 实际在运行（`cnn_feat = self.cnn(events)`），但注释表明曾经尝试过禁用 CNN。如果 CNN 输出被置零，模型完全依赖 PoseMLP + GRU，相当于只用 tilt 信息预测速度，事件帧信息被丢弃。这与项目目标（从事件相机预测速度）矛盾。
+
+### 3.2 CNN 架构对事件帧不友好
+
+`model.py:28-36` 的 CNN 使用标准 Conv-BN-ReLU-Pool 结构，但事件帧是稀疏的（大部分像素为 0，只有边缘处为 ±1）：
+
+- **BatchNorm** 在稀疏输入上统计量不稳定（均值和方差被大量零像素拉偏）。
+- **MaxPool** 对稀疏信号不友好：如果事件只占几个像素，Pool 窗口内最大值可能始终为 0。
+- 没有使用空洞卷积或更大的感受野来捕获事件的空间结构。
+
+### 3.3 PoseMLP 容量可能不足
+
+`config.py:69-71`：
+
+```python
+input_dim: int = 3
+hidden_dim: int = 32
+output_dim: int = 64
+```
+
+PoseMLP 只有 2 层线性层（3→32→64），隐藏层仅 32 维。对于编码 tilt 信息（pitch/roll 的非线性映射），容量可能偏低。
+
+### 3.4 GRU 单层且无 dropout
+
+`config.py:77-79`：
+
+```python
+hidden_size: int = 128
+num_layers: int = 1
+dropout: float = 0.0
+```
+
+- 单层 GRU 表达能力有限，难以建模长时间依赖。
+- 无 dropout，容易过拟合（特别是训练数据量不大时）。
+
+### 3.5 Head 输出层初始化过小
+
+`model.py:122-123`：
+
+```python
+self.head[-1].weight.data.mul_(0.01)
+self.head[-1].bias.data.zero_()
+```
+
+输出层权重缩小 100 倍，初始输出接近零。这有助于训练初期避免大梯度，但也意味着模型需要更多迭代才能"激活"输出层。如果训练 epoch 不够，模型可能一直输出接近零的值。
+
+### 3.6 缺少序列级损失监督
+
+模型只对最后一帧计算 loss（`target[:, -1, :]`），GRU 的中间时间步输出被完全忽略。可以添加辅助损失：
+
+- 每个时间步的预测与对应 GT 的损失（teacher forcing 风格）
+- 速度变化量的平滑性约束（相邻帧速度差的正则化）
+
+### 3.7 没有位置编码或时间信息
+
+模型输入不包含时间戳或帧间间隔信息。事件帧之间的时间间隔可能不均匀（实际 DAVIS 相机帧率有波动），但模型假设所有帧等间隔。
+
+---
+
+## 四、数据管线问题
+
+### 4.1 SimulateEvents 跨 shard 状态残留（已知问题）
+
+`AGENTS.md` 已记录：`EventProcessor` 的 `_prev_frame` 在 shard 边界不重置，导致新 shard 第一帧的事件帧错误。
+
+### 4.2 滑窗跨 shard 边界（已知问题）
+
+`AGENTS.md` 已记录：滑窗不感知 shard 边界，序列可能跨越两个 shard。
+
+### 4.3 验证集使用滑窗而非逐帧
+
+`dataset.py:139` 中验证集也使用滑窗（`stride=32`），与评估时的逐帧 stateful 推理不一致。验证 loss 不能反映实际推理时的性能。
+
+---
+
+## 五、总结优先级
+
+| 优先级 | 问题 | 影响 |
+|--------|------|------|
+| P0 | 归一化被注释 + 评估假设归一化 | 评估结果完全错误 |
+| P0 | 验证集 = 测试集 | 无法评估泛化 |
+| P1 | CNN 可能被禁用 | 事件信息被丢弃 |
+| P1 | 只监督最后一帧 | GRU 中间状态无梯度 |
+| P1 | 训练滑窗 vs 评估 stateful 不一致 | 训练/评估分布偏移 |
+| P2 | 输出维度命名混乱 | 代码可维护性差 |
+| P2 | 学习率衰减过快 | 后期训练无效 |
+| P2 | 滑窗无重叠 | 数据利用率低 |
+| P3 | CNN 架构对稀疏事件不友好 | 特征提取效率低 |
+| P3 | 缺少时间编码 | 忽略帧间间隔变化 |

benchmark/benchmark.py

@@ -1,309 +0,0 @@
-"""
-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

@@ -1,98 +0,0 @@
-"""
-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

@@ -1,458 +0,0 @@
-"""
-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()}")

src/velocity_prediction/config.py

@@ -39,7 +39,8 @@ VAL_SCENES = [
     # "indoor_forward_3", "indoor_forward_9", "indoor_forward_10",  # Easy
 ]
 TEST_SCENES = [
-    "indoor_forward_9","indoor_forward_3",
+    "indoor_forward_9",
+    # "indoor_forward_9","indoor_forward_3",
     # "indoor_forward_7",      # Hard 室内
     # "outdoor_forward_1",     # Easy 室外
     # "outdoor_forward_5"      # Hard 室外
@@ -92,8 +93,8 @@ class ModelConfig:
 
 @dataclass
 class TrainConfig:
-    seq_len: int = 8                    # frames per training sequence
-    batch_size: int = 32
+    seq_len: int = 128                    # frames per training sequence
+    batch_size: int = 4
     epochs: int = 300
     lr: float = 1e-3
     weight_decay: float = 1e-5

src/velocity_prediction/dataset.py

@@ -145,3 +145,48 @@ def create_val_loader(
         shuffle=False,
     )
     return loader
+
+
+def create_tbptt_loader(
+    scene_names: Optional[List[str]] = None,
+    seq_len: int = 8,
+    batch_size: int = 32,
+    event_threshold: float = 0.1,
+    event_use_log: bool = True,
+):
+    """Create a DataLoader for TBPTT training.
+
+    Non-overlapping windows (stride=seq_len), no shuffle, single worker
+    for strict temporal order.  Scene order is shuffled per epoch but
+    frames within each scene are always in temporal sequence, so GRU
+    hidden state can be carried across consecutive batches.
+
+    Frames per scene = floor(scene_frames / seq_len) * seq_len;
+    trailing frames that don't fill a full window are dropped.
+    """
+    import random
+    if scene_names is None:
+        from src.velocity_prediction.config import TRAIN_SCENES
+        scene_names = TRAIN_SCENES
+
+    urls = _scene_urls(scene_names)
+    # Scene-level shuffle: randomize scene order, but each scene's frames
+    # are strictly in temporal order (essential for cross-batch TBPTT).
+    random.shuffle(urls)
+
+    transform = build_train_transform(
+        event_threshold=event_threshold,
+        event_use_log=event_use_log,
+    )
+    pipeline = _build_pipeline(
+        urls, transform, seq_len=seq_len, stride=seq_len,
+        shuffle=0, deterministic=True,
+    )
+
+    loader = wds.WebLoader(
+        pipeline,
+        batch_size=batch_size,
+        num_workers=0,          # strict temporal ordering
+        shuffle=False,
+    )
+    return loader

src/velocity_prediction/evaluate.py

@@ -15,54 +15,52 @@ 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
+from src.velocity_prediction.config import train_cfg
 
 
 @torch.no_grad()
-def evaluate(
+def evaluate_stateful(
     model: nn.Module,
     loader,
     device: torch.device,
 ) -> dict:
     """
-    Run evaluation on a dataloader.
+    Stateful evaluation: process frames one-by-one, maintaining GRU
+    hidden state across timesteps for full temporal context.
+
+    The loader must yield single-frame samples (seq_len=1, stride=1)
+    in strict temporal order (num_workers=0).
 
     Returns:
         dict with keys:
-            preds:  np.ndarray (N, 2)  predicted [vx, vy]
-            targets: np.ndarray (N, 2) ground truth [vx, vy]
+            preds:  np.ndarray (N, 2)  predicted [v_right, v_forward]
+            targets: np.ndarray (N, 2) ground truth
     """
     model.eval()
     all_preds = []
     all_targets = []
+    h = None
 
     for batch in loader:
-        events = batch["events"].to(device)
-        tilt = batch["tilt"].to(device)
-        target = batch["v_body_target"].to(device)  # (B, S, 2)
+        events = batch["events"].to(device)          # (B=1, S=1, 1, H, W)
+        tilt = batch["tilt"].to(device)              # (B=1, S=1, 3)
+        target = batch["v_body_target"].to(device)   # (B=1, S=1, 2)
 
-        pred = model(events, tilt)                   # (B, 2)
-        target_last = target[:, -1, :]               # (B, 2)
+        pred, h = model.step(events, tilt, h)       # (1, 2), (1, 1, 128)
 
         all_preds.append(pred.cpu().numpy())
-        all_targets.append(target_last.cpu().numpy())
+        all_targets.append(target[:, -1, :].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) ==========")
+    # ── Diagnostics (in model output space — original m/s) ─────────
+    print("\n========== Evaluation Diagnostics (model output space) ==========")
     print(f"Total samples: {len(preds)}")
-    print(f"\n--- Targets (normalized) ---")
+    print(f"\n--- Targets ---")
     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"\n--- Predictions ---")
     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}, "
@@ -79,14 +77,14 @@ def evaluate(
     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)))
+    # Per-axis and overall RMSE (in original m/s space)
+    rmse_x = np.sqrt(np.mean((preds[:, 0] - targets[:, 0]) ** 2))
+    rmse_y = np.sqrt(np.mean((preds[:, 1] - targets[:, 1]) ** 2))
+    rmse_xy = np.sqrt(np.mean(np.sum((preds - targets) ** 2, axis=1)))
 
     return {
-        "preds": preds_denorm,      # denormalized for plotting
-        "targets": targets_denorm,  # denormalized for plotting
+        "preds": preds,      # original m/s space
+        "targets": targets,  # original m/s space
         "rmse_x": rmse_x,
         "rmse_y": rmse_y,
         "rmse_xy": rmse_xy,
@@ -178,13 +176,14 @@ def main():
     for scene in TEST_SCENES:
         loader = create_val_loader(
             scene_names=[scene],
-            seq_len=train_cfg.seq_len,
-            batch_size=train_cfg.batch_size,
-            num_workers=2,
+            seq_len=1,
+            stride=1,
+            batch_size=1,
+            num_workers=0,               # strict temporal order
             event_threshold=train_cfg.event_threshold,
             event_use_log=train_cfg.event_use_log,
         )
-        results = evaluate(model, loader, device)
+        results = evaluate_stateful(model, loader, device)
         n = len(results["preds"])
         print(f"  [{scene}] RMSE vx={results['rmse_x']:.4f}  vy={results['rmse_y']:.4f}  "
               f"xy={results['rmse_xy']:.4f}  samples={n}")

src/velocity_prediction/model.py

@@ -86,13 +86,13 @@ class PoseMLP(nn.Module):
 
 class VelocityPredictionModel(nn.Module):
     """
-    Full model: CNN + PoseMLP → concat → GRU → Head → [vx, vy].
+    Full model: CNN + PoseMLP → concat → GRU → Head → [v_right, v_forward].
 
     Input:
         events: (B, S, 1, H, W)
         tilt:   (B, S, 3)
     Output:
-        v_body: (B, 2)  — body-frame [v_forward, v_lateral] for the last frame in the sequence
+        v_body: (B, S, 2)  — body-frame [v_right, v_forward] for each frame in the sequence
     """
 
     def __init__(self, cnn_cfg=model_cfg.cnn, pose_cfg=model_cfg.pose_mlp,
@@ -124,15 +124,20 @@ class VelocityPredictionModel(nn.Module):
         # nn.init.uniform_(self.head[-1].weight, -0.001, 0.001)
         # nn.init.zeros_(self.head[-1].bias)
 
-    def forward(self, events: torch.Tensor, tilt: torch.Tensor) -> torch.Tensor:
+    def forward(self, events: torch.Tensor, tilt: torch.Tensor,
+                h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Args:
             events: (B, S, 1, H, W)
             tilt:   (B, S, 3)
+            h:      (num_layers, B, hidden_size) or None → zeros (TBPTT state)
 
         Returns:
-            v_body: (B, 2)  predicted body-frame [v_forward, v_lateral] at the last timestep
+            v_body: (B, S, 2)  predicted body-frame [v_right, v_forward] at every timestep
+            h_new:  (num_layers, B, hidden_size)  final GRU hidden state (detached by caller)
         """
+        B, S = events.shape[:2]
+
         # Per-frame encoding
         cnn_feat = self.cnn(events)           # (B, S, 256)
         # B, S = events.shape[:2]
@@ -143,15 +148,39 @@ class VelocityPredictionModel(nn.Module):
         # 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)
+        # GRU temporal modelling — accepts external hidden state for TBPTT
+        gru_out, h_new = self.gru(fused, h)    # (B, S, 128), (num_layers, B, 128)
 
-        # Use last hidden state
-        last_hidden = h_n[-1]                 # (B, 128)
+        # Head regression over every timestep (single-layer GRU → gru_out == h_n[-1] at t=-1)
+        v_body = self.head(gru_out)           # (B, S, 128) → (B, S, 2)
+        return v_body, h_new
 
-        # Head regression
+    @torch.no_grad()
+    def step(self, events: torch.Tensor, tilt: torch.Tensor,
+             h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Single-frame stateful forward pass.
+
+        Processes one timestep and returns the prediction plus updated
+        GRU hidden state.  Call repeatedly with the output ``h`` to
+        accumulate temporal context across frames.
+
+        Args:
+            events: (B, 1, 1, H, W)  single-frame event map
+            tilt:   (B, 1, 3)         single-frame tilt vector
+            h:      (num_layers, B, hidden_size) or None  → zeros
+
+        Returns:
+            v_body: (B, 2)  body-frame [v_right, v_forward]
+            h_new:  (num_layers, B, hidden_size)
+        """
+        cnn_feat = self.cnn(events)          # (B, 1, 256)
+        pose_feat = self.pose_mlp(tilt)      # (B, 1, 64)
+        fused = torch.cat([cnn_feat, pose_feat], dim=-1)  # (B, 1, 320)
+        _, h_new = self.gru(fused, h)        # (num_layers, B, 128)
+        last_hidden = h_new[-1]              # (B, 128)
         v_body = self.head(last_hidden)      # (B, 2)
-        return v_body
+        return v_body, h_new
 
 
 def count_parameters(model: nn.Module) -> int:
@@ -169,7 +198,7 @@ if __name__ == "__main__":
     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)
+    out, h = model(events, tilt)
     print(f"Input events: {events.shape}")
     print(f"Input tilt:   {tilt.shape}")
-    print(f"Output:       {out.shape}  (should be [4, 2])")
+    print(f"Output:       {out.shape}  (should be [4, 8, 2])")

src/velocity_prediction/train.py

@@ -18,7 +18,7 @@ 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
+from src.velocity_prediction.dataset import create_val_loader, create_tbptt_loader
 
 
 def set_seed(seed: int):
@@ -39,8 +39,9 @@ def train_one_epoch(
     log_interval: int = 50,
     global_step: int = 0,
     use_amp: bool = True,
-) -> tuple[float, int]:
-    """Train for one epoch. Returns (avg_loss, updated_global_step)."""
+    h_state: torch.Tensor = None,
+) -> tuple[float, int, torch.Tensor]:
+    """Train for one epoch. Returns (avg_loss, updated_global_step, h_state_final)."""
     model.train()
     total_loss = 0.0
     num_batches = 0
@@ -51,17 +52,25 @@ def train_one_epoch(
         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
+        # Drop carried state if batch size changed (partial last batch / new scene)
+        if h_state is not None and h_state.shape[1] != events.shape[0]:
+            h_state = None
+
+        # Per-step supervision over the whole sequence — TBPTT hidden state
         with torch.amp.autocast(device.type, enabled=use_amp):
-            pred = model(events, tilt)               # (B, 2)
-            target_last = target[:, -1, :]           # (B, 2)
-            loss = criterion(pred, target_last)
+            pred_seq, h_new = model(events, tilt, h_state)   # (B, S, 2), (L, B, H)
+            loss_per_step = criterion(pred_seq, target)      # (B, S, 2)
+            loss_per_step = loss_per_step.mean(-1)           # (B, S)
+            loss = loss_per_step.mean()
 
         optimizer.zero_grad()
         scaler.scale(loss).backward()
         scaler.step(optimizer)
         scaler.update()
 
+        # Detach hidden state for next batch — gradient only back to current seq_len
+        h_state = h_new.detach()
+
         total_loss += loss.item()
         num_batches += 1
         global_step += 1
@@ -70,10 +79,15 @@ def train_one_epoch(
             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(), global_step)
+            # Per-step loss quantiles — observe GRU stabilization across timesteps
+            S = loss_per_step.shape[1]
+            for i in range(0, S, max(1, S // 8)):
+                writer.add_scalar(f"train/loss_step_{i:02d}",
+                                  loss_per_step[:, i].mean().item(), global_step)
 
     avg_loss = total_loss / max(num_batches, 1)
     print(f"  Epoch {epoch} | Avg Loss: {avg_loss:.6f}")
-    return avg_loss, global_step
+    return avg_loss, global_step, h_state
 
 
 @torch.no_grad()
@@ -95,9 +109,9 @@ def validate(
         target = batch["v_body_target"].to(device)
 
         with torch.amp.autocast(device.type, enabled=use_amp):
-            pred = model(events, tilt)
-            target_last = target[:, -1, :]
-            loss = criterion(pred, target_last)
+            pred_seq, _ = model(events, tilt)          # (B, S, 2), h=None stateless
+            loss_per_step = criterion(pred_seq, target)  # (B, S, 2)
+            loss = loss_per_step.mean(-1).mean()
 
         total_loss += loss.item()
         num_batches += 1
@@ -114,6 +128,8 @@ def main():
                         help="Path to checkpoint .pt file to resume training from")
     parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True,
                         help="Enable Automatic Mixed Precision (default: True)")
+    parser.add_argument("--threshold", type=float, default=None,
+                        help="Event brightness change threshold (overrides config, default: same as train_cfg)")
     args = parser.parse_args()
     use_amp = args.amp
 
@@ -121,6 +137,11 @@ def main():
     device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
     print(f"Device: {device}")
 
+    # Override threshold from CLI if provided
+    event_threshold = args.threshold if args.threshold is not None else train_cfg.event_threshold
+    if args.threshold is not None:
+        print(f"Event threshold overridden: {train_cfg.event_threshold} → {event_threshold}")
+
     # Create model
     model = VelocityPredictionModel()
     model.to(device)
@@ -129,13 +150,11 @@ def main():
     print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
     print(f"AMP: {'enabled' if use_amp else 'disabled'}")
 
-    # Data loaders
-    train_loader = create_train_loader(
+    # Data loaders — TBPTT training requires strict temporal order
+    train_loader = create_tbptt_loader(
         seq_len=train_cfg.seq_len,
-        stride=train_cfg.sliding_window_stride,
         batch_size=train_cfg.batch_size,
-        num_workers=train_cfg.num_workers,
-        event_threshold=train_cfg.event_threshold,
+        event_threshold=event_threshold,
         event_use_log=train_cfg.event_use_log,
     )
     val_loader = create_val_loader(
@@ -143,7 +162,7 @@ def main():
         stride=train_cfg.sliding_window_stride,
         batch_size=train_cfg.batch_size,
         num_workers=train_cfg.num_workers,
-        event_threshold=train_cfg.event_threshold,
+        event_threshold=event_threshold,
         event_use_log=train_cfg.event_use_log,
     )
 
@@ -158,11 +177,8 @@ def main():
         step_size=train_cfg.lr_scheduler_step,
         gamma=train_cfg.lr_scheduler_gamma,
     )
-    # criterion = nn.SmoothL1Loss()
-    criterion = nn.MSELoss()
-
-    ckpt_dir = Path(train_cfg.checkpoint_dir)
-    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    # criterion = nn.SmoothL1Loss(reduction='none')
+    criterion = nn.MSELoss(reduction='none')
 
     # ── Resume from checkpoint ────────────────────────────────────
     start_epoch = 1
@@ -191,11 +207,13 @@ def main():
     else:
         print(f"\nStarting training for {train_cfg.epochs} epochs...")
 
-    # Logging — run-specific subdirectory for isolation + resume continuity
+    # Logging & checkpoint — run-specific subdirectories for isolation + resume continuity
     if run_id is None:
         run_id = time.strftime("run_%Y%m%d_%H%M%S")
     log_dir = Path(train_cfg.log_dir) / run_id
     log_dir.mkdir(parents=True, exist_ok=True)
+    ckpt_dir = Path(train_cfg.checkpoint_dir) / run_id
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
     writer = SummaryWriter(log_dir=str(log_dir))
 
     print(f"  seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}")
@@ -205,11 +223,13 @@ def main():
     for epoch in range(start_epoch, train_cfg.epochs + 1):
         epoch_start = time.time()
 
-        train_loss, global_step = train_one_epoch(
+        # Reset GRU hidden state at epoch start — each epoch begins with h=0
+        train_loss, global_step, _ = train_one_epoch(
             model, train_loader, optimizer, criterion, scaler, device, epoch, writer,
             log_interval=train_cfg.log_interval,
             global_step=global_step,
             use_amp=use_amp,
+            h_state=None,
         )
         val_loss = validate(model, val_loader, criterion, device, use_amp=use_amp)
         scheduler.step()