Skip to content

Latest commit

 

History

History
149 lines (119 loc) · 9.12 KB

td_outer_loop.md

File metadata and controls

149 lines (119 loc) · 9.12 KB

TD Outer Loop — Fitted Policy Iteration Driver

scripts/td_outer_loop.py runs many (prepare, fit) iterations in a row, refreshing V_snap from the latest passing checkpoint each iter. The single-pass mechanics (Stage 1–4) live in td_training.md; this doc owns orchestration, validation, retention, and failure handling.

Run layout

runs/td_loop_<tag>/
  manifest.json              # hyperparams (initial + history), per-iter records, latest_good, best_val
  latest_good.npz            # symlink → iterK/weights.npz of last promoted iter (V_snap source)
  best_val.npz               # symlink → iter with the best val_mc_nll
  iter0/
    V_snap.npz               # copy of --bootstrap (iter 0 only)
    branch_data.npz          # prepare output
    weights.npz              # fit output
    weights.training.json    # fit history + LR schedule + abort marker
    validation.json          # stage 3 + stage 4 check results
  iter1/ ...
  iter{K}_retry1/ ...        # only if --on-fail retry-then-abort triggered

Retention: branch_data.npz is large; the driver prunes iter{i-2}/branch_data.npz (and its retry siblings) at the start of iter i. The KL-drift comparison (S3.7) keeps one prior cache so the comparison window is exactly one iter. Weight and training-history files are small and always kept.

Bootstrap: iter 0 copies --bootstrap (default models/value_spline_v2.npz) into iter0/V_snap.npz and seeds latest_good. Subsequent iters use latest_good.npz as V_snap.

Resume: a re-invoked driver loads the existing manifest.json, refreshes hyperparams_current, appends a new entry to hyperparams_history if it changed, trims trailing non-promoted iters, and continues from the next index.

Driver shell

python scripts/td_outer_loop.py \
  --tag default \
  --iters 20 \
  --n 80000 \
  --k 0.3 --budget-s 0.2 --workers 24 \
  --partial-input-p 0.5 --mc-target-p 0.5 \
  --epochs 10 --batches-per-epoch 64 --batch-size 128 \
  --lambda-mc 0.3 --lambda-td 0.7 \
  --lr-peak 0.03 --lr-warmup-frac 0.10 --lr-hold-frac 0.40 --lr-floor-frac 0.10 \
  --on-fail abort \
  --val-blowup-tau 0.25 --val-regression-tau 0.05 --seg-uniform-tau 0.05

Per iter the driver shells out to:

  1. prepare_td_data.py --weights V_snap --out iter{i}/branch_data.npz --seed args.seed --iter-index i .... The prepare RNG uses seed + iter_index, so each iter draws an independent row sample without changing the driver's --seed. Retries (see below) bump the seed by +1000.
  2. Stage-3 validation on the cache (see checks). Hard failures abort the loop.
  3. fit_value_spline.py --resume V_snap --td-cache iter{i}/branch_data.npz --output iter{i}/weights.npz --lr <peak> --lr-warmup-frac ... --val-blowup-tau .... Fit writes weights.training.json and aborts mid-run with exit code 2 if val_mc_nll > val_mc_nll_baseline * (1 + tau).
  4. Stage-4 validation on the training history.
  5. Promote or revert. Promoted iter: latest_good.npz symlink moves to the new weights; best_val.npz updates if the iter's final val_mc_nll is the best so far. Reverted iter: latest_good stays; the loop's next behavior depends on --on-fail.

Per-iter LR schedule

Every iter does its own warmup → hold → decay around --lr-peak because each iter starts with a freshly refreshed V_snap, so targets shift relative to the last fit's end state. Defaults: linear warmup over the first 10% of optimizer steps from lr_peak * 0.1 to lr_peak, hold for 40%, cosine decay over the remaining 50% down to lr_peak * 0.1. The schedule is computed against total_steps = epochs * batches_per_epoch and applied per optimizer param group's base_lr. It is orthogonal to the existing --warmup-steps flag (the knot-requantize warmup; --resume keeps that at 0).

Stage-3 checks (after prepare)

Hard checks abort the loop with status = "hard_aborted". Soft checks log to validation.json and the manifest record; they do not fail the iter on their own.

# Check Severity
S3.1 I5: ((branch_{r,a}_target > 0) & ~alive_at_s_for_rows_{r,a}).sum() == 0 on both branches hard
S3.2 n_r + n_a == n hard
S3.3 branch_{r,a}_target is finite and per-row sum over alive seats ≈ 1.0 (atol 1e-5) hard
S3.4 Empirical Branch-A fraction within ±3σ of Binom(n, k) soft
S3.5 branch_a_engine_fallback.mean() (degenerate Branch-A proxy) soft
S3.6 Per-seat target mean per player-count segment (2P → 0.5; 4P → 0.25). Warn if any deviates from uniform by > --seg-uniform-tau soft
S3.7 Mean KL of branch_r_target_i vs branch_r_target_{i-1} on the frame-id × focal-seat intersection, with a 1e-6 floor on the prior. Flagged as "plateau" if near-zero for ≥3 consecutive iters. soft

