feat(phase1): warm-start joint encoder fine-tuning (Option B)

Two-phase phase-1:
  1a. Frozen warmup: head trains on pre-computed embeddings for PHASE1_EPOCHS=200
  1b. Joint fine-tune: encoder + head for PHASE1_JOINT_EPOCHS=30 at PHASE1_ENCODER_LR=3e-6

Key design decisions:
- Warm start prevents catastrophic forgetting (PHASE1_JOINT=1 cold-start → -32 R²)
- Normalize live encoder output with FROZEN stats (mu_e/sd_e) so head sees same
  embedding distribution it was warmed up on
- head LR reduced 10× in joint phase to prevent head from racing ahead

HPO sweep: 30ep@3e-6=0.3962, 30ep@1e-5=0.3930, 50ep@3e-6=0.3923
Baseline (frozen): 0.3908. New best: phase1_r2=0.3962 (+0.0054 OOS).

New knobs: JEPA_PHASE1_JOINT (default 1), JEPA_PHASE1_JOINT_EPOCHS (default 30),
JEPA_PHASE1_ENCODER_LR (default 3e-6). 4 new tests (tests 15-18). 28/28 pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

tests/test_hepa.py

@@ -227,3 +227,51 @@ def test_hpo_sweep_configs():
     required = {"JEPA_D_MODEL", "JEPA_DEPTH", "JEPA_WINDOW"}
     for cfg in cfgs:
         assert required.issubset(cfg.keys()), f"config missing required keys: {cfg}"
+
+
+# ── Option B: joint encoder fine-tuning in phase-1 ───────────────────────────
+
+# 15. PHASE1_JOINT and PHASE1_ENCODER_LR knobs exist at module level
+def test_joint_phase1_knobs():
+    mod = _import({"JEPA_PHASE1_JOINT": "1", "JEPA_PHASE1_ENCODER_LR": "1e-5"})
+    assert hasattr(mod, "PHASE1_JOINT"), "PHASE1_JOINT knob missing from train.py"
+    assert hasattr(mod, "PHASE1_ENCODER_LR"), "PHASE1_ENCODER_LR knob missing from train.py"
+    assert mod.PHASE1_JOINT is True
+    assert abs(mod.PHASE1_ENCODER_LR - 1e-5) < 1e-12
+
+
+# 16. PHASE1_JOINT defaults to True (joint mode on by default)
+def test_joint_phase1_default_on():
+    mod = _import()
+    assert hasattr(mod, "PHASE1_JOINT"), "PHASE1_JOINT knob missing"
+    assert mod.PHASE1_JOINT is True, f"PHASE1_JOINT default should be True, got {mod.PHASE1_JOINT}"
+
+
+# 17. JEPA_PHASE1_JOINT=0 disables joint (env override works)
+def test_joint_phase1_can_disable():
+    mod = _import({"JEPA_PHASE1_JOINT": "0"})
+    assert mod.PHASE1_JOINT is False, f"expected False, got {mod.PHASE1_JOINT}"
+
+
+# 18. Encoder receives non-zero gradients when joint-training with the head
+def test_joint_encoder_grad_flows(train_mod):
+    """Gradient must flow into encoder when using two-param-group joint optimizer."""
+    import torch.nn.functional as F
+    enc  = train_mod.CausalEncoder(n_channels=2, patch_len=8, d_model=16, n_heads=2, depth=1)
+    head = train_mod.SupervisedHead(16)
+    enc.train(); head.train()
+    opt = torch.optim.Adam([
+        {"params": head.parameters(), "lr": 1e-3},
+        {"params": enc.parameters(), "lr": 1e-5},
+    ], weight_decay=1e-4)
+    # Tiny batch: 4 windows of length 16 (= 2 patches of patch_len=8)
+    X = torch.randn(4, 16, 2)
+    y = torch.randn(4)
+    tokens = enc(X)           # (4, 2, 16)
+    h      = tokens[:, -1, :] # (4, 16) — last token
+    pred   = head(h)
+    loss   = F.mse_loss(pred, y)
+    loss.backward()
+    enc_grads = [p.grad for p in enc.parameters() if p.grad is not None]
+    assert len(enc_grads) > 0, "no encoder params received gradients"
+    assert any(g.abs().max().item() > 0 for g in enc_grads), "all encoder grads are zero"