Buckets:

AsianPlayer
/

TriDrive

Files

xet

AsianPlayer/TriDrive / code /aide /train /run_kv2_program.sh

AsianPlayer

17 days ago

download

raw

14.4 kB

#!/usr/bin/env bash
# K-v2 programme (docs/decisions.md 2026-09-14 "K-v2 programme — S0 pre-registration", aide/docs/protocol.md C23) — detached lane runner.
#
#   setsid nohup bash aide/train/run_kv2_program.sh <A|B|C|D> > aide/results/kv2/lane_<X>.log 2>&1 &
#
# Lane A  S1 val screen: 9 K-v2 configs at seed 42 (cwm3d_v2[_<cfg>]_s42_v2) -> `kv2_gate.py s1` (needs aide/models/cwm3d_v2.py + its tests green; exit 4 if not)
# Lane B  S3 (i): 9 MotionBERT-kp grid points at seed 42 -> `kv2_gate.py s3` -> MB-best at seeds 123/456/789/1011 if != default -> MB-best + trunk-only init
#         at 5 seeds (existing motionbert_kp_batonpre_pretrain_s123 checkpoint; a d/L change needs a new pretraining = follow-up, exit 5)
# Lane C  motionbert_kp_batonpre_trunkp123_s{789,1011}_v2 and cwm3d_kp_batonpre_trunkp123_s{789,1011}_v2 (recipes / inits of the existing s42/123/456 runs verbatim)
# Lane D  S1b (seed 123) per S1 survivor -> `kv2_gate.py s1b` -> per continuing survivor: architecture-matched trunk-only pretraining
#         (cwm3d_v2[_nospatial]_batonpre_pretrain_s123, survivor flags, --zeroshot 0; verified by pretrain-check) -> seeds 456/789/1011 and +pre at 5 seeds
# Every python launch passes the project GPU gate first (re-checked every 30 s): free VRAM >= need + 4 GB, MemAvailable >= 28 GB, fewer than 4 trainer
# processes (`pgrep -fc "python.*(train_cwm3|train_kp|pretrain)"`), OMP/MKL threads 4.  One trainer per lane.  Existing runs are never touched: a run whose
# run json exists is skipped; new run dirs only.  PIDs -> aide/results/kv2/pids.txt.  Selection is val-only inside train_kp / the pretraining script; the
# official test split is evaluated once per run by the existing harness.  Config flags / run names come from aide/eval/kv2_gate.py (single source of truth).
set -u
LANE=${1:?lane A|B|C|D}
ROOT=/home/USER/cwm3
# Run from a per-launch snapshot so that editing this file while a lane is alive can never change the running lane (bash reads scripts
# incrementally; lane B of 2026-09-14 hit `flags: unbound variable` after its last run when the lane-D block was edited underneath it).
if [ -z "${KV2_SNAPSHOT:-}" ]; then
  mkdir -p "$ROOT/aide/results/kv2/launcher_snapshots"
  SNAP="$ROOT/aide/results/kv2/launcher_snapshots/run_kv2_program_${LANE}_$(date -u +%Y%m%dT%H%M%SZ).sh"
  cp "${BASH_SOURCE[0]}" "$SNAP" && KV2_SNAPSHOT=$SNAP exec bash "$SNAP" "$@"
fi
PY=python
RUNS=$ROOT/aide/results/runs
KV2=$ROOT/aide/results/kv2
PIDS=$KV2/pids.txt
GATE=aide/eval/kv2_gate.py
GATE_OUT=${GATE_OUT:-$KV2}          # decision json dir (tests / dry runs point it elsewhere)
DRY_RUN=${DRY_RUN:-0}               # 1: print every assembled command instead of launching
K_CKPT=aide/results/runs/cwm3d_kp_batonpre_pretrain_s123/ckpt_best.pt
MB_CKPT=aide/results/runs/motionbert_kp_batonpre_pretrain_s123/ckpt_best.pt
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 NUMEXPR_NUM_THREADS=4
mkdir -p "$KV2"
cd "$ROOT" || exit 1
# Stage-A recipe verbatim (aide/train/run_arch_attribution.sh:24 without --loss euc = K's loss defaults euc + 0.1 NLL + 0.05 bone + 0.01 jerk);
# tensor-output models (MotionBERT-kp) add --loss euc exactly as motionbert_kp_s*_v2.
RECIPE_K="--lr 3e-4 --wd 0.05 --warmup 10 --epochs 200 --bs 32 --ema 0.999 --clip 1.0 --patience 40 --aug-window 1 --noise-px 2.0 --joint-drop 0.1 --translate 20 --scale 0.05 --flip-p 0"
RECIPE="$RECIPE_K --loss euc"
COMMON="--no-lock --max-minutes 45"
KV2_CFGS="full nospatial khead nohaz hazaux plain q060 q090 lam03"
MB_CFGS="d192 L4 lr1e-4 lr1e-3 wd0.01 wd0.1 do0.1 d192_L4 d192_L4_do0.1"
SEEDS5="42 123 456 789 1011"

