ks-equation-1d-forecast

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

cd /app && echo "===CANDIDATE MODEL===" && cat candidate_model.py && echo "===DATASET===" && cat dataset.py && echo "===TRAIN===" && 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))
print('cpu count', __import__('os').cpu_count())
" && echo "---" && nproc && free -h 2>/dev/null | head -3

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

cd /app && python3 -c "
import torch
t = torch.load('data/train/trajectories.pt', weights_only=True)['trajectories'].float()
# temporal autocorrelation: how fast frames decorrelate
tr = t[0]  # [100,128]
f0 = tr[0]
for dt in [1,2,4,8,16,20,40]:
    # relative L2 between frame 0 and frame dt (persistence error proxy)
    rel = (tr[dt]-tr[0]).norm()/tr[0].norm()
    print(f'frame0 vs frame{dt}: relL2={rel:.3f}')
print('---- step-to-step rel change ----')
for dt in [1,2,4]:
    rels=[]
    for i in range(t.shape[0]):
        tr=t[i]
        for s in range(0,96,10):
            rels.append(((tr[s+dt]-tr[s]).norm()/tr[s].norm()).item())
    import numpy as np
    print(f'dt={dt}: mean persistence relL2={np.mean(rels):.3f}')
"

command: create
path: /app/rollout_eval.py
file_text: """Mimic the hidden grader: autoregressive rollout for H steps, score rel-L2."""
from __future__ import annotations
import json, os, sys
from pathlib import Path
import torch

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


@torch.no_grad()
def rollout_metrics(model, traj, feat_mean, feat_std, k=4, horizon=20, device="cuda"):
    """traj: [N, T, X] ground truth. Start from first k frames, roll horizon steps."""
    model.eval()
    N, T, X = traj.shape
    assert T >= k + horizon
    traj = traj.to(device)
    window = ((traj[:, :k] - feat_mean) / feat_std).clone()  # normalized [N,k,X]
    preds = []
    for step in range(horizon):
        out = model(window)                      # [N,1,X] normalized
        preds.append(out)
        window = torch.cat([window[:, 1:], out], dim=1)
    pred = torch.cat(preds, dim=1) * feat_std + feat_mean     # [N,horizon,X] physical
    true = traj[:, k:k + horizon]                             # [N,horizon,X]

    # Per-frame relative L2 then averaged (most common convention)
    num = (pred - true).pow(2).sum(-1).sqrt()                 # [N,horizon]
    den = true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
    per_frame_rel = (num / den)                               # [N,horizon]
    rollout_rel_l2 = per_frame_rel.mean().item()
    final_rel_l2 = per_frame_rel[:, -1].mean().item()

    # Global (whole spatiotemporal block) relative L2
    g_num = (pred - true).pow(2).sum(dim=(1, 2)).sqrt()
    g_den = true.pow(2).sum(dim=(1, 2)).sqrt().clamp_min(1e-8)
    global_rel_l2 = (g_num / g_den).mean().item()

    rmse = (pred - true).pow(2).mean().sqrt().item()
    # per-step curve
    curve = per_frame_rel.mean(0).tolist()
    return {
        "rollout_rel_l2": rollout_rel_l2,
        "final_rel_l2": final_rel_l2,
        "global_rel_l2": global_rel_l2,
        "rmse": rmse,
        "curve": curve,
    }


def load_model_from_ckpt(path, cfg, device="cuda"):
    ckpt = torch.load(path, map_location=device, weights_only=True)
    model = build_model(cfg).to(device)
    model.load_state_dict(ckpt["state_dict"])
    return model, float(ckpt["feat_mean"]), float(ckpt["feat_std"])


def main():
    with (APP_DIR / "task_config.json").open() as fh:
        cfg = json.load(fh)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    ckpt_path = sys.argv[1] if len(sys.argv) > 1 else str(APP_DIR / "outputs" / "model.pt")
    model, fm, fs = load_model_from_ckpt(ckpt_path, cfg, device)
    print(f"ckpt={ckpt_path} feat_mean={fm:.5f} feat_std={fs:.5f}")
    for split in ["val", "public_test"]:
        traj = torch.load(APP_DIR / "data" / split / "trajectories.pt", weights_only=True)["trajectories"].float()
        m = rollout_metrics(model, traj, fm, fs, k=int(cfg["in_channels"]), horizon=int(cfg["rollout_horizon"]), device=device)
        print(f"[{split}] rollout_rel_l2={m['rollout_rel_l2']:.4f}  final={m['final_rel_l2']:.4f}  global={m['global_rel_l2']:.4f}  rmse={m['rmse']:.4f}")
        c = m["curve"]
        print("   curve:", " ".join(f"{v:.3f}" for v in c))


