feat(trainer): replace eager _augment with streaming TrajectoryDataset

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

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()