feat(hpo): env-var knob overrides + sweep script (18 configs)

- train.py knobs all readable from JEPA_* env vars (JEPA_WINDOW, JEPA_D_MODEL,
  JEPA_DEPTH, etc.) so hpo_sweep.py can override without touching source
- scripts/hpo_sweep.py: 3×2×3 grid over D_MODEL × DEPTH × WINDOW,
  logs to results/hpo/hpo_results.jsonl with leaderboard at end
- 3 new tests: env override correctness, configs() schema validation
- 19/19 tests pass

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

scripts/hpo_sweep.py

@@ -0,0 +1,97 @@
+"""HPO sweep for jepa-fx-risk HEPA backbone.
+
+Runs train.py with different JEPA_* env overrides, logs results to
+results/hpo/hpo_results.jsonl. Each config writes its metrics.json then
+the result is appended to the JSONL.
+
+Usage:
+    python scripts/hpo_sweep.py
+    python scripts/hpo_sweep.py --dry-run   # print configs, don't train
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+from datetime import datetime
+from itertools import product
+from pathlib import Path
+
+# ── Search space ──────────────────────────────────────────────────────────────
+SEARCH_SPACE = {
+    "JEPA_D_MODEL": [64, 128, 256],
+    "JEPA_DEPTH":   [2, 4],
+    "JEPA_WINDOW":  [120, 240, 480],
+}
+# Fixed: PATCH_LEN=24 (1-day patches), N_HEADS=4, EPOCHS=300, PHASE1_EPOCHS=200
+
+PYTHON = str(Path(sys.executable))
+OUT_DIR = Path("results/hpo")
+
+
+def configs():
+    """Yield all configs as dicts of JEPA_* env overrides."""
+    keys = list(SEARCH_SPACE.keys())
+    for vals in product(*SEARCH_SPACE.values()):
+        yield dict(zip(keys, vals))
+
+
+def run_config(cfg: dict, metrics_path: str = "metrics.json") -> dict:
+    env = {**os.environ, **{k: str(v) for k, v in cfg.items()}}
+    result = subprocess.run(
+        [PYTHON, "train.py"],
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        return {"config": cfg, "error": result.stderr[-500:]}
+    stdout_last = result.stdout.strip().split("\n")[-1]
+    with open(metrics_path) as f:
+        m = json.load(f)
+    return {
+        "config": cfg,
+        "val_vol_r2":  m.get("val_vol_r2"),
+        "phase1_r2":   m.get("phase1_r2"),
+        "stdout_last": stdout_last,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+    out_file = OUT_DIR / "hpo_results.jsonl"
+
+    all_cfgs = list(configs())
+    print(f"HPO sweep: {len(all_cfgs)} configs")
+    for i, cfg in enumerate(all_cfgs):
+        label = " ".join(f"{k.replace('JEPA_','')}={v}" for k, v in cfg.items())
+        print(f"\n[{i+1}/{len(all_cfgs)}] {label}")
+        if args.dry_run:
+            continue
+        ts = datetime.utcnow().isoformat()
+        row = run_config(cfg)
+        row["ts"] = ts
+        with open(out_file, "a") as f:
+            f.write(json.dumps(row) + "\n")
+        if "error" in row:
+            print(f"  ERROR: {row['error'][:200]}")
+        else:
+            print(f"  val_vol_r2={row['val_vol_r2']:.4f}  phase1_r2={row['phase1_r2']:.4f}")
+
+    if not args.dry_run:
+        # Print leaderboard
+        rows = [json.loads(l) for l in open(out_file) if l.strip()]
+        rows = [r for r in rows if "error" not in r]
+        rows.sort(key=lambda r: r.get("phase1_r2", -999), reverse=True)
+        print("\n── Leaderboard (by phase1_r2) ─────────────────────────")
+        for r in rows[:5]:
+            cfg_str = " ".join(f"{k.replace('JEPA_','')}={v}" for k,v in r["config"].items())
+            print(f"  {r['phase1_r2']:.4f}  {cfg_str}")
+
+
+if __name__ == "__main__":
+    main()

tests/test_hepa.py

@@ -1,9 +1,10 @@
-"""Failing tests for HEPA backbone + Phase-1 supervised head in train.py.
+"""Failing tests for HEPA backbone + Phase-1 supervised head + HPO in train.py.
 
 Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_hepa.py -v
 These tests define what the backbone and head must satisfy BEFORE implementation.
 """
 import math
+import os
 import torch
 import torch.nn as nn
 import pytest
@@ -12,11 +13,24 @@ import pytest
 # They will fail until train.py implements: CausalEncoder, HorizonPredictor, vicreg_loss
 
 
-def _import():
+def _import(env_overrides=None):
     import importlib.util, sys
-    spec = importlib.util.spec_from_file_location("train", "train.py")
+    saved = {}
+    if env_overrides:
+        for k, v in env_overrides.items():
+            saved[k] = os.environ.get(k)
+            os.environ[k] = str(v)
+    # Force fresh module load (env vars must be read at import time)
+    name = f"train_{id(env_overrides)}"
+    spec = importlib.util.spec_from_file_location(name, "train.py")
     mod = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(mod)
