ks-equation-1d-forecast

cd /app && ls -la && echo "---- task_config.json ----" && cat task_config.json && echo "---- feature_schema.json ----" && 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, time
print('torch version:', torch.__version__)
print('cuda available:', torch.cuda.is_available())
if torch.cuda.is_available():
    print('device:', torch.cuda.get_device_name(0))
import os
print('CPU count:', os.cpu_count())
" 2>&1 | head -20

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, '== keys:', list(payload.keys()))
    t = payload['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()
# frame-to-frame relative L2
d = t[:,1:] - t[:,:-1]
rel = d.pow(2).sum(-1).sqrt() / t[:,1:].pow(2).sum(-1).sqrt()
print('mean frame-to-frame rel L2 (persistence 1-step):', rel.mean().item())
# persistence rollout: predict next=last over 20 steps from frame idx 3
def persistence_rollout(traj, start=3, horizon=20):
    # window last frame = traj[:,start]; predict constant
    errs=[]
    for h in range(1,horizon+1):
        pred = traj[:,start]  # constant
        true = traj[:,start+h]
        rl2 = (pred-true).pow(2).sum(-1).sqrt()/true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
        errs.append(rl2.mean().item())
    return errs
errs = persistence_rollout(t)
print('persistence rollout rel L2 per step (1,5,10,20):', errs[0], errs[4], errs[9], errs[19])
print('persistence mean over 20:', sum(errs)/len(errs))
# how correlated across time - decorrelation
import math
print('std per frame approx constant:', t.std(dim=(0,2))[:5])
"

command: create
path: /app/eval_rollout.py
file_text: """Local evaluation harness mimicking the hidden grader's autoregressive rollout."""
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


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


@torch.no_grad()
def rollout_metrics(model, traj, feat_mean, feat_std, k=4, horizon=20, start=0, device="cpu"):
    """traj: [N, T, X]. Start from frames [start:start+k], roll horizon steps."""
    model.eval()
    traj = traj.to(device)
    N, T, X = traj.shape
    window = traj[:, start:start + k].clone()  # [N,k,X] real units
    per_step_rel = []
    per_step_sq = []  # for rmse
    preds = []
    for h in range(1, horizon + 1):
        xin = (window - feat_mean) / feat_std
        out = model(xin)  # [N,1,X] normalized
        pred = out[:, 0] * feat_std + feat_mean  # [N,X]
        true = traj[:, start + k + h - 1]
        rl2 = (pred - true).pow(2).sum(-1).sqrt() / true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
        per_step_rel.append(rl2.mean().item())
        per_step_sq.append((pred - true).pow(2).mean().item())
        preds.append(pred)
        window = torch.cat([window[:, 1:], pred.unsqueeze(1)], dim=1)
    mean_rel = sum(per_step_rel) / len(per_step_rel)
    final_rel = per_step_rel[-1]
    rmse = (sum(per_step_sq) / len(per_step_sq)) ** 0.5
    return {"mean_rel_l2": mean_rel, "final_rel_l2": final_rel, "rmse": rmse,
            "per_step": per_step_rel}


def evaluate_all(ckpt_path="/app/outputs/model.pt", horizon=20, starts=(0,)):
    with (APP_DIR / "task_config.json").open() as fh:
        cfg = json.load(fh)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    k = int(cfg["in_channels"])
    model, fm, fs = load_model(ckpt_path, cfg, device)
    results = {}
    for split in ["val", "public_test"]:
        traj = torch.load(APP_DIR / "data" / split / "trajectories.pt", weights_only=True)["trajectories"].float()
        # average over multiple start windows for a robust estimate
        agg = {"mean_rel_l2": [], "final_rel_l2": [], "rmse": []}
        for s in starts:
            if s + k + horizon > traj.shape[1]:
                continue
            m = rollout_metrics(model, traj, fm, fs, k=k, horizon=horizon, start=s, device=device)
            for key in agg:
                agg[key].append(m[key])
        results[split] = {key: sum(v) / len(v) for key, v in agg.items()}
        # also store the start=0 per-step curve
        results[split]["per_step_start0"] = rollout_metrics(
            model, traj, fm, fs, k=k, horizon=horizon, start=0, device=device)["per_step"]
    return results


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="/app/outputs/model.pt")
    ap.add_argument("--multi", action="store_true", help="average over many start windows")
    args = ap.parse_args()
    starts = tuple(range(0, 76, 4)) if args.multi else (0,)
    res = evaluate_all(args.ckpt, starts=starts)
    for split, m in res.items():
        ps = m.pop("per_step_start0")
        print(f"[{split}] mean_rel_l2={m['mean_rel_l2']:.4f} final_rel_l2={m['final_rel_l2']:.4f} rmse={m['rmse']:.4f}")
        print(f"   per-step(start0) steps[1,5,10,15,20]={ps[0]:.3f},{ps[4]:.3f},{ps[9]:.3f},{ps[14]:.3f},{ps[19]:.3f}")

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

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

