ks-equation-1d-forecast

cd /app && ls -la && echo "---CONFIG---" && cat task_config.json && echo "---SCHEMA---" && cat feature_schema.json

cd /app && echo "===candidate_model.py===" && cat candidate_model.py && echo "===dataset.py===" && cat dataset.py && echo "===train.py===" && cat train.py

cd /app && python3 -c "
import torch, sys
print('torch', torch.__version__)
print('cuda available', torch.cuda.is_available())
if torch.cuda.is_available():
    print('device', torch.cuda.get_device_name(0))
import subprocess
print('cpus', __import__('os').cpu_count())
"
echo "---nvidia-smi---"
nvidia-smi 2>/dev/null | head -15 || echo "no nvidia-smi"
echo "---mem---"
free -h 2>/dev/null | head -3

cd /app && python3 -c "
import torch
for split in ['train','val','public_test']:
    p = torch.load(f'data/{split}/trajectories.pt', map_location='cpu', weights_only=True)
    print(split, type(p), list(p.keys()) if isinstance(p,dict) else '')
    t = p['trajectories']
    print('  shape', t.shape, 'dtype', t.dtype)
    print('  mean %.5f std %.5f min %.4f max %.4f' % (t.mean(), t.std(), t.min(), t.max()))
"

cd /app && time TASK_EPOCHS_OVERRIDE=2 python3 train.py 2>&1 | tail -20

cd /app && cat > /tmp/rollout_eval.py << 'EOF'
import torch, json
from pathlib import Path
from candidate_model import build_model

cfg = json.load(open('task_config.json'))
k = int(cfg['in_channels']); H = int(cfg['rollout_horizon'])
dev = 'cuda'

def load_model(ckpt_path):
    ck = torch.load(ckpt_path, map_location='cpu', weights_only=True)
    m = build_model(cfg); m.load_state_dict(ck['state_dict']); m.to(dev).eval()
    return m, ck['feat_mean'], ck['feat_std']

@torch.no_grad()
def rollout_metrics(model, mean, std, traj, starts=None):
    # traj: [N, T, 128]
    N, T, G = traj.shape
    if starts is None:
        starts = [k]  # predict frames k..k+H-1 from window [0,k)
    per_step_all = []
    glob_all = []
    for s in starts:
        if s + H > T: continue
        win = ((traj[:, s-k:s] - mean)/std).to(dev)  # [N,k,G]
        preds = []
        for i in range(H):
            p = model(win)  # [N,1,G]
            preds.append(p)
            win = torch.cat([win[:,1:], p], dim=1)
        pred = torch.cat(preds, dim=1)*std + mean  # [N,H,G] physical
        true = traj[:, s:s+H].to(dev)
        # per-step rel l2 (mean over N, steps)
        num = (pred-true).pow(2).sum(-1).sqrt()
        den = true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
        per_step = (num/den)  # [N,H]
        per_step_all.append(per_step)
        # global rel l2 over whole rollout per trajectory
        gnum = (pred-true).reshape(N,-1).pow(2).sum(-1).sqrt()
        gden = true.reshape(N,-1).pow(2).sum(-1).sqrt().clamp_min(1e-8)
        glob_all.append(gnum/gden)
    per_step = torch.cat(per_step_all,0)  # [*,H]
    glob = torch.cat(glob_all,0)
    return per_step.mean().item(), per_step.mean(0).cpu(), glob.mean().item()

if __name__ == '__main__':
    import sys
    ckpt = sys.argv[1] if len(sys.argv)>1 else 'outputs/model.pt'
    val = torch.load('data/val/trajectories.pt', weights_only=True)['trajectories'].float()
    test = torch.load('data/public_test/trajectories.pt', weights_only=True)['trajectories'].float()
    m, mean, std = load_model(ckpt)
    for name, tr in [('val', val), ('test', test)]:
        mstep, per, glob = rollout_metrics(m, mean, std, tr, starts=[k])
        print(f'{name}: mean_per_step_relL2={mstep:.4f} global_relL2={glob:.4f}')
        print('   per-step:', ' '.join(f'{v:.3f}' for v in per.tolist()))
EOF
python3 /tmp/rollout_eval.py outputs/model.pt

cd /app && cp /tmp/rollout_eval.py /app/rollout_eval.py && python3 /app/rollout_eval.py outputs/model.pt