+    if env_overrides:
+        for k, orig in saved.items():
+            if orig is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = orig
     return mod
 
 
@@ -171,7 +185,7 @@ def test_phase1_beats_linear_on_nonlinear(train_mod):
 
 # 11. main() returns phase1_r2 in metrics.json (integration — needs real data)
 def test_metrics_json_has_phase1_r2(train_mod):
-    import os, json
+    import json
     if not os.path.exists("metrics.json"):
         pytest.skip("metrics.json not present — run train.py first")
     with open("metrics.json") as f:
@@ -181,3 +195,35 @@ def test_metrics_json_has_phase1_r2(train_mod):
         f"MLP head phase1_r2={m['phase1_r2']:.4f} should beat linear probe "
         f"val_vol_r2={m['val_vol_r2']:.4f}"
     )
+
+
+# ── HPO: env-var knob overrides ───────────────────────────────────────────────
+
+# 12. JEPA_WINDOW env var overrides WINDOW at import time
+def test_env_override_window():
+    mod = _import({"JEPA_WINDOW": "48"})
+    assert mod.WINDOW == 48, f"expected WINDOW=48, got {mod.WINDOW}"
+
+
+# 13. JEPA_D_MODEL and JEPA_DEPTH env vars work
+def test_env_override_d_model_depth():
+    mod = _import({"JEPA_D_MODEL": "64", "JEPA_DEPTH": "4"})
+    assert mod.D_MODEL == 64, f"expected D_MODEL=64, got {mod.D_MODEL}"
+    assert mod.DEPTH == 4, f"expected DEPTH=4, got {mod.DEPTH}"
+
+
+# 14. hpo_sweep.py exists and generates correct config list
+def test_hpo_sweep_configs():
+    import importlib.util
+    sweep_path = "scripts/hpo_sweep.py"
+    if not os.path.exists(sweep_path):
+        pytest.fail(f"{sweep_path} not found — implement it")
+    spec = importlib.util.spec_from_file_location("hpo_sweep", sweep_path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    cfgs = list(mod.configs())
+    assert len(cfgs) > 0, "configs() returned empty list"
+    # Every config must have at least D_MODEL, DEPTH, WINDOW keys
+    required = {"JEPA_D_MODEL", "JEPA_DEPTH", "JEPA_WINDOW"}
+    for cfg in cfgs:
+        assert required.issubset(cfg.keys()), f"config missing required keys: {cfg}"

train.py

@@ -17,21 +17,22 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-# --- agent-tunable knobs ---
-USE_HOURLY  = True   # prefer eurusd_hourly.parquet when available
-WINDOW      = 240    # hourly: 10 trading days; if USE_HOURLY=False reset to 60
-PATCH_LEN   = 24     # hourly: 1-day patches (10 tokens); if USE_HOURLY=False reset to 10
-D_MODEL     = 128
-DEPTH       = 2
-N_HEADS     = 4
-ALPHA       = 0.1   # VICReg mixing weight (fixed at 0.1 in HEPA paper)
-DELTA_T_MAX = 3     # max prediction horizon in patches (1..min(DELTA_T_MAX, N-1-c))
-BATCH_SIZE    = 512   # mini-batch per step (hourly dataset is too large for full-batch)
-EPOCHS        = 300
-LR            = 3e-4
-PHASE1_EPOCHS = 200   # supervised head epochs (encoder frozen)
-PHASE1_LR     = 1e-3
-SEED          = 0
+# --- agent-tunable knobs (all overridable via JEPA_* env vars for HPO) ---
+import os as _os
+USE_HOURLY    = True
+WINDOW        = int(_os.environ.get("JEPA_WINDOW",        240))
+PATCH_LEN     = int(_os.environ.get("JEPA_PATCH_LEN",     24))
+D_MODEL       = int(_os.environ.get("JEPA_D_MODEL",       128))
+DEPTH         = int(_os.environ.get("JEPA_DEPTH",         2))
+N_HEADS       = int(_os.environ.get("JEPA_N_HEADS",       4))
+ALPHA         = float(_os.environ.get("JEPA_ALPHA",       0.1))
+DELTA_T_MAX   = int(_os.environ.get("JEPA_DELTA_T_MAX",   3))
+BATCH_SIZE    = int(_os.environ.get("JEPA_BATCH_SIZE",    512))
+EPOCHS        = int(_os.environ.get("JEPA_EPOCHS",        300))
+LR            = float(_os.environ.get("JEPA_LR",          3e-4))
+PHASE1_EPOCHS = int(_os.environ.get("JEPA_PHASE1_EPOCHS", 200))
+PHASE1_LR     = float(_os.environ.get("JEPA_PHASE1_LR",   1e-3))
+SEED          = int(_os.environ.get("JEPA_SEED",          0))
 # ---------------------------
 
 torch.manual_seed(SEED)