From e992823aefb3d588f7ee31b7827b963e0f9f22c1 Mon Sep 17 00:00:00 2001
From: Huang Qi <yellowdog016@gmail.com>
Date: Sun, 10 May 2026 13:08:37 +0800
Subject: [PATCH] feat(trainer): add resume_from for two-stage training

---
 ai_mouse/trainer.py   | 16 ++++++++++++
 tests/test_trainer.py | 57 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/ai_mouse/trainer.py b/ai_mouse/trainer.py
index 9e5864c..73cdc41 100644
--- a/ai_mouse/trainer.py
+++ b/ai_mouse/trainer.py
@@ -353,6 +353,7 @@ def train(
     seq_len: int = 64,
     progress_callback: Callable[[dict], None] | None = None,
     config: TrainConfig | None = None,
+    resume_from: Path | None = None,
 ) -> None:
     """Train TrajectoryFlowModel with OT-Conditional Flow Matching.
 
@@ -367,6 +368,9 @@ def train(
                             {"epoch": n, "total": N, "loss": f}.
                             Called once at the end with {"done": True, "mu", "sigma"}.
         config:             optional TrainConfig for model hyperparameters
+        resume_from:        if given, load model weights from this checkpoint
+                            directory (must contain flow_model.pt). Used for
+                            two-stage training (pretrain → fine-tune).
     """
     data_path = Path(data_path)
     output_dir = Path(output_dir)
@@ -402,6 +406,18 @@ def train(
         dropout=config.dropout,
         cond_dim=3,
     )
+
+    # ---- Resume from checkpoint if requested ----
+    if resume_from is not None:
+        resume_path = Path(resume_from) / "flow_model.pt"
+        if not resume_path.exists():
+            raise FileNotFoundError(
+                f"resume_from checkpoint not found: {resume_path}"
+            )
+        logger.info("Resuming from checkpoint: %s", resume_path)
+        state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict)
+
     optimiser = torch.optim.AdamW(model.parameters(), lr=lr)
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)
 
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
index 3e61c21..d3f83fe 100644
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -170,3 +170,60 @@ class TestTrajectoryDataset:
         # idx=1 → base_idx=0, aug_id=1 → flip
         s1, _ = ds[1]
         assert (s1[:, 1] < 0).all()
+
+
+class TestResumeFrom:
+    def test_resume_from_loads_checkpoint(self, synthetic_traces_file, tmp_path):
+        """train() with resume_from should load weights from given checkpoint dir."""
+        import torch
+        from ai_mouse.trainer import train
+        from ai_mouse.models import TrajectoryFlowModel
+
+        # First, train an initial model and save it
+        ckpt_dir = tmp_path / "pretrain"
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=ckpt_dir,
+            epochs=2,
+            batch_size=8,
+            seq_len=64,
+        )
+        assert (ckpt_dir / "flow_model.pt").exists()
+
+        # Read its weights to compare later
+        m_pretrain = TrajectoryFlowModel(seq_len=64)
+        m_pretrain.load_state_dict(torch.load(ckpt_dir / "flow_model.pt", weights_only=True))
+        first_param_pre = next(m_pretrain.parameters()).clone()
+
+        # Now train with resume_from for 1 epoch — weights should still be loaded
+        out_dir = tmp_path / "finetune"
+        train(
+            data_path=synthetic_traces_file,
+            output_dir=out_dir,
+            epochs=1,
+            batch_size=8,
+            seq_len=64,
+            resume_from=ckpt_dir,
+        )
+
+        m_after = TrajectoryFlowModel(seq_len=64)
+        m_after.load_state_dict(torch.load(out_dir / "flow_model.pt", weights_only=True))
+        first_param_after = next(m_after.parameters())
+
+        # After 1 epoch, weights should be close to pre-train, not random init
+        # (random init would be O(1) magnitude apart; 1 epoch on small data shifts O(0.1))
+        diff = (first_param_pre - first_param_after).abs().mean().item()
+        assert diff < 0.5, f"Resume_from weights diverged too much: {diff}"
+
+    def test_resume_from_missing_path_raises(self, synthetic_traces_file, tmp_path):
+        from ai_mouse.trainer import train
+
+        with pytest.raises(FileNotFoundError):
+            train(
+                data_path=synthetic_traces_file,
+                output_dir=tmp_path / "out",
+                epochs=1,
+                batch_size=8,
+                seq_len=64,
+                resume_from=tmp_path / "nonexistent",
+            )