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/train.py

@@ -51,11 +51,12 @@ 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
+        # Per-step supervision over the whole sequence
         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 = model(events, tilt)               # (B, S, 2)
+            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()
@@ -70,6 +71,11 @@ 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}")
@@ -95,9 +101,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)
+            loss_per_step = criterion(pred_seq, target)  # (B, S, 2)
+            loss = loss_per_step.mean(-1).mean()
 
         total_loss += loss.item()
         num_batches += 1
@@ -165,8 +171,8 @@ def main():
         step_size=train_cfg.lr_scheduler_step,
         gamma=train_cfg.lr_scheduler_gamma,
     )
-    # criterion = nn.SmoothL1Loss()
-    criterion = nn.MSELoss()
+    # criterion = nn.SmoothL1Loss(reduction='none')
+    criterion = nn.MSELoss(reduction='none')
 
     # ── Resume from checkpoint ────────────────────────────────────
     start_epoch = 1