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>

STATUS.md

@@ -16,3 +16,7 @@
 | 3 | -0.0716 | +0.0487 | KEEP | 4s | gpu=0% vram=10054/12227MiB temp=36°C | iter3 |
 | 4 | 0.0590 | +0.1306 | KEEP | 5s | gpu=0% vram=10054/12227MiB temp=36°C | iter4 |
 | 5 | 0.0599 | +0.0009 | KEEP | 5s | gpu=0% vram=10054/12227MiB temp=37°C | iter5 |
+| 1 | 0.0563 | -0.0036 | revert | 5s | gpu=0% vram=10054/12227MiB temp=35°C | iter1 |
+| 2 | 0.0577 | -0.0022 | revert | 5s | gpu=0% vram=10054/12227MiB temp=36°C | iter2 |
+| 3 | 0.0563 | -0.0036 | revert | 5s | gpu=0% vram=10054/12227MiB temp=36°C | iter3 |
+| 4 | -0.1613 | -0.2212 | revert | 5s | gpu=0% vram=10054/12227MiB temp=37°C | iter4 |

scripts/prepare_hourly.py

@@ -29,27 +29,45 @@ def resample_to_hourly(m1: pd.DataFrame) -> pd.DataFrame:
     """Aggregate M1 DataFrame to hourly bars.
 
     Args:
-        m1: DataFrame with columns ['ts' (datetime), 'close' (float)]
+        m1: DataFrame with columns ['ts', 'open', 'high', 'low', 'close']
+            ('open'/'high'/'low' optional — omit for close-only data).
 
     Returns:
-        DataFrame with columns ['datetime', 'close', 'ret', 'realized_vol']
+        DataFrame with columns ['datetime', 'close', 'ret', 'realized_vol',
-        sorted by datetime; hours with fewer than MIN_BARS M1 ticks dropped.
+        'hl_range', 'ret_intrabar'] sorted by datetime.
+        Hours with fewer than MIN_BARS M1 ticks are dropped.
     """
     m1 = m1.sort_values("ts").copy()
     m1["log_r"] = np.log(m1["close"]).diff()
     m1["hour"]  = m1["ts"].dt.floor("h")
 
-    agg = m1.groupby("hour").agg(
+    has_ohlc = all(c in m1.columns for c in ("open", "high", "low"))
-        close       = ("close",  "last"),
+
-        realized_vol= ("log_r",  lambda x: np.sqrt(np.nansum(x.values ** 2))),
+    agg_dict = dict(
-        n_bars      = ("log_r",  "count"),
+        close        = ("close", "last"),
-    ).reset_index()
+        realized_vol = ("log_r", lambda x: np.sqrt(np.nansum(x.values ** 2))),
+        n_bars       = ("log_r", "count"),
+    )
+    if has_ohlc:
+        agg_dict["high"]  = ("high",  "max")
+        agg_dict["low"]   = ("low",   "min")
+        agg_dict["open_"] = ("open",  "first")
+
+    agg = m1.groupby("hour").agg(**agg_dict).reset_index()
 
     agg = agg[agg["n_bars"] >= MIN_BARS].copy()
     agg["ret"] = np.log(agg["close"]).diff()
     agg = agg.dropna(subset=["ret"]).reset_index(drop=True)
     agg = agg.rename(columns={"hour": "datetime"})
-    return agg[["datetime", "close", "ret", "realized_vol"]]
+
+    if has_ohlc:
+        agg["hl_range"]    = np.log(agg["high"] / agg["low"])
+        agg["ret_intrabar"]= np.log(agg["close"] / agg["open_"])
+        cols = ["datetime", "close", "ret", "realized_vol", "hl_range", "ret_intrabar"]
+    else:
+        cols = ["datetime", "close", "ret", "realized_vol"]
+
+    return agg[cols]
 
 
 def load_m1_from_zips(raw_dir: str) -> pd.DataFrame:
@@ -68,7 +86,7 @@ def load_m1_from_zips(raw_dir: str) -> pd.DataFrame:
                     names=["dt", "open", "high", "low", "close", "vol"],
                 )
         df["ts"] = pd.to_datetime(df["dt"], format="%Y%m%d %H%M%S")