cd /app && cp outputs/model.pt /tmp/baseline2.pt && python3 train.py 2>&1 | tail -12 && echo "===ROLLOUT===" && python3 /app/rollout_eval.py outputs/model.pt && cp outputs/model.pt /tmp/baseline10.pt

cd /app && cat > /app/train_improved.py << 'EOF'
from __future__ import annotations
import json, os, sys, time, copy, math
from pathlib import Path
import torch
import torch.nn.functional as F

APP_DIR = Path(os.environ.get("APP_DIR", "/app"))
sys.path.insert(0, str(APP_DIR))
from candidate_model import build_model
from dataset import KSForecast, compute_standardization

def getenv(name, default, cast):
    v = os.environ.get(name)
    return cast(v) if v is not None else default

# ---------------- config knobs (env-overridable) ----------------
EPOCHS      = getenv("EPOCHS", 80, int)
BATCH       = getenv("BATCH", 128, int)
LR          = getenv("LR", 2e-3, float)
WD          = getenv("WD", 1e-4, float)
WARMUP      = getenv("WARMUP", 5, int)
MMAX        = getenv("MMAX", 10, int)        # max rollout length in training
RAMP        = getenv("RAMP_EPOCHS", 30, int) # epoch by which M reaches MMAX
DETACH      = getenv("DETACH", 1, int)       # 1 = pushforward (detach between steps)
BPTT        = getenv("BPTT", 4, int)         # keep graph through last BPTT steps
NOISE       = getenv("NOISE", 0.0, float)    # input noise std (normalized units)
EMA_DECAY   = getenv("EMA", 0.999, float)
MSE_W       = getenv("MSE_W", 0.0, float)    # extra mse weight
GRAD_CLIP   = getenv("GRAD_CLIP", 1.0, float)
SEED_EXTRA  = getenv("SEED_EXTRA", 0, int)
TAG         = os.environ.get("TAG", "run")
SAVE        = os.environ.get("SAVE_PATH", "")
SPEC_W      = getenv("SPEC_W", 0.0, float)   # spectral (gradient) loss weight

def main():
    cfg = json.load(open(APP_DIR/"task_config.json"))
    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(int(cfg["seed"]) + SEED_EXTRA)
    k = int(cfg["in_channels"]); H = int(cfg["rollout_horizon"])

    train_ds = KSForecast(APP_DIR/"data"/"train", k=k)
    val_tr   = torch.load(APP_DIR/"data"/"val"/"trajectories.pt", weights_only=True)["trajectories"].float()
    feat_mean, feat_std = compute_standardization(train_ds.trajectories)
    traj = train_ds.trajectories.to(dev)                 # [N,T,G]
    trajn = (traj - feat_mean)/feat_std                  # normalized
    N, T, G = trajn.shape
    val_tr = val_tr.to(dev)

    model = build_model(cfg).to(dev)
    ema = copy.deepcopy(model)
    for p in ema.parameters(): p.requires_grad_(False)

    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WD)
    steps_per_epoch = math.ceil(N*(T-k)/BATCH)
    total_steps = EPOCHS*steps_per_epoch
    warmup_steps = WARMUP*steps_per_epoch
    def lr_at(step):
        if step < warmup_steps:
            return LR*step/max(1,warmup_steps)
        prog = (step-warmup_steps)/max(1,total_steps-warmup_steps)
        return 0.5*LR*(1+math.cos(math.pi*prog))
    def M_at(epoch):
        if epoch < WARMUP: return 1
        f = min(1.0, (epoch-WARMUP)/max(1,RAMP-WARMUP))
        return max(1, int(round(1 + f*(MMAX-1))))

    arangeW = torch.arange(-k,0,device=dev)
    def rel_l2(pred, true):  # [B,.,G] physical
        num = (pred-true).pow(2).sum(-1).sqrt()
        den = true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
        return (num/den).mean()

    @torch.no_grad()
    def val_rollout(m, starts):
        m.eval(); outs=[]
        for s in starts:
            if s+H>T: continue
            win = ((val_tr[:, s-k:s]-feat_mean)/feat_std)
            preds=[]
            for _ in range(H):
                p=m(win); preds.append(p); win=torch.cat([win[:,1:],p],1)
            pred=torch.cat(preds,1)*feat_std+feat_mean
            true=val_tr[:, s:s+H]
            gn=(pred-true).reshape(val_tr.shape[0],-1).pow(2).sum(-1).sqrt()
            gd=true.reshape(val_tr.shape[0],-1).pow(2).sum(-1).sqrt().clamp_min(1e-8)
            outs.append((gn/gd))
        return torch.cat(outs).mean().item()

    def ema_update(decay):
        with torch.no_grad():
            for pe,pm in zip(ema.parameters(), model.parameters()):
                pe.mul_(decay).add_(pm, alpha=1-decay)
            for be,bm in zip
