ks-equation-1d-forecast

cd /app && ls -la && echo "---- nvidia ----" && nvidia-smi 2>/dev/null || echo "no gpu"; echo "---- python ----"; python3 --version; python3 -c "import torch; print('torch', torch.__version__, 'cuda', torch.cuda.is_available())"

cd /app && echo "==== task_config.json ====" && cat task_config.json && echo && echo "==== feature_schema.json ====" && cat feature_schema.json && echo && echo "==== candidate_model.py ====" && cat candidate_model.py

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

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

cd /app && python3 -c "
import torch
tr = torch.load('data/val/trajectories.pt', map_location='cpu', weights_only=True)['trajectories']
# relative L2 between frame t and t+lag, averaged
def rel_l2(a,b):
    return ((a-b).pow(2).sum(-1).sqrt() / b.pow(2).sum(-1).sqrt().clamp_min(1e-8)).mean().item()
print('rel_l2 frame t vs t+1 :', rel_l2(tr[:,:-1], tr[:,1:]))
print('rel_l2 frame t vs t+5 :', rel_l2(tr[:,:-5], tr[:,5:]))
print('rel_l2 frame t vs t+20:', rel_l2(tr[:,:-20], tr[:,20:]))
# persistence rollout: predict next = last frame, 20 steps from t0=4
# window last frame is frame index 3, predict frames 4..23 all equal frame 3
N,T,X = tr.shape
import numpy as np
scores=[]
for start in range(4, T-20):
    last = tr[:, start-1:start]  # frame start-1
    gt = tr[:, start:start+20]
    pred = last.expand(-1,20,-1)
    r = ((pred-gt).pow(2).sum(-1).sqrt()/gt.pow(2).sum(-1).sqrt().clamp_min(1e-8)).mean().item()
    scores.append(r)
print('persistence 20-step rollout mean rel_l2:', np.mean(scores))
"

cd /app && mkdir -p outputs && TASK_EPOCHS_OVERRIDE=3 python3 train.py 2>&1 | tail -8

cd /app && cat > rollout_eval.py << 'EOF'
from __future__ import annotations
import json, os, sys
from pathlib import Path
import torch
sys.path.insert(0, '/app')
from candidate_model import build_model

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

@torch.no_grad()
def rollout_eval(model, feat_mean, feat_std, traj, device, k=4, horizon=20, starts=None):
    # traj: [N, T, X]. For each start s: window = frames [s-k:s], predict s..s+horizon-1
    N, T, X = traj.shape
    if starts is None:
        starts = list(range(k, T - horizon + 1))
    all_step_rl2 = []  # per-step rel l2 averaged over traj and starts
    final_rl2 = []
    rmse_all = []
    for s in starts:
        window = traj[:, s-k:s].to(device)  # [N,k,X]
        w = (window - feat_mean) / feat_std
        preds = []
        for h in range(horizon):
            p = model(w)  # [N,1,X] normalized
            preds.append(p)
            w = torch.cat([w[:, 1:], p], dim=1)
        pred = torch.cat(preds, dim=1) * feat_std + feat_mean  # [N,horizon,X]
        gt = traj[:, s:s+horizon].to(device)
        num = (pred - gt).pow(2).sum(-1).sqrt()
        den = gt.pow(2).sum(-1).sqrt().clamp_min(1e-8)
        rl2 = (num / den)  # [N, horizon]
        all_step_rl2.append(rl2.mean(0).cpu())  # [horizon]
        final_rl2.append(rl2[:, -1].mean().cpu())
        rmse = (pred - gt).pow(2).mean(-1).sqrt().mean(0).cpu()
        rmse_all.append(rmse)
    step_rl2 = torch.stack(all_step_rl2).mean(0)  # [horizon]
    rmse_step = torch.stack(rmse_all).mean(0)
    return {
        'mean_rollout_rel_l2': step_rl2.mean().item(),
        'final_step_rel_l2': step_rl2[-1].item(),
        'mean_rollout_rmse': rmse_step.mean().item(),
        'step_rl2': step_rl2.tolist(),
        'n_starts': len(starts),
    }