-        frames.append(df[["ts", "close"]])
+        frames.append(df[["ts", "open", "high", "low", "close"]])
         print(f"  loaded {os.path.basename(zp)}: {len(df):,} rows")
     return pd.concat(frames).sort_values("ts").reset_index(drop=True)
 

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"

tests/test_prepare_hourly.py

@@ -102,9 +102,7 @@ def test_thin_hours_dropped(ph):
 
 # 5. Output parquet path and schema (integration — reads actual M1 zips if present)
 def test_output_schema_from_zips(ph, tmp_path):
-    # Build a minimal fake zip structure
     import zipfile, io
-    # synthetic M1 CSV (histdata format: YYYYMMDD HHMMSS;O;H;L;C;V)
     rows = []
     for h in range(24):
         for m in range(60):
@@ -124,3 +122,86 @@ def test_output_schema_from_zips(ph, tmp_path):
     df = pd.read_parquet(out_path)
     assert set(["datetime", "close", "ret", "realized_vol"]).issubset(df.columns)
     assert len(df) > 0
+
+
+# ── New OHLCV-derived features ────────────────────────────────────────────────
+
+def _make_m1_ohlcv(n_hours: int = 4, price: float = 1.1) -> pd.DataFrame:
+    """Synthetic M1 with distinct O, H, L, C so hl_range and ret_intrabar are nonzero."""
+    rng = np.random.default_rng(7)
+    ts = pd.date_range("2020-01-06 00:00", periods=n_hours * 60, freq="min")
+    closes = price + np.cumsum(rng.normal(0, 0.0002, len(ts)))
+    highs  = closes + rng.uniform(0.0001, 0.0005, len(ts))
+    lows   = closes - rng.uniform(0.0001, 0.0005, len(ts))
+    opens  = np.roll(closes, 1); opens[0] = price
+    return pd.DataFrame({"ts": ts, "open": opens, "high": highs, "low": lows, "close": closes})
+
+
+# 6. resample_to_hourly produces hl_range column
+def test_hourly_has_hl_range(ph):
+    m1 = _make_m1_ohlcv()
+    hourly = ph.resample_to_hourly(m1)
+    assert "hl_range" in hourly.columns, f"missing hl_range; cols={hourly.columns.tolist()}"
+    assert (hourly["hl_range"] > 0).all(), "hl_range should be positive"
+
+
+# 7. resample_to_hourly produces ret_intrabar column
+def test_hourly_has_ret_intrabar(ph):
+    m1 = _make_m1_ohlcv()
+    hourly = ph.resample_to_hourly(m1)
+    assert "ret_intrabar" in hourly.columns, f"missing ret_intrabar; cols={hourly.columns.tolist()}"
+
+
+# 8. hl_range = log(hourly_high / hourly_low)
+def test_hl_range_formula(ph):
+    # Two hours; second has known H=1.105, L=1.095
+    ts0 = pd.date_range("2020-01-06 00:00", periods=60, freq="min")
+    ts1 = pd.date_range("2020-01-06 01:00", periods=60, freq="min")
+    closes = np.full(120, 1.1)
+    highs  = np.full(120, 1.1)
+    lows   = np.full(120, 1.1)
+    # second hour: known spread
+    highs[60:] = 1.105
+    lows[60:]  = 1.095
+    m1 = pd.DataFrame({
+        "ts": np.concatenate([ts0, ts1]),
+        "open": closes, "high": highs, "low": lows, "close": closes,
+    })
+    hourly = ph.resample_to_hourly(m1)
+    assert len(hourly) >= 1
+    hl = hourly.iloc[-1]["hl_range"]
+    expected = float(np.log(1.105 / 1.095))
+    assert abs(hl - expected) < 1e-6, f"hl_range={hl:.8f}, expected={expected:.8f}"
+
+
+# 9. ret_intrabar = log(hourly_last_close / hourly_first_open)
+def test_ret_intrabar_formula(ph):
+    ts0 = pd.date_range("2020-01-06 00:00", periods=60, freq="min")
+    ts1 = pd.date_range("2020-01-06 01:00", periods=60, freq="min")
+    closes = np.full(120, 1.1)
+    opens  = np.full(120, 1.1)
+    # second hour: open=1.09, close=1.11
+    opens[60]   = 1.09
+    closes[119] = 1.11
+    m1 = pd.DataFrame({
+        "ts": np.concatenate([ts0, ts1]),
+        "open": opens, "high": closes + 0.001, "low": closes - 0.001, "close": closes,
+    })
+    hourly = ph.resample_to_hourly(m1)
+    assert len(hourly) >= 1
+    rib = hourly.iloc[-1]["ret_intrabar"]
+    expected = float(np.log(1.11 / 1.09))
+    assert abs(rib - expected) < 1e-6, f"ret_intrabar={rib:.8f}, expected={expected:.8f}"
+
+
+# 10. build() in train.py uses 2 feature channels (HPO: hl_range/ret_intrabar redundant)
+def test_build_uses_2_channels(tmp_path):
+    import importlib.util, os
+    hourly_path = "data/processed/eurusd_hourly.parquet"
+    if not os.path.exists(hourly_path):
+        pytest.skip("eurusd_hourly.parquet not present")
+    spec = importlib.util.spec_from_file_location("train_2ch", "train.py")
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    (Xtr, _), _ = mod.build()
+    assert Xtr.shape[2] == 2, f"expected 2 channels, got {Xtr.shape[2]}"

