Skip to content

fix(worldmodel): persist the EMA target encoder; pin the random baseline - #54

Merged
aurascoper merged 1 commit into
mainfrom
fix/persist-target-and-pin-baseline
Aug 2, 2026
Merged

aurascoper merged 1 commit into
mainfrom
fix/persist-target-and-pin-baseline

Conversation

@aurascoper

Copy link
Copy Markdown
Owner

First two items of the benchmark repair — plus a correction to node 14 that doing the work surfaced.

1. Checkpoint format v2 — persist target_encoder_state_dict

The documented failure mode was wrong, and the truth is worse. Node 14, frame_diagnostic.py and this file all said a restored model has target == encoder, so it reports zero mismatch. But __init__ deep-copies the target from the encoder at construction, before any load — so a v1 restore leaves the target at its random initialization.

Measured, after loading encoder_state_dict into a fresh module:

target compared against max param diff
that module's own random init 0.00e+00 — byte-identical
the loaded online encoder 2.04
the trained target it should have been 0.39

target == encoder holds only if you construct and never load, which is not what using a checkpoint means. A v1-restored model plans toward goal latents from an untrained network — not the one-frame remedy arm, and not any condition that was ever run. v2 restore verified exact (0.00e+00 against the original target).

This is not tidiness. The registered acceptance criterion says the repair must separate zero from random from MSE on existing data. Run against v1 checkpoints, it would have been scoring an untrained goal encoder — and it would have looked like it worked.

train.py has no equivalent defect: it saves model.state_dict(), which already covers every submodule. So export_coreml.py's GoalEncoder.mlpackage is unaffected — checked, not assumed.

2. Pin the random baseline — baseline_reps, default 100

The baseline touches no encoder, no predictor and no CEM, only _step with uniform actions, so precision costs env steps alone. It was estimated at one draw per episode and re-drawn per seed, injecting its own sampling noise into every arm comparison as if it were model variance.

Decomposing node 14's 0.15 → 0.45 swing — it is not episode-limited binomial noise alone:

range over 8 draws sd
same episodes, reps=1 (old) 0.2500 0.0750
same episodes, reps=100 (new) 0.0195 0.0054
same episodes, reps=500 0.0069 0.0022
different episodes, reps=100 0.2440 0.0853

Roughly half sampling noise, half genuine episode-set difficulty, since starts/goal_pos are redrawn from random.Random(seed). Reps removes the first half; decoupling the episode seed from the model seed is required for the second.

Kept episode-paired rather than moved to an independent set — a paired comparison needs P(random succeeds | THIS episode). Per-episode rates are returned so a later paired analysis can use them as a difficulty covariate, and random_baseline_stderr is returned so a reader can see the pin held.

The model arm is bit-identical

Verified against node 14's recorded ld=32 success_rate at seeds 0–4 — 0.30, 0.55, 0.20, 0.35, 0.35 — exact match on all five.

starts/goal_pos are drawn before the loop and the CEM generator is separate, so moving baseline estimation after the model loop cannot perturb it. But that is an argument; the check is the evidence.

3. Ledger correction, appended not retro-edited

A dated ## Corrections to node 14 block, matching the file's existing ## Review clarifications precedent. Node 14's measurements stand — forward_eval.py has no torch.load and trains in-process, so nothing it reported came from a restored checkpoint. Only the parenthetical account of the hazard was wrong.

Verification

smoke test        3 consecutive runs, exit 0
ledger gate       PASS (0 warnings)      ledger tests   OK
ADR gate          PASS                   ADR tests      OK
frame tests       OK (2 skipped under the shim, as designed)

Remaining in the ordering

Decouple episode seed from model seed → per-episode CEM generator → paired continuous primary → raise n to ~200. Not in this PR.

First two items of the benchmark repair, plus a correction to node 14 that the
work surfaced.

1. CHECKPOINT FORMAT v2 — persist target_encoder_state_dict.

v1 saved encoder_state_dict and predictor_state_dict only. The failure mode was
documented here, in frame_diagnostic.py and in ledger node 14 as "a restored
model has target == encoder, so it reports zero mismatch". That is WRONG, and
the truth is worse: __init__ deep-copies the target from the encoder AT
CONSTRUCTION, before any load, so a v1 restore leaves the target at its RANDOM
INITIALIZATION. Measured after loading encoder_state_dict into a fresh module:

  vs that module's own random init   0.00e+00   <- byte-identical
  vs the loaded online encoder       2.04
  vs the trained target              0.39

target == encoder holds only if you construct and never load, which is not what
using a checkpoint means. A v1-restored model plans toward goal latents from an
UNTRAINED network -- not the one-frame remedy arm, and not any condition that
was ever run. v2 restore verified exact (max parameter diff 0.00e+00).

This matters beyond tidiness: the registered acceptance criterion for the
benchmark repair says it must separate zero from random from MSE ON EXISTING
DATA. Run against v1 checkpoints that would have been scoring an untrained goal
encoder, and it would have looked like it worked.

train.py has no equivalent defect -- it saves model.state_dict(), which already
covers every submodule. export_coreml.py's GoalEncoder export is therefore also
unaffected; checked rather than assumed.

2. PIN THE RANDOM BASELINE — baseline_reps, default 100.

The baseline touches no encoder, no predictor and no CEM, only _step with
uniform actions, so precision costs env steps alone. It was estimated at one
draw per episode -- a scalar with binomial SE ~0.10 at n=20 -- and re-drawn per
seed, injecting its own sampling noise into every arm comparison as if it were
model variance.

Decomposing node 14's "0.15 to 0.45" swing, it is NOT episode-limited binomial
noise alone. Same episode set, varying only the draw: sd 0.0750 at reps=1,
0.0054 at reps=100, 0.0022 at reps=500. Different episode sets with reps pinned
at 100: sd 0.0853. So it is roughly half sampling noise and half genuine
episode-set difficulty. Reps removes the first half; decoupling the episode seed
from the model seed (next item) is required for the second.

Kept episode-paired rather than moved to an independent episode set: a paired
comparison needs P(random succeeds | THIS episode). Per-episode rates are
returned so a later paired analysis can use them as a difficulty covariate.
random_baseline_stderr is returned so a reader can see the pin held.

THE MODEL ARM IS BIT-IDENTICAL. Verified against node 14's recorded ld=32
success_rate at seeds 0-4 (0.30, 0.55, 0.20, 0.35, 0.35): exact match on all
five. starts/goal_pos are drawn before the loop and the CEM generator is
separate, so moving the baseline estimation after the model loop cannot
perturb it -- but that is an argument, and the check is the evidence.

3. Ledger gains a dated "Corrections to node 14" block rather than a retro-edit
of node 14 itself, matching the file's existing "Review clarifications"
precedent. Node 14's measurements stand: forward_eval.py has no torch.load and
trains in-process, so nothing it reported came from a restored checkpoint.

Verified: smoke test passes on three consecutive runs; ledger gate, ADR gate and
both test suites green.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DJJ8y1zo7WSqEmGNW8DQap
@aurascoper
aurascoper merged commit 0e0bd81 into main Aug 2, 2026
2 checks passed
@aurascoper
aurascoper deleted the fix/persist-target-and-pin-baseline branch August 2, 2026 04:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant