feat: per-step supervision over full sequence

- model outputs (B, S, 2) instead of (B, 2) — GRU output at every timestep
- train/val loss computed over all S timesteps with reduction=none
- benchmark/evaluate.py takes pred[:, -1, :] for final-step evaluation
- added per-step loss logging (8 evenly spaced steps) to TensorBoard

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

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,
@@ -131,8 +131,10 @@ class VelocityPredictionModel(nn.Module):
             tilt:   (B, S, 3)
 
         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
         """
+        B, S = events.shape[:2]
+
         # Per-frame encoding
         cnn_feat = self.cnn(events)           # (B, S, 256)
         # B, S = events.shape[:2]
@@ -144,13 +146,10 @@ class VelocityPredictionModel(nn.Module):
         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_out, _ = self.gru(fused)           # (B, S, 128)
 
-        # Use last hidden state
-        last_hidden = h_n[-1]                 # (B, 128)
-
-        # Head regression
-        v_body = self.head(last_hidden)       # (B, 2)
+        # 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
 
     @torch.no_grad()
@@ -199,4 +198,4 @@ if __name__ == "__main__":
     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])")
+    print(f"Output:       {out.shape}  (should be [4, 8, 2])")