…[truncated]

cd /app && EPOCHS=8 WARMUP=2 RAMP=6 MMAX=6 TAG=smoke python3 train_improved.py 2>&1 | tail -15

cd /app && for cfg in \
  "TAG=A LR=1e-3 MMAX=10 RAMP=22 DETACH=1 BPTT=1" \
  "TAG=B LR=1e-3 MMAX=10 RAMP=22 DETACH=1 BPTT=4" \
  "TAG=C LR=2e-3 MMAX=10 RAMP=22 DETACH=1 BPTT=4" \
  ; do
  env EPOCHS=40 WARMUP=4 $cfg python3 train_improved.py 2>&1 | tail -2
done

restart: true

cd /app && mkdir -p /tmp/logs
nohup env EPOCHS=40 WARMUP=4 TAG=B LR=1e-3 MMAX=10 RAMP=22 DETACH=1 BPTT=4 python3 train_improved.py > /tmp/logs/B.log 2>&1 &
echo "B pid $!"
nohup env EPOCHS=40 WARMUP=4 TAG=C LR=2e-3 MMAX=10 RAMP=22 DETACH=1 BPTT=4 python3 train_improved.py > /tmp/logs/C.log 2>&1 &
echo "C pid $!"
sleep 1 && echo "launched"

cd /app && sleep 90 && echo "===B===" && tail -4 /tmp/logs/B.log && echo "===C===" && tail -4 /tmp/logs/C.log

cd /app && cp train_improved.py train_improved.py.bak
python3 - << 'EOF'
import re
s = open('/app/train_improved.py').read()

# Track best over both model and EMA, using multi-start selection metric.
s = s.replace(
'''        vr_m = val_rollout(model, [k])
        vr_e = val_rollout(ema, [k])
        vr_e_multi = val_rollout(ema, [k, 20, 40, 60])
        sel = vr_e
        tag=""
        if sel<best:
            best=sel; best_state=copy.deepcopy(ema.state_dict()); tag=" *"
        if epoch%5==0 or epoch==EPOCHS-1 or tag:
            print(f"[{TAG}] ep{epoch+1}/{EPOCHS} M={M} loss={run/max(nb,1):.4e} "
                  f"val(model)={vr_m:.4f} val(ema)={vr_e:.4f} ema_multi={vr_e_multi:.4f}{tag}", flush=True)''',
'''        starts_sel = [k, 20, 40, 60]
        vr_m = val_rollout(model, [k]); vr_m_multi = val_rollout(model, starts_sel)
        vr_e = val_rollout(ema, [k]);   vr_e_multi = val_rollout(ema, starts_sel)
        tag=""
        if vr_e_multi < best:
            best=vr_e_multi; best_state=copy.deepcopy(ema.state_dict()); best_src="ema"; tag=" *E"
        if vr_m_multi < best:
            best=vr_m_multi; best_state=copy.deepcopy(model.state_dict()); best_src="model"; tag=" *M"
        if epoch%5==0 or epoch==EPOCHS-1 or tag:
            print(f"[{TAG}] ep{epoch+1}/{EPOCHS} M={M} loss={run/max(nb,1):.4e} "
                  f"m={vr_m:.4f}/{vr_m_multi:.4f} e={vr_e:.4f}/{vr_e_multi:.4f} best={best:.4f}{tag}", flush=True)''')

s = s.replace('gstep=0; best=1e9; best_state=None; t0=time.time()',
              'gstep=0; best=1e9; best_state=None; best_src="ema"; t0=time.time()')
s = s.replace('print(f"[{TAG}] done best_val_ema={best:.4f} time={dt:.1f}s", flush=True)',
              'print(f"[{TAG}] done best_val={best:.4f} src={best_src} time={dt:.1f}s", flush=True)')
open('/app/train_improved.py','w').write(s)
print("patched OK")
EOF
python3 -c "import ast; ast.parse(open('/app/train_improved.py').read()); print('syntax OK')"

