World-model probe — does full-transcript SFT build a better terminal model?

Phase 0 measures the agent (pass@1). This probe measures the world model: how well each checkpoint predicts terminal output it has never been trained to generate. The two questions turn out to point in opposite directions, which is the most interesting result in the project so far.

Setup

• Held-out set: 100 SWE-ZERO trajectories from 100 unseen repos (zero repo overlap and zero instance-id overlap with the 1M training subset). Sampled from the full 12M-trajectory corpus.
• Task: precondition on the first 14 (command, observation) pairs + the 15th command, then score the perplexity of the 15th observation (the real terminal output). Intuition: with 14 turns of context the repo state is heavily constrained, so this is the most determined, least-ambiguous version of “predict the next terminal output.”
• Scoring: teacher-forced HF forward pass, per-token NLL over the target observation only, prompts rendered with each checkpoint’s own chat template (byte-identical framing to its SFT). All seven checkpoints score the same 100 targets (33,520 tokens), so only the weights differ.
• Note: vllm-tpu (our pinned 0.13.x) does not implement prompt_logprobs — the feature PR is still open — so scoring runs through transformers on CPU, not vLLM.

Headline

PPL on held-out env tokens (lower = better world model). Two aggregations: token-weighted (every target token counts equally — dominated by long outputs) and median per-trajectory (the typical observation).

Rows are grouped by training arm (base anchor, then assistant-only, full-transcript, user-only). All four metrics are lower = better (↓); within each arm, the best cell per column is bolded.

Three takeaways

1. (a) forgets the world model with scale; (b) sharpens it — and only (b) has no catastrophic failures. Anchor on the pretrained base (median per-traj PPL 2.27, token-weighted 1.81 — it saw plenty of terminal text in pretraining):

• Median (the typical observation): (a) 10K (1.81) beats base, then forgets past it — 100K → 2.45, 1M → 3.33, both worse than the untrained model. (b) is ~2× better than base at every scale and still improving: 1.17 → 1.11 → 1.07.
• Token-weighted: base (1.81) is already strong; (a) diverges far above base and keeps worsening (2.01 → 2.21 → 2.58), while (b) tracks base closely (1.74 → 1.79 → 1.87 — better at 10K/100K, a hair above base only at 1M; see the length-effect nuance below).
• The tail is the sharpest tell — and only (b) has none. base carries a heavy tail: 17 of 100 held-out trajectories score PPL > 10 — all 17 are short outputs (≤ 20 tokens, where a few surprising tokens spike the per-traj PPL) — dragging the mean to 16.25; the worst is 343× on a 5-token output from jbetancur/react-data-table-component. (a) doesn’t fix this tail — it grows it: 8 → 15 → 21 trajectories PPL > 10 from 10K → 1M (worst 203 → 325). (b) has zero catastrophic trajectories at any scale (worst ~5–8, mean ~1.5). On base’s worst case (where (a) 1M is also catastrophic at 324×), (b) scores 1.07.

The mechanism is catastrophic forgetting: (a)‘s loss is on action tokens only, so env-token prediction is protected by nothing and erodes as the model specializes on actions over more optimizer steps — by 1M it degrades on 96 of 100 trajectories relative to its own 10K checkpoint and beats base on only 32 of 100. (b) keeps env tokens in the loss, so the world model is actively maintained — sharpening the median, holding token-weighted near base, and erasing the catastrophic tail that base and (a) both carry. (b) beats base on 83 of 100 trajectories at 1M; it is systematic, not outliers.

2. Pass@1 and world-model PPL point in opposite directions. From Phase 0, arm (a) is the slightly better agent (e.g. 11 vs 6 / 100 at 1M, (b) a 3-rep mean). Here, arm (b) is the dramatically better world model (3× lower median env PPL at 1M). So the two arms make opposite trades: (a) reallocates capacity to the action policy at the cost of its world model; (b) keeps a usable world model at a modest cost to the policy. For TerminalWorld’s actual goal — a model that can serve as a training environment / be used for model-based RL or planning — this is a stronger argument for ECHO-style unmasking than the pass@1 numbers ever were.

3. The pure simulator (c) confirms it’s the mask, not the data — and the predicted ordering (c) ≤ (b) < (a) holds. Arm (c) = user-only loss (loss on task + observation tokens, the exact complement of (a)) is the apples-to-apples “pure simulator”: same agent-format data and same scoring framing as (a)/(b), differing only in which tokens get loss. It lands at median 1.14 (10K) / 1.10 (100K) / 1.05 (1M) — the best world model at every scale, edging or tying (b) (1.17 / 1.11 / 1.07) and crushing (a) (1.81 / 2.45 / 3.33), with zero catastrophic trajectories like (b) at all three scales. Since (a) ⊕ (c) = (b) exactly at the token level, the three arms are a clean decomposition: putting loss on observation tokens is necessary and sufficient for the world model — (a) (no observation loss) forgets it, (c) (only observation loss) and (b) (both) build it.

The (b) ≈ (c) tie is itself the headline, and it holds across all three scales (10K/100K/1M): adding the command-loss (going (c) → (b)) costs the world model essentially nothing — at 1M (c) is even a hair ahead (1.05 vs 1.07). So (b) full-transcript buys you (c)‘s simulator quality for free while keeping a usable agent policy — the strongest case for ECHO-style unmasking. If you only want a frozen simulator (e.g. for model-based RL), (c) is the marginally-cleaner, policy-free choice. The dropped dedicated env-rewritten simulator (token-weighted 1.79) wasn’t comparable — it was scored in a different framing (observation in the assistant role) — and (c) now supersedes it. The (a)/(b)/(c) × 10K/100K/1M grid is now complete.

Nuance: (b)‘s token-weighted PPL rises slightly — why?

The token-weighted column shows (b) ticking up 1.74 → 1.87 even though the median falls. This is a length effect, not a broad regression. Bucketing (b)‘s per-trajectory change from 10K → 1M by observation length:

Short and mid outputs (the typical observation) improve with scale. The long bucket holds 92% of all scored tokens, so it determines the token-weighted headline — and within it the median is flat while the mean rises (heavy right tail: the worst trajectory goes 4.35 → 7.85, and ~5 trajectories account for 60% of the entire aggregate increase).

The mechanism is format generalizes, content memorizes:

• Short outputs are format-dominated — Observation: + a couple of grep hits, an exit code, a short listing. Structure is shared across all repos, so more training sharpens it and it transfers to unseen repos. These improve.
• Long outputs are content-dominated — a full cat file.py, a long sed range, a test log: literal repo-specific source. This is only predictable if you’ve seen that repo. On held-out unseen repos there is nothing to ground on, and sharpening content priors toward SWE-ZERO’s training repos doesn’t transfer. These degrade — but it’s a heavy tail of the most idiosyncratic outputs, not a uniform shift.

This is an inherent ceiling, not an objective failure: no volume of SWE-ZERO data lets any model emit the literal contents of a file in a repo it has never seen. The only fixes are to put the real content in context (retrieval-grounding, exactly R-WoM’s move) or accept a floor on long-output fidelity for novel repos. It is the same root cause as the Phase 0.5 simulator hallucinating plausible-but-wrong file dumps.

Next steps

The loss mask is the clean axis to push on. Partition every non-system token: assistant (commands) vs user (task + observations). Then the three arms are an exact decomposition — (a) = assistant-only, (b) = both, and the missing third point (c) = user-only:

with (a) ⊕ (c) = (b) exactly. (c) is a one-line chat-template change (the full-transcript template with the assistant markers removed) on the same agent-format data — no rewrite pipeline, no sim system prompt. It is the apples-to-apples “pure simulator” this page’s takeaway 3 is missing: scored in the same agent framing as (a)/(b), it isolates the mask from the format and tests whether the dedicated simulator’s rewrite buys anything. Predicted env-PPL ordering: (c) ≤ (b) < (a). Note (c) is a simulator, so it is graded here (env-PPL), not on SWE-bench pass@1 (which would be ~0 — it never learns to emit commands). The env-substitution / closed-loop work that previously held the “(c)” label moves to its own Phase 0.5 thread.

Reproduce

Code pinned at commit da8d12e on the terminalworld branch.

1. Build the held-out set — scripts/sample_heldout_unseen_repos.py: enumerates the 1M training subset’s repos/instance-ids, then samples 100 trajectories (≥15 turns, distinct unseen repos) from the 12M corpus (seed=42). Writes gs://marin-us-east5/heldout/sim_ppl_heldout_100.jsonl.
2. Score each checkpoint — scripts/sim_ppl_eval_hf.py --model-path <gcs-or-hf> --mode {agent,sim} [--chat-template-from <ckpt>]. Teacher-forced HF forward pass on CPU; per-token NLL over the 15th observation only. The base model uses --mode agent --chat-template-from an SFT ckpt so prompts render identically.

The scoring had to run through transformers because vllm-tpu (0.13.x) returns [None] for prompt_logprobs — confirmed by scripts/probe_vllm_logprobs.py; the feature PR is still open. The vLLM and Levanter attempts (sim_ppl_eval.py, sim_ppl_eval_levanter.py) are kept to document the dead-ends.

Artifacts: held-out set gs://marin-us-east5/heldout/sim_ppl_heldout_100.jsonl; per-checkpoint scores gs://marin-us-east5/sim-ppl-eval/sim_ppl_*.jsonl (one row per trajectory + a summary row). Checkpoint paths: see the Phase 0 reproduce table.