From 1d9dff4f2a9858590ec895c97479738766012021 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Sun, 10 May 2026 13:11:27 +0800
Subject: [PATCH] =?UTF-8?q?fix(trainer):=20rename=20=5FLOG=5F1=5F1=20?=
 =?UTF-8?q?=E2=86=92=20=5FLOG=5FINV=5F0=5F9,=20add=20tests=20for=20aug=5Fi?=
 =?UTF-8?q?ds=202-5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ai_mouse/trainer.py   |  6 ++---
 tests/test_trainer.py | 57 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/ai_mouse/trainer.py b/ai_mouse/trainer.py
index 73cdc41..cc5fdfd 100644
--- a/ai_mouse/trainer.py
+++ b/ai_mouse/trainer.py
@@ -245,7 +245,7 @@ class TrajectoryDataset(torch.utils.data.Dataset):
 
     _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)
+    _LOG_INV_0_9 = math.log(1.0 / 0.9)
 
     def __init__(self, seq: np.ndarray, cond: np.ndarray, augment: bool = True):
         self.seq = seq
@@ -276,8 +276,8 @@ class TrajectoryDataset(torch.utils.data.Dataset):
             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
+            s[1:, 2] += self._LOG_INV_0_9
+            c[2] += self._LOG_INV_0_9
 
         return torch.from_numpy(s), torch.from_numpy(c)
 
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index d3f83fe..8cbf79b 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -171,6 +171,63 @@ class TestTrajectoryDataset:
         s1, _ = ds[1]
         assert (s1[:, 1] < 0).all()
 
+    def test_aug_id_two_slows_speed(self):
+        """Aug id 2 should add log(1.25) to log_dt channel and cond[2]."""
+        import math
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s2, c2 = ds[2]  # idx=2 → base=0, aug_id=2
+        expected_delta = math.log(1.25)
+        np.testing.assert_allclose(s2[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c2[2].item() - expected_delta) < 1e-5
+
+    def test_aug_id_three_speeds_up(self):
+        """Aug id 3 should add log(1/1.2) to log_dt channel and cond[2]."""
+        import math
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s3, c3 = ds[3]  # idx=3 → base=0, aug_id=3
+        expected_delta = math.log(1.0 / 1.2)
+        np.testing.assert_allclose(s3[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c3[2].item() - expected_delta) < 1e-5
+
+    def test_aug_id_four_adds_temporal_noise(self):
+        """Aug id 4 should add Gaussian noise to log_dt (channel 2), leaving other channels unchanged."""
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s4, c4 = ds[4]  # idx=4 → base=0, aug_id=4
+        # Channels 0 and 1 should be unchanged (still 0)
+        np.testing.assert_allclose(s4[:, 0].numpy(), 0.0, atol=1e-6)
+        np.testing.assert_allclose(s4[:, 1].numpy(), 0.0, atol=1e-6)
+        # Channel 2 (log_dt) at position 0 is unchanged; positions 1+ have noise
+        assert s4[0, 2].item() == 0.0
+        # Noise should be non-zero (with overwhelming probability for 63 samples)
+        assert not np.allclose(s4[1:, 2].numpy(), 0.0)
+        # cond is unchanged for aug_id 4
+        np.testing.assert_allclose(c4.numpy(), cond[0], atol=1e-6)
+
+    def test_aug_id_five_flips_and_slows(self):
+        """Aug id 5 should flip lateral and add log(1/0.9) to log_dt and cond[2]."""
+        import math
+        from ai_mouse.trainer import TrajectoryDataset
+        seq = np.zeros((1, 64, 3), dtype=np.float32)
+        seq[0, :, 1] = 0.5  # lateral positive
+        cond = np.zeros((1, 3), dtype=np.float32)
+        ds = TrajectoryDataset(seq, cond, augment=True)
+        s5, c5 = ds[5]  # idx=5 → base=0, aug_id=5
+        # Lateral flipped
+        assert (s5[:, 1] < 0).all()
+        # log_dt shifted
+        expected_delta = math.log(1.0 / 0.9)
+        np.testing.assert_allclose(s5[1:, 2].numpy(), expected_delta, rtol=1e-5)
+        assert abs(c5[2].item() - expected_delta) < 1e-5
+
 
 class TestResumeFrom:
     def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):