log() { echo "$(date -u +%FT%TZ) [lane $LANE] $*"; }

gate() {   # gate <need_gb>: block until the project launch gate is satisfied
  local need=$1 free_mib avail_kb ntr ok
  while :; do
    ok=1
    free_mib=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | head -1 | tr -d ' ')
    avail_kb=$(awk '/MemAvailable/{print $2}' /proc/meminfo)
    ntr=$(pgrep -fc "python.*(train_cwm3|train_kp|pretrain)")
    [ "${free_mib:-0}" -lt $(( (need + 4) * 1024 )) ] && ok=0
    [ "${avail_kb:-0}" -lt $(( 28 * 1024 * 1024 )) ] && ok=0
    [ "${ntr:-9}" -ge 4 ] && ok=0
    if [ "$ok" = 1 ]; then
      log "gate OK: free VRAM ${free_mib} MiB (need $((need + 4)) GB), MemAvailable $((avail_kb / 1048576)) GB (>= 28), trainers ${ntr} (< 4)"
      return 0
    fi
    log "gate WAIT: free VRAM ${free_mib} MiB (need $((need + 4)) GB), MemAvailable $((avail_kb / 1048576)) GB, trainers ${ntr}"
    sleep 30
  done
}

run() {    # run <name> <need_gb> <cmd...>: gate, launch, record pid, wait
  local name=$1 need=$2; shift 2
  if ls "$RUNS/$name"/run_"$name"_2*.json >/dev/null 2>&1; then log "skip $name (run json exists; existing runs are never re-run)"; return 0; fi
  if [ "$DRY_RUN" = 1 ]; then log "DRY $name (need ${need} GB): $*"; return 0; fi
  gate "$need"
  "$@" > "$RUNS/${name}_train.log" 2>&1 &
  local pid=$!
  echo "$(date -u +%FT%TZ) lane=$LANE name=$name pid=$pid START" >> "$PIDS"
  log "launched $name pid $pid"
  wait "$pid"; local rc=$?
  echo "$(date -u +%FT%TZ) lane=$LANE name=$name pid=$pid END rc=$rc" >> "$PIDS"
  log "finished $name rc=$rc"
  return $rc
}

refs_for() {   # <seed>: evaluation-only paired references (simlpe exists for 42/123/456 only; 789/1011 use the s42 run as lane C of arch_attribution did)
  local s=$1 simlpe=simlpe_kp_s42_v2
  case $s in 42|123|456) simlpe=simlpe_kp_s${s}_v2 ;; esac
  echo "zvshrink=zv_shrink_val,simlpe=$simlpe,kp=cwm3d_kp_s${s}_v2,mb=motionbert_kp_s${s}_v2"
}
read_dec() { "$PY" $GATE read "$1" "$2" --out "$GATE_OUT"; }                      # <json-stem> <key>

v2_ready() {   # lanes A / D: model file + its tests + trainer flags must exist (exit 4 otherwise)
  [ -f aide/models/cwm3d_v2.py ] || { log "NOT READY: aide/models/cwm3d_v2.py missing"; return 1; }
  [ -f aide/tests/test_cwm3d_v2.py ] || { log "NOT READY: aide/tests/test_cwm3d_v2.py missing"; return 1; }
  "$PY" -m pytest -q aide/tests/test_cwm3d_v2.py > "$KV2/pytest_cwm3d_v2_lane${LANE}.log" 2>&1 || { log "NOT READY: pytest aide/tests/test_cwm3d_v2.py failed (see $KV2/pytest_cwm3d_v2_lane${LANE}.log)"; return 1; }
  local help; help=$("$PY" aide/train/train_kp.py --help 2>&1)
  local fl; for fl in --no-v2-spatial --no-lean-head --no-hazard --haz-mod --haz-quantile --haz-weight --no-physics --w-nll; do
    echo "$help" | grep -q -- "$fl" || { log "NOT READY: train_kp.py lacks flag $fl"; return 1; }
  done
  log "v2 ready: model file, tests and trainer flags present"
  return 0
}

run_kv2() {   # <cfg> <seed> <pre:0|1> <stage-tag>
  local cfg=$1 s=$2 pre=$3 stage=$4 name flags refs init=""
  if [ "$pre" = 1 ]; then name=$("$PY" $GATE name kv2 "$cfg" "$s" pre); else name=$("$PY" $GATE name kv2 "$cfg" "$s"); fi
  flags=$("$PY" $GATE flags kv2 "$cfg")
  refs=$(refs_for "$s")
  if [ "$pre" = 1 ]; then
    [ -n "${V2_CKPT:-}" ] || { log "ABORT: V2_CKPT unset for $name"; return 1; }
    init="--init-ckpt $V2_CKPT --init-parts trunk"
    refs="$refs,seedinit=$("$PY" $GATE name kv2 "$cfg" "$s")"
  fi
  # shellcheck disable=SC2086
  run "$name" 3 "$PY" aide/train/train_kp.py $flags --seed "$s" --name "$name" $RECIPE_K $init --refs "$refs" $COMMON \
      --notes "kv2_program;${stage};cfg=${cfg}${init:+;init=$(basename "$(dirname "$V2_CKPT")"):trunk}"
}

run_mb() {    # <cfg> <seed> <pre:0|1> <stage-tag>
  local cfg=$1 s=$2 pre=$3 stage=$4 name flags refs init=""
  if [ "$pre" = 1 ]; then name=$("$PY" $GATE name mb "$cfg" "$s" pre); else name=$("$PY" $GATE name mb "$cfg" "$s"); fi
  flags=$("$PY" $GATE flags mb "$cfg")
  refs=$(refs_for "$s")
  if [ "$pre" = 1 ]; then
    [ -n "${MB_PRE_CKPT:-}" ] || { log "ABORT: MB_PRE_CKPT unset for $name"; return 1; }
    init="--init-ckpt $MB_PRE_CKPT --init-parts trunk"
    refs="$refs,seedinit=$("$PY" $GATE name mb "$cfg" "$s")"
  fi
  # shellcheck disable=SC2086
  run "$name" 3 "$PY" aide/train/train_kp.py $flags --seed "$s" --name "$name" $RECIPE $init --refs "$refs" $COMMON \
      --notes "kv2_program;${stage};cfg=${cfg}${init:+;init=$(basename "$(dirname "$MB_PRE_CKPT")"):trunk}"
}

case "$LANE" in
  A)
    v2_ready || exit 4
    for cfg in $KV2_CFGS; do run_kv2 "$cfg" 42 0 "S1" || log "WARN: $cfg seed 42 failed (rc $?), continuing the grid"; done
    "$PY" $GATE s1 --out "$GATE_OUT"; rc=$?
    log "S1 gate rc=$rc survivors='$(read_dec s1_decision survivors)' complete=$(read_dec s1_decision complete)"
    [ "$rc" = 0 ] || { log "S1 incomplete — inspect pending runs, then re-run lane A (finished runs are skipped)"; exit 3; }
    ;;
  B)
    for cfg in $MB_CFGS; do run_mb "$cfg" 42 0 "S3i" || log "WARN: MB $cfg seed 42 failed (rc $?), continuing the grid"; done
    "$PY" $GATE s3 --out "$GATE_OUT"; rc=$?
    best=$(read_dec s3_decision mb_best)
    log "S3 gate rc=$rc mb_best='$best' changes_dL=$(read_dec s3_decision mb_best_changes_dL) complete=$(read_dec s3_decision complete)"
    [ "$rc" = 0 ] && [ -n "$best" ] || { log "S3 incomplete — re-run lane B after the pending grid runs finish"; exit 3; }
    if [ "$best" = default ]; then
      log "MB-best = default: 5-seed and +pre arms are the existing runs + lane C (motionbert_kp_batonpre_trunkp123_s{789,1011}_v2); nothing more to run"
    else
      for s in 123 456 789 1011; do run_mb "$best" "$s" 0 "S3i:mb_best_5seed" || log "WARN: MB-best $best seed $s failed"; done
      if [ "$(read_dec s3_decision mb_best_changes_dL)" = 1 ]; then
        pre=$(read_dec s3_decision mb_best_pretrain_run); model=$(read_dec s3_decision mb_best_model)
        log "MB-best changes d/L -> new pretraining $pre (needs pretrain_cwm3d_baton.py --model $model; FOLLOW-UP if the script rejects the name)"
        run "$pre" 10 "$PY" aide/train/pretrain_cwm3d_baton.py --model "$model" --seed 123 --zeroshot 0 ; rc=$?
        if [ "$rc" != 0 ]; then log "FOLLOW-UP: pretraining $pre failed rc=$rc (pretrain_cwm3d_baton.py must accept --model $model); S3 (ii) for MB-best not run"; exit 5; fi
        [ -f "$RUNS/$pre/ckpt_best.pt" ] && [ -f "$RUNS/$pre/TRAIN_DONE" ] || { log "ABORT: $pre has no ckpt_best.pt / TRAIN_DONE"; exit 5; }
        MB_PRE_CKPT="aide/results/runs/$pre/ckpt_best.pt"
      else
        MB_PRE_CKPT=$MB_CKPT
      fi
      export MB_PRE_CKPT
      for s in $SEEDS5; do run_mb "$best" "$s" 1 "S3ii:mb_best_pre_5seed" || log "WARN: MB-best+pre $best seed $s failed"; done
    fi
    ;;
  C)
    export MB_PRE_CKPT=$MB_CKPT
    for s in 789 1011; do run_mb default "$s" 1 "S3ii:mb_pre_5seed_completion" || log "WARN: MB+pre seed $s failed"; done
    for s in 789 1011; do
      name="cwm3d_kp_batonpre_trunkp123_s${s}_v2"
      run "$name" 3 "$PY" aide/train/train_kp.py --model cwm3d --seed "$s" --name "$name" $RECIPE_K \
          --init-ckpt "$K_CKPT" --init-parts trunk --refs "$(refs_for "$s"),seedinit=cwm3d_kp_s${s}_v2" $COMMON \
          --notes "kv2_program;S2:k_pre_5seed_completion;cfg=k;init=cwm3d_kp_batonpre_pretrain_s123:trunk" || log "WARN: K+pre seed $s failed"
    done
    ;;
  D)
    v2_ready || exit 4
    [ "$(read_dec s1_decision complete)" = 1 ] || { log "NOT READY: s1_decision.json missing or incomplete (run lane A first)"; exit 4; }
    survivors=$(read_dec s1_decision survivors)
    [ -n "$survivors" ] && [ "$survivors" != "None" ] || { log "S1: no survivors — programme ends at S1 (report via report_kv2_program.py)"; exit 0; }
    log "S1 survivors: $survivors"
    # (1) S1b FIRST (coordinator 2026-09-14): seed 123 for every survivor, then the S1b gate
    for cfg in $survivors; do run_kv2 "$cfg" 123 0 "S1b" || log "WARN: $cfg seed 123 failed"; done
    cont=""
    for cfg in $survivors; do
      "$PY" $GATE s1b "$cfg" --out "$GATE_OUT"; rc=$?
      dec=$(read_dec "s1b_${cfg}_decision" decision); log "S1b $cfg rc=$rc -> $dec: $(read_dec "s1b_${cfg}_decision" reason)"
      [ "$dec" = CONTINUE ] && cont="$cont $cfg"
    done
    [ -n "$cont" ] || { log "S1b: no survivor CONTINUEs (decisions above: KILLED = val-killed at seed 123, PENDING = seed-123 run missing) — S2 not run"; exit 0; }
    # (2) per continuing survivor: trunk-only pretraining PER ARCHITECTURE with the survivor's flags (nospatial -> cwm3d_v2_nospatial_batonpre_pretrain_s123,
    #     full trunk -> cwm3d_v2_batonpre_pretrain_s123), --zeroshot 0; K's checkpoint only if it loads the whole trunk; the own checkpoint must load
    #     every trunk tensor (0 shape-skipped / 0 missing) before any +pre fine-tune; then seeds 456/789/1011 (no pre) and 5 seeds +pre
    for cfg in $cont; do
      pre=$("$PY" $GATE name kv2pre "$cfg")
      "$PY" $GATE pretrain-check "$cfg" --out "$GATE_OUT" || { log "pretrain-check not ready for $cfg"; exit 4; }
      if [ "$(read_dec "pretrain_check_${cfg}" decision)" = REUSE_K_CKPT ]; then
        V2_CKPT=$K_CKPT; log "$cfg: K's checkpoint loads the whole trunk -> reused"
      else
        flags=$("$PY" $GATE flags kv2 "$cfg")
        # shellcheck disable=SC2086
        run "$pre" 10 "$PY" aide/train/pretrain_cwm3d_baton.py $flags --seed 123 --name "$pre" --zeroshot 0 ; rc=$?
        [ "$rc" = 0 ] || { log "FOLLOW-UP: pretraining $pre failed rc=$rc"; exit 5; }
        [ -f "$RUNS/$pre/ckpt_best.pt" ] && [ -f "$RUNS/$pre/TRAIN_DONE" ] || { log "ABORT: $pre has no ckpt_best.pt / TRAIN_DONE"; exit 5; }
        V2_CKPT="aide/results/runs/$pre/ckpt_best.pt"
        "$PY" $GATE pretrain-check "$cfg" --ckpt "$V2_CKPT" --stem "pretrain_check_${cfg}_own" --out "$GATE_OUT" || { log "pretrain-check ($V2_CKPT) not ready"; exit 4; }
        [ "$(read_dec "pretrain_check_${cfg}_own" trunk_complete)" = 1 ] || { log "ABORT: $V2_CKPT does not load every trunk tensor of $cfg (shape-skipped $(read_dec "pretrain_check_${cfg}_own" n_skipped_shape), missing $(read_dec "pretrain_check_${cfg}_own" n_missing_trunk))"; exit 5; }
        log "$cfg: own pretraining $pre verified (loaded $(read_dec "pretrain_check_${cfg}_own" n_loaded) / $(read_dec "pretrain_check_${cfg}_own" n_trunk_tensors) trunk tensors, 0 skipped, 0 missing)"
      fi
      export V2_CKPT; log "$cfg: trunk-only init checkpoint $V2_CKPT"
      for s in 456 789 1011; do run_kv2 "$cfg" "$s" 0 "S2" || log "WARN: $cfg seed $s failed"; done
      for s in $SEEDS5; do run_kv2 "$cfg" "$s" 1 "S2:pre" || log "WARN: $cfg+pre seed $s failed"; done
    done
    ;;
  *) echo "unknown lane $LANE"; exit 2 ;;
esac
log "lane $LANE done"
[ "$DRY_RUN" = 1 ] || echo "$(date -u +%FT%TZ) lane=$LANE LANE_DONE" >> "$PIDS"

Xet Storage Details

Size:: 14.4 kB
Xet hash:: e1e6fb481961c734918459ffa95caacf17aa4cd554bb001115c08e27ed2ba06f

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#!/usr/bin/env bash
	# K-v2 programme (docs/decisions.md 2026-09-14 "K-v2 programme — S0 pre-registration", aide/docs/protocol.md C23) — detached lane runner.
	#
	# setsid nohup bash aide/train/run_kv2_program.sh <A\|B\|C\|D> > aide/results/kv2/lane_<X>.log 2>&1 &
	#
	# Lane A S1 val screen: 9 K-v2 configs at seed 42 (cwm3d_v2[_<cfg>]_s42_v2) -> `kv2_gate.py s1` (needs aide/models/cwm3d_v2.py + its tests green; exit 4 if not)
	# Lane B S3 (i): 9 MotionBERT-kp grid points at seed 42 -> `kv2_gate.py s3` -> MB-best at seeds 123/456/789/1011 if != default -> MB-best + trunk-only init
	# at 5 seeds (existing motionbert_kp_batonpre_pretrain_s123 checkpoint; a d/L change needs a new pretraining = follow-up, exit 5)
	# Lane C motionbert_kp_batonpre_trunkp123_s{789,1011}_v2 and cwm3d_kp_batonpre_trunkp123_s{789,1011}_v2 (recipes / inits of the existing s42/123/456 runs verbatim)
	# Lane D S1b (seed 123) per S1 survivor -> `kv2_gate.py s1b` -> per continuing survivor: architecture-matched trunk-only pretraining
	# (cwm3d_v2[_nospatial]_batonpre_pretrain_s123, survivor flags, --zeroshot 0; verified by pretrain-check) -> seeds 456/789/1011 and +pre at 5 seeds
	# Every python launch passes the project GPU gate first (re-checked every 30 s): free VRAM >= need + 4 GB, MemAvailable >= 28 GB, fewer than 4 trainer
	# processes (`pgrep -fc "python.*(train_cwm3\|train_kp\|pretrain)"`), OMP/MKL threads 4. One trainer per lane. Existing runs are never touched: a run whose
	# run json exists is skipped; new run dirs only. PIDs -> aide/results/kv2/pids.txt. Selection is val-only inside train_kp / the pretraining script; the
	# official test split is evaluated once per run by the existing harness. Config flags / run names come from aide/eval/kv2_gate.py (single source of truth).
	set -u
	LANE=${1:?lane A\|B\|C\|D}
	ROOT=/home/USER/cwm3
	# Run from a per-launch snapshot so that editing this file while a lane is alive can never change the running lane (bash reads scripts
	# incrementally; lane B of 2026-09-14 hit `flags: unbound variable` after its last run when the lane-D block was edited underneath it).
	if [ -z "${KV2_SNAPSHOT:-}" ]; then
	mkdir -p "$ROOT/aide/results/kv2/launcher_snapshots"
	SNAP="$ROOT/aide/results/kv2/launcher_snapshots/run_kv2_program_${LANE}_$(date -u +%Y%m%dT%H%M%SZ).sh"
	cp "${BASH_SOURCE[0]}" "$SNAP" && KV2_SNAPSHOT=$SNAP exec bash "$SNAP" "$@"
	fi
	PY=python
	RUNS=$ROOT/aide/results/runs
	KV2=$ROOT/aide/results/kv2
	PIDS=$KV2/pids.txt
	GATE=aide/eval/kv2_gate.py
	GATE_OUT=${GATE_OUT:-$KV2} # decision json dir (tests / dry runs point it elsewhere)
	DRY_RUN=${DRY_RUN:-0} # 1: print every assembled command instead of launching
	K_CKPT=aide/results/runs/cwm3d_kp_batonpre_pretrain_s123/ckpt_best.pt
	MB_CKPT=aide/results/runs/motionbert_kp_batonpre_pretrain_s123/ckpt_best.pt
	export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 NUMEXPR_NUM_THREADS=4
	mkdir -p "$KV2"
	cd "$ROOT" \|\| exit 1
	# Stage-A recipe verbatim (aide/train/run_arch_attribution.sh:24 without --loss euc = K's loss defaults euc + 0.1 NLL + 0.05 bone + 0.01 jerk);
	# tensor-output models (MotionBERT-kp) add --loss euc exactly as motionbert_kp_s*_v2.
	RECIPE_K="--lr 3e-4 --wd 0.05 --warmup 10 --epochs 200 --bs 32 --ema 0.999 --clip 1.0 --patience 40 --aug-window 1 --noise-px 2.0 --joint-drop 0.1 --translate 20 --scale 0.05 --flip-p 0"
	RECIPE="$RECIPE_K --loss euc"
	COMMON="--no-lock --max-minutes 45"
	KV2_CFGS="full nospatial khead nohaz hazaux plain q060 q090 lam03"
	MB_CFGS="d192 L4 lr1e-4 lr1e-3 wd0.01 wd0.1 do0.1 d192_L4 d192_L4_do0.1"
	SEEDS5="42 123 456 789 1011"

	log() { echo "$(date -u +%FT%TZ) [lane $LANE] $*"; }

	gate() { # gate <need_gb>: block until the project launch gate is satisfied
	local need=$1 free_mib avail_kb ntr ok
	while :; do
	ok=1
	free_mib=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits \| head -1 \| tr -d ' ')
	avail_kb=$(awk '/MemAvailable/{print $2}' /proc/meminfo)
	ntr=$(pgrep -fc "python.*(train_cwm3\|train_kp\|pretrain)")
	[ "${free_mib:-0}" -lt $(( (need + 4) * 1024 )) ] && ok=0
	[ "${avail_kb:-0}" -lt $(( 28 * 1024 * 1024 )) ] && ok=0
	[ "${ntr:-9}" -ge 4 ] && ok=0
	if [ "$ok" = 1 ]; then
	log "gate OK: free VRAM ${free_mib} MiB (need $((need + 4)) GB), MemAvailable $((avail_kb / 1048576)) GB (>= 28), trainers ${ntr} (< 4)"
	return 0
	fi
	log "gate WAIT: free VRAM ${free_mib} MiB (need $((need + 4)) GB), MemAvailable $((avail_kb / 1048576)) GB, trainers ${ntr}"
	sleep 30
	done
	}

	run() { # run <name> <need_gb> <cmd...>: gate, launch, record pid, wait
	local name=$1 need=$2; shift 2
	if ls "$RUNS/$name"/run_"$name"_2*.json >/dev/null 2>&1; then log "skip $name (run json exists; existing runs are never re-run)"; return 0; fi
	if [ "$DRY_RUN" = 1 ]; then log "DRY $name (need ${need} GB): $*"; return 0; fi
	gate "$need"
	"$@" > "$RUNS/${name}_train.log" 2>&1 &
	local pid=$!
	echo "$(date -u +%FT%TZ) lane=$LANE name=$name pid=$pid START" >> "$PIDS"
	log "launched $name pid $pid"
	wait "$pid"; local rc=$?
	echo "$(date -u +%FT%TZ) lane=$LANE name=$name pid=$pid END rc=$rc" >> "$PIDS"
	log "finished $name rc=$rc"
	return $rc
	}

	refs_for() { # <seed>: evaluation-only paired references (simlpe exists for 42/123/456 only; 789/1011 use the s42 run as lane C of arch_attribution did)
	local s=$1 simlpe=simlpe_kp_s42_v2
	case $s in 42\|123\|456) simlpe=simlpe_kp_s${s}_v2 ;; esac
	echo "zvshrink=zv_shrink_val,simlpe=$simlpe,kp=cwm3d_kp_s${s}_v2,mb=motionbert_kp_s${s}_v2"
	}
	read_dec() { "$PY" $GATE read "$1" "$2" --out "$GATE_OUT"; } # <json-stem> <key>

	v2_ready() { # lanes A / D: model file + its tests + trainer flags must exist (exit 4 otherwise)
	[ -f aide/models/cwm3d_v2.py ] \|\| { log "NOT READY: aide/models/cwm3d_v2.py missing"; return 1; }
	[ -f aide/tests/test_cwm3d_v2.py ] \|\| { log "NOT READY: aide/tests/test_cwm3d_v2.py missing"; return 1; }
	"$PY" -m pytest -q aide/tests/test_cwm3d_v2.py > "$KV2/pytest_cwm3d_v2_lane${LANE}.log" 2>&1 \|\| { log "NOT READY: pytest aide/tests/test_cwm3d_v2.py failed (see $KV2/pytest_cwm3d_v2_lane${LANE}.log)"; return 1; }
	local help; help=$("$PY" aide/train/train_kp.py --help 2>&1)
	local fl; for fl in --no-v2-spatial --no-lean-head --no-hazard --haz-mod --haz-quantile --haz-weight --no-physics --w-nll; do
	echo "$help" \| grep -q -- "$fl" \|\| { log "NOT READY: train_kp.py lacks flag $fl"; return 1; }
	done
	log "v2 ready: model file, tests and trainer flags present"
	return 0
	}

	run_kv2() { # <cfg> <seed> <pre:0\|1> <stage-tag>
	local cfg=$1 s=$2 pre=$3 stage=$4 name flags refs init=""
	if [ "$pre" = 1 ]; then name=$("$PY" $GATE name kv2 "$cfg" "$s" pre); else name=$("$PY" $GATE name kv2 "$cfg" "$s"); fi
	flags=$("$PY" $GATE flags kv2 "$cfg")
	refs=$(refs_for "$s")
	if [ "$pre" = 1 ]; then
	[ -n "${V2_CKPT:-}" ] \|\| { log "ABORT: V2_CKPT unset for $name"; return 1; }
	init="--init-ckpt $V2_CKPT --init-parts trunk"
	refs="$refs,seedinit=$("$PY" $GATE name kv2 "$cfg" "$s")"
	fi
	# shellcheck disable=SC2086
	run "$name" 3 "$PY" aide/train/train_kp.py $flags --seed "$s" --name "$name" $RECIPE_K $init --refs "$refs" $COMMON \
	--notes "kv2_program;${stage};cfg=${cfg}${init:+;init=$(basename "$(dirname "$V2_CKPT")"):trunk}"
	}

	run_mb() { # <cfg> <seed> <pre:0\|1> <stage-tag>
	local cfg=$1 s=$2 pre=$3 stage=$4 name flags refs init=""
	if [ "$pre" = 1 ]; then name=$("$PY" $GATE name mb "$cfg" "$s" pre); else name=$("$PY" $GATE name mb "$cfg" "$s"); fi
	flags=$("$PY" $GATE flags mb "$cfg")
	refs=$(refs_for "$s")
	if [ "$pre" = 1 ]; then
	[ -n "${MB_PRE_CKPT:-}" ] \|\| { log "ABORT: MB_PRE_CKPT unset for $name"; return 1; }
	init="--init-ckpt $MB_PRE_CKPT --init-parts trunk"
	refs="$refs,seedinit=$("$PY" $GATE name mb "$cfg" "$s")"
	fi
	# shellcheck disable=SC2086
	run "$name" 3 "$PY" aide/train/train_kp.py $flags --seed "$s" --name "$name" $RECIPE $init --refs "$refs" $COMMON \
	--notes "kv2_program;${stage};cfg=${cfg}${init:+;init=$(basename "$(dirname "$MB_PRE_CKPT")"):trunk}"
	}

	case "$LANE" in
	A)
	v2_ready \|\| exit 4
	for cfg in $KV2_CFGS; do run_kv2 "$cfg" 42 0 "S1" \|\| log "WARN: $cfg seed 42 failed (rc $?), continuing the grid"; done
	"$PY" $GATE s1 --out "$GATE_OUT"; rc=$?
	log "S1 gate rc=$rc survivors='$(read_dec s1_decision survivors)' complete=$(read_dec s1_decision complete)"
	[ "$rc" = 0 ] \|\| { log "S1 incomplete — inspect pending runs, then re-run lane A (finished runs are skipped)"; exit 3; }
	;;
	B)
	for cfg in $MB_CFGS; do run_mb "$cfg" 42 0 "S3i" \|\| log "WARN: MB $cfg seed 42 failed (rc $?), continuing the grid"; done
	"$PY" $GATE s3 --out "$GATE_OUT"; rc=$?
	best=$(read_dec s3_decision mb_best)
	log "S3 gate rc=$rc mb_best='$best' changes_dL=$(read_dec s3_decision mb_best_changes_dL) complete=$(read_dec s3_decision complete)"
	[ "$rc" = 0 ] && [ -n "$best" ] \|\| { log "S3 incomplete — re-run lane B after the pending grid runs finish"; exit 3; }
	if [ "$best" = default ]; then
	log "MB-best = default: 5-seed and +pre arms are the existing runs + lane C (motionbert_kp_batonpre_trunkp123_s{789,1011}_v2); nothing more to run"
	else
	for s in 123 456 789 1011; do run_mb "$best" "$s" 0 "S3i:mb_best_5seed" \|\| log "WARN: MB-best $best seed $s failed"; done
	if [ "$(read_dec s3_decision mb_best_changes_dL)" = 1 ]; then
	pre=$(read_dec s3_decision mb_best_pretrain_run); model=$(read_dec s3_decision mb_best_model)
	log "MB-best changes d/L -> new pretraining $pre (needs pretrain_cwm3d_baton.py --model $model; FOLLOW-UP if the script rejects the name)"
	run "$pre" 10 "$PY" aide/train/pretrain_cwm3d_baton.py --model "$model" --seed 123 --zeroshot 0 ; rc=$?
	if [ "$rc" != 0 ]; then log "FOLLOW-UP: pretraining $pre failed rc=$rc (pretrain_cwm3d_baton.py must accept --model $model); S3 (ii) for MB-best not run"; exit 5; fi
	[ -f "$RUNS/$pre/ckpt_best.pt" ] && [ -f "$RUNS/$pre/TRAIN_DONE" ] \|\| { log "ABORT: $pre has no ckpt_best.pt / TRAIN_DONE"; exit 5; }
	MB_PRE_CKPT="aide/results/runs/$pre/ckpt_best.pt"
	else
	MB_PRE_CKPT=$MB_CKPT
	fi
	export MB_PRE_CKPT
	for s in $SEEDS5; do run_mb "$best" "$s" 1 "S3ii:mb_best_pre_5seed" \|\| log "WARN: MB-best+pre $best seed $s failed"; done
	fi
	;;
	C)
	export MB_PRE_CKPT=$MB_CKPT
	for s in 789 1011; do run_mb default "$s" 1 "S3ii:mb_pre_5seed_completion" \|\| log "WARN: MB+pre seed $s failed"; done
	for s in 789 1011; do
	name="cwm3d_kp_batonpre_trunkp123_s${s}_v2"
	run "$name" 3 "$PY" aide/train/train_kp.py --model cwm3d --seed "$s" --name "$name" $RECIPE_K \
	--init-ckpt "$K_CKPT" --init-parts trunk --refs "$(refs_for "$s"),seedinit=cwm3d_kp_s${s}_v2" $COMMON \
	--notes "kv2_program;S2:k_pre_5seed_completion;cfg=k;init=cwm3d_kp_batonpre_pretrain_s123:trunk" \|\| log "WARN: K+pre seed $s failed"
	done
	;;
	D)
	v2_ready \|\| exit 4
	[ "$(read_dec s1_decision complete)" = 1 ] \|\| { log "NOT READY: s1_decision.json missing or incomplete (run lane A first)"; exit 4; }
	survivors=$(read_dec s1_decision survivors)
	[ -n "$survivors" ] && [ "$survivors" != "None" ] \|\| { log "S1: no survivors — programme ends at S1 (report via report_kv2_program.py)"; exit 0; }
	log "S1 survivors: $survivors"
	# (1) S1b FIRST (coordinator 2026-09-14): seed 123 for every survivor, then the S1b gate
	for cfg in $survivors; do run_kv2 "$cfg" 123 0 "S1b" \|\| log "WARN: $cfg seed 123 failed"; done
	cont=""
	for cfg in $survivors; do
	"$PY" $GATE s1b "$cfg" --out "$GATE_OUT"; rc=$?
	dec=$(read_dec "s1b_${cfg}_decision" decision); log "S1b $cfg rc=$rc -> $dec: $(read_dec "s1b_${cfg}_decision" reason)"
	[ "$dec" = CONTINUE ] && cont="$cont $cfg"
	done
	[ -n "$cont" ] \|\| { log "S1b: no survivor CONTINUEs (decisions above: KILLED = val-killed at seed 123, PENDING = seed-123 run missing) — S2 not run"; exit 0; }
	# (2) per continuing survivor: trunk-only pretraining PER ARCHITECTURE with the survivor's flags (nospatial -> cwm3d_v2_nospatial_batonpre_pretrain_s123,
	# full trunk -> cwm3d_v2_batonpre_pretrain_s123), --zeroshot 0; K's checkpoint only if it loads the whole trunk; the own checkpoint must load
	# every trunk tensor (0 shape-skipped / 0 missing) before any +pre fine-tune; then seeds 456/789/1011 (no pre) and 5 seeds +pre
	for cfg in $cont; do
	pre=$("$PY" $GATE name kv2pre "$cfg")
	"$PY" $GATE pretrain-check "$cfg" --out "$GATE_OUT" \|\| { log "pretrain-check not ready for $cfg"; exit 4; }
	if [ "$(read_dec "pretrain_check_${cfg}" decision)" = REUSE_K_CKPT ]; then
	V2_CKPT=$K_CKPT; log "$cfg: K's checkpoint loads the whole trunk -> reused"
	else
	flags=$("$PY" $GATE flags kv2 "$cfg")
	# shellcheck disable=SC2086
	run "$pre" 10 "$PY" aide/train/pretrain_cwm3d_baton.py $flags --seed 123 --name "$pre" --zeroshot 0 ; rc=$?
	[ "$rc" = 0 ] \|\| { log "FOLLOW-UP: pretraining $pre failed rc=$rc"; exit 5; }
	[ -f "$RUNS/$pre/ckpt_best.pt" ] && [ -f "$RUNS/$pre/TRAIN_DONE" ] \|\| { log "ABORT: $pre has no ckpt_best.pt / TRAIN_DONE"; exit 5; }
	V2_CKPT="aide/results/runs/$pre/ckpt_best.pt"
	"$PY" $GATE pretrain-check "$cfg" --ckpt "$V2_CKPT" --stem "pretrain_check_${cfg}_own" --out "$GATE_OUT" \|\| { log "pretrain-check ($V2_CKPT) not ready"; exit 4; }
	[ "$(read_dec "pretrain_check_${cfg}_own" trunk_complete)" = 1 ] \|\| { log "ABORT: $V2_CKPT does not load every trunk tensor of $cfg (shape-skipped $(read_dec "pretrain_check_${cfg}_own" n_skipped_shape), missing $(read_dec "pretrain_check_${cfg}_own" n_missing_trunk))"; exit 5; }
	log "$cfg: own pretraining $pre verified (loaded $(read_dec "pretrain_check_${cfg}_own" n_loaded) / $(read_dec "pretrain_check_${cfg}_own" n_trunk_tensors) trunk tensors, 0 skipped, 0 missing)"
	fi
	export V2_CKPT; log "$cfg: trunk-only init checkpoint $V2_CKPT"
	for s in 456 789 1011; do run_kv2 "$cfg" "$s" 0 "S2" \|\| log "WARN: $cfg seed $s failed"; done
	for s in $SEEDS5; do run_kv2 "$cfg" "$s" 1 "S2:pre" \|\| log "WARN: $cfg+pre seed $s failed"; done
	done
	;;
	*) echo "unknown lane $LANE"; exit 2 ;;
	esac
	log "lane $LANE done"
	[ "$DRY_RUN" = 1 ] \|\| echo "$(date -u +%FT%TZ) lane=$LANE LANE_DONE" >> "$PIDS"