feat(trainer): add resume_from for two-stage training

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",
+            )