if __name__ == "__main__":
    main()

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

cd /app && python3 rollout_eval.py 2>&1 | tail -10

cd /app && python3 -c "
import torch
t = torch.load('data/train/trajectories.pt', weights_only=True)['trajectories'].float()
# Reflection symmetry: v(x,t) = -u(-x,t). Under periodic discrete flip.
# If original (u[t0-4:t0] -> u[t0]) is a valid transition, so is the reflected one.
# We test statistical invariance: distribution of reflected frames should match original stats.
tr = t[0,10]
flip_neg = -torch.flip(tr, dims=[-1])
print('orig  mean/std', tr.mean().item(), tr.std().item())
print('flipN mean/std', flip_neg.mean().item(), flip_neg.std().item())
# Roll invariance is exact for periodic grid by construction.
# Verify spectrum preserved under flip+neg (KS reflection):
import numpy as np
s1 = torch.fft.rfft(tr).abs()
s2 = torch.fft.rfft(flip_neg).abs()
print('spectrum max abs diff', (s1-s2).abs().max().item())
"

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


def augment(frames, do_roll=True, do_reflect=True):
    """frames: [B, F, X] physical. Apply exact KS symmetries jointly across F."""
    B, Fr, X = frames.shape
    if do_roll:
        shift = torch.randint(0, X, (B,), device=frames.device)
        base = torch.arange(X, device=frames.device).unsqueeze(0)      # [1,X]
        idx = (base - shift.unsqueeze(1)) % X                          # [B,X]
        idx = idx.unsqueeze(1).expand(B, Fr, X)
        frames = torch.gather(frames, 2, idx)
    if do_reflect:
        m = (torch.rand(B, 1, 1, device=frames.device) < 0.5).float()
        flipped = -torch.flip(frames, dims=[-1])
        frames = m * flipped + (1 - m) * frames
    return frames


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--wd", type=float, default=1e-4)
    ap.add_argument("--warmup", type=int, default=3)
    ap.add_argument("--rollout_max", type=int, default=8)
    ap.add_argument("--rollout_start", type=int, default=2)
    ap.add_argument("--curr_epochs", type=int, default=20, help="epoch by which rollout reaches max")
    ap.add_argument("--noise", type=float, default=0.01, help="input noise std (normalized units)")
    ap.add_argument("--pf", type=int, default=0, help="pushforward no-grad warmup steps (0=full BPTT)")
    ap.add_argument("--loss", type=str, default="rel", choices=["rel", "mse", "mix"])
    ap.add_argument("--no_aug", action="store_true")
    ap.add_argument("--ema", type=float, default=0.999)
    ap.add_argument("--steps_per_epoch", type=int, default=400)
    ap.add_argument("--clip", type=float, default=1.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", type=str, default=str(APP_DIR / "outputs" / "model.pt"))
    ap.add_argument("--tag", type=str, default="exp")
    args = ap.parse_args()

    with (APP_DIR / "task_config.json").open() as fh:
        cfg = json.load(fh)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    seed = int(cfg["seed"]) + args.seed
    torch.manual_seed(seed)
    k = int(cfg["in_channels"])
    H = int(cfg["rollout_horizon"])

    train_ds = KSForecast(APP_DIR / "data" / "train", k=k)
    feat_mean, feat_std = compute_standardization(train_ds.trajectories)
    traj = train_ds.trajectories.to(device)                 # [N,T,X] physical
    N, T, X = traj.shape
    val_traj = torch.load(APP_DIR / "data" / "val" / "trajectories.pt", weights_only=True)["trajectories"].float()

    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=args.lr, weight_decay=args.wd)

    def lr_at(ep):
        if ep < args.warmup:
            return args.lr * (ep + 1) / max(1, args.warmup)
        import math
        prog = (ep - args.warmup) / max(1, args.epochs - args.warmup)
        return 0.5 * args.lr * (1 + math.cos(math.pi * prog))

    def rollout_len(ep):
        if ep >= args.curr_epochs:
            return args.rollout_max
        frac = ep / max(1, args.curr_epochs)
        return int(round(args.rollout_start + frac * (args.rollout_max - args.rollout_start)))

    fm, fs = feat_mean, feat_std
    best = {"rollout_rel_l2": 1e9}
    gstep = 0
    for ep in range(args.epochs):
        model.train()
        for g in opt.param_groups:
            g["lr"] = lr_at(ep)
        R = rollout_len(ep)
        running = 0.0