Stage-4 checks (after fit)

# Check Severity
S4.1 No inf / nan / ≥ 1e9 in any *_nll history field hard
S4.2 Mid-run val blow-up (exit code 2 from fit, written to training.json["aborted"]) hard
S4.3 Final val_mc_nll ≤ min(best_val_so_far, this_iter's_first_epoch_val) * (1 + --val-regression-tau) soft → iter reverted
S4.4 First-quartile vs last-quartile mean of sample_td_{r,a}_nll — informational only, since TD NLLs aren't comparable when the target itself moves log

Knot-drift monitor (after fit, non-blocking)

Score knots are frozen on --resume, but the feature distribution (mainly c_prod) drifts as the policy/gamma evolve. The fit writes a knot_coverage block to training.json (knot_coverage_stats); _knot_coverage_verdict reads it and prints a [knot-drift] warning when frac_above_top_knot > --knot-above-top-tau (default 0.30) or frac_below_first_interior_knot > --knot-below-first-tau (default 0.70). The verdict is stored on the iter record (knot_coverage) but never aborts or reverts — a sustained warning is the signal to manually re-baseline the knots (re-derive + reproject + reset V_snap); see value_init.md for why auto-adjust is avoided and how to re-baseline safely.

Failure handling

S3.1 / S3.2 / S3.3 / S4.1 / S4.2     →  hard_aborted (loop halts; latest_good unchanged)
S4.3                                  →  reverted    (iter fails soft check)

--on-fail controls how the loop reacts to a soft failure:

  • abort (default): the loop stops after the reverted iter.
  • retry-then-abort: re-run the same iter once with lr_peak *= 0.5 and seed += 1000. If the retry also reverts, the loop aborts. The retry writes to iter{i}_retry1/ and its record's previous_attempt field carries the original.

A hard-aborted iter never updates latest_good; a reverted iter doesn't either. Per the project memory, val_mc_nll is a "did we break something" floor only — judge play strength by tournament (use scripts/value_net_tournament.py against the promoted checkpoints) and pick the iter to deploy based on that, not on the best val_mc_nll.

manifest.json shape

{
  "tag": "...",
  "bootstrap_sha256": "...",
  "hyperparams": { ... current CLI values ... },
  "hyperparams_initial": { ... },
  "hyperparams_current": { ... },
  "hyperparams_history": [{"resume_index": 0, "hyperparams": {...}}, ...],
  "iters": [
    {
      "i": 0,
      "attempt": 0,
      "iter_dir": "iter0",
      "weights": "iter0/weights.npz",
      "v_snap": "...",
      "v_snap_sha256": "...",
      "prepare": {"wall_s": 1320.4, "exit_code": 0, "n_r": 56000, "n_a": 24000, "engine_fallback_rate": 0.02},
      "stage3_validation": { ... S3.* records and config ... },
      "fit":     {"wall_s": 540.1, "exit_code": 0, "final_val_mc_nll": 0.812},
      "stage4_validation": { ... S4.* records ... },
      "knot_coverage": { "available": true, "drift_flag": false, "warnings": [], "coverage": { ... } },
      "status": "promoted",
      "promoted_to_latest_good": true,
      "weights_sha256": "..."
    }
  ],
  "latest_good": "iter0/weights.npz",
  "best_val": {"iter": 0, "val_mc_nll": 0.812, "path": "iter0/weights.npz"}
}

Inner-loop cost model

Stages 1 + 2 are pure functions of the replay set, so the outer loop touches only stages 3 + 4. Inside stage 3, Branch A is the only meaningful compute: roughly n_a * budget_s / workers + per-task overhead. Stage 4 is GPU-bound on the spline forward/backward.

Levers, in rough order of leverage:

  • k shifts rows R↔A; the lowest-cost setting is k=0 (pure replay-TD), at the cost of dropping the on-policy partial-state correction. k is a learning hyperparameter, not a speed knob — change deliberately.
  • budget-s linearly scales Branch A wall time. The deploy budget is 1.0 s; preparing at 0.2 is a quality choice that yields sharper Branch-A targets without affecting inference cost.
  • workers scales Branch A linearly until memory-bound. The scaling MHz ceiling on this box clamps the useful range to roughly the P-core count; past that, tasks land on weaker E-cores under a throttled clock.
  • n = epochs * batches_per_epoch * batch_size scales both prepare and fit linearly.
  • epochs per iter trades fit cost for V_snap refresh rate. Small epochs = fresher targets, more (cheaper) iters.