if __name__ == '__main__':
    cfg = json.load(open('/app/task_config.json'))
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    ckpt = sys.argv[1] if len(sys.argv) > 1 else '/app/outputs/model.pt'
    split = sys.argv[2] if len(sys.argv) > 2 else 'val'
    traj = torch.load(f'/app/data/{split}/trajectories.pt', map_location='cpu', weights_only=True)['trajectories'].float()
    model, fm, fs = load_model(ckpt, cfg, device)
    # Evaluation mimicking grader: use early start (window=first 4 frames)
    res_single = rollout_eval(model, fm, fs, traj, device, starts=[4])
    res_multi = rollout_eval(model, fm, fs, traj, device)
    print(f'== {split} ckpt={ckpt} ==')
    print('start=4  : mean_rl2=%.4f final=%.4f rmse=%.4f' % (res_single['mean_rollout_rel_l2'], res_single['final_step_rel_l2'], res_single['mean_rollout_rmse']))
    print('all starts(%d): mean_rl2=%.4f final=%.4f rmse=%.4f' % (res_multi['n_starts'], res_multi['mean_rollout_rel_l2'], res_multi['final_step_rel_l2'], res_multi['mean_rollout_rmse']))
    sr = res_multi['step_rl2']
    print('per-step rl2 [1,2,3,5,10,15,20]:', [round(sr[i],4) for i in [0,1,2,4,9,14,19]])
EOF
python3 rollout_eval.py /app/outputs/model.pt val

cd /app && python3 train.py 2>&1 | tail -12 && cp outputs/model.pt outputs/baseline_10ep.pt && python3 rollout_eval.py outputs/baseline_10ep.pt val

cd /app && python3 -c "
import torch
tr = torch.load('data/train/trajectories.pt', map_location='cpu', weights_only=True)['trajectories'].float()
# energy per fourier mode (rfft along x)
ft = torch.fft.rfft(tr, dim=-1, norm='ortho')  # [N,T,65]
energy = ft.abs().pow(2).mean(dim=(0,1))  # [65]
cum = torch.cumsum(energy, 0) / energy.sum()
print('num rfft modes:', energy.shape[0])
print('cumulative energy fraction at modes [4,8,12,16,20,24,32]:')
for m in [4,8,12,16,20,24,32,64]:
    print('  modes<=%2d : %.5f' % (m, cum[m].item() if m < len(cum) else 1.0))
print('energy beyond mode 16 fraction: %.5e' % (1 - cum[16].item()))
# mean frame norm
print('mean ||frame||_2:', tr.pow(2).sum(-1).sqrt().mean().item())
"

cd /app && python3 -c "
import torch
# Verify reflect+negate symmetry: if u(x,t) solves KS, so does -u(-x,t).
# Test via spectral time-derivative residual on real data consistency:
# Instead, check that transformed trajectory is dynamically consistent by
# comparing one-step finite-difference operators' statistics (sanity proxy).
tr = torch.load('data/train/trajectories.pt', map_location='cpu', weights_only=True)['trajectories'].float()
L=22.0; N=128; dx=L/N
u = tr[0,50]  # a frame
# build spatial derivatives via spectral
k = 2*torch.pi*torch.fft.rfftfreq(N, d=dx)
def deriv(u, order):
    uf = torch.fft.rfft(u)
    return torch.fft.irfft((1j*k)**order * uf, n=N)
# KS RHS: -u*u_x - u_xx - u_xxxx
def rhs(u):
    return -(u*deriv(u,1) + deriv(u,2) + deriv(u,4))
r1 = rhs(u)
# transform w(x) = -u(-x): grid reflection about origin = -roll(flip(u),1)
w = -torch.roll(torch.flip(u, dims=[-1]), shifts=1, dims=-1)
rw = rhs(w)
# w's rhs should equal the transform of r1: (R r1)(x) = -r1(-x) = -roll(flip(r1),1)
r1_transformed = -torch.roll(torch.flip(r1, dims=[-1]), shifts=1, dims=-1)
print('reflect-negate symmetry residual (should be ~0):', (rw - r1_transformed).abs().max().item())
# translation: shift by s cells -> rhs shifts by s cells
s=7
us = torch.roll(u, shifts=s, dims=-1)
print('translation symmetry residual (should be ~0):', (rhs(us) - torch.roll(r1, shifts=s, dims=-1)).abs().max().item())
"

cd /app && python3 -c "
import torch
tr = torch.load('data/train/trajectories.pt', map_location='cpu', weights_only=True)['trajectories'].float()
L=22.0; N=128; dx=L/N
u = tr[0,50]
k = 2*torch.pi*torch.fft.rfftfreq(N, d=dx)
def deriv(u, order):
    return torch.fft.irfft((1j*k)**order * torch.fft.rfft(u), n=N)