…[truncated]

cd /app && time python3 train_exp.py --epochs 3 --steps_per_epoch 150 --rollout_start 2 --rollout_max 4 --curr_epochs 3 --tag smoke --out /app/outputs/smoke.pt 2>&1 | tail -10

cd /app && mkdir -p logs exp && nohup python3 train_exp.py \
  --epochs 45 --steps_per_epoch 300 --lr 1e-3 --warmup 3 \
  --rollout_start 2 --rollout_max 8 --curr_epochs 22 \
  --noise 0.01 --ema 0.999 --loss rel --clip 1.0 \
  --tag rollout --out /app/exp/rollout.pt > logs/rollout.log 2>&1 &
echo "started PID $!"; sleep 60; echo "--- after 60s ---"; cat logs/rollout.log

cd /app && kill 167 2>/dev/null; sleep 2; echo "killed"; nvidia-smi --query-gpu=memory.used --format=csv,noheader 2>/dev/null

cd /app && cat > run_ablation.sh <<'EOF'
#!/bin/bash
cd /app
# A: one-step MSE + aug + EMA + long (tests aug/EMA/long vs baseline)
python3 train_exp.py --epochs 30 --steps_per_epoch 250 --lr 1e-3 --warmup 3 \
  --rollout_start 1 --rollout_max 1 --curr_epochs 1 --noise 0 --loss mse \
  --ema 0.999 --tag A_1step --out /app/exp/A.pt > logs/A.log 2>&1
# B: one-step MSE + small noise
python3 train_exp.py --epochs 30 --steps_per_epoch 250 --lr 1e-3 --warmup 3 \
  --rollout_start 1 --rollout_max 1 --curr_epochs 1 --noise 0.005 --loss mse \
  --ema 0.999 --tag B_noise --out /app/exp/B.pt > logs/B.log 2>&1
# C: rollout MSE + aug + EMA, no noise
python3 train_exp.py --epochs 30 --steps_per_epoch 250 --lr 1e-3 --warmup 3 \
  --rollout_start 2 --rollout_max 6 --curr_epochs 15 --noise 0 --loss mse \
  --ema 0.999 --tag C_rollout --out /app/exp/C.pt > logs/C.log 2>&1
echo "ALL DONE" > logs/ablation_done.txt
EOF
chmod +x run_ablation.sh && nohup ./run_ablation.sh > logs/driver.log 2>&1 &
echo "driver PID $!"

cd /app && sleep 75 && echo "=== A (1-step MSE+aug+EMA) ===" && tail -6 logs/A.log 2>/dev/null; echo "=== B (1-step +noise) ===" && tail -6 logs/B.log 2>/dev/null; echo "=== C (rollout MSE) ===" && tail -6 logs/C.log 2>/dev/null

cd /app && sleep 90 && echo "=== B (1-step +noise 0.005) ===" && tail -4 logs/B.log 2>/dev/null; echo "=== C (rollout MSE 2->6) ===" && tail -6 logs/C.log 2>/dev/null; ls logs/ablation_done.txt 2>/dev/null && echo DONE

cd /app && sleep 100 && echo "=== C (rollout MSE 2->6) ===" && tail -8 logs/C.log 2>/dev/null; ls logs/ablation_done.txt 2>/dev/null && echo ALLDONE

cd /app && while [ ! -f logs/ablation_done.txt ]; do sleep 5; done; echo "=== C final ==="; tail -3 logs/C.log
echo "=== SUMMARY (best lines) ==="; grep BEST logs/A.log logs/B.log logs/C.log

cd /app && for c in A B C; do echo "=== $c ==="; python3 rollout_eval.py /app/exp/$c.pt 2>&1 | grep -E "rollout_rel_l2"; done

cd /app && rm -f logs/ablation_done.txt && cat > run_tune.sh <<'EOF'
#!/bin/bash
cd /app
# D: no-aug control (isolate augmentation's effect) - same as C but --no_aug
python3 train_exp.py --epochs 25 --steps_per_epoch 250 --lr 1e-3 --warmup 3 \
  --rollout_start 2 --rollout_max 6 --curr_epochs 15 --noise 0 --loss mse \
  --ema 0.999 --no_aug --tag D_noaug --out /app/exp/D.pt > logs/D.log 2>&1
# E: main candidate - rollout MSE + aug, 60 epochs, max 6
python3 train_exp.py --epochs 60 --steps_per_epoch 300 --lr 1e-3 --warmup 4 \
  --rollout_start 2 --rollout_max 6 --curr_epochs 28 --noise 0 --loss mse \
  --ema 0.999 --tag E_main --out /app/exp/E.pt > logs/E.log 2>&1
# F: longer rollout horizon, max 10
python3 train_exp.py --epochs 60 --steps_per_epoch 300 --lr 1e-3 --warmup 4 \
  --rollout_start 2 --rollout_max 10 --curr_epochs 30 --noise 0 --loss mse \
  --ema 0.999 --tag F_r10 --out /app/exp/F.pt > logs/F.log 2>&1
echo "DONE" > logs/tune_done.txt
EOF
chmod +x run_tune.sh && nohup ./run_tune.sh > logs/driver2.log 2>&1 &
echo "driver2 PID $!"; sleep 70; echo "--- D so far ---"; tail -4 logs/D.log

cd /app && sleep 180 && echo "=== D final ===" && grep BEST logs/D.log; echo "=== E so far ===" && tail -3 logs/E.log; echo "=== F so far ===" && tail -3 logs/F.log 2>/dev/null

cd /app && sleep 240 && echo "=== E so far ===" && tail -3 logs/E.log; echo "=== F so far ===" && tail -3 logs/F.log 2>/dev/null; ls logs/tune_done.txt 2>/dev/null && echo ALLDONE

cd /app && while [ ! -f logs/tune_done.txt ]; do sleep 10; done; echo "=== E final ==="; grep BEST logs/E.log; tail -3 logs/E.log; echo "=== F final ==="; grep BEST logs/F.log; tail -3 logs/F.log

restart: true

cd /app && echo "=== E final ==="; grep BEST logs/E.log; echo "=== F progress ==="; tail -3 logs/F.log; ls logs/tune_done.txt 2>/dev/null && echo ALLDONE || echo "F still running"

cd /app && python3 rollout_eval.py /app/exp/E.pt 2>&1

cd /app && sleep 200; echo "=== F progress ==="; tail -4 logs/F.log; ls logs/tune_done.txt 2>/dev/null && echo ALLDONE || echo "F still running"

cd /app && ps aux | grep -E "train_exp|run_tune" | grep -v grep; echo "---GPU---"; nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader 2>/dev/null; echo "---driver2 log---"; tail -5 logs/driver2.log

cd /app && ps aux | grep python | grep -v grep; echo "--- all procs w/ cuda ---"; nvidia-smi 2>/dev/null | tail -12

cd /app && mkdir -p outputs && cp /app/exp/E.pt /app/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)   # must work with weights_only=True
assert set(['state_dict','feat_mean','feat_std']).issubset(ck.keys()), ck.keys()
print('keys:', list(ck.keys()))
print('feat_mean', ck['feat_mean'], 'feat_std', ck['feat_std'], 'types', type(ck['feat_mean']).__name__, type(ck['feat_std']).__name__)
m = build_model(cfg)
missing, unexpected = m.load_state_dict(ck['state_dict'], strict=True)
print('load_state_dict strict OK; missing', missing, 'unexpected', unexpected)
# sanity forward
x = torch.randn(2,4,128)
print('forward out shape', tuple(m(x).shape))
"

cd /app && rm -f logs/G_done.txt && setsid bash -c '
python3 train_exp.py --epochs 90 --steps_per_epoch 300 --lr 1e-3 --warmup 5 \
  --rollout_start 2 --rollout_max 8 --curr_epochs 40 --noise 0 --loss mse \
  --ema 0.999 --tag G_long --out /app/exp/G.pt > /app/logs/G.log 2>&1
