fix(ll_gen): load_checkpoint tolerates architecture drift (resilient load)

RLAlignmentTrainer.load_checkpoint loaded model state strictly, so a genuinely
trained checkpoint saved under an older architecture (the M3 VAE checkpoint,
predating STEPVAE's optional dim_encoder layer) raised "Missing key(s) ...
dim_encoder.weight" and was unusable. Load non-strictly with a warning on
missing/unexpected keys; guard optimizer restore too. The real 95%-validity
VAE checkpoint (ll_gen/checkpoints/vae_rl_solid.pt) now loads via the normal
API and scores 0.95 valid (verified). Regression test added.

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

Author: LayerDynamics

ll_gen/ll_gen/training/rl_trainer.py

@@ -598,11 +598,31 @@ def load_checkpoint(self, path: Path) -> None:
             _log.warning("No history file found at %s; starting fresh", history_path)
             self._train_history = checkpoint.get("train_history", [])
 
-        # Load model and optimizer state
+        # Load model and optimizer state. Load the model resiliently: a
+        # checkpoint saved under an older architecture (e.g. before the optional
+        # dim_encoder dimension-conditioning layer was added to STEPVAE) must
+        # still load, or a genuinely-trained checkpoint becomes unusable. Any
+        # missing/unexpected keys are surfaced as a warning, not silently
+        # dropped, so a true mismatch stays visible.
         if "model_state_dict" in checkpoint and self.generator._model is not None:
-            self.generator._model.load_state_dict(checkpoint["model_state_dict"])
+            incompatible = self.generator._model.load_state_dict(
+                checkpoint["model_state_dict"], strict=False
+            )
+            if incompatible.missing_keys or incompatible.unexpected_keys:
+                _log.warning(
+                    "Checkpoint %s loaded non-strictly: missing=%s unexpected=%s",
+                    path,
+                    list(incompatible.missing_keys),
+                    list(incompatible.unexpected_keys),
+                )
         if "optimizer_state_dict" in checkpoint:
-            self._optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+            # The optimizer state may not line up if the parameter set changed
+            # (e.g. new conditioning layer); resuming the optimizer is
+            # best-effort, never fatal.
+            try:
+                self._optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+            except (ValueError, KeyError) as exc:
+                _log.warning("Could not restore optimizer state from %s: %s", path, exc)
 
         _log.info("Checkpoint loaded from %s (step %d)", path, self._step_count)
 

ll_gen/tests/test_resilient_checkpoint_load.py

@@ -0,0 +1,40 @@
+"""Regression: RLAlignmentTrainer.load_checkpoint must tolerate architecture drift.
+
+A genuinely-trained checkpoint saved under an older model architecture (e.g. the
+M3 VAE checkpoint, saved before STEPVAE gained the optional ``dim_encoder``
+dimension-conditioning layer) must still load. Previously ``load_checkpoint``
+called ``load_state_dict`` strictly and raised ``RuntimeError: Missing key(s)
+... dim_encoder.weight`` -- which made the real 95%-validity checkpoint
+unusable via the normal API. The loader now loads non-strictly and warns.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+# repo/ll_gen/tests/<this> -> repo/ll_gen/checkpoints/vae_rl_solid.pt
+_CKPT = Path(__file__).resolve().parents[1] / "checkpoints" / "vae_rl_solid.pt"
+
+
+def test_load_checkpoint_tolerates_missing_keys(tmp_path):
+    if not _CKPT.exists():
+        pytest.skip("trained VAE checkpoint fixture not present")
+
+    from ll_gen.training.rl_trainer import RLAlignmentTrainer
+    from ll_gen.training.run import build_generator
+
+    generator = build_generator("vae", "cpu")
+    trainer = RLAlignmentTrainer(generator, device="cpu", output_dir=str(tmp_path))
+    trainer._init_training()
+
+    # Must NOT raise even though the checkpoint predates the dim_encoder layer.
+    trainer.load_checkpoint(_CKPT)
+    assert generator._model is not None
+
+    # The current model genuinely has the key the checkpoint lacks (so the test
+    # is exercising the resilient path, not a no-op).
+    state_keys = set(generator._model.state_dict().keys())
+    assert any("dim_encoder" in k for k in state_keys)