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

WorldModel/EXPERIMENT_LEDGER.md

@@ -265,3 +265,10 @@ Two reviewers verified the delivered synthetic A/Bs are sound (symlog threaded i
   - **The benchmark cannot resolve this.** `random_baseline_success` ranges **0.15 to 0.45 across five seeds** at `mpc_episodes=20`, while every arm sits at ~0.35. A chance baseline swinging +/-0.15 cannot resolve a 0.07 effect. This is measurement, not preference.
 - Decision: **REJECT frame mismatch as the explanation for the control null, and KEEP the diagnostic.** The mismatch is real and large, so `_encode_states`' two-encoder split is a genuine defect in the planner's geometry — but removing it entirely leaves control at chance, so it is not the binding constraint. Also **DROP the isotropy/gauge line as a route to control**: the disagreement is a translation, and no isotropic regularizer addresses translations. This is the fourth time (nodes 9, 11, 12, 13) a representation-side change has been measured and found not to reach control, and the first where the mechanism was measured precisely enough to be excluded rather than merely not-supported.
 - Next question: NOT another objective. **The benchmark's resolving power is now the blocker, and it is measured rather than asserted** — a random baseline varying 0.15–0.45 at n=20 episodes cannot adjudicate any regularizer. Repair that first, specified blind (using only the existing MSE baseline and the zero/random policies, never a candidate arm) and required to separate zero from random from MSE on existing data before adoption. A scale-only vs SIGReg vs VISReg comparison run on this benchmark as it stands would be unfalsifiable in exactly the way this node was almost allowed to be.
+
+## Corrections to node 14 (2026-08-02)
+
+- **Correction A (factual, does not change node 14's result).** Node 14's Implementation section says a checkpoint-loaded model "reports zero mismatch: a false negative reading as a clean result", on the grounds that `target_encoder == encoder`. That is wrong, and the truth is worse. `EEGJEPAModule.__init__` deep-copies the target from the encoder **at construction, before any load**, so a format-v1 restore leaves the target at its RANDOM INITIALIZATION. Measured: after loading `encoder_state_dict` into a fresh module, the target is byte-identical to that module's own random init (0.00e+00), **2.04** away from the loaded online encoder and **0.39** from the trained target it should have been. `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 produced by an **untrained** network — not the one-frame remedy arm, and not any condition that was ever run.
+- **Node 14's measurements stand.** `forward_eval.py` contains no `torch.load`; it trains in-process via `train_jepa`, so nothing it reported came from a restored checkpoint. Only the parenthetical account of the hazard was wrong.
+- **Consequence for the benchmark repair.** The registered acceptance criterion — that the repair must separate zero from random from MSE *on existing data* — cannot be run against v1 checkpoints at all. It would be scoring an untrained goal encoder. `eeg_jepa.py` now writes format **v2** with `target_encoder_state_dict`; a v2 restore is exact (verified: max parameter difference 0.00e+00 against the original target). Acceptance runs must use v2 checkpoints or run in-process.
+- **Correction B (measurement, refines node 14's "0.15 to 0.45" reading).** That swing was attributed to the episode-limited binomial noise of estimating the baseline at one draw per episode. Decomposed, it is roughly **half** that and half genuine episode-set difficulty, because `starts` and `goal_pos` are redrawn from `random.Random(seed)` so each seed scores a different episode set. Holding the episode set fixed and varying only the random draw: sd **0.0750** at reps=1, **0.0054** at reps=100, **0.0022** at reps=500. Varying the episode set with reps pinned at 100: sd **0.0853**, essentially the original spread. So raising reps removes one half for the price of env steps — the baseline touches no encoder, predictor or CEM — and decoupling the episode seed from the model seed is required for the other half.

WorldModel/eeg_jepa.py

@@ -592,7 +592,21 @@ def main() -> None:
     args.output.parent.mkdir(parents=True, exist_ok=True)
     torch.save(
         {
-            "format": "neuralcompose.eeg-jepa.v1",
+            # v2 adds target_encoder_state_dict. v1 omitted it, and __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, it is byte-identical to the fresh init
+            # (0.00e+00), 2.04 away from the loaded online encoder and 0.39 from
+            # the trained target it should have been.
+            #
+            # That is worse than the "target == online" this was previously
+            # documented as, here and in frame_diagnostic.py and ledger node 14.
+            # target == online 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 a lossy
+            # restore, and not the one-frame arm either, but a third thing that
+            # was never an experimental condition at all.
+            "format": "neuralcompose.eeg-jepa.v2",
             "config": asdict(config),
             "state_feature_names": dataset.state_feature_names,
             "action_dim": dataset.action_dim,
@@ -604,6 +618,12 @@ def main() -> None:
             "normalization_log_epsilon": dataset.log_epsilon,
             "encoder_state_dict": model.encoder.state_dict(),
             "predictor_state_dict": model.predictor.state_dict(),
+            # The EMA target is a distinct set of weights, not a derivable one:
+            # update_target_ema mixes at tau=0.99 every step, so it lags the
+            # online encoder by a trajectory that cannot be recovered from the
+            # final online weights. (train.py has no equivalent defect — it
+            # saves model.state_dict(), which already covers every submodule.)
+            "target_encoder_state_dict": model.target_encoder.state_dict(),
             "final_loss": history[-1],
         },
         args.output,

WorldModel/forward_eval.py

@@ -45,6 +45,7 @@
 
 import argparse
 import json
+import math
 import sys
 import tempfile
 from pathlib import Path
@@ -319,7 +320,8 @@ def _cem_plan(model, z0_lat, zg_lat, device, *, horizon, cem_iters, n_samples, e
 @torch.no_grad()
 def _mpc_success(model, train_mean, train_std, device, *, n_episodes=20, horizon=6,
                  cem_iters=3, n_samples=64, elite_frac=0.2, mode="signal", seed=2,
-                 goal_offset=0.4, goal_tol=0.15, log_features: bool = False,
+                 goal_offset=0.4, goal_tol=0.15, baseline_reps=100,
+                 log_features: bool = False,
                  log_epsilon: float = 1e-6, symlog: bool = False) -> dict[str, Any]:
     """Goal-conditioned latent MPC/CEM planning success. Per episode: sample a start
     state + a goal `pos`, plan with CEM (the JEPA predictor as forward model), then
@@ -348,7 +350,8 @@ def _enc(states, use_target: bool = False):
 
     frame = frame_diagnostic(_enc, starts, goals)
 
-    successes, rand_successes, dists = 0, 0, []
+    successes, dists = 0, []
+    rand_per_episode: list[float] = []
     for i in range(n_episodes):
         plan = _cem_plan(model, z0[i], zg[i], device, horizon=horizon, cem_iters=cem_iters,
                          n_samples=n_samples, elite_frac=elite_frac, gen=gen)
@@ -359,15 +362,44 @@ def _enc(states, use_target: bool = False):
         dists.append(d)
         if d < goal_tol:
             successes += 1
-        z = dict(starts[i])
-        for h in range(horizon):
-            z = _step(z, [rng.uniform(-1.0, 1.0), rng.uniform(-1.0, 1.0), float(h % 2)], mode)
-        if abs(z["pos"] - goal_pos[i]) < goal_tol:
-            rand_successes += 1
+    # The random baseline is estimated AFTER the model arm, over
+    # `baseline_reps` rollouts per episode rather than one. It touches no
+    # encoder, no predictor and no CEM — only `_step` with uniform actions —
+    # so its cost is env steps alone and precision here is nearly free.
+    #
+    # Estimating it at one draw per episode made a scalar with a binomial
+    # standard error of ~0.10 at n=20 the thing every arm was compared
+    # against, and re-drawing it per seed injected that sampling noise into
+    # every comparison as if it were model variance. The 0.15-0.45 swing
+    # recorded in ledger node 14 needs no model explanation: five draws from
+    # Binomial(20, 0.30) have a median range of 0.25.
+    #
+    # Kept episode-paired rather than moved to an independent episode set:
+    # what a paired comparison needs is P(random succeeds | THIS episode),
+    # and per-episode rates are returned so a later paired analysis can use
+    # them as a difficulty covariate instead of re-deriving them.
+    for i in range(n_episodes):
+        hits = 0
+        for _ in range(baseline_reps):
+            z = dict(starts[i])
+            for h in range(horizon):
+                z = _step(z, [rng.uniform(-1.0, 1.0), rng.uniform(-1.0, 1.0), float(h % 2)], mode)
+            if abs(z["pos"] - goal_pos[i]) < goal_tol:
+                hits += 1
+        rand_per_episode.append(hits / baseline_reps)
+    rand_mean = float(np.mean(rand_per_episode))
+    # SE of the mean over episodes, so a reader can see the pin held rather
+    # than take it on faith.
+    rand_se = float(np.std(rand_per_episode, ddof=1) / math.sqrt(n_episodes)) \
+        if n_episodes > 1 else float("nan")
+
     return {
         "success_rate": successes / n_episodes,
         "mean_final_distance": float(np.mean(dists)),
-        "random_baseline_success": rand_successes / n_episodes,
+        "random_baseline_success": rand_mean,
+        "random_baseline_stderr": rand_se,
+        "random_baseline_reps": baseline_reps,
+        "random_baseline_per_episode": rand_per_episode,
         "frame": frame,
     }
 

WorldModel/frame_diagnostic.py

@@ -72,14 +72,24 @@
 ratio at or above VICReg arms, because better isotropy means a freer rotation.
 A LOWER SIGReg ratio falsifies the mechanism.
 
-NOTE ON CHECKPOINTS: this cannot be run from a saved checkpoint.
-`eeg_jepa.py`'s torch.save (:593) persists `encoder_state_dict` and
-`predictor_state_dict` only -- `target_encoder_state_dict` is written nowhere in
-this repo -- and `EEGJEPAModule.__init__` rebuilds the target as
-`copy.deepcopy(self.encoder)` (:327). A checkpoint-loaded model therefore has
-target_encoder == encoder EXACTLY and would report zero mismatch: a false
-negative that reads like a clean result. It must run against a live in-memory
-model, i.e. inside the process that trained it.
+NOTE ON CHECKPOINTS: run this against a live in-memory model, i.e. inside the
+process that trained it.
+
+`eeg_jepa.py` checkpoints are now format v2 and DO persist
+`target_encoder_state_dict`, so a v2 restore is exact. Format **v1** did not,
+and `EEGJEPAModule.__init__` deep-copies the target from the encoder AT
+CONSTRUCTION, before any load — so a v1 restore leaves the target at its RANDOM
+INITIALIZATION.
+
+CORRECTION (2026-08-02): an earlier version of this note, and ledger node 14,
+said a checkpoint-loaded model has `target_encoder == encoder` exactly and would
+report zero mismatch. That is wrong. Measured on a v1 restore, the target is
+byte-identical to the fresh random init (0.00e+00), 2.04 from the loaded online
+encoder and 0.39 from the trained target. `target == encoder` holds only if you
+construct and never load. A v1-restored model plans toward goal latents from an
+UNTRAINED network — not the one-frame arm, and not any condition that was ever
+run. Node 14's numbers are unaffected: `forward_eval.py` contains no
+`torch.load` and trains in-process.
 
 PROVENANCE NOTE: an earlier commit of this file attributed the motivating null
 ("prediction error 4.6x lower, rollout norm ratio 1.20 -> 0.98, control 0/35")