echo DONE > /app/logs/G_done.txt
' < /dev/null > /dev/null 2>&1 &
echo "launched G (detached)"; sleep 25; echo "--- G start ---"; tail -3 logs/G.log 2>/dev/null

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

# ----------------------------------------------------------------------------
# Improved training recipe for KS-1D autoregressive forecasting.
#
# Key ideas (the baseline only did teacher-forced single-step MSE):
#   1. Exact KS symmetry augmentation: periodic spatial translation (roll) and
#      reflection u(x) -> -u(-x). These are exact symmetries of the PDE, so they
#      are "free" extra data and enforce translation equivariance.
#   2. Rollout curriculum: unroll the model autoregressively for R steps and
#      backprop through the rollout (R grows 2 -> ROLLOUT_MAX). This directly
#      optimises the quantity the grader scores and exposes the model to its own
#      predictions, stopping errors from compounding during the 20-step rollout.
#   3. MSE loss in normalised space (pushes the single-step error floor very low)
#      -- the rollout error here grows additively, so single-step precision wins.
#   4. EMA of weights (decay 0.999) -- averages out the noisy rollout-training
#      updates and was consistently the best model late in training.
#   5. Warmup + cosine LR, longer training.
#
# The best of {EMA, raw} is selected each epoch by the true 20-step rollout
# relative-L2 on the validation split (exactly the grader's metric).
# ----------------------------------------------------------------------------

DEFAULT_EPOCHS = 80
ROLLOUT_MAX = 6
ROLLOUT_START = 2
BASE_LR = 1e-3
EMA_DECAY = 0.999
STEPS_PER_EPOCH = 300
GRAD_CLIP = 1.0


def augment(frames: torch.Tensor) -> torch.Tensor:
    """Apply exact KS symmetries to a [B, F, X] block (same transform per frame)."""
    B, Fr, X = frames.shape
    # periodic translation
    shift = torch.randint(0, X, (B,), device=frames.device)
    base = torch.arange(X, device=frames.device).unsqueeze(0)
    idx = ((base - shift.unsqueeze(1)) % X).unsqueeze(1).expand(B, Fr, X)
    frames = torch.gather(frames, 2, idx)
    # reflection u(x) -> -u(-x)
    m = (torch.rand(B, 1, 1, device=frames.device) < 0.5).float()
    frames = m * (-torch.flip(frames, dims=[-1])) + (1 - m) * frames
    return frames


@torch.no_grad()
def rollout_rel_l2(model, traj, fm, fs, k, horizon, device):
    model.eval()
    traj = traj.to(device)
    window = ((traj[:, :k] - fm) / fs).clone()
    preds = []
    for _ in range(horizon):
        out = model(window)
        preds.append(out)
        window = torch.cat([window[:, 1:], out], dim=1)
    pred = torch.cat(preds, dim=1) * fs + fm
    true = traj[:, k:k + horizon]
    num = (pred - true).pow(2).sum(-1).sqrt()
    den = true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
    return (num / den).mean().item()


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"])
    horizon = int(cfg["rollout_horizon"])

    epochs = int(os.environ.get("TASK_EPOCHS_OVERRIDE", DEFAULT_EPOCHS))
    batch_size = int(cfg["training"]["batch_size"])
    warmup = max(1, round(epochs * 0.06))
    curr_epochs = max(1, round(epochs * 0.45))  # epoch at which rollout hits max
    rollout_max = max(1, min(ROLLOUT_MAX, horizon))

    train_ds = KSForecast(APP_DIR / "data" / "train", k=k)
    feat_mean, feat_std = compute_standardization(train_ds.trajectories)
    fm, fs = feat_mean, feat_std
    traj = train_ds.trajectories.to(device)
    N, T, X = traj.shape
    val_traj = torch.load(APP_DIR / "data" / "val" / "trajectories.pt", weights_only=True)["trajectories"].float()

    model = build
…[truncated]

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

cd /app && cp /app/exp/best_so_far.pt /app/outputs/model.pt && echo "restored E to outputs"; echo "=== G progress ==="; tail -4 logs/G.log; ls logs/G_done.txt 2>/dev/null && echo GDONE || echo "G running"

cd /app && sleep 260; echo "=== G progress ==="; tail -3 logs/G.log; ls logs/G_done.txt 2>/dev/null && echo GDONE || echo "G running"

cd /app && sleep 260; echo "=== G progress ==="; tail -3 logs/G.log; ls logs/G_done.txt 2>/dev/null && echo GDONE || echo "G running"

cd /app && sleep 260; echo "=== G tail ==="; tail -4 logs/G.log; echo "=== G BEST ==="; grep BEST logs/G.log; ls logs/G_done.txt 2>/dev/null && echo GDONE || echo "G running"

cd /app && for i in 1 2 3 4 5 6; do [ -f logs/G_done.txt ] && break; sleep 10; done; echo "=== G BEST ==="; grep BEST logs/G.log; ls logs/G_done.txt 2>/dev/null && echo GDONE || echo "still running"

cd /app && echo "=== G ===" && python3 rollout_eval.py /app/exp/G.pt 2>&1; echo; echo "=== E (prev best) ===" && python3 rollout_eval.py /app/exp/E.pt 2>&1 | grep rollout_rel_l2

cd /app && sed -i 's/^DEFAULT_EPOCHS = 80/DEFAULT_EPOCHS = 90/; s/^ROLLOUT_MAX = 6/ROLLOUT_MAX = 8/' train.py && grep -E "^DEFAULT_EPOCHS|^ROLLOUT_MAX|^ROLLOUT_START|^EMA_DECAY|^BASE_LR" train.py
# Promote G to final output
cp /app/exp/G.pt /app/outputs/model.pt && echo "promoted G -> outputs/model.pt"

cd /app && python3 -c "
# Confirm train.py derived schedule for default 90 epochs matches G
epochs=90
print('warmup', max(1, round(epochs*0.06)), '| curr_epochs', max(1, round(epochs*0.45)))
" && echo "--- FINAL CONTRACT CHECK ---" && 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)
assert isinstance(ck, dict) and {'state_dict','feat_mean','feat_std'} <= set(ck), 'missing keys'
assert isinstance(ck['feat_mean'], float) and isinstance(ck['feat_std'], float), 'stats must be float'
m = build_model(cfg); m.load_state_dict(ck['state_dict'], strict=True); m.eval()
x = torch.randn(3, cfg['in_channels'], cfg['grid_size'])
y = m(x); assert y.shape == (3, cfg['out_channels'], cfg['grid_size']), y.shape
n_params = sum(p.numel() for p in m.parameters())
print('OK: weights_only load, strict state_dict, forward', tuple(y.shape), '| params', n_params)
print('feat_mean=%.3e feat_std=%.6f' % (ck['feat_mean'], ck['feat_std']))
"

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
from rollout_eval import rollout_metrics
cfg = json.load(open('task_config.json'))
dev='cuda'
ck = torch.load('outputs/model.pt', weights_only=True)
m = build_model(cfg).to(dev); m.load_state_dict(ck['state_dict']); m.eval()
fm, fs = ck['feat_mean'], ck['feat_std']
k, H = cfg['in_channels'], cfg['rollout_horizon']
for split in ['val','public_test']:
    tr = torch.load(f'data/{split}/trajectories.pt', weights_only=True)['trajectories'].float()
    print(f'[{split}] T={tr.shape[1]}')
    for start in [0, 20, 40, 60, 76]:   # need start+k+H <= T(=100)
        sub = tr[:, start:start+k+H]
        met = rollout_metrics(m, sub, fm, fs, k=k, horizon=H, device=dev)
        print(f'   start={start:3d}: rollout_rel_l2={met[\"rollout_rel_l2\"]:.4f}  final={met[\"final_rel_l2\"]:.4f}')
"

cd /app && echo "=== FINAL outputs/model.pt rollout ===" && python3 rollout_eval.py /app/outputs/model.pt 2>&1 | grep -E "feat|rollout_rel_l2"
echo "=== outputs dir ===" && ls -la outputs/
echo "=== cleaning temp smoke files ===" && rm -f outputs/smoke.pt /app/exp/best_so_far.pt logs/*_done.txt 2>/dev/null; ls -la outputs/

cd /app && echo "=== lingering train procs? ===" && (ps aux | grep train_exp | grep -v grep || echo "none"); echo "=== train.py defaults ===" && sed -n '40,48p' train.py; echo "=== files ===" && ls outputs/ && echo "model.pt is G (best)"