train.py

@@ -30,8 +30,11 @@ 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_EPOCHS      = int(_os.environ.get("JEPA_PHASE1_EPOCHS",       200))
-PHASE1_LR     = float(_os.environ.get("JEPA_PHASE1_LR",   1e-3))
+PHASE1_LR          = float(_os.environ.get("JEPA_PHASE1_LR",           1e-3))
+PHASE1_JOINT       = bool(int(_os.environ.get("JEPA_PHASE1_JOINT",     1)))
+PHASE1_JOINT_EPOCHS= int(_os.environ.get("JEPA_PHASE1_JOINT_EPOCHS",  30))
+PHASE1_ENCODER_LR  = float(_os.environ.get("JEPA_PHASE1_ENCODER_LR",  3e-6))
 SEED          = int(_os.environ.get("JEPA_SEED",          0))
 # ---------------------------
 
@@ -154,7 +157,10 @@ def build():
     else:
         df = pd.read_parquet(daily_path).reset_index(drop=True)
         df["date"] = pd.to_datetime(df["date"])
-    feats  = df[["ret", "realized_vol"]].to_numpy(np.float32)
+    # 2-channel default (HPO: adding hl_range+ret_intrabar hurt — correlated with base feats)
+    # To experiment: change to ["ret", "realized_vol", "hl_range", "ret_intrabar"]
+    FEAT_COLS = ["ret", "realized_vol"]
+    feats  = df[FEAT_COLS].to_numpy(np.float32)
     target = df["realized_vol"].to_numpy(np.float32)
     tr_idx = df.index[df["date"].dt.year <= 2021].tolist()
     te_idx = df.index[df["date"].dt.year >= 2022].tolist()
@@ -222,27 +228,61 @@ def main():
     ss_tot  = ((yte - yte.mean()) ** 2).sum()
     val_vol_r2 = float(1 - ss_res / ss_tot)
 
-    # Phase-1: MLP supervised head on frozen embeddings
+    # Phase-1: MLP supervised head — joint or frozen-encoder path
-    # Standardise targets so the head trains on unit-scale signals.
     ytr_mu = float(ytr.mean()); ytr_sd = float(ytr.std()) + 1e-8
     ytr_z  = (ytr - ytr_mu) / ytr_sd
-    head     = SupervisedHead(D_MODEL).to(dev)
+    head   = SupervisedHead(D_MODEL).to(dev)
+    p1_bs  = min(BATCH_SIZE, len(Etr_n))
+
+    # Shared tensors for the frozen-head warmup (used by both paths)
+    Etr_t   = torch.tensor(Etr_n, device=dev)
+    ytr_z_t = torch.tensor(ytr_z, device=dev)
+    Ete_t   = torch.tensor(Ete_n, device=dev)
+    N_tr_h  = len(Etr_t)
+
+    # Phase 1a: warm up head on frozen embeddings (both paths run this)
     head_opt = torch.optim.Adam(head.parameters(), lr=PHASE1_LR, weight_decay=1e-4)