cd /app && python3 -c "
import torch
t = torch.load('data/train/trajectories.pt', weights_only=True)['trajectories'].float()
# check spectral content: energy per fourier mode to justify modes=16
uf = torch.fft.rfft(t, dim=-1).abs().mean(dim=(0,1))
import numpy as np
e = uf.numpy()
print('rfft bins:', len(e))
tot = e.sum()
c = np.cumsum(e)/tot
for m in [4,8,12,16,20,24,32,64]:
    print(f'  cum energy up to mode {m}: {c[m-1]:.4f}')
"

command: create
path: /app/exp.py
file_text: """Experimentation harness for KS rollout training recipes."""
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("/app")
sys.path.insert(0, str(APP_DIR))
from candidate_model import build_model

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
with (APP_DIR / "task_config.json").open() as fh:
    CFG = json.load(fh)
K = int(CFG["in_channels"])


def load_split(split):
    return torch.load(APP_DIR / "data" / split / "trajectories.pt", weights_only=True)["trajectories"].float()


TRAIN = load_split("train").to(DEVICE)
VAL = load_split("val").to(DEVICE)
FEAT_MEAN = float(TRAIN.mean().item())
FEAT_STD = float(TRAIN.std().clamp_min(1e-6).item())


def make_index(T, M):
    """valid t0 so that [t0, t0+k+M) fits in T."""
    return list(range(0, T - K - M + 1))


def rel_l2(pred, true):
    # pred,true: [B, X]
    return (pred - true).pow(2).sum(-1).sqrt() / true.pow(2).sum(-1).sqrt().clamp_min(1e-8)


@torch.no_grad()
def eval_rollout(model, traj, horizon=20, start=0):
    model.eval()
    N, T, X = traj.shape
    window = (traj[:, start:start + K] - FEAT_MEAN) / FEAT_STD
    errs = []
    for h in range(1, horizon + 1):
        out = model(window)
        pred = out[:, 0]
        true = (traj[:, start + K + h - 1] - FEAT_MEAN) / FEAT_STD
        errs.append(rel_l2(pred, true).mean().item())
        window = torch.cat([window[:, 1:], pred.unsqueeze(1)], dim=1)
    return sum(errs) / len(errs), errs


@torch.no_grad()
def eval_rollout_multi(model, traj, horizon=20, starts=range(0, 76, 4)):
    vals = []
    finals = []
    for s in starts:
        m, errs = eval_rollout(model, traj, horizon, s)
        vals.append(m); finals.append(errs[-1])
    return sum(vals) / len(vals), sum(finals) / len(finals)


class EMA:
    def __init__(self, model, decay):
        self.decay = decay
        self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()}

    def update(self, model):
        for k, v in model.state_dict().items():
            if v.dtype.is_floating_point:
                self.shadow[k].mul_(self.decay).add_(v.detach(), alpha=1 - self.decay)
            else:
                self.shadow[k].copy_(v)

    def copy_to(self, model):
        model.load_state_dict(self.shadow)


def train(cfg):
    torch.manual_seed(int(CFG["seed"]))
    model = build_model(CFG).to(DEVICE)
    epochs = cfg["epochs"]
    bs = cfg.get("batch_size", 64)
    lr = cfg.get("lr", 1e-3)
    wd = cfg.get("wd", 1e-4)
    noise = cfg.get("noise", 0.0)
    clip = cfg.get("clip", 1.0)
    warmup = cfg.get("warmup", 0.05)
    ema_decay = cfg.get("ema", 0.0)
    loss_type = cfg.get("loss", "rel_l2")
    pushforward = cfg.get("pushforward", False)
    get_M = cfg.get("rollout_schedule")  # function(epoch)->M

    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd)
    ema = EMA(model, ema_decay) if ema_decay > 0 else None

    # normalize trajectories once
    trn = (TRAIN - FEAT_MEAN) / FEAT_STD  # [N,T,X]
    N, T, X = trn.shape

    total_steps_est = 0
    for e in range(epochs):
        M = get_M(e) if get_M else cfg.get("M", 1)
        idx_list = make_index(T, M)
        total_steps_est += len(idx_list) * N // bs
    steps_done = 0

    best = (1e9, None)
    t0 = time.time()
    for epoch in range(epochs):
        model.train()
        M = get_M(epoch) if get_M else cfg.get("M", 1)
        idx_list = make_index(T, M)
        # build all (traj, t0) pairs
        pairs = torch.tensor([(i, t) for i in range(N) for t in idx_list], device=DEVICE)
        perm = torch.randperm(pairs.shape[0], device=DEVICE)
        pairs = pairs[perm]
        running = 0.0; nb = 0
        for s in range(0, pairs.shape[0], bs):
            pb = pairs[s:s + bs]
            ti = pb[:, 0]; t0i = pb[:, 1]
            # gather window [B,k,X]
            ar = torch.arange(K, device=DE
…[truncated]

cd /app && python3 -c "
import exp
print('--- A: baseline M=1 mse (repro) ---')
exp.train(dict(epochs=8, M=1, loss='mse', noise=0.0, ema=0.0, warmup=0.05))
print('--- B: M=1 rel_l2 ---')
exp.train(dict(epochs=8, M=1, loss='rel_l2', noise=0.0, ema=0.0, warmup=0.05))
" 2>&1 | tail -40

cd /app && python3 -c "
import exp
for noise in [0.01, 0.03]:
    print(f'--- noise={noise} M=1 rel_l2 ---')
    exp.train(dict(epochs=10, M=1, loss='rel_l2', noise=noise, ema=0.0, warmup=0.05))
print('--- M=4 rel_l2 full-BPTT ---')
exp.train(dict(epochs=10, M=4, loss='rel_l2', noise=0.0, ema=0.0, warmup=0.05))
print('--- M=4 rel_l2 pushforward ---')
exp.train(dict(epochs=10, M=4, loss='rel_l2', noise=0.0, ema=0.0, warmup=0.05, pushforward=True))
" 2>&1 | grep -E "(---|ep10/|ep9/)"

cd /app && python3 -c "
import exp
print('--- F: M=4 pushforward + noise=0.01 ---')
exp.train(dict(epochs=10, M=4, loss='rel_l2', noise=0.01, pushforward=True, warmup=0.05))
print('--- G: M=6 pushforward ---')
exp.train(dict(epochs=10, M=6, loss='rel_l2', noise=0.0, pushforward=True, warmup=0.05))
print('--- H: M=8 pushforward ---')
exp.train(dict(epochs=10, M=8, loss='rel_l2', noise=0.0, pushforward=True, warmup=0.05))
print('--- I: M=4 pushforward + ema=0.999 ---')
exp.train(dict(epochs=10, M=4, loss='rel_l2', noise=0.0, pushforward=True, ema=0.999, warmup=0.05))
" 2>&1 | grep -E "(---|ep10/|ep9/)"

restart: true

cd /app && mkdir -p logs && cat > run_exp.py << 'EOF'
import exp, sys, json
name = sys.argv[1]
cfg = json.loads(sys.argv[2])
print(f'=== {name} cfg={cfg} ===', flush=True)
best = exp.train(cfg)
print(f'=== {name} BEST val_roll={best[0]:.5f} ===', flush=True)
EOF
echo "created run_exp.py"

cd /app && nohup python3 run_exp.py F '{"epochs":10,"M":4,"loss":"rel_l2","noise":0.01,"pushforward":true,"warmup":0.05}' > logs/F.log 2>&1 &
echo "started F, pid $!"

cd /app && cat > chain.sh << 'EOF'
#!/bin/bash
cd /app
python3 run_exp.py G '{"epochs":10,"M":6,"loss":"rel_l2","noise":0.0,"pushforward":true,"warmup":0.05}' > logs/G.log 2>&1
python3 run_exp.py H '{"epochs":12,"M":8,"loss":"rel_l2","noise":0.0,"pushforward":true,"warmup":0.05}' > logs/H.log 2>&1
python3 run_exp.py I '{"epochs":10,"M":4,"loss":"rel_l2","noise":0.0,"pushforward":true,"ema":0.999,"warmup":0.05}' > logs/I.log 2>&1
echo "chain done"
EOF
chmod +x chain.sh
# wait for F to finish first before starting chain to avoid GPU contention
( while kill -0 238 2>/dev/null; do sleep 5; done; ./chain.sh ) > logs/chain_driver.log 2>&1 &
echo "chain driver started, pid $!"
sleep 60 && echo "--- F.log so far ---" && cat logs/F.log

cd /app && sleep 30 && tail -3 logs/F.log

cd /app && sleep 240 && echo "=== G ===" && tail -2 logs/G.log && echo "=== H ===" && tail -2 logs/H.log && echo "=== I ===" && tail -2 logs/I.log

cd /app && sleep 45 && echo "=== I final ===" && tail -3 logs/I.log

cd /app && pkill -f run_exp.py; pkill -f chain.sh; sleep 3
# Add a curriculum helper to exp via a small wrapper script
cat > run_curric.py << 'EOF'
import exp, sys, json
name = sys.argv[1]
cfg = json.loads(sys.argv[2])
sched = cfg.pop("Msched", None)
if sched:
    def get_M(e, s=sched):
        return s[min(e, len(s)-1)]
    cfg["rollout_schedule"] = get_M
print(f'=== {name} cfg keys={list(cfg.keys())} sched={sched} ===', flush=True)
best = exp.train(cfg)
print(f'=== {name} BEST val_roll={best[0]:.5f} ===', flush=True)
EOF
echo "ok"

cd /app && cat > chain2.sh << 'EOF'
#!/bin/bash
cd /app
python3 run_exp.py J '{"epochs":24,"M":4,"loss":"rel_l2","noise":0.005,"pushforward":true,"warmup":0.05}' > logs/J.log 2>&1
python3 run_curric.py K '{"epochs":18,"loss":"rel_l2","noise":0.005,"pushforward":true,"warmup":0.05,"Msched":[2,2,2,4,4,4,4,6,6,6,8,8]}' > logs/K.log 2>&1
echo "chain2 done" > logs/chain2_done.txt
EOF
chmod +x chain2.sh
nohup ./chain2.sh > logs/chain2_driver.log 2>&1 &
echo "started chain2, pid $!"
sleep 120 && echo "=== J so far ===" && tail -6 logs/J.log

cd /app && sleep 90 && echo "=== J so far ===" && tail -5 logs/J.log

cd /app && sleep 150 && echo "=== K result ===" && tail -5 logs/K.log && ls logs/chain2_done.txt 2>/dev/null && echo "chain2 finished"

cd /app && sleep 20 && echo "=== K final ===" && tail -2 logs/K.log

cd /app && cat > chain3.sh << 'EOF'
#!/bin/bash
cd /app
python3 run_exp.py L '{"epochs":40,"M":4,"loss":"rel_l2","noise":0.005,"pushforward":true,"warmup":0.04}' > logs/L.log 2>&1
python3 run_curric.py N '{"epochs":40,"loss":"rel_l2","noise":0.005,"pushforward":true,"warmup":0.04,"Msched":[2,2,3,3,4,4,4,4,4,4,4,4,5,5,5,6,6,6]}' > logs/N.log 2>&1
echo "done" > logs/chain3_done.txt
EOF
chmod +x chain3.sh
nohup ./chain3.sh > logs/chain3_driver.log 2>&1 &
echo "started chain3 pid $!"
sleep 5 && echo "launched"

cd /app && sleep 180 && echo "=== L progress ===" && tail -4 logs/L.log

cd /app && sleep 120 && echo "=== L final ===" && tail -4 logs/L.log

cd /app && sleep 200 && echo "=== N progress ===" && tail -4 logs/N.log

cd /app && sleep 130 && echo "=== N final ===" && tail -3 logs/N.log && ls logs/chain3_done.txt 2>/dev/null && echo DONE

command: create
path: /app/train.py
file_text: """Training recipe for KS-1D autoregressive forecasting.

Key ideas (vs. the weak single-step/MSE baseline), all targeting the fact that
the grader rolls the model out autoregressively for 20 steps on a chaotic
system where one-step errors compound:

  * Relative-L2 loss that matches the evaluation metric (scale-invariant per
    sample, so every window contributes equally).
  * Multi-step "pushforward" rollout training (Brandstetter et al. 2022): each
    sample is unrolled M=4 steps feeding the model its OWN predictions, with
    intermediate states detached. This exposes the network to its own error
    distribution at train time, curing the train/test distribution shift that
    makes naive single-step models drift.
  * Small input-noise augmentation for extra robustness to accumulated error.
  * Warmup + cosine LR, gradient clipping, long training, and checkpoint
    selection by a multi-start 20-step rollout metric on the val split.

The architecture (candidate_model.build_model) and checkpoint contract are
unchanged: we save {state_dict, feat_mean, feat_std}.
"""
from __future__ import annotations
import json
import math
import os
import sys
import time
import copy
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

# ----------------------------- recipe knobs --------------------------------
ROLLOUT_M = 4          # unrolled steps per training sample (pushforward)
NOISE_STD = 0.005      # stddev of input noise (in normalized units)
GRAD_CLIP = 1.0
WARMUP_FRAC = 0.04
DEFAULT_EPOCHS = 60    # strong model in a few minutes on a modern GPU
EVAL_HORIZON = 20      # matches grader rollout_horizon
EVAL_STARTS = tuple(range(0, 76, 4))  # multi-start robust val metric


def rel_l2(pred: torch.Tensor, true: torch.Tensor) -> torch.Tensor:
    """Per-sample relative L2 over the spatial axis. pred/true: [B, X]."""
    num = (pred - true).pow(2).sum(-1).sqrt()
    den = true.pow(2).sum(-1).sqrt().clamp_min(1e-8)
    return num / den


@torch.no_grad()
def rollout_rel_l2(model, traj_norm, k, horizon, starts):
    """Autoregressive rollout rel-L2 averaged over steps and start windows.

    traj_norm: [N, T, X] already standardized. Mirrors the hidden grader's
    procedure (feed predictions back in, compare each step to ground truth)."""
    model.eval()
    T = traj_norm.shape[1]
    mean_vals, final_vals = [], []
    for s in starts:
        if s + k + horizon > T:
            continue
        window = traj_norm[:, s:s + k].clone()
        errs = []
        for h in range(1, horizon + 1):
            pred = model(window)[:, 0]
            true = traj_norm[:, s + k + h - 1]
            errs.append(rel_l2(pred, true).mean().item())
            window = torch.cat([window[:, 1:], pred.unsqueeze(1)], dim=1)
        mean_vals.append(sum(errs) / len(errs))
        final_vals.append(errs[-1])
    return sum(mean_vals) / len(mean_vals), sum(final_vals) / len(final_vals)


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"])

    train_ds = KSForecast(APP_DIR / "data" / "train", k=k)
    val_ds = KSForecast(APP_DIR / "data" / "val", k=k)
    feat_mean, feat_std = compute_standardization(train_ds.trajectories)

    # Standardize whole trajectories once and keep on device.
    trn = ((train_ds.trajectories - feat_mean) / feat_std).to(device)   # [N,T,X]
    val = ((val_ds.trajectories - feat_mean) / feat_std).to(device)
    N, T, X = trn.shape

    model = build_model(cfg).to(device)
    epochs = int(os.environ.get("TASK_EPOCHS_OVERRIDE", DEFAULT_EPOCHS))
    batch_size = int(cfg["train
…[truncated]

cd /app && TASK_EPOCHS_OVERRIDE=2 timeout 120 python3 train.py 2>&1 | tail -6

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
ckpt = torch.load('outputs/model.pt', weights_only=True)
print('keys:', list(ckpt.keys()))
print('feat_mean:', ckpt['feat_mean'], 'feat_std:', ckpt['feat_std'])
print('types:', type(ckpt['feat_mean']).__name__, type(ckpt['feat_std']).__name__)
cfg = json.load(open('task_config.json'))
m = build_model(cfg)
missing, unexpected = m.load_state_dict(ckpt['state_dict'], strict=True)
print('load_state_dict strict OK; missing=', missing, 'unexpected=', unexpected)
print('num params:', sum(p.numel() for p in m.parameters()))
"

cd /app && nohup python3 train.py > logs/final60.log 2>&1 &
echo "started final training pid $!"
sleep 150 && echo "--- progress ---" && tail -6 logs/final60.log

cd /app && sleep 240 && echo "--- final epochs ---" && tail -8 logs/final60.log

cd /app && python3 eval_rollout.py --multi 2>&1 | tail -6

cd /app && cp outputs/model.pt outputs/model_ep60.pt && echo "backed up ep60 model"
nohup env TASK_EPOCHS_OVERRIDE=100 python3 train.py > logs/final100.log 2>&1 &
echo "started 100-epoch run pid $!"

cd /app && python3 -c "
import torch, json
from candidate_model import build_model
cfg = json.load(open('task_config.json'))
ckpt = torch.load('outputs/model_ep60.pt', weights_only=True)
# confirm _grid not in state_dict (non-persistent buffer)
print('_grid in state_dict:', any('_grid' in kk for kk in ckpt['state_dict']))
m = build_model(cfg); m.load_state_dict(ckpt['state_dict']); m.eval()
x = torch.randn(3,4,128)
with torch.no_grad():
    y1 = m(x)
# reload fresh and compare
m2 = build_model(cfg); m2.load_state_dict(ckpt['state_dict']); m2.eval()
with torch.no_grad():
    y2 = m2(x)
print('fresh-reload identical:', torch.allclose(y1,y2), 'out shape:', tuple(y1.shape))
" 2>&1 | tail

cd /app && sleep 280 && tail -4 logs/final100.log

cd /app && sleep 280 && tail -5 logs/final100.log

cd /app && sleep 110 && tail -4 logs/final100.log

cd /app && echo "===== EP60 model =====" && python3 eval_rollout.py --ckpt outputs/model_ep60.pt --multi 2>&1 | tail -4
echo "===== EP100 model =====" && python3 eval_rollout.py --ckpt outputs/model.pt --multi 2>&1 | tail -4

cd /app && cp outputs/model_ep60.pt outputs/model.pt && echo "restored EP60 as final model.pt"
python3 -c "
import torch
a=torch.load('outputs/model.pt',weights_only=True); b=torch.load('outputs/model_ep60.pt',weights_only=True)
same=all(torch.equal(a['state_dict'][k],b['state_dict'][k]) for k in a['state_dict'])
print('final model.pt == ep60:', same, '| feat_mean=%.3e feat_std=%.4f'%(a['feat_mean'],a['feat_std']))
"

cd /app && python3 -c "
import json, torch
from eval_rollout import load_model, rollout_metrics
cfg = json.load(open('task_config.json'))
dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model, fm, fs = load_model('outputs/model.pt', cfg, dev)
print('FINAL MODEL (outputs/model.pt)  feat_mean=%.2e feat_std=%.4f'%(fm,fs))
for split in ['val','public_test']:
    traj = torch.load(f'data/{split}/trajectories.pt', weights_only=True)['trajectories'].float()
    # start=0 (most likely grader protocol)
    m0 = rollout_metrics(model, traj, fm, fs, horizon=20, start=0, device=dev)
    # robust average over all valid starts
    starts = range(0, traj.shape[1]-4-20+1)
    ms = [rollout_metrics(model, traj, fm, fs, horizon=20, start=s, device=dev) for s in starts]
    mean_all = sum(x['mean_rel_l2'] for x in ms)/len(ms)
    fin_all  = sum(x['final_rel_l2'] for x in ms)/len(ms)
    print(f'[{split}] start0: mean={m0[\"mean_rel_l2\"]:.4f} final={m0[\"final_rel_l2\"]:.4f} rmse={m0[\"rmse\"]:.4f}'
          f'  | all-starts({len(ms)}): mean={mean_all:.4f} final={fin_all:.4f}')
"

cd /app && rm -f outputs/model_ep60.pt outputs/model_ep100.pt && ls -la outputs/
echo "----- GRADER SIMULATION (fresh process) -----"
python3 - << 'EOF'
import json, torch
from candidate_model import build_model
cfg = json.load(open('/app/task_config.json'))
# 1) load exactly per the contract
ckpt = torch.load('/app/outputs/model.pt', weights_only=True)
assert set(['state_dict','feat_mean','feat_std']).issubset(ckpt), "missing keys"
assert isinstance(ckpt['feat_mean'], float) and isinstance(ckpt['feat_std'], float)
model = build_model(cfg); model.load_state_dict(ckpt['state_dict'], strict=True); model.eval()
fm, fs = ckpt['feat_mean'], ckpt['feat_std']
k, H = int(cfg['in_channels']), int(cfg['rollout_horizon'])
# 2) autoregressive rollout exactly like the grader, input shape [B,4,128]
traj = torch.load('/app/data/public_test/trajectories.pt', weights_only=True)['trajectories'].float()
win = ((traj[:, :k] - fm)/fs)                       # [B,4,128] normalized
assert tuple(win.shape) == (traj.shape[0], 4, 128)
errs=[]
with torch.no_grad():
    for h in range(1, H+1):
        out = model(win)                            # [B,1,128] normalized
        assert out.shape[1]==1 and out.shape[2]==128
        pred = out[:,0]*fs + fm                     # de-normalize to real units
        true = traj[:, k+h-1]
        errs.append(((pred-true).pow(2).sum(-1).sqrt()/true.pow(2).sum(-1).sqrt()).mean().item())
        win = torch.cat([win[:,1:], ((pred-fm)/fs).unsqueeze(1)], dim=1)
print('rollout OK | 20-step mean rel_l2 = %.4f | final = %.4f'%(sum(errs)/len(errs), errs[-1]))
print('CONTRACT: keys/dtypes/shapes all valid; strict load OK')
EOF

cd /app && pkill -f run_exp.py 2>/dev/null; pkill -f train.py 2>/dev/null; pkill -f chain 2>/dev/null; sleep 2
echo "running python procs:"; pgrep -af python3 || echo "  (none)"
echo "--- confirming the saved model.pt came from a default 'python3 train.py' run ---"
head -1 logs/final60.log
grep -c "epoch=" logs/final60.log | xargs echo "epochs logged:"
grep "wrote" logs/final60.log