feat: add stateful evaluation with model.step() for per-frame GRU inference

- Add VelocityPredictionModel.step() for single-frame forward pass
  with external GRU hidden state management
- Add evaluate_stateful() that processes frames sequentially,
  maintaining hidden state across timesteps (seq_len=1, stride=1, batch=1)
- Evaluation uses strict temporal ordering (num_workers=0)

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

src/velocity_prediction/evaluate.py

@@ -19,33 +19,37 @@ from src.velocity_prediction.config import train_cfg, VELOCITY_MEAN, VELOCITY_ST
 
 
 @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]
+            preds:  np.ndarray (N, 2)  predicted [v_right, v_forward]
-            targets: np.ndarray (N, 2) ground truth [vx, vy]
+            targets: np.ndarray (N, 2) ground truth
     """
     model.eval()
     all_preds = []
     all_targets = []
+    h = None
 
     for batch in loader:
-        events = batch["events"].to(device)
+        events = batch["events"].to(device)          # (B=1, S=1, 1, H, W)
-        tilt = batch["tilt"].to(device)
+        tilt = batch["tilt"].to(device)              # (B=1, S=1, 3)
-        target = batch["v_body_target"].to(device)  # (B, S, 2)
+        target = batch["v_body_target"].to(device)   # (B=1, S=1, 2)
 
-        pred = model(events, tilt)                   # (B, 2)
+        pred, h = model.step(events, tilt, h)       # (1, 2), (1, 1, 128)
-        target_last = target[:, -1, :]               # (B, 2)
 
         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)
@@ -178,13 +182,14 @@ def main():
     for scene in TEST_SCENES:
         loader = create_val_loader(
             scene_names=[scene],
-            seq_len=train_cfg.seq_len,
+            seq_len=1,
-            batch_size=train_cfg.batch_size,
+            stride=1,
-            num_workers=2,
+            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

@@ -153,6 +153,33 @@ class VelocityPredictionModel(nn.Module):
         v_body = self.head(last_hidden)       # (B, 2)
         return v_body
 
+    @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, h_new
+
 
 def count_parameters(model: nn.Module) -> int:
     """Count trainable parameters."""