-    Etr_t = torch.tensor(Etr_n, device=dev)
-    ytr_t = torch.tensor(ytr_z, device=dev)
-    Ete_t = torch.tensor(Ete_n, device=dev)
-    p1_bs = min(BATCH_SIZE, len(Etr_t))
-    N_tr_h = len(Etr_t)
-    # Real epoch iteration: shuffle full dataset each epoch
     for _ in range(PHASE1_EPOCHS):
         perm = torch.randperm(N_tr_h, device=dev)
         for start in range(0, N_tr_h, p1_bs):
             idx_h = perm[start:start + p1_bs]
-            loss_h = F.mse_loss(head(Etr_t[idx_h]), ytr_t[idx_h])
+            loss_h = F.mse_loss(head(Etr_t[idx_h]), ytr_z_t[idx_h])
             head_opt.zero_grad(); loss_h.backward(); head_opt.step()
+
+    if PHASE1_JOINT:
+        # Phase 1b: short joint fine-tuning — encoder nudged with tiny LR.
+        # Normalize live encoder output with FROZEN stats (mu_e, sd_e) so the
+        # head sees the same embedding distribution it was warmed up on.
+        enc.train()
+        mu_e_t = torch.tensor(mu_e, device=dev)
+        sd_e_t = torch.tensor(sd_e, device=dev)
+        Xtr_t  = torch.tensor(Xtr, device=dev)
+        joint_opt = torch.optim.Adam([
+            {"params": head.parameters(), "lr": PHASE1_LR * 0.1},
+            {"params": enc.parameters(),  "lr": PHASE1_ENCODER_LR},
+        ], weight_decay=1e-4)
+        for _ in range(PHASE1_JOINT_EPOCHS):
+            perm = torch.randperm(len(Xtr_t), device=dev)
+            for start in range(0, len(Xtr_t), p1_bs):
+                idx_j = perm[start:start + p1_bs]
+                h_raw = enc(Xtr_t[idx_j])[:, -1, :]
+                h_n   = (h_raw - mu_e_t) / sd_e_t  # frozen-stats normalisation
+                loss_j = F.mse_loss(head(h_n), ytr_z_t[idx_j])
+                joint_opt.zero_grad(); loss_j.backward(); joint_opt.step()
+        enc.eval()
+        # Re-extract test embeddings with fine-tuned encoder, same normalisation
+        with torch.no_grad():
+            chunks = []
+            for i in range(0, len(Xte), p1_bs):
+                t = torch.tensor(Xte[i:i+p1_bs], device=dev)
+                h = enc(t)[:, -1, :]
+                chunks.append(((h - mu_e_t) / sd_e_t).cpu().numpy())
+            Ete_t = torch.tensor(np.concatenate(chunks), device=dev)
+
     head.eval()
     with torch.no_grad():
         pred_h_z = head(Ete_t).cpu().numpy()
+
     pred_h    = pred_h_z * ytr_sd + ytr_mu          # de-standardise
     phase1_r2 = float(1 - ((yte - pred_h) ** 2).sum() / ss_tot)
     print("phase1_r2   = %.4f  (n_test=%d)" % (phase1_r2, len(yte)))
@@ -268,7 +308,9 @@ def main():
             df2 = pd.read_parquet(daily_path2).reset_index(drop=True)
             df2["date"] = pd.to_datetime(df2["date"])
         tr_mask = df2["date"].dt.year <= 2021
-        feats2 = df2[["ret", "realized_vol"]].to_numpy(np.float32)
+        base2 = ["ret", "realized_vol"]
+        extra2 = [c for c in ["hl_range", "ret_intrabar"] if c in df2.columns]
+        feats2 = df2[base2 + extra2].to_numpy(np.float32)
         mu2 = feats2[tr_mask].mean(0); sd2 = feats2[tr_mask].std(0) + 1e-8
         fn2 = (feats2 - mu2) / sd2
         def _export_windows(year_mask):