feat(trainer): replace eager _augment with streaming TrajectoryDataset

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ai_mouse/trainer.py

@@ -21,7 +21,7 @@ from pathlib import Path
 
 import numpy as np
 import torch
-from torch.utils.data import DataLoader, TensorDataset
+from torch.utils.data import DataLoader
 
 from ai_mouse.config import TrainConfig
 from ai_mouse.coord import encode_trajectory
@@ -226,6 +226,62 @@ def _augment(
     return np.concatenate(seqs, axis=0), np.concatenate(conds, axis=0)
 
 
+class TrajectoryDataset(torch.utils.data.Dataset):
+    """Trajectory dataset with on-the-fly 6× augmentation.
+
+    Replaces the old eager `_augment(seq, cond)` which expanded the dataset
+    6× in memory before training. With this class, the original (N, T, 3)
+    arrays stay as-is and each `__getitem__` call computes one of the 6
+    augmentation variants on demand.
+
+    Augmentation variants (matching legacy `_augment` semantics):
+      0 — original
+      1 — lateral flip (lateral → −lateral)
+      2 — speed ×0.8  (log_dt[1:] += log(1.25),  cond[2] += log(1.25))
+      3 — speed ×1.2  (log_dt[1:] += log(1/1.2), cond[2] += log(1/1.2))
+      4 — temporal noise (log_dt[1:] += N(0, 0.05))
+      5 — flip + speed ×0.9 (lateral flip, log_dt[1:] += log(1/0.9), cond[2] += log(1/0.9))
+    """
+
+    _LOG_1_25 = math.log(1.25)
+    _LOG_INV_1_2 = math.log(1.0 / 1.2)
+    _LOG_1_1 = math.log(1.0 / 0.9)
+
+    def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
+        self.seq = seq
+        self.cond = cond
+        self.augment = augment
+        self._n_aug = 6 if augment else 1
+
+    def __len__(self) -> int:
+        return len(self.seq) * self._n_aug
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        base = idx // self._n_aug
+        aug_id = idx % self._n_aug
+
+        s = self.seq[base].copy()
+        c = self.cond[base].copy()
+
+        if aug_id == 1:
+            s[:, 1] = -s[:, 1]
+        elif aug_id == 2:
+            s[1:, 2] += self._LOG_1_25
+            c[2] += self._LOG_1_25
+        elif aug_id == 3:
+            s[1:, 2] += self._LOG_INV_1_2
+            c[2] += self._LOG_INV_1_2
+        elif aug_id == 4:
+            noise = np.random.normal(0.0, 0.05, size=s[1:, 2].shape).astype(np.float32)
+            s[1:, 2] += noise
+        elif aug_id == 5:
+            s[:, 1] = -s[:, 1]
+            s[1:, 2] += self._LOG_1_1
+            c[2] += self._LOG_1_1
+
+        return torch.from_numpy(s), torch.from_numpy(c)
+
+
 # ---------------------------------------------------------------------------
 # Duration distribution (per distance bin)
 # ---------------------------------------------------------------------------
@@ -328,12 +384,11 @@ def train(
     seq_np, cond_np, click_durs = load_and_prepare_data(data_path, seq_len=seq_len)
     logger.info("Loaded %d traces", len(seq_np))
 
-    # ---- Augment ----
+    # ---- Build streaming dataset (on-the-fly 6× augmentation) ----
-    seq_np, cond_np = _augment(seq_np, cond_np)
+    if config.augment:
-    logger.info("After augmentation: %d samples", len(seq_np))
+        logger.info("Using on-the-fly 6× augmentation, base samples: %d", len(seq_np))
-
+    ds = TrajectoryDataset(seq_np, cond_np, augment=config.augment)
-    seq_t = torch.from_numpy(seq_np)    # (N, seq_len, 3)
+    logger.info("Effective dataset size: %d", len(ds))
-    cond_t = torch.from_numpy(cond_np)  # (N, 3)
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -350,7 +405,6 @@ def train(
     optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)
 
-    ds = TensorDataset(seq_t, cond_t)
     loader = DataLoader(ds, batch_size=batch_size, shuffle=True, drop_last=False)
 
     # ---- Training loop: OT-Conditional Flow Matching ----

tests/test_trainer.py

@@ -119,3 +119,54 @@ class TestTrain:
         first_half = np.mean(losses[:10])
         second_half = np.mean(losses[10:])
         assert second_half < first_half
+
+
+class TestTrajectoryDataset:
+    def test_dataset_length_with_augmentation(self):
+        """Dataset length = N * 6 when augment=True."""
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((10, 64, 3), dtype=np.float32)
+        cond = np.zeros((10, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        assert len(ds) == 60
+
+    def test_dataset_length_without_augmentation(self):
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((10, 64, 3), dtype=np.float32)
+        cond = np.zeros((10, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=False)
+        assert len(ds) == 10
+
+    def test_getitem_returns_tensors(self):
+        from ai_mouse.trainer import TrajectoryDataset
+        import torch
+        seq = np.random.randn(5, 64, 3).astype(np.float32)
+        cond = np.random.randn(5, 3).astype(np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s, c = ds[0]
+        assert isinstance(s, torch.Tensor)
+        assert isinstance(c, torch.Tensor)
+        assert s.shape == (64, 3)
+        assert c.shape == (3,)
+
+    def test_aug_id_zero_returns_original(self):
+        """Aug id 0 (idx=0 % 6 == 0) should return the original sample unchanged."""
+        from ai_mouse.trainer import TrajectoryDataset
+        import torch
+        seq = np.array([[[0.5, 0.7, 0.3]] * 64] * 3, dtype=np.float32)
+        cond = np.array([[1.0, 2.0, 3.0]] * 3, dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s0, c0 = ds[0]
+        np.testing.assert_allclose(s0.numpy(), seq[0], rtol=1e-5)
+        np.testing.assert_allclose(c0.numpy(), cond[0], rtol=1e-5)
+
+    def test_aug_id_one_flips_lateral(self):
+        """Aug id 1 should flip the sign of the lateral channel (index 1)."""
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        seq[0, :, 1] = 0.5  # lateral all positive
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        # idx=1 → base_idx=0, aug_id=1 → flip
+        s1, _ = ds[1]
+        assert (s1[:, 1] < 0).all()