cd /app && sleep 60 && echo "===B tail===" && tail -3 /tmp/logs/B.log && echo "===C tail===" && tail -3 /tmp/logs/C.log && echo "---running?---" && (pgrep -af train_improved | grep -v grep || echo none)

cd /app && sleep 40 && echo "===B final===" && tail -2 /tmp/logs/B.log && echo "===C final===" && tail -2 /tmp/logs/C.log && (pgrep -af train_improved | grep -v grep || echo "all done")

cd /app && mkdir -p /tmp/ck
nohup env EPOCHS=90 WARMUP=5 TAG=D LR=1.5e-3 MMAX=16 RAMP=45 DETACH=1 BPTT=1 EMA=0.999 SAVE_PATH=/tmp/ck/D.pt python3 train_improved.py > /tmp/logs/D.log 2>&1 &
echo "D pid $!"
nohup env EPOCHS=90 WARMUP=5 TAG=E LR=1.5e-3 MMAX=20 RAMP=50 DETACH=1 BPTT=2 EMA=0.999 SAVE_PATH=/tmp/ck/E.pt python3 train_improved.py > /tmp/logs/E.log 2>&1 &
echo "E pid $!"
sleep 1 && echo launched

cd /app && sleep 150 && echo "===D===" && tail -3 /tmp/logs/D.log && echo "===E===" && tail -3 /tmp/logs/E.log

cd /app && sleep 230 && echo "===D===" && tail -3 /tmp/logs/D.log && echo "===E===" && tail -3 /tmp/logs/E.log

cd /app && sleep 270 && echo "===D===" && tail -3 /tmp/logs/D.log && echo "===E===" && tail -3 /tmp/logs/E.log

cd /app && sleep 200 && echo "===D===" && tail -2 /tmp/logs/D.log && echo "===E===" && tail -3 /tmp/logs/E.log && (pgrep -af train_improved | grep -v grep || echo "all done")

cd /app && sleep 60 && echo "===E final===" && tail -2 /tmp/logs/E.log && (pgrep -af train_improved | grep -v grep || echo "all done")
echo "=== Eval D.pt ===" && python3 /app/rollout_eval.py /tmp/ck/D.pt
echo "=== Eval E.pt ===" && python3 /app/rollout_eval.py /tmp/ck/E.pt

cd /app && mkdir -p outputs && cp /tmp/ck/D.pt outputs/model.pt
python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json'))
ck = torch.load('outputs/model.pt', weights_only=True)  # must work with weights_only
assert isinstance(ck, dict), type(ck)
assert set(['state_dict','feat_mean','feat_std']).issubset(ck.keys()), ck.keys()
assert isinstance(ck['feat_mean'], float) and isinstance(ck['feat_std'], float), (type(ck['feat_mean']), type(ck['feat_std']))
m = build_model(cfg)
missing, unexpected = m.load_state_dict(ck['state_dict'], strict=True)
print('strict load OK; missing', missing, 'unexpected', unexpected)
print('feat_mean=%.6f feat_std=%.6f' % (ck['feat_mean'], ck['feat_std']))
print('num params', sum(p.numel() for p in m.parameters()))
# sanity forward
x = torch.randn(2,4,128)
print('forward out shape', m(x).shape)
print('CONTRACT OK')
EOF

cd /app && nohup env EPOCHS=90 WARMUP=5 TAG=F LR=1.5e-3 MMAX=20 RAMP=50 DETACH=1 BPTT=1 EMA=0.999 SAVE_PATH=/tmp/ck/F.pt python3 train_improved.py > /tmp/logs/F.log 2>&1 &
echo "F pid $!"
nohup env EPOCHS=130 WARMUP=6 TAG=G LR=1.5e-3 MMAX=16 RAMP=55 DETACH=1 BPTT=1 EMA=0.9995 SAVE_PATH=/tmp/ck/G.pt python3 train_improved.py > /tmp/logs/G.log 2>&1 &
echo "G pid $!"
sleep 1 && echo launched

cd /app && sleep 285 && echo "===F===" && tail -2 /tmp/logs/F.log && echo "===G===" && tail -2 /tmp/logs/G.log

cd /app && sleep 290 && echo "===F===" && tail -2 /tmp/logs/F.log && echo "===G===" && tail -2 /tmp/logs/G.log

cd /app && sleep 290 && echo "===F===" && tail -2 /tmp/logs/F.log && echo "===G===" && tail -2 /tmp/logs/G.log && (pgrep -af train_improved | grep -v grep || echo "all done")

cd /app && sleep 245 && echo "===F final===" && tail -1 /tmp/logs/F.log && echo "===G final===" && tail -2 /tmp/logs/G.log && (pgrep -af train_improved | grep -v grep || echo "all done")

cd /app && sleep 55 && echo "===G final===" && tail -2 /tmp/logs/G.log && (pgrep -af train_improved | grep -v grep || echo "all done")
echo "=== F.pt ===" && python3 /app/rollout_eval.py /tmp/ck/F.pt
echo "=== G.pt ===" && python3 /app/rollout_eval.py /tmp/ck/G.pt

cd /app && cat > /tmp/compare.py << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json')); k=cfg['in_channels']; H=cfg['rollout_horizon']; dev='cuda'
def load(p):
    c=torch.load(p,map_location='cpu',weights_only=True); m=build_model(cfg); m.load_state_dict(c['state_dict']); m.to(dev).eval(); return m,c['feat_mean'],c['feat_std']
@torch.no_grad()
def ev(m,mean,std,tr,starts):
    N,T,G=tr.shape; tr=tr.to(dev); gl=[]; fs=[]; ps=[]
    for s in starts:
        if s+H>T: continue
        win=((tr[:,s-k:s]-mean)/std); P=[]
        for _ in range(H):
            p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
        pred=torch.cat(P,1)*std+mean; true=tr[:,s:s+H]
        gn=(pred-true).reshape(N,-1).pow(2).sum(-1).sqrt(); gd=true.reshape(N,-1).pow(2).sum(-1).sqrt().clamp_min(1e-8)
        gl.append(gn/gd)
        # final step rel l2
        fn=(pred[:,-1]-true[:,-1]).pow(2).sum(-1).sqrt(); fd=true[:,-1].pow(2).sum(-1).sqrt().clamp_min(1e-8)
        fs.append(fn/fd)
        ps.append((pred-true).pow(2).sum(-1).sqrt()/true.pow(2).sum(-1).sqrt().clamp_min(1e-8))
    return torch.cat(gl).mean().item(), torch.cat(fs).mean().item(), torch.cat(ps,0).mean().item()
val=torch.load('data/val/trajectories.pt',weights_only=True)['trajectories'].float()
test=torch.load('data/public_test/trajectories.pt',weights_only=True)['trajectories'].float()
starts=list(range(k,81,4))  # many starts
print(f"{'ckpt':>6} | {'val_glob':>8} {'val_fin':>8} {'val_ps':>7} | {'tst_glob':>8} {'tst_fin':>8} {'tst_ps':>7}")
for name in ['D','F','G']:
    m,me,sd=load(f'/tmp/ck/{name}.pt')
    vg,vf,vp=ev(m,me,sd,val,starts); tg,tf,tp=ev(m,me,sd,test,starts)
    print(f"{name:>6} | {vg:8.4f} {vf:8.4f} {vp:7.4f} | {tg:8.4f} {tf:8.4f} {tp:7.4f}")
EOF
python3 /tmp/compare.py

cd /app && cp /tmp/compare.py /app/compare.py && python3 /app/compare.py

cd /app && python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json')); k=cfg['in_channels']; dev='cuda'
def load(p):
    c=torch.load(p,weights_only=True); m=build_model(cfg); m.load_state_dict(c['state_dict']); m.to(dev).eval(); return m,c['feat_mean'],c['feat_std']
test=torch.load('data/public_test/trajectories.pt',weights_only=True)['trajectories'].float().to(dev)
N,T,G=test.shape
for name in ['D','G']:
    m,mean,std=load(f'/tmp/ck/{name}.pt')
    win=((test[:,0:k]-mean)/std); P=[]
    with torch.no_grad():
        for _ in range(T-k):  # roll to end of trajectory (96 steps)
            p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
    pred=torch.cat(P,1)*std+mean; true=test[:,k:T]
    rel=((pred-true).pow(2).sum(-1).sqrt()/true.pow(2).sum(-1).sqrt().clamp_min(1e-8)).mean(0)
    print(f"{name}: relL2 @step 20={rel[19]:.4f} @40={rel[39]:.4f} @60={rel[59]:.4f} @96={rel[-1]:.4f} maxval={pred.abs().max():.2f}")
EOF

cd /app && cp /tmp/ck/D.pt outputs/model.pt && echo "copied D -> outputs/model.pt"
md5sum /tmp/ck/D.pt outputs/model.pt
echo "=== contract re-verify + 20-step eval on outputs/model.pt ===" && python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json')); k=cfg['in_channels']; H=cfg['rollout_horizon']; dev='cuda'
ck=torch.load('outputs/model.pt', weights_only=True)
assert isinstance(ck,dict) and {'state_dict','feat_mean','feat_std'}<=set(ck)
assert isinstance(ck['feat_mean'],float) and isinstance(ck['feat_std'],float)
m=build_model(cfg); m.load_state_dict(ck['state_dict'],strict=True); m.to(dev).eval()
print('contract OK; feat_mean=%.6f feat_std=%.6f params=%d'%(ck['feat_mean'],ck['feat_std'],sum(p.numel() for p in m.parameters())))
def ev(tr):
    tr=tr.to(dev); N,T,_=tr.shape; g=[]
    for s in range(k, T-H+1, 4):
        win=((tr[:,s-k:s]-ck['feat_mean'])/ck['feat_std']); P=[]
        with torch.no_grad():
            for _ in range(H): p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
        pred=torch.cat(P,1)*ck['feat_std']+ck['feat_mean']; true=tr[:,s:s+H]
        gn=(pred-true).reshape(N,-1).pow(2).sum(-1).sqrt(); gd=true.reshape(N,-1).pow(2).sum(-1).sqrt().clamp_min(1e-8)
        g.append(gn/gd)
    return torch.cat(g).mean().item()
for nm in ['val','public_test']:
    tr=torch.load(f'data/{nm}/trajectories.pt',weights_only=True)['trajectories'].float()
    print(f'{nm}: 20-step multi-start global relL2 = {ev(tr):.4f}')
EOF

cd /app && cp outputs/model.pt /tmp/ck/D_final.pt && cp train.py /tmp/train_orig_baseline.py && echo "backed up deliverable and original baseline train.py"

command: create
path: /app/train.py
file_text: from __future__ import annotations
"""Improved training recipe for KS-1D autoregressive forecasting.

Key ideas over the weak baseline (single-step teacher forcing + MSE):
  * Rollout-aware training (scheduled-sampling / pushforward): the model is
    unrolled for M steps feeding its OWN predictions back as input, so it learns
    to correct the error distribution it actually sees at eval time. This is the
    single biggest win -- it flattens the compounding-error curve of the chaotic
    rollout.
  * Curriculum on the rollout length M (1 -> MMAX) so early training learns a
    good one-step map before being asked to stay stable over long horizons.
  * Relative-L2 loss in physical units -- matches the evaluation metric exactly.
  * Warmup + cosine LR, gradient clipping, and an EMA of the weights.
  * Stability-guarded checkpoint selection: a candidate is only accepted if it
    does NOT diverge at 2x the eval horizon (prevents picking a model that is
    marginally better at 20 steps but blows up on longer rollouts).

All behaviour is config-driven with sensible defaults; env vars allow overrides.
The final best checkpoint is written to /app/outputs/model.pt in the required
format: {"state_dict", "feat_mean", "feat_std"} (torch.load weights_only=True).
"""
import json
import os
import sys
import time
import copy
import math
from pathlib import Path
import torch
import torch.nn.functional as F

APP_DIR = Path(os.environ.get("APP_DIR", "/app"))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from candidate_model import build_model
from dataset import KSForecast, compute_standardization


def _env(name, default, cast):
    v = os.environ.get(name)
    return cast(v) if v is not None else default


def main() -> None:
    with (APP_DIR / "task_config.json").open() as fh:
        cfg = json.load(fh)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(int(cfg["seed"]))
    k = int(cfg["in_channels"])
    H = int(cfg["rollout_horizon"])

    # --- recipe hyper-parameters (the defaults reproduce the delivered model) ---
    epochs    = int(os.environ.get("TASK_EPOCHS_OVERRIDE", _env("EPOCHS", 90, int)))
    batch     = _env("BATCH", 128, int)
    lr        = _env("LR", 1.5e-3, float)
    wd        = _env("WD", float(cfg["training"]["weight_decay"]), float)
    mmax      = _env("MMAX", 16, int)            # max training rollout length
    grad_clip = _env("GRAD_CLIP", 1.0, float)
    ema_decay = _env("EMA", 0.999, float)
    stab_thr  = _env("STAB_THRESH", 0.1, float)  # reject models diverging at 2H
    out_path  = Path(os.environ.get("OUT_PATH", str(APP_DIR / "outputs" / "model.pt")))

    # curriculum / schedule fractions (scale with the number of epochs so that a
    # small TASK_EPOCHS_OVERRIDE still runs and produces a valid checkpoint).
    warmup = max(1, int(round(epochs * 0.06)))
    ramp   = max(warmup + 1, int(round(epochs * 0.5)))

    train_ds = KSForecast(APP_DIR / "data" / "train", k=k)
    val_tr = torch.load(APP_DIR / "data" / "val" / "trajectories.pt",
                        weights_only=True)["trajectories"].float().to(device)
    feat_mean, feat_std = compute_standardization(train_ds.trajectories)
    trajn = ((train_ds.trajectories.to(device) - feat_mean) / feat_std)
    N, T, G = trajn.shape

    model = build_model(cfg).to(device)
    ema = copy.deepcopy(model)
    for p in ema.parameters():
        p.requires_grad_(False)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)

    steps_per_epoch = math.ceil(N * (T - k) / batch)
    total_steps = max(1, epochs * steps_per_epoch)
    warmup_steps = warmup * steps_per_epoch

    def lr_at(step):
        if step < warmup_steps:
            return lr * step / max(1, warmup_steps)
        prog = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * lr * (1 + math.cos(math.pi * min(1.0, prog)))

    def M_at(epoch):
       
…[truncated]

cd /app && OUT_PATH=/tmp/ck/smoke.pt TASK_EPOCHS_OVERRIDE=3 python3 train.py 2>&1 | tail -8
echo "=== verify smoke checkpoint contract ===" && python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json'))
ck=torch.load('/tmp/ck/smoke.pt', weights_only=True)
assert {'state_dict','feat_mean','feat_std'}<=set(ck) and isinstance(ck['feat_mean'],float)
build_model(cfg).load_state_dict(ck['state_dict'], strict=True)
print('smoke checkpoint: contract OK')
EOF

cd /app && echo "outputs/model.pt md5:" && md5sum outputs/model.pt && echo "D_final md5:" && md5sum /tmp/ck/D_final.pt
nohup env OUT_PATH=/tmp/ck/canon.pt python3 train.py > /tmp/logs/canon.log 2>&1 &
echo "canon pid $!  (running full 90-epoch canonical train.py in background)"

cd /app && ls -la outputs/ && echo "---" && du -sh outputs/model.pt && sleep 120 && echo "=== canon progress ===" && tail -3 /tmp/logs/canon.log

cd /app && sleep 295 && echo "=== canon progress ===" && tail -4 /tmp/logs/canon.log

cd /app && sleep 260 && echo "=== canon final ===" && tail -4 /tmp/logs/canon.log && (pgrep -af 'train.py' | grep -v grep || echo "done")

cd /app && python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json')); k=cfg['in_channels']; H=cfg['rollout_horizon']; dev='cuda'
def load(p):
    c=torch.load(p,map_location='cpu',weights_only=True); m=build_model(cfg); m.load_state_dict(c['state_dict']); m.to(dev).eval(); return m,c['feat_mean'],c['feat_std']
@torch.no_grad()
def metrics(m,mean,std,tr):
    tr=tr.to(dev); N,T,_=tr.shape
    # 20-step multi-start global + final-step
    g=[]; f=[]
    for s in range(k,T-H+1,4):
        win=((tr[:,s-k:s]-mean)/std); P=[]
        for _ in range(H): p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
        pred=torch.cat(P,1)*std+mean; true=tr[:,s:s+H]
        g.append((pred-true).reshape(N,-1).pow(2).sum(-1).sqrt()/true.reshape(N,-1).pow(2).sum(-1).sqrt().clamp_min(1e-8))
        f.append((pred[:,-1]-true[:,-1]).pow(2).sum(-1).sqrt()/true[:,-1].pow(2).sum(-1).sqrt().clamp_min(1e-8))
    # long-horizon stability: roll to end from t=k
    win=((tr[:,0:k]-mean)/std); P=[]
    for _ in range(T-k): p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
    pl=torch.cat(P,1)*std+mean
    rel_long=((pl-tr[:,k:]).pow(2).sum(-1).sqrt()/tr[:,k:].pow(2).sum(-1).sqrt().clamp_min(1e-8)).mean(0)
    return torch.cat(g).mean().item(), torch.cat(f).mean().item(), rel_long[min(39,len(rel_long)-1)].item(), pl.abs().max().item()
val=torch.load('data/val/trajectories.pt',weights_only=True)['trajectories'].float()
test=torch.load('data/public_test/trajectories.pt',weights_only=True)['trajectories'].float()
print(f"{'ckpt':>12} | {'val_g20':>8} {'val_fin':>8} | {'tst_g20':>8} {'tst_fin':>8} | {'tst_rel@40':>10} {'tst_max':>8}")
for name,path in [('D_final','/tmp/ck/D_final.pt'),('canon','/tmp/ck/canon.pt')]:
    m,me,sd=load(path)
    vg,vf,_,_=metrics(m,me,sd,val); tg,tf,tl,tmx=metrics(m,me,sd,test)
    print(f"{name:>12} | {vg:8.4f} {vf:8.4f} | {tg:8.4f} {tf:8.4f} | {tl:10.4f} {tmx:8.2f}")
EOF

cd /app && cp /tmp/ck/canon.pt outputs/model.pt && echo "promoted canon.pt -> outputs/model.pt"
python3 - << 'EOF'
import torch, json
from candidate_model import build_model
cfg=json.load(open('task_config.json')); k=cfg['in_channels']; H=cfg['rollout_horizon']; dev='cuda'
# ---- contract ----
ck=torch.load('outputs/model.pt', weights_only=True)
assert isinstance(ck,dict), "top-level must be dict"
assert {'state_dict','feat_mean','feat_std'}<=set(ck), f"missing keys: {set(ck)}"
assert isinstance(ck['feat_mean'],float) and isinstance(ck['feat_std'],float), "feat_* must be float scalars"
m=build_model(cfg); miss,unexp=m.load_state_dict(ck['state_dict'], strict=True); m.to(dev).eval()
assert miss==[] and unexp==[], (miss,unexp)
assert m(torch.randn(3,4,128,device=dev)).shape==(3,1,128)
print("[contract] OK: dict keys + float scalars + strict state_dict load + forward shape")
print(f"[contract] feat_mean={ck['feat_mean']:.6f} feat_std={ck['feat_std']:.6f} params={sum(p.numel() for p in m.parameters())}")
# ---- final metric (the grader protocol: 20-step autoregressive rollout) ----
@torch.no_grad()
def rollout_global(tr):
    tr=tr.to(dev); N,T,_=tr.shape; win=((tr[:,0:k]-ck['feat_mean'])/ck['feat_std']); P=[]
    for _ in range(H): p=m(win); P.append(p); win=torch.cat([win[:,1:],p],1)
    pred=torch.cat(P,1)*ck['feat_std']+ck['feat_mean']; true=tr[:,k:k+H]
    return ((pred-true).reshape(N,-1).pow(2).sum(-1).sqrt()/true.reshape(N,-1).pow(2).sum(-1).sqrt().clamp_min(1e-8)).mean().item()
for nm in ['val','public_test']:
    tr=torch.load(f'data/{nm}/trajectories.pt',weights_only=True)['trajectories'].float()
    print(f"[metric] {nm}: 20-step rollout global rel_L2 (from first window) = {rollout_global(tr):.4f}")
EOF
echo "=== final outputs/ ===" && ls -la outputs/

cd /app && echo "=== BASELINE (stock train.py, 10 epochs) ===" && python3 rollout_eval.py /tmp/baseline10.pt 2>/dev/null | grep -E 'val|test'
echo "" && echo "=== FINAL (improved recipe) ===" && python3 rollout_eval.py outputs/model.pt 2>/dev/null | grep -E 'val|test'

cd /app && rm -f rollout_eval.py compare.py train_improved.py train_improved.py.bak
echo "=== final /app contents ===" && ls -la /app
echo "" && echo "=== sanity: improved train.py still imports & is valid ===" && python3 -c "import ast; ast.parse(open('/app/train.py').read()); print('train.py syntax OK')"
echo "" && echo "=== final deliverable ===" && ls -la /app/outputs/model.pt && python3 -c "
import torch,json; from candidate_model import build_model
ck=torch.load('/app/outputs/model.pt',weights_only=True)
build_model(json.load(open('/app/task_config.json'))).load_state_dict(ck['state_dict'],strict=True)
print('outputs/model.pt loads & matches build_model(task_config) , contract satisfied')
"