rhs = -(u*deriv(u,1)+deriv(u,2)+deriv(u,4))
print('RHS magnitude (max abs): %.3f, symmetry residual ~0.016 => relative %.2e' % (rhs.abs().max().item(), 0.016/rhs.abs().max().item()))
"

command: create
path: /app/train_improved.py
file_text: from __future__ import annotations
import json, os, sys, math, copy, time
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 compute_standardization
from rollout_eval import rollout_eval

# ----------------------------- config -----------------------------
def getenv_f(name, default):
    return float(os.environ.get(name, default))
def getenv_i(name, default):
    return int(os.environ.get(name, default))

CFG = json.load(open(APP_DIR / "task_config.json"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
K = int(CFG["in_channels"])          # 4 input frames
HORIZON = int(CFG["rollout_horizon"])  # 20 (eval)

EPOCHS      = getenv_i("EPOCHS", 80)
BATCH       = getenv_i("BATCH", 128)
LR          = getenv_f("LR", 1.5e-3)
WD          = getenv_f("WD", 1e-4)
WARMUP_FRAC = getenv_f("WARMUP_FRAC", 0.05)
EMA_DECAY   = getenv_f("EMA_DECAY", 0.9995)
GRAD_CLIP   = getenv_f("GRAD_CLIP", 1.0)
AUG         = getenv_i("AUG", 1)          # symmetry augmentation on/off
NOISE_STD   = getenv_f("NOISE_STD", 0.0)  # input noise (in normalized units)
MAX_M       = getenv_i("MAX_M", 10)       # max rollout length in curriculum
SEED        = getenv_i("SEED", int(CFG["seed"]))
TAG         = os.environ.get("TAG", "run")
LOSS        = os.environ.get("LOSS", "rel_l2")  # rel_l2 | mse
PUSHFWD     = getenv_i("PUSHFWD", 0)     # pushforward: unroll no-grad then grad on tail
SAVE_PATH   = os.environ.get("SAVE_PATH", str(APP_DIR / "outputs" / f"model_{TAG}.pt"))

torch.manual_seed(SEED)

# ----------------------------- data -----------------------------
def load_split(name):
    return torch.load(APP_DIR / "data" / name / "trajectories.pt",
                      map_location="cpu", weights_only=True)["trajectories"].float()

train_traj = load_split("train").to(DEVICE)   # [200,100,128]
val_traj   = load_split("val")                # keep on cpu; eval moves as needed
feat_mean, feat_std = compute_standardization(train_traj.cpu())
print(f"[{TAG}] feat_mean={feat_mean:.5f} feat_std={feat_std:.5f} device={DEVICE}", flush=True)

N, T, X = train_traj.shape
train_norm = (train_traj - feat_mean) / feat_std   # normalized, on GPU

# ----------------------------- augmentation -----------------------------
def augment(seq):
    """seq: [B, L, X] normalized frames. Apply KS symmetries.
    - reflect+negate: u(x) -> -u(-x)  (valid: -u(-x,t) solves KS)
    - random circular shift (translation on periodic domain)
    Same transform applied to all L frames of a sample."""
    B, L, Xl = seq.shape
    # reflect+negate with prob 0.5 (normalized mean~0 so negation is well-defined)
    do_ref = (torch.rand(B, device=seq.device) < 0.5)
    reflected = -torch.roll(torch.flip(seq, dims=[-1]), shifts=1, dims=-1)
    seq = torch.where(do_ref.view(B, 1, 1), reflected, seq)
    # random circular shift
    shifts = torch.randint(0, Xl, (B,), device=seq.device)
    ar = torch.arange(Xl, device=seq.device)
    idx = (ar.view(1, Xl) - shifts.view(B, 1)) % Xl           # [B,X]
    idx = idx.view(B, 1, Xl).expand(B, L, Xl)
    seq = torch.gather(seq, 2, idx)
    return seq

# ----------------------------- curriculum -----------------------------
def rollout_len_for_epoch(ep):
    # gradually grow the training rollout horizon
    frac = ep / max(1, EPOCHS - 1)
    if frac < 0.12:   return 1
    if frac < 0.25:   return 2
    if frac < 0.40:   return 4
    if frac < 0.55:   return 6
    if frac < 0.72:   return 8
    return MAX_M

# ----------------------------- model / optim -----------------------------
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, betas=(0.9, 0.99))

steps_per_epoch = math.ceil((N * (T - K - 1)) / BATC
…[truncated]

cd /app && EPOCHS=6 TAG=smoke SAVE_PATH=/app/outputs/model_smoke.pt python3 train_improved.py 2>&1 | tail -15

cd /app && nohup env EPOCHS=80 TAG=main SAVE_PATH=/app/outputs/model_main.pt python3 train_improved.py > /app/outputs/log_main.txt 2>&1 &
echo "launched PID $!"; sleep 60; echo "---- progress ----"; tail -5 /app/outputs/log_main.txt

cd /app && sleep 150; tail -8 /app/outputs/log_main.txt

cd /app && sleep 180; tail -10 /app/outputs/log_main.txt

cd /app && sleep 130; tail -6 /app/outputs/log_main.txt; echo "--- still running? ---"; pgrep -f train_improved.py || echo "DONE"

cd /app && sleep 20; tail -2 /app/outputs/log_main.txt; pgrep -f train_improved.py >/dev/null && echo RUNNING || echo DONE

cd /app && for split in val public_test; do
echo "######## $split ########"
echo "-- baseline --"; python3 rollout_eval.py outputs/baseline_10ep.pt $split | sed -n '2,4p'
echo "-- main --";     python3 rollout_eval.py outputs/model_main.pt   $split | sed -n '2,4p'
done

command: str_replace
path: /app/train_improved.py
old_str:     if frac < 0.55:   return 6
    if frac < 0.72:   return 8
    return MAX_M
new_str:     if frac < 0.55:   return 6
    if frac < 0.72:   return 8
    if frac < 0.85:   return min(MAX_M, 10)
    return MAX_M

cd /app && nohup env EPOCHS=90 MAX_M=16 TAG=longM SAVE_PATH=/app/outputs/model_longM.pt python3 train_improved.py > /app/outputs/log_longM.txt 2>&1 &
echo "longM PID $!"
nohup env EPOCHS=80 MAX_M=10 AUG=0 TAG=noaug SAVE_PATH=/app/outputs/model_noaug.pt python3 train_improved.py > /app/outputs/log_noaug.txt 2>&1 &
echo "noaug PID $!"
sleep 5; echo "both launched"; nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json'))
ck = torch.load('outputs/model_main.pt', weights_only=True)
print('top-level keys:', list(ck.keys()))
print('feat_mean type/val:', type(ck['feat_mean']).__name__, ck['feat_mean'])
print('feat_std  type/val:', type(ck['feat_std']).__name__, ck['feat_std'])
devs = set(str(v.device) for v in ck['state_dict'].values())
print('state_dict tensor devices:', devs)
m = build_model(cfg)
ret = m.load_state_dict(ck['state_dict'], strict=True)
print('load_state_dict strict=True OK:', ret)
# confirm a forward pass on cpu
import torch
x = torch.randn(2,4,128)
print('forward output shape:', tuple(m(x).shape))
"

command: str_replace
path: /app/train_improved.py
old_str:             if sc < best["score"]:
                best = {"score": sc, "final": r["final_step_rel_l2"], "which": name, "epoch": ep,
                        "state": copy.deepcopy(mdl.state_dict())}
new_str:             if sc < best["score"]:
                cpu_state = {k: v.detach().cpu().clone() for k, v in mdl.state_dict().items()}
                best = {"score": sc, "final": r["final_step_rel_l2"], "which": name, "epoch": ep,
                        "state": cpu_state}

cd /app && sleep 210; echo "=== longM ==="; tail -3 /app/outputs/log_longM.txt; echo "=== noaug ==="; tail -3 /app/outputs/log_noaug.txt

cd /app && sleep 220; echo "=== longM ==="; tail -3 /app/outputs/log_longM.txt; echo "=== noaug ==="; tail -3 /app/outputs/log_noaug.txt; echo "--- running ---"; pgrep -f train_improved.py || echo NONE

cd /app && sleep 150; echo "=== noaug tail ==="; tail -2 /app/outputs/log_noaug.txt; echo "=== longM tail ==="; tail -2 /app/outputs/log_longM.txt; pgrep -f train_improved.py || echo NONE

cd /app && sleep 120; echo "=== noaug ==="; grep SAVED /app/outputs/log_noaug.txt || tail -1 /app/outputs/log_noaug.txt; echo "=== longM ==="; grep SAVED /app/outputs/log_longM.txt || tail -2 /app/outputs/log_longM.txt; pgrep -f train_improved.py || echo ALL_DONE

cd /app && for ck in baseline_10ep model_noaug model_main model_longM; do
  echo "### $ck ###"
  for split in val public_test; do
    line=$(python3 rollout_eval.py outputs/$ck.pt $split | sed -n '3p')
    echo "  $split  $line"
  done
done

cd /app && python3 -c "
import json, torch
from rollout_eval import load_model, rollout_eval
cfg = json.load(open('task_config.json')); dev='cuda'
model, fm, fs = load_model('outputs/model_longM.pt', cfg, dev)
for split in ['val','public_test']:
    traj = torch.load(f'data/{split}/trajectories.pt', weights_only=True)['trajectories'].float()
    r4 = rollout_eval(model, fm, fs, traj, dev, horizon=20, starts=[4])
    r50 = rollout_eval(model, fm, fs, traj, dev, horizon=50, starts=[4])
    print(f'{split}: start=4 20-step mean={r4[\"mean_rollout_rel_l2\"]:.4f} final={r4[\"final_step_rel_l2\"]:.4f} | 50-step final={r50[\"final_step_rel_l2\"]:.4f} (stability)')
"

cd /app && cp outputs/model_longM.pt outputs/model.pt && python3 -c "
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)
build_model(cfg).load_state_dict(ck['state_dict'],strict=True)
print('provisional /app/outputs/model.pt OK; keys', list(ck.keys()), 'devices', set(str(v.device) for v in ck['state_dict'].values()))
"

command: str_replace
path: /app/train_improved.py
old_str: def rollout_len_for_epoch(ep):
    # gradually grow the training rollout horizon
    frac = ep / max(1, EPOCHS - 1)
    if frac < 0.12:   return 1
    if frac < 0.25:   return 2
    if frac < 0.40:   return 4
    if frac < 0.55:   return 6
    if frac < 0.72:   return 8
    if frac < 0.85:   return min(MAX_M, 10)
    return MAX_M
new_str: CURRICULUM = os.environ.get("CURRICULUM", "stager")  # stager | ramp
def rollout_len_for_epoch(ep):
    # gradually grow the training rollout horizon
    frac = ep / max(1, EPOCHS - 1)
    if CURRICULUM == "ramp":
        # single-step warmup, smooth ramp 2..MAX_M by frac=0.80, then hold MAX_M
        if frac < 0.12:  return 1
        if frac >= 0.80: return MAX_M
        r = (frac - 0.12) / (0.80 - 0.12)
        return int(round(2 + r * (MAX_M - 2)))
    if frac < 0.12:   return 1
    if frac < 0.25:   return 2
    if frac < 0.40:   return 4
    if frac < 0.55:   return 6
    if frac < 0.72:   return 8
    if frac < 0.85:   return min(MAX_M, 10)
    return MAX_M

cd /app && EPOCHS=110 python3 -c "
import os
os.environ['EPOCHS']='110'; os.environ['MAX_M']='20'; os.environ['CURRICULUM']='ramp'
# mimic the schedule function
EPOCHS=110; MAX_M=20
def sched(ep):
    frac=ep/(EPOCHS-1)
    if frac<0.12: return 1
    if frac>=0.80: return MAX_M
    r=(frac-0.12)/(0.80-0.12); return int(round(2+r*(MAX_M-2)))
vals=[sched(e) for e in range(EPOCHS)]
from itertools import groupby
print('M schedule (value:count):', [(k,len(list(g))) for k,g in groupby(vals)])
"

cd /app && nohup env EPOCHS=110 MAX_M=20 CURRICULUM=ramp TAG=m20 SAVE_PATH=/app/outputs/model_m20.pt python3 train_improved.py > /app/outputs/log_m20.txt 2>&1 &
echo "m20 PID $!"; sleep 90; tail -4 /app/outputs/log_m20.txt

cd /app && sleep 230; tail -5 /app/outputs/log_m20.txt

cd /app && sleep 240; tail -5 /app/outputs/log_m20.txt

cd /app && sleep 260; tail -5 /app/outputs/log_m20.txt; pgrep -f train_improved.py >/dev/null && echo RUNNING || echo DONE

cd /app && sleep 90; grep SAVED /app/outputs/log_m20.txt || tail -2 /app/outputs/log_m20.txt; pgrep -f train_improved.py >/dev/null && echo RUNNING || echo DONE

cd /app && sleep 60; grep SAVED /app/outputs/log_m20.txt; pgrep -f train_improved.py >/dev/null && echo RUNNING || echo DONE

cd /app && for ck in model_main model_longM model_m20; do
  v=$(python3 rollout_eval.py outputs/$ck.pt val | sed -n '3p')
  p=$(python3 rollout_eval.py outputs/$ck.pt public_test | sed -n '3p')
  echo "$ck"; echo "   val: $v"; echo "   pub: $p"
done

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
ck = torch.load('outputs/model_longM.pt', weights_only=True)
cpu_state = {k: v.detach().cpu().clone() for k, v in ck['state_dict'].items()}
out = {'state_dict': cpu_state, 'feat_mean': float(ck['feat_mean']), 'feat_std': float(ck['feat_std'])}
torch.save(out, 'outputs/model.pt')
# verify
cfg = json.load(open('task_config.json'))
ck2 = torch.load('outputs/model.pt', weights_only=True)
assert set(ck2.keys()) == {'state_dict','feat_mean','feat_std'}
assert isinstance(ck2['feat_mean'], float) and isinstance(ck2['feat_std'], float)
devs = set(str(v.device) for v in ck2['state_dict'].values())
m = build_model(cfg); m.load_state_dict(ck2['state_dict'], strict=True); m.eval()
print('OK: keys', list(ck2.keys()), '| devices', devs, '| feat_mean', ck2['feat_mean'], 'feat_std', ck2['feat_std'])
"

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json'))
k = cfg['in_channels']; H = cfg['rollout_horizon']
# Fresh load exactly per contract
ck = torch.load('outputs/model.pt', weights_only=True)
m = build_model(cfg); m.load_state_dict(ck['state_dict']); m.eval()
fm, fs = ck['feat_mean'], ck['feat_std']

@torch.no_grad()
def grade(split, start=4):
    tr = torch.load(f'data/{split}/trajectories.pt', weights_only=True)['trajectories'].float()
    N,T,X = tr.shape
    win = (tr[:, start-k:start] - fm)/fs      # normalized window [N,k,X]
    preds=[]
    for _ in range(H):
        p = m(win)                            # [N,1,X]
        preds.append(p)
        win = torch.cat([win[:,1:], p], dim=1)
    pred = torch.cat(preds,1)*fs+fm           # [N,H,X]
    gt = tr[:, start:start+H]
    rl2 = ((pred-gt).pow(2).sum(-1).sqrt()/gt.pow(2).sum(-1).sqrt().clamp_min(1e-8))  # [N,H]
    return rl2.mean().item(), rl2[:,-1].mean().item(), rl2.mean(0)
for sp in ['val','public_test']:
    mean, fin, steps = grade(sp)
    print(f'{sp:12s} rollout_rel_l2(mean over 20)={mean:.5f}  final_step={fin:.5f}')
    print('            per-step:', ' '.join(f'{v:.4f}' for v in steps.tolist()))
"

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

Key ideas (all target the graded 20-step autoregressive rollout, where errors
compound), relative to the weak single-step MSE baseline:

  * Multi-step rollout training with a curriculum: the model is unrolled for M
    steps feeding its own predictions back in, and the loss is accumulated over
    all M steps against ground truth. M grows over training (1 -> MAX_M). This
    directly teaches the model to suppress compounding error. This is the single
    biggest lever and flattens the per-step error-growth curve.
  * Symmetry data augmentation from the KS equation on a periodic domain:
      - translation  : random circular shift (exact for band-limited fields)
      - reflect+negate: u(x,t) -> -u(-x,t) is also a KS solution
    Applied identically to every frame of a sample; improves generalization.
  * Relative-L2 loss (matches the evaluation metric) in normalized space.
  * AdamW + linear warmup + cosine decay, gradient clipping.
  * EMA of weights; best checkpoint chosen by a 20-step val rollout (raw & EMA).

Output contract (torch.load(..., weights_only=True)):
  {"state_dict": <matches build_model(task_config)>,
   "feat_mean": float, "feat_std": float}

Env overrides: TASK_EPOCHS_OVERRIDE, EPOCHS, BATCH, LR, WD, MAX_M, AUG, SEED ...
"""
import json, os, sys, math, copy, time
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 compute_standardization


def ei(name, d):  # env int
    return int(os.environ.get(name, d))
def ef(name, d):  # env float
    return float(os.environ.get(name, d))


def main() -> None:
    cfg = json.load((APP_DIR / "task_config.json").open())
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    K = int(cfg["in_channels"])              # input window length (4)
    HORIZON = int(cfg["rollout_horizon"])    # eval rollout length (20)
    SEED = ei("SEED", int(cfg["seed"]))
    torch.manual_seed(SEED)

    # ---- hyperparameters (defaults reproduce the best M=16 recipe) ----
    EPOCHS = ei("TASK_EPOCHS_OVERRIDE", ei("EPOCHS", 90))
    BATCH = ei("BATCH", 128)
    LR = ef("LR", 1.5e-3)
    WD = ef("WD", 1e-4)
    WARMUP_FRAC = ef("WARMUP_FRAC", 0.05)
    EMA_DECAY = ef("EMA_DECAY", 0.9995)
    GRAD_CLIP = ef("GRAD_CLIP", 1.0)
    AUG = ei("AUG", 1)
    MAX_M = ei("MAX_M", 16)

    # ---- data (kept on-device; trajectories are tiny) ----
    def load(name):
        return torch.load(APP_DIR / "data" / name / "trajectories.pt",
                          map_location="cpu", weights_only=True)["trajectories"].float()
    train_traj = load("train")
    val_traj = load("val")
    feat_mean, feat_std = compute_standardization(train_traj)
    N, T, X = train_traj.shape
    train_norm = ((train_traj - feat_mean) / feat_std).to(device)
    print(f"[train] device={device} epochs={EPOCHS} traj={N} feat_mean={feat_mean:.4g} "
          f"feat_std={feat_std:.4g} MAX_M={MAX_M} aug={AUG}", flush=True)

    # ---- symmetry augmentation ----
    def augment(seq):  # seq: [B, L, X] normalized
        B, L, Xl = seq.shape
        do_ref = (torch.rand(B, device=seq.device) < 0.5).view(B, 1, 1)
        reflected = -torch.roll(torch.flip(seq, dims=[-1]), shifts=1, dims=-1)
        seq = torch.where(do_ref, reflected, seq)
        shifts = torch.randint(0, Xl, (B,), device=seq.device)
        ar = torch.arange(Xl, device=seq.device)
        idx = ((ar.view(1, Xl) - shifts.view(B, 1)) % Xl).view(B, 1, Xl).expand(B, L, Xl)
        return torch.gather(seq, 2, idx)

    # ---- rollout-length curriculum ----
    def rollout_len(ep):
        frac = ep / max(1, EPOCHS - 1)
        if frac < 0.12: return 1
        if frac < 0.25: return 2
        if frac < 0.40: return 4
   
…[truncated]

cd /app && cp outputs/model.pt outputs/model_BEST.pt && echo "backed up" && TASK_EPOCHS_OVERRIDE=2 python3 train.py 2>&1 | tail -6

cd /app && python3 -c "
import torch,json; from candidate_model import build_model
ck=torch.load('outputs/model.pt',weights_only=True)
build_model(json.load(open('task_config.json'))).load_state_dict(ck['state_dict'],strict=True)
print('2-epoch smoke output valid & loadable:', list(ck.keys()))
" && cp outputs/model_BEST.pt outputs/model.pt && echo "restored best -> model.pt" && python3 -c "
import torch,json; from candidate_model import build_model
ck=torch.load('outputs/model.pt',weights_only=True)
m=build_model(json.load(open('task_config.json'))); m.load_state_dict(ck['state_dict'],strict=True)
print('final model.pt devices:', set(str(v.device) for v in ck['state_dict'].values()), '| feat_mean',ck['feat_mean'],'feat_std',ck['feat_std'])
"

cd /app && nohup python3 train.py > /app/outputs/log_repro.txt 2>&1 &
echo "repro PID $!"; sleep 5; echo "launched; best is safe in model_BEST.pt"; ls -la outputs/model_BEST.pt

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
ck = torch.load('outputs/model_BEST.pt', weights_only=True)
bad = {k: (torch.isnan(v).any().item() or torch.isinf(v).any().item()) for k,v in ck['state_dict'].items()}
print('any NaN/Inf in weights:', any(bad.values()))
m = build_model(json.load(open('task_config.json'))); m.load_state_dict(ck['state_dict']); m.eval()
import torch
x = torch.randn(8,4,128)*1.2
with torch.no_grad(): y = m((x - ck['feat_mean'])/ck['feat_std'])
print('output shape', tuple(y.shape), 'finite:', torch.isfinite(y).all().item(), 'range [%.3f, %.3f]' % (y.min(), y.max()))
print('num params:', sum(p.numel() for p in m.parameters()))
"

cd /app && sleep 250; tail -4 /app/outputs/log_repro.txt

cd /app && sleep 260; tail -4 /app/outputs/log_repro.txt

cd /app && sleep 60; grep "wrote" /app/outputs/log_repro.txt || tail -2 /app/outputs/log_repro.txt; pgrep -f "python3 train.py" >/dev/null && echo RUNNING || echo DONE

cd /app && echo "=== reproduced (current outputs/model.pt) ===" && python3 rollout_eval.py outputs/model.pt val | sed -n '3p' && python3 rollout_eval.py outputs/model.pt public_test | sed -n '3p'
echo "=== model_BEST (longM) ===" && python3 rollout_eval.py outputs/model_BEST.pt val | sed -n '3p' && python3 rollout_eval.py outputs/model_BEST.pt public_test | sed -n '3p'

cd /app && python3 -c "
import json, torch
from rollout_eval import load_model, rollout_eval
cfg=json.load(open('task_config.json')); dev='cuda'
for tag,path in [('reproduced','outputs/model.pt'),('longM_BEST','outputs/model_BEST.pt')]:
    m,fm,fs=load_model(path,cfg,dev)
    row=[]
    for sp in ['val','public_test']:
        tr=torch.load(f'data/{sp}/trajectories.pt',weights_only=True)['trajectories'].float()
        r=rollout_eval(m,fm,fs,tr,dev)
        row.append((sp,r['mean_rollout_rel_l2'],r['final_step_rel_l2']))
    print(f'{tag:12s}', ' | '.join(f'{sp}: mean={a:.5f} final={b:.5f}' for sp,a,b in row))
"

cd /app && cp outputs/model_BEST.pt outputs/model.pt && python3 -c "
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json'))
# Exact contract load
ck = torch.load('outputs/model.pt', weights_only=True)
assert set(ck.keys()) == {'state_dict','feat_mean','feat_std'}, ck.keys()
assert isinstance(ck['feat_mean'], float) and isinstance(ck['feat_std'], float)
m = build_model(cfg); res = m.load_state_dict(ck['state_dict'], strict=True); m.eval()
print('FINAL model.pt:')
print('  keys         :', list(ck.keys()))
print('  strict load  :', res)
print('  devices      :', set(str(v.device) for v in ck['state_dict'].values()))
print('  feat_mean/std: %.6g / %.6g' % (ck['feat_mean'], ck['feat_std']))
print('  #params      :', sum(p.numel() for p in m.parameters()))
import os; print('  file size    : %d bytes' % os.path.getsize('outputs/model.pt'))
"

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json')); k=cfg['in_channels']; H=cfg['rollout_horizon']
def evalck(path, split, start=4):
    ck = torch.load(path, weights_only=True)
    m = build_model(cfg); m.load_state_dict(ck['state_dict']); m.eval()
    fm, fs = ck['feat_mean'], ck['feat_std']
    tr = torch.load(f'data/{split}/trajectories.pt', weights_only=True)['trajectories'].float()
    with torch.no_grad():
        w=(tr[:,start-k:start]-fm)/fs; preds=[]
        for _ in range(H):
            p=m(w); preds.append(p); w=torch.cat([w[:,1:],p],1)
        pred=torch.cat(preds,1)*fs+fm; gt=tr[:,start:start+H]
        rl2=((pred-gt).pow(2).sum(-1).sqrt()/gt.pow(2).sum(-1).sqrt().clamp_min(1e-8))
    return rl2.mean().item(), rl2[:,-1].mean().item()
print('%-10s %-12s %10s %12s' % ('model','split','mean_rl2','final_rl2'))
for split in ['val','public_test']:
    bm,bf=evalck('outputs/baseline_10ep.pt',split)
    fm_,ff=evalck('outputs/model.pt',split)
    print('%-10s %-12s %10.5f %12.5f' % ('baseline',split,bm,bf))
    print('%-10s %-12s %10.5f %12.5f   (%.1fx mean, %.1fx final better)' % ('FINAL',split,fm_,ff,bm/fm_,bf/ff))
"

cd /app && rm -f outputs/model_BEST.pt outputs/model_main.pt outputs/model_longM.pt outputs/model_noaug.pt outputs/model_m20.pt outputs/model_smoke.pt outputs/baseline_10ep.pt && ls -la outputs/ && echo "--- final contract check ---" && python3 -c "
import torch, json; from candidate_model import build_model
ck=torch.load('outputs/model.pt', weights_only=True)
build_model(json.load(open('task_config.json'))).load_state_dict(ck['state_dict'], strict=True)
print('model.pt OK | keys:', list(ck.keys()), '| feat_mean=%.4g feat_std=%.4g'%(ck['feat_mean'],ck['feat_std']))
"