Add model save/load and --resume for training continuation

- AsymmetricEncoder.save_checkpoint(): saves model_state + metadata
- AsymmetricEncoder.load_checkpoint(): rebuilds model with frozen backbones,
  then loads trainable weights from checkpoint
- --resume flag restores optimizer, loss (learnable tau), and scheduler state
- Training continues from the saved epoch

Usage:
  python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/training/train_gtauav.py

@@ -65,6 +65,7 @@ class TrainConfigGTAUAV:
     baseline_mode: bool = False
 
     # Training.
+    resume_from: str | None = None  # path to checkpoint for resuming
     output_dir: str = "out/gtauav/with_text"
     epochs: int = 10
     batch_size: int = 64
@@ -211,17 +212,31 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         json.dump(vars(cfg), f, indent=2)
 
     # Model.
-    mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
-    LOGGER.info("🏗️  Building model — %s", mode_str)
-    model = AsymmetricEncoder(
-        dino_web_path=cfg.dino_web_path,
-        dino_sat_path=cfg.dino_sat_path,
-        lrsclip_path=cfg.lrsclip_path,
-        proj_dim=cfg.proj_dim,
-        init_gate=cfg.init_gate,
-        baseline_mode=cfg.baseline_mode,
-        device=cfg.device,
-    ).to(cfg.device)
+    start_epoch = 0
+    resume_ckpt = None
+
+    if cfg.resume_from is not None:
+        LOGGER.info("🔄 Resuming from %s", cfg.resume_from)
+        model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
+            cfg.resume_from,
+            dino_web_path=cfg.dino_web_path,
+            dino_sat_path=cfg.dino_sat_path,
+            lrsclip_path=cfg.lrsclip_path,
+            device=cfg.device,
+        )
+        start_epoch = resume_ckpt.get("epoch", -1) + 1
+    else:
+        mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
+        LOGGER.info("🏗️  Building model — %s", mode_str)
+        model = AsymmetricEncoder(
+            dino_web_path=cfg.dino_web_path,
+            dino_sat_path=cfg.dino_sat_path,
+            lrsclip_path=cfg.lrsclip_path,
+            proj_dim=cfg.proj_dim,
+            init_gate=cfg.init_gate,
+            baseline_mode=cfg.baseline_mode,
+            device=cfg.device,
+        ).to(cfg.device)
 
     n_trainable = sum(p.numel() for p in model.trainable_parameters())
     n_total = sum(p.numel() for p in model.parameters())
@@ -308,11 +323,24 @@ def train(cfg: TrainConfigGTAUAV) -> None:
     )
     scaler = GradScaler(enabled=cfg.use_amp)
 
+    # Restore optimizer/scheduler/loss state on resume.
+    if resume_ckpt is not None:
+        if "optimizer_state" in resume_ckpt:
+            optimizer.load_state_dict(resume_ckpt["optimizer_state"])
+            LOGGER.info("🔄 Optimizer state restored")
+        if "loss_state" in resume_ckpt:
+            loss_fn.load_state_dict(resume_ckpt["loss_state"])
+            LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature)
+        # Advance scheduler to the correct step.
+        for _ in range(start_epoch * steps_per_epoch):
+            scheduler.step()
+        LOGGER.info("🔄 Resuming from epoch %d", start_epoch)
+
     history: list[dict] = []
 
-    LOGGER.info("🚀 Starting training for %d epochs", cfg.epochs)
+    LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
 
-    for epoch in range(cfg.epochs):
+    for epoch in range(start_epoch, cfg.epochs):
         model.train()
         epoch_start = time.time()
         agg: dict[str, float] = {}
@@ -448,6 +476,10 @@ def main() -> None:
         "--baseline", action="store_true",
         help="Run baseline mode (no text).",
     )
+    parser.add_argument(
+        "--resume", type=str, default=None,
+        help="Path to checkpoint to resume training from.",
+    )
     parser.add_argument(
         "--output-dir", type=str, default=None,
         help="Override output directory.",
@@ -484,6 +516,7 @@ def main() -> None:
 
     cfg = TrainConfigGTAUAV()
     cfg.baseline_mode = args.baseline
+    cfg.resume_from = args.resume
     cfg.batch_size = args.batch_size
     cfg.epochs = args.epochs
     cfg.learning_rate = args.lr