generated from mathias/template-go-web
loop.py: Karpathy-style keep/revert loop. Agent (berget/gemma4-31b, iguana model, NOT koala GPU) proposes one change to train.py per iter → train.py runs on koala GPU (<2s) → read val_vol_r2 from metrics.json → keep if improved, else restore original content. STATUS.md tracks per-iter metric + delta + GPU snap. Iter 1 kept: improved EMBED_DIM/capacity, +9.3% on OOS R². Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
"""train.py — the ONLY file the autoresearch agent may edit (Phase-1 contract).
|
|
|
|
Toy slice: a tiny self-supervised encoder (masked reconstruction of windowed
|
|
daily [return, realized_vol]) → FROZEN → linear probe predicts NEXT-day realized
|
|
vol → val_vol_r2 = OOS R². The agent improves val_vol_r2 by editing the encoder /
|
|
objective / masking below. Writes metrics.json (the scalar the loop reads).
|
|
|
|
python train.py
|
|
"""
|
|
import json
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
# --- agent-tunable knobs ---
|
|
WINDOW = 20
|
|
EMBED_DIM = 32
|
|
MASK_FRAC = 0.30
|
|
EPOCHS = 200
|
|
LR = 1e-3
|
|
SEED = 0
|
|
# ---------------------------
|
|
|
|
torch.manual_seed(SEED)
|
|
np.random.seed(SEED)
|
|
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def build():
|
|
df = pd.read_parquet("data/processed/eurusd_daily.parquet").reset_index(drop=True)
|
|
feats = df[["ret", "realized_vol"]].to_numpy(np.float32)
|
|
target = df["realized_vol"].to_numpy(np.float32) # predict NEXT-day RV
|
|
X, y = [], []
|
|
for t in range(WINDOW, len(df) - 1):
|
|
X.append(feats[t - WINDOW:t])
|
|
y.append(target[t + 1])
|
|
X = np.stack(X); y = np.array(y, np.float32)
|
|
n_tr = int(0.7 * len(X)) # time-ordered OOS split
|
|
mu, sd = X[:n_tr].mean((0, 1)), X[:n_tr].std((0, 1)) + 1e-8 # train-only stats
|
|
X = (X - mu) / sd
|
|
return (X[:n_tr], y[:n_tr]), (X[n_tr:], y[n_tr:])
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
def __init__(self, win, emb):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Flatten(),
|
|
nn.Linear(win * 2, 128),
|
|
nn.LayerNorm(128),
|
|
nn.GELU(),
|
|
nn.Linear(128, emb)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
def main():
|
|
(Xtr, ytr), (Xte, yte) = build()
|
|
Xtr_t = torch.tensor(Xtr, device=dev)
|
|
enc = Encoder(WINDOW, EMBED_DIM).to(dev)
|
|
dec = nn.Sequential(nn.Linear(EMBED_DIM, 128), nn.GELU(), nn.Linear(128, WINDOW * 2)).to(dev)
|
|
opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=LR)
|
|
|
|
for _ in range(EPOCHS): # SSL: masked reconstruction of the window
|
|
mask = (torch.rand_like(Xtr_t) > MASK_FRAC).float()
|
|
rec = dec(enc((Xtr_t * mask)))
|
|
loss = (((rec - Xtr_t.flatten(1)) ** 2) * (1 - mask.flatten(1))).mean()
|
|
opt.zero_grad(); loss.backward(); opt.step()
|
|
|
|
enc.eval()
|
|
with torch.no_grad(): # FROZEN embeddings
|
|
Etr = enc(Xtr_t).cpu().numpy()
|
|
Ete = enc(torch.tensor(Xte, device=dev)).cpu().numpy()
|
|
|
|
# linear probe (ridge, closed form) on frozen embeddings → val_vol_r2 (OOS R²)
|
|
A = np.hstack([Etr, np.ones((len(Etr), 1))])
|
|
w = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
|
|
pred = np.hstack([Ete, np.ones((len(Ete), 1))]) @ w
|
|
ss_res = ((yte - pred) ** 2).sum()
|
|
ss_tot = ((yte - yte.mean()) ** 2).sum()
|
|
val_vol_r2 = float(1 - ss_res / ss_tot)
|
|
|
|
json.dump({"val_vol_r2": val_vol_r2, "n_test": len(yte),
|
|
"knobs": {"WINDOW": WINDOW, "EMBED_DIM": EMBED_DIM, "MASK_FRAC": MASK_FRAC, "EPOCHS": EPOCHS}},
|
|
open("metrics.json", "w"), indent=2)
|
|
print("val_vol_r2 = %.4f (n_test=%d, dev=%s)" % (val_vol_r2, len(yte), dev))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |