Merge pull request #55 from aurascoper/fix/pin-episodes-and-paired-endpoint

fix(worldmodel): pin the episode set, per-episode CEM noise, paired endpoint

Author: aurascoper

WorldModel/EXPERIMENT_LEDGER.md

@@ -272,3 +272,11 @@ Two reviewers verified the delivered synthetic A/Bs are sound (symlog threaded i
 - **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.
+
+## Corrections to node 14, part 2 (2026-08-02)
+
+- **Node 14's `success_rate` values are superseded, deliberately.** Per-episode CEM generators change the planner's sampling noise, so the recorded 32d/64d rates no longer reproduce. The EPISODE SET for model-seed 0 is unchanged — `evaluate()` previously passed `seed = model_seed + 2` and the new fixed `episode_seed` is 2 — but every seed now scores that same set, where before each scored its own. Node 14's *conclusions* are untouched: the frame ratio, the Procrustes residual and the one-frame remedy were all measured within a run, not across seeds.
+- **The baseline is now a constant.** Across 10 runs (2 latent dims x 5 seeds) `random_baseline_success` is 0.303 for every one, against 0.254-0.483 before. It is a property of the episode set, as it should be, and no longer contributes anything to arm-vs-baseline comparisons.
+- **The per-episode CEM generator is not cosmetic.** A single generator consumed sequentially made the planning noise on episode i depend on the draws taken by episodes 0..i-1, i.e. on `n_samples * cem_iters * horizon`. Any arm touching a CEM knob got different noise on the same episode. Nodes 10 and 11 swept exactly those knobs. That sweep was therefore not the paired comparison it appeared to be — it is not necessarily wrong, but its "within seed noise" reading rested on an assumption that did not hold.
+- **The continuous endpoint helps, by less than predicted.** Measured paired, ld32 vs ld64 on identical episodes and identical planner noise: mean |t| of 0.62 for the per-episode binary against 0.88 for per-episode `log(d_final/d_start)`, a factor of **1.43**, implying median n for |t|=2 of ~248 versus ~154. A power simulation had suggested nearer 3-4x. Only 3-8 of 20 pairs are discordant, which is why the binary endpoint is weak and is the mechanism the simulation got roughly right while overstating the size. Both endpoints are now returned per episode so this can be recomputed rather than re-derived.
+- **`log_distance_ratio_mean` is reported, NOT adopted as the adjudicator.** Switching the primary endpoint in the same change that repins the episode set would confound the two. Adoption is a separate pre-registered decision, and it still has to clear the blind criterion: separate zero from random from MSE on existing data.

WorldModel/forward_eval.py

@@ -319,8 +319,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, baseline_reps=100,
+                 cem_iters=3, n_samples=64, elite_frac=0.2, mode="signal",
+                 episode_seed=2, 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
@@ -329,9 +329,13 @@ def _mpc_success(model, train_mean, train_std, device, *, n_episodes=20, horizon
     true final pos reached the goal. A random-action policy is the baseline — the
     JEPA-planned success rate beating it is the signal that the latent dynamics are
     useful for control. Returns success_rate, mean_final_distance, random_baseline."""
-    rng = random.Random(seed)
-    gen = torch.Generator()
-    gen.manual_seed(seed + 7)
+    # The episode set, the baseline draws and the planner's sampling noise all
+    # derive from `episode_seed` and NOTHING else. Previously `evaluate()` passed
+    # `seed = model_seed + 2`, so every model seed scored a DIFFERENT episode set:
+    # "which episodes" was confounded with "which model", and the between-set
+    # difficulty spread (sd 0.0853, measured) was being read as model variance.
+    # The benchmark's episodes are a property of the benchmark, not of the run.
+    rng = random.Random(episode_seed)
     starts = [_sample_latent(rng) for _ in range(n_episodes)]
     goal_pos = [starts[i]["pos"] + rng.uniform(-goal_offset, goal_offset) for i in range(n_episodes)]
     goals = [{**starts[i], "pos": goal_pos[i]} for i in range(n_episodes)]
@@ -352,7 +356,18 @@ def _enc(states, use_target: bool = False):
 
     successes, dists = 0, []
     rand_per_episode: list[float] = []
+    start_dists: list[float] = []
+    log_ratios: list[float] = []
     for i in range(n_episodes):
+        # A generator PER EPISODE, seeded from the episode identity alone. A
+        # single generator consumed sequentially made the planning noise on
+        # episode i depend on how many draws episodes 0..i-1 had taken, which is
+        # a function of n_samples * cem_iters * horizon. Any arm that touched a
+        # CEM knob therefore got different noise on the SAME episode, silently
+        # defeating pairing even with the episode set fixed. Nothing in the panel
+        # would have revealed it.
+        gen = torch.Generator()
+        gen.manual_seed(episode_seed * 1_000_003 + i)
         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)
         z = dict(starts[i])
@@ -362,6 +377,24 @@ def _enc(states, use_target: bool = False):
         dists.append(d)
         if d < goal_tol:
             successes += 1
+        # Paired continuous endpoint: log(d_final / d_start), per episode.
+        #
+        # `success_rate` thresholds a continuous distance at goal_tol=0.15 and
+        # throws away everything else. Simulating a 10% distance improvement
+        # with shared episode difficulty, power at n=20 is 17.7% for a paired
+        # continuous statistic against 0.5% for paired binary; at n=200 it is
+        # 85.1% against 23.4%. The distance is already computed and returned --
+        # it simply was not the adjudicator.
+        #
+        # The RATIO, not the difference: episodes differ enormously in how far
+        # the goal starts, so a raw difference is dominated by episode
+        # difficulty. The log makes "halved the distance" the same effect size
+        # wherever it happens, which is what pairing needs. Both terms are
+        # clamped because goal_offset can draw a start essentially on top of
+        # its goal, and log(0) would take the whole panel with it.
+        d0 = max(abs(starts[i]["pos"] - goal_pos[i]), 1e-6)
+        start_dists.append(d0)
+        log_ratios.append(math.log(max(d, 1e-6) / d0))
     # 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 —
@@ -400,6 +433,18 @@ def _enc(states, use_target: bool = False):
         "random_baseline_stderr": rand_se,
         "random_baseline_reps": baseline_reps,
         "random_baseline_per_episode": rand_per_episode,
+        # Primary candidate for the repaired benchmark. Reported alongside
+        # success_rate rather than replacing it -- adopting it as THE adjudicator
+        # is a separate, pre-registered decision, and switching endpoints while
+        # also changing the episode set would confound the two.
+        # Per-episode finals, so a downstream paired analysis can compute EITHER
+        # endpoint exactly instead of approximating the binary one from the two
+        # marginal rates. Without this the paired binary sd has to be guessed.
+        "final_distance_per_episode": dists,
+        "log_distance_ratio_mean": float(np.mean(log_ratios)),
+        "log_distance_ratio_per_episode": log_ratios,
+        "start_distance_per_episode": start_dists,
+        "episode_seed": episode_seed,
         "frame": frame,
     }
 
@@ -463,7 +508,7 @@ def evaluate(
         )
         mpc = _mpc_success(
             model, dataset.mean, dataset.std, device,
-            n_episodes=mpc_episodes, horizon=mpc_horizon, mode=mode, seed=seed + 2,
+            n_episodes=mpc_episodes, horizon=mpc_horizon, mode=mode,
             cem_iters=mpc_cem_iters, n_samples=mpc_n_samples, elite_frac=mpc_elite_frac,
             log_features=log_features, log_epsilon=log_epsilon, symlog=symlog,
         )