refactor: stateless training with randomized seq_len per epoch

- stateless: each batch starts from zero GRU hidden state
- curriculum: random seq_len from [64,96,128,160,192] per epoch
- val fixed at seq_len=128 for consistent cross-epoch comparison
- remove 256 from choices (OOM at B=4 on 24GB VRAM)

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

src/velocity_prediction/train.py

@@ -8,6 +8,7 @@ Usage:
 
 import argparse
 import os
+import random
 import time
 import numpy as np
 from pathlib import Path
@@ -39,9 +40,9 @@ def train_one_epoch(
     log_interval: int = 50,
     global_step: int = 0,
     use_amp: bool = True,
-    h_state: torch.Tensor = None,
+) -> tuple[float, int]:
-) -> tuple[float, int, torch.Tensor]:
+    """Train for one epoch (stateless — each batch starts from zero hidden state).
-    """Train for one epoch. Returns (avg_loss, updated_global_step, h_state_final)."""
+    Returns (avg_loss, updated_global_step)."""
     model.train()
     total_loss = 0.0
     num_batches = 0
@@ -52,15 +53,11 @@ def train_one_epoch(
         tilt = batch["tilt"].to(device)              # (B, S, 3)
         target = batch["v_body_target"].to(device)   # (B, S, 2)
 
-        # Drop carried state if batch size changed (partial last batch / new scene)
+        # Stateless: each batch starts from zero hidden state
-        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_seq, h_new = model(events, tilt, h_state)   # (B, S, 2), (L, B, H)
+            pred_seq, _ = model(events, tilt, None)     # (B, S, 2), h=None
-            loss_per_step = criterion(pred_seq, target)      # (B, S, 2)
+            loss_per_step = criterion(pred_seq, target)  # (B, S, 2)
-            loss_per_step = loss_per_step.mean(-1)           # (B, S)
+            loss_per_step = loss_per_step.mean(-1)       # (B, S)
             loss = loss_per_step.mean()
 
         optimizer.zero_grad()
@@ -68,9 +65,6 @@ def train_one_epoch(
         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
@@ -87,7 +81,7 @@ def train_one_epoch(
 
     avg_loss = total_loss / max(num_batches, 1)
     print(f"  Epoch {epoch} | Avg Loss: {avg_loss:.6f}")
-    return avg_loss, global_step, h_state
+    return avg_loss, global_step
 
 
 @torch.no_grad()
@@ -150,15 +144,12 @@ def main():
     print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
     print(f"AMP: {'enabled' if use_amp else 'disabled'}")
 
-    # Data loaders — TBPTT training requires strict temporal order
+    # ── Randomised sequence lengths for curriculum ──────────────
-    train_loader = create_tbptt_loader(
+    SEQ_LEN_CHOICES = [64, 96, 128, 160, 192]
-        seq_len=train_cfg.seq_len,
+
-        batch_size=train_cfg.batch_size,
+    # Validation uses a fixed seq_len for consistent comparison
-        event_threshold=event_threshold,
-        event_use_log=train_cfg.event_use_log,
-    )
     val_loader = create_val_loader(
-        seq_len=train_cfg.seq_len,
+        seq_len=128,
         stride=train_cfg.sliding_window_stride,
         batch_size=train_cfg.batch_size,
         num_workers=train_cfg.num_workers,
@@ -216,20 +207,28 @@ def main():
     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}")
+    print(f"  seq_len=random({min(SEQ_LEN_CHOICES)}~{max(SEQ_LEN_CHOICES)}), "
+          f"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(start_epoch, train_cfg.epochs + 1):
         epoch_start = time.time()
 
-        # Reset GRU hidden state at epoch start — each epoch begins with h=0
+        # Random seq_len per epoch — GRU learns to handle varying temporal horizons
-        train_loss, global_step, _ = train_one_epoch(
+        seq_len = random.choice(SEQ_LEN_CHOICES)
+        train_loader = create_tbptt_loader(
+            seq_len=seq_len,
+            batch_size=train_cfg.batch_size,
+            event_threshold=event_threshold,
+            event_use_log=train_cfg.event_use_log,
+        )
+
+        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()
@@ -238,6 +237,7 @@ def main():
         current_lr = scheduler.get_last_lr()[0]
 
         print(f"Epoch {epoch:3d}/{train_cfg.epochs} | "
+              f"seq_len={seq_len:3d} | "
               